Coverage for src/utils/data.py: 69%

2# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.

3# Direct all questions to legal@cloudzero.com

5"""

6Data loading and schema utilities for CloudZero Data Tool.

7Handles CSV loading, schema detection, caching, and optimization.

8"""

10import json

11import sys

12import threading

13import time

14from datetime import datetime

15from pathlib import Path

16from typing import Dict, Any

18import polars as pl

19from rich.console import Console

20from rich.progress import (

21 Progress,

22 SpinnerColumn,

23 TextColumn,

24 BarColumn,

25 MofNCompleteColumn,

26 TimeElapsedColumn,

27)

29console = Console()

32def polars_dtype_to_string(dtype) -> str:

33 """Convert Polars dtype to string representation for JSON serialization."""

34 if dtype == pl.String:

35 return "String"

36 elif dtype == pl.Int64:

37 return "Int64"

38 elif dtype == pl.Float64:

39 return "Float64"

40 elif dtype == pl.Boolean:

41 return "Boolean"

42 elif dtype == pl.Date:

43 return "Date"

44 elif hasattr(dtype, "time_unit"): # Datetime type

45 return "Datetime"

46 else:

47 return str(dtype)

50def string_to_polars_dtype(dtype_str: str):

51 """Convert string representation back to Polars dtype."""

52 mapping = {

53 "String": pl.String,

54 "Utf8": pl.String, # Backwards compatibility

55 "Int64": pl.Int64,

56 "Float64": pl.Float64,

57 "Boolean": pl.Boolean,

58 "Date": pl.Date,

59 "Datetime": pl.Datetime(

60 time_unit="us"

61 ), # Use microsecond precision for datetime parsing

62 }

63 return mapping.get(dtype_str, pl.String) # Default to String if unknown

66def detect_csv_schema(file_path: Path) -> Dict[str, Any]:

67 """Detect the complete schema of a CSV file by reading all data."""

68 console.print(f"🔍 [cyan]Detecting schema for {file_path.name}...[/cyan]")

70 file_size = file_path.stat().st_size

71 file_size_mb = file_size / (1024 * 1024)

73 try:

74 # Read with comprehensive schema detection

75 df = pl.read_csv(

76 file_path,

77 infer_schema_length=None, # Read all rows for complete schema detection

78 null_values=["", "null", "NULL", "None"],

79 try_parse_dates=True,

80 ignore_errors=True,

81 low_memory=False,

82 )

84 # Extract schema information with special handling for datetime columns

85 schema = {}

86 for col_name in df.columns:

87 dtype = df[col_name].dtype

88 dtype_string = polars_dtype_to_string(dtype)

90 # Special check for columns that should be datetime but were detected as date

91 if dtype_string == "Date":

92 # Check if the column actually contains time information

93 sample_values = df[col_name].drop_nulls().head(10)

94 has_time_component = False

95 for value in sample_values:

96 value_str = str(value)

97 if " " in value_str and ":" in value_str: # Contains time component

98 has_time_component = True

99 break

100

101 # Also check column name patterns that indicate datetime

102 datetime_patterns = ["CREATED_AT", "_AT_", "TIMESTAMP", "DATETIME"]

103 if has_time_component or any(

104 pattern in col_name.upper() for pattern in datetime_patterns

105 ):

106 dtype_string = "Datetime"

107

108 schema[col_name] = dtype_string

109

110 return {

111 "file_name": file_path.name,

112 "file_size_bytes": file_size,

113 "file_size_mb": round(file_size_mb, 2),

114 "num_rows": len(df),

115 "num_columns": len(df.columns),

116 "schema": schema,

117 "generated_at": datetime.now().isoformat(),

118 }

119

120 except Exception as e:

121 console.print(f"❌ [red]Error detecting schema for {file_path.name}: {e}[/red]")

122 return None

123

124

125def generate_schema_manifest(folder_path: Path) -> None:

126 """Generate a schema manifest file for all CSV files in the given folder."""

127 if not folder_path.exists() or not folder_path.is_dir():

128 console.print(f"❌ [red]Error: {folder_path} is not a valid directory[/red]")

129 return

130

131 # Find all CSV files in the folder

132 csv_files = list(folder_path.glob("*.csv"))

133

134 if not csv_files:

135 console.print(f"📁 [yellow]No CSV files found in {folder_path}[/yellow]")

136 return

137

138 console.print(f"📁 [cyan]Found {len(csv_files)} CSV files in {folder_path}[/cyan]")

139

140 manifest = {

141 "generated_at": datetime.now().isoformat(),

142 "folder_path": str(folder_path),

143 "schemas": {},

144 }

145

146 # Process each CSV file

147 for csv_file in csv_files:

148 # Skip files that start with "customer-" prefix as they might be large/sensitive

149 if csv_file.name.startswith("customer-"):

150 console.print(

151 f"⚠️ [yellow]Skipping customer data file: {csv_file.name}[/yellow]"

152 )

153 continue

154

155 schema_info = detect_csv_schema(csv_file)

156 if schema_info:

157 manifest["schemas"][csv_file.name] = schema_info

158 console.print(f"✅ [green]Processed schema for {csv_file.name}[/green]")

159 else:

160 console.print(f"❌ [red]Failed to process {csv_file.name}[/red]")

161

162 # Write manifest to file

163 manifest_path = folder_path / "manifest.json"

164 try:

165 with open(manifest_path, "w") as f:

166 json.dump(manifest, f, indent=2)

167 console.print(f"📄 [green]Schema manifest saved to {manifest_path}[/green]")

168 console.print(f"📊 [cyan]Processed {len(manifest['schemas'])} files[/cyan]")

169 except Exception as e:

170 console.print(f"❌ [red]Error saving manifest: {e}[/red]")

171

172

173def load_manifest_schema(file_path: Path) -> Dict[str, Any]:

174 """Load schema from manifest file if available, generate if missing."""

175 manifest_path = file_path.parent / "manifest.json"

176

177 if not manifest_path.exists():

178 console.print(

179 "\nℹ️ [yellow]Schema manifest not found. Schema detection is required before analysis.[/yellow]"

180 )

181 console.print(

182 "🔍 [cyan]Generating schema manifest to optimize future data loading...[/cyan]"

183 )

184

185 # Auto-generate schema for the directory

186 generate_schema_manifest(file_path.parent)

187

188 # Try loading again after generation

189 if not manifest_path.exists():

190 console.print(

191 "⚠️ [red]Schema generation failed. Continuing with automatic schema detection.[/red]"

192 )

193 return None

194

195 try:

196 with open(manifest_path, "r") as f:

197 manifest = json.load(f)

198

199 file_name = file_path.name

200 if file_name in manifest.get("schemas", {}):

201 schema_info = manifest["schemas"][file_name]

202

203 # Convert string dtypes back to Polars dtypes

204 polars_schema = {}

205 for col_name, dtype_str in schema_info["schema"].items():

206 polars_schema[col_name] = string_to_polars_dtype(dtype_str)

207

208 console.print(

209 f"📋 [green]Using cached schema for {file_name} from manifest[/green]"

210 )

211 return polars_schema

212

213 except Exception as e:

214 console.print(f"⚠️ [yellow]Error loading manifest schema: {e}[/yellow]")

215

216 return None

217

218

219def get_required_columns(

220 rollup_spec: str = None,

221 column: str = None,

222 analysis: bool = False,

223 group_by_time: str = None,

224) -> list:

225 """Determine the minimum set of columns needed for the requested analysis."""

226 required_cols = set()

227

228 if column:

229 required_cols.add(column)

230

231 if rollup_spec:

232 # Parse rollup spec to extract column names

233 if ":" in rollup_spec:

234 group_part, agg_part = rollup_spec.split(":", 1)

235 else:

236 group_part = rollup_spec

237 agg_part = ""

238

239 # Add group columns

240 group_cols = [col.strip() for col in group_part.split(",")]

241 required_cols.update(group_cols)

242

243 # Extract columns from aggregation functions

244 if agg_part:

245 import re

246

247 agg_pattern = r"(count|sum|avg|min|max)\(([^)]+)\)"

248 matches = re.findall(agg_pattern, agg_part)

249 for _, col_name in matches:

250 col_name = col_name.strip()

251 if col_name != "*": # count(*) doesn't require specific column

252 required_cols.add(col_name)

253

254 if group_by_time:

255 # Add common datetime columns that might be needed

256 datetime_cols = [

257 "CREATED_AT_UTC",

258 "CREATED_AT_PT",

259 "DATE_PT",

260 "TIMESTAMP",

261 "CREATED_AT",

262 ]

263 required_cols.update(datetime_cols)

264

265 if analysis:

266 # For full analysis, we need all columns, so return None to indicate no filtering

267 return None

268

269 return list(required_cols) if required_cols else None

270

271

272def load_data_optimized(

273 file_path: Path, required_cols: list = None, use_lazy: bool = False

274) -> pl.DataFrame:

275 """Optimized data loading that only reads required columns."""

276 file_size = file_path.stat().st_size

277 file_size_mb = file_size / (1024 * 1024)

278

279 # Show optimization info

280 if required_cols:

281 console.print(

282 f"🚀 [green]Optimizing for {len(required_cols)} columns: {', '.join(required_cols[:5])}{'...' if len(required_cols) > 5 else ''}[/green]"

283 )

284

285 # Auto-enable optimizations for large files

286 if file_size_mb > 1000: # > 1GB

287 console.print(

288 f"📈 [yellow]Large file detected ({file_size_mb:.1f}MB), using optimized loading...[/yellow]"

289 )

290 use_lazy = True

291

292 if use_lazy:

293 # Try to load schema from manifest first

294 manifest_schema = load_manifest_schema(file_path)

295

296 # Use lazy evaluation for large files

297 df = pl.scan_csv(

298 file_path,

299 null_values=["", "null", "NULL", "None"],

300 try_parse_dates=True,

301 low_memory=False,

302 schema_overrides=manifest_schema if manifest_schema else None,

303 infer_schema_length=10000,

304 )

305

306 # Select only required columns before materializing

307 if required_cols:

308 # Use collect_schema() to avoid performance warning

309 schema_names = df.collect_schema().names()

310 available_cols = [col for col in required_cols if col in schema_names]

311 if available_cols:

312 df = df.select(available_cols)

313 console.print(

314 f"🎯 [green]Selected {len(available_cols)} of {len(required_cols)} requested columns[/green]"

315 )

316

317 return df.collect()

318 else:

319 return load_data(file_path, use_lazy=use_lazy, columns=required_cols)

320

321

322def load_data(

323 file_path: Path, use_lazy: bool = False, columns: list = None

324) -> pl.DataFrame:

325 """Load CSV data with progress tracking and error handling."""

326 if not file_path.exists():

327 console.print(f"❌ [red]Error: File {file_path} does not exist[/red]")

328 sys.exit(1)

329

330 file_size = file_path.stat().st_size

331 file_size_mb = file_size / (1024 * 1024)

332

333 # Progress tracking

334 progress_thread = None

335

336 with Progress(

337 SpinnerColumn(),

338 TextColumn(f"[green]Loading {file_size_mb:.1f}MB CSV from {file_path.name}..."),

339 BarColumn(),

340 MofNCompleteColumn(),

341 TimeElapsedColumn(),

342 console=console,

343 ) as progress:

344 task = progress.add_task("Loading...", total=100)

345

346 # Start progress animation in background

347 def update_progress():

348 while not progress.tasks[task].finished:

349 progress.update(task, advance=1)

350 time.sleep(0.1)

351

352 if file_size_mb > 100: # Only show progress for large files

353 progress_thread = threading.Thread(target=update_progress)

354 progress_thread.daemon = True

355 progress_thread.start()

356

357 try:

358 # Try to load schema from manifest first

359 manifest_schema = load_manifest_schema(file_path)

360

361 if manifest_schema:

362 # Use cached schema from manifest with optional column selection

363 read_columns = columns if columns else None

364

365 if use_lazy:

366 df = pl.scan_csv(

367 file_path,

368 null_values=["", "null", "NULL", "None"],

369 try_parse_dates=True,

370 low_memory=False,

371 schema_overrides=manifest_schema,

372 infer_schema_length=10000,

373 ignore_errors=True,

374 )

375 if read_columns:

376 df = df.select(read_columns)

377 df = df.collect()

378 else:

379 df = pl.read_csv(

380 file_path,

381 null_values=["", "null", "NULL", "None"],

382 try_parse_dates=True,

383 low_memory=False,

384 rechunk=True,

385 schema_overrides=manifest_schema,

386 columns=read_columns,

387 infer_schema_length=10000,

388 ignore_errors=True,

389 )

390 else:

391 # No cached schema available, try optimized parsing with fallback

392 read_columns = columns if columns else None

393

394 try:

395 if use_lazy:

396 df = pl.scan_csv(

397 file_path,

398 null_values=["", "null", "NULL", "None"],

399 try_parse_dates=True,

400 low_memory=False,

401 columns=read_columns,

402 ).collect()

403 else:

404 df = pl.read_csv(

405 file_path,

406 null_values=["", "null", "NULL", "None"],

407 try_parse_dates=True,

408 low_memory=False,

409 rechunk=True,

410 columns=read_columns,

411 )

412 except Exception:

413 # Fallback to more robust parsing if initial attempt fails

414 console.print(

415 "⚠️ [yellow]Initial parsing failed, using robust mode...[/yellow]"

416 )

417 if use_lazy:

418 df = pl.scan_csv(

419 file_path,

420 null_values=["", "null", "NULL", "None"],

421 try_parse_dates=True,

422 infer_schema_length=10000,

423 ignore_errors=True,

424 low_memory=False,

425 columns=read_columns,

426 ).collect()

427 else:

428 df = pl.read_csv(

429 file_path,

430 null_values=["", "null", "NULL", "None"],

431 try_parse_dates=True,

432 infer_schema_length=10000,

433 ignore_errors=True,

434 low_memory=False,

435 rechunk=True,

436 columns=read_columns,

437 )

438

439 progress.update(task, completed=True)

440

441 except Exception as e:

442 progress.update(task, completed=True)

443 console.print(f"❌ [red]Error loading data: {e}[/red]")

444 sys.exit(1)

445

446 console.print(

447 f"✅ [green]Loaded {len(df):,} rows, {len(df.columns)} columns[/green]"

448 )

449 return df