Coverage for src/main_old.py: 0%

612 statements  

« prev     ^ index     » next       coverage.py v7.8.2, created at 2025-05-25 13:51 -0400

1# Copyright (c) CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL 

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

3# Direct all questions to legal@cloudzero.com 

4 

5#!/usr/bin/env python3 

6""" 

7Data analysis script for Perplexity CSV data using Polars. 

8Provides comprehensive analysis including business insights, performance metrics, and cost patterns. 

9""" 

10 

11import argparse 

12import sys 

13from pathlib import Path 

14import polars as pl 

15from rich.console import Console 

16from rich.table import Table 

17from rich.panel import Panel 

18from rich.progress import ( 

19 Progress, 

20 SpinnerColumn, 

21 TextColumn, 

22 BarColumn, 

23 MofNCompleteColumn, 

24 TimeElapsedColumn, 

25) 

26from rich import box 

27import re 

28import threading 

29import time 

30import json 

31from datetime import datetime 

32from typing import Dict, Any 

33 

34 

35console = Console() 

36 

37 

38def detect_datetime_column(df: pl.DataFrame, datetime_col: str = None) -> str: 

39 """Detect the most appropriate datetime column in the DataFrame. 

40 

41 Args: 

42 df: DataFrame to analyze 

43 datetime_col: Specific column name to use, or None for auto-detection 

44 

45 Returns: 

46 Column name to use for datetime grouping 

47 """ 

48 if datetime_col: 

49 if datetime_col in df.columns: 

50 return datetime_col 

51 else: 

52 console.print( 

53 f"ERROR [red]Error: Specified datetime column '{datetime_col}' not found[/red]" 

54 ) 

55 console.print( 

56 f"INFO [cyan]Available columns: {', '.join(df.columns)}[/cyan]" 

57 ) 

58 return None 

59 

60 # Auto-detect datetime column by analyzing column types and names 

61 datetime_candidates = [] 

62 

63 # First, prioritize columns with Datetime type that have time information 

64 for col in df.columns: 

65 col_dtype = df[col].dtype 

66 if ( 

67 str(col_dtype).startswith("Datetime") 

68 or "datetime" in str(col_dtype).lower() 

69 ): 

70 # Check if it actually contains time information (not just date) 

71 sample_values = df[col].drop_nulls().head(5) 

72 has_time_info = False 

73 for value in sample_values: 

74 value_str = str(value) 

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

76 has_time_info = True 

77 break 

78 

79 if has_time_info: 

80 datetime_candidates.append( 

81 (col, "datetime_type", 100) 

82 ) # Highest priority 

83 else: 

84 datetime_candidates.append( 

85 (col, "datetime_type", 50) 

86 ) # Medium priority 

87 

88 # Then look for columns with datetime-like names 

89 datetime_name_patterns = [ 

90 (r".*created.*at.*utc.*", 90), # created_at_utc gets high priority 

91 (r".*created.*at.*", 80), # created_at gets high priority 

92 (r".*timestamp.*", 70), # timestamp columns 

93 (r".*time.*", 60), # time columns 

94 (r".*date.*", 40), # date columns (lower priority) 

95 ] 

96 

97 for col in df.columns: 

98 col_lower = col.lower() 

99 for pattern, priority in datetime_name_patterns: 

100 if re.match(pattern, col_lower): 

101 datetime_candidates.append((col, "name_match", priority)) 

102 break 

103 

104 if not datetime_candidates: 

105 console.print("WARNING [yellow]Warning: No datetime columns detected[/yellow]") 

106 return None 

107 

108 # Sort by priority (highest first) and select the best candidate 

109 datetime_candidates.sort(key=lambda x: x[2], reverse=True) 

110 selected_col = datetime_candidates[0][0] 

111 

112 console.print(f"TIME [green]Auto-detected datetime column: {selected_col}[/green]") 

113 return selected_col 

114 

115 

116def polars_dtype_to_string(dtype) -> str: 

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

118 dtype_str = str(dtype) 

119 if dtype_str.startswith("String") or dtype_str.startswith("Utf8"): 

120 return "String" 

121 elif dtype_str.startswith("Int"): 

122 return "Int64" 

123 elif dtype_str.startswith("Float"): 

124 return "Float64" 

125 elif dtype_str.startswith("Boolean"): 

126 return "Boolean" 

127 elif dtype_str.startswith("Datetime") or "datetime" in dtype_str.lower(): 

128 return "Datetime" 

129 elif dtype_str.startswith("Date") and "datetime" not in dtype_str.lower(): 

130 return "Date" 

131 else: 

132 return dtype_str 

133 

134 

135def string_to_polars_dtype(dtype_str: str): 

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

137 mapping = { 

138 "String": pl.String, 

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

140 "Int64": pl.Int64, 

141 "Float64": pl.Float64, 

142 "Boolean": pl.Boolean, 

143 "Date": pl.Date, 

144 "Datetime": pl.Datetime( 

145 time_unit="us" 

146 ), # Use microsecond precision for datetime parsing 

147 } 

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

149 

150 

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

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

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

154 

155 file_size = file_path.stat().st_size 

156 file_size_mb = file_size / (1024 * 1024) 

157 

158 try: 

159 # Read with comprehensive schema detection 

160 df = pl.read_csv( 

161 file_path, 

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

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

164 try_parse_dates=True, 

165 ignore_errors=True, 

166 low_memory=False, 

167 ) 

168 

169 # Extract schema information with special handling for datetime columns 

170 schema = {} 

171 for col_name in df.columns: 

172 dtype = df[col_name].dtype 

173 dtype_string = polars_dtype_to_string(dtype) 

174 

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

176 if dtype_string == "Date": 

177 # Check if the column actually contains time information 

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

179 has_time_component = False 

180 for value in sample_values: 

181 value_str = str(value) 

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

183 has_time_component = True 

184 break 

185 

186 # Also check column name patterns that indicate datetime 

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

188 if has_time_component or any( 

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

190 ): 

191 dtype_string = "Datetime" 

192 

193 schema[col_name] = dtype_string 

194 

195 return { 

196 "file_name": file_path.name, 

197 "file_size_bytes": file_size, 

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

199 "num_rows": len(df), 

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

201 "schema": schema, 

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

203 } 

204 

205 except Exception as e: 

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

207 return None 

208 

209 

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

211 """Generate schema manifest for all CSV files in the specified folder.""" 

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

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

214 return 

215 

216 # Find all CSV files 

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

218 if not csv_files: 

219 console.print(f"⚠️ [yellow]No CSV files found in {folder_path}[/yellow]") 

220 return 

221 

222 console.print( 

223 f"📁 [green]Found {len(csv_files)} CSV files in {folder_path}[/green]" 

224 ) 

225 

226 manifest = { 

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

228 "folder_path": str(folder_path), 

229 "schemas": {}, 

230 } 

231 

232 with Progress( 

233 SpinnerColumn(), 

234 TextColumn("[progress.description]{task.description}"), 

235 BarColumn(), 

236 TextColumn("{task.completed}/{task.total}"), 

237 console=console, 

238 ) as progress: 

239 task = progress.add_task("Generating schemas...", total=len(csv_files)) 

240 

241 for csv_file in csv_files: 

242 schema_info = detect_csv_schema(csv_file) 

243 if schema_info: 

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

245 console.print( 

246 f"✅ [green]{csv_file.name}: {schema_info['num_rows']:,} rows, {schema_info['num_columns']} cols[/green]" 

247 ) 

248 else: 

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

250 

251 progress.advance(task) 

252 

253 # Write manifest file 

254 manifest_path = folder_path / "manifest.json" 

255 try: 

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

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

258 

259 console.print(f"\n🎉 [green]Schema manifest written to {manifest_path}[/green]") 

260 console.print( 

261 f"📊 [cyan]Processed {len(manifest['schemas'])} files successfully[/cyan]" 

262 ) 

263 

264 except Exception as e: 

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

266 

267 

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

269 """Load schema from manifest file if available.""" 

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

271 

272 if not manifest_path.exists(): 

273 return None 

274 

275 try: 

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

277 manifest = json.load(f) 

278 

279 file_name = file_path.name 

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

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

282 

283 # Convert string dtypes back to Polars dtypes 

284 polars_schema = {} 

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

286 polars_schema[col_name] = string_to_polars_dtype(dtype_str) 

287 

288 console.print( 

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

290 ) 

291 return polars_schema 

292 

293 except Exception as e: 

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

295 

296 return None 

297 

298 

299def get_required_columns( 

300 rollup_spec: str = None, 

301 column: str = None, 

302 full_analysis: bool = False, 

303 group_by_time: str = None, 

304) -> list: 

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

306 required_cols = set() 

307 

308 if rollup_spec: 

309 group_cols, agg_specs = parse_rollup_spec(rollup_spec) 

310 required_cols.update(group_cols) 

311 if agg_specs: 

312 required_cols.update(agg_specs.keys()) 

313 else: 

314 # Default aggregation columns 

315 default_cols = [ 

316 "TOTAL_COST", 

317 "CT_PROMPT_TOKENS", 

318 "CT_TOTAL_COMPLETION_TOKENS", 

319 "NUM_SEARCH_QUERIES", 

320 ] 

321 required_cols.update(default_cols) 

322 

323 if column: 

324 required_cols.add(column) 

325 

326 if group_by_time: 

327 # Add common datetime columns when time grouping is needed 

328 datetime_cols = ["DATE_PT", "CREATED_AT_UTC", "CREATED_AT_PT"] 

329 required_cols.update(datetime_cols) 

330 

331 if full_analysis: 

332 # Add columns needed for comprehensive analysis 

333 analysis_cols = [ 

334 "RUNTIME_MODEL_NAME", 

335 "TOTAL_TIME", 

336 "TOTAL_DECODING_LATENCY", 

337 "PRICE_PER_INPUT_TOKEN", 

338 "PRICE_PER_OUTPUT_TOKEN", 

339 "TOTAL_COST", 

340 "API_KEY_SUFFIX", 

341 "MODEL", 

342 "NUM_SEARCH_QUERIES", 

343 "IS_PERPLEXITY_INTERNAL", 

344 "CT_PROMPT_TOKENS", 

345 "CT_TOTAL_COMPLETION_TOKENS", 

346 ] 

347 required_cols.update(analysis_cols) 

348 

349 return list(required_cols) if required_cols else None 

350 

351 

352def load_data_optimized( 

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

354) -> pl.DataFrame: 

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

356 # Get file size for optimization decisions 

357 file_size = file_path.stat().st_size 

358 file_size_mb = file_size / (1024 * 1024) 

359 

360 # For small files (< 100MB), just load everything normally 

361 if file_size_mb < 100: 

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

363 

364 # For large files, use optimized strategies 

365 console.print( 

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

367 ) 

368 

369 # Try to load manifest schema first for better performance 

370 manifest_schema = load_manifest_schema(file_path) 

371 

372 if use_lazy or file_size_mb > 1000: # Use lazy for very large files 

373 if manifest_schema: 

374 df = pl.scan_csv( 

375 file_path, 

376 schema_overrides=manifest_schema, 

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

378 try_parse_dates=True, 

379 low_memory=False, 

380 ) 

381 else: 

382 df = pl.scan_csv( 

383 file_path, 

384 infer_schema_length=1000, 

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

386 try_parse_dates=True, 

387 low_memory=False, 

388 ) 

389 

390 # Select only required columns before materializing 

391 if required_cols: 

392 # Use collect_schema() to avoid performance warning 

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

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

395 if available_cols: 

396 df = df.select(available_cols) 

397 console.print( 

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

399 ) 

400 

401 return df.collect() 

402 else: 

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

404 

405 

406def load_data( 

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

408) -> pl.DataFrame: 

409 """Load CSV data using Polars with optimizations for large files.""" 

410 try: 

411 # Get file size for progress tracking 

412 file_size = file_path.stat().st_size 

413 file_size_mb = file_size / (1024 * 1024) 

414 

415 with Progress( 

416 SpinnerColumn(), 

417 TextColumn("[progress.description]{task.description}"), 

418 BarColumn(), 

419 MofNCompleteColumn(), 

420 TextColumn("•"), 

421 TimeElapsedColumn(), 

422 console=console, 

423 ) as progress: 

424 # Create progress task with file size 

425 task = progress.add_task( 

426 f"Loading {file_size_mb:.1f}MB CSV from {file_path.name}...", total=100 

427 ) 

428 

429 # Start a background thread to simulate progress based on time 

430 # Since Polars doesn't provide read callbacks, we estimate based on typical read speeds 

431 stop_progress = threading.Event() 

432 

433 def update_progress(): 

434 """Update progress based on estimated read time.""" 

435 # Estimate read time based on file size (rough heuristic: ~50MB/sec for CSV parsing) 

436 estimated_seconds = max(1.0, file_size_mb / 50.0) 

437 update_interval = estimated_seconds / 90 # Update 90 times during read 

438 

439 progress_value = 0 

440 while not stop_progress.is_set() and progress_value < 90: 

441 time.sleep(update_interval) 

442 progress_value += 1 

443 progress.update(task, completed=progress_value) 

444 

445 # Start progress thread 

446 progress_thread = threading.Thread(target=update_progress, daemon=True) 

447 progress_thread.start() 

448 

449 try: 

450 # Try to load schema from manifest first 

451 manifest_schema = load_manifest_schema(file_path) 

452 

453 if manifest_schema: 

454 # Use cached schema from manifest with optional column selection 

455 read_columns = columns if columns else None 

456 

457 if use_lazy: 

458 df = pl.scan_csv( 

459 file_path, 

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

461 try_parse_dates=True, 

462 low_memory=False, 

463 schema_overrides=manifest_schema, 

464 infer_schema_length=10000, 

465 ignore_errors=True, 

466 ) 

467 if read_columns: 

468 df = df.select(read_columns) 

469 df = df.collect() 

470 else: 

471 df = pl.read_csv( 

472 file_path, 

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

474 try_parse_dates=True, 

475 low_memory=False, 

476 rechunk=True, 

477 schema_overrides=manifest_schema, 

478 columns=read_columns, 

479 infer_schema_length=10000, 

480 ignore_errors=True, 

481 ) 

482 else: 

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

484 read_columns = columns if columns else None 

485 

486 try: 

487 if use_lazy: 

488 # Lazy loading - only materializes when needed 

489 df = pl.scan_csv( 

490 file_path, 

491 infer_schema_length=1000, # Reduced from 10000 

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

493 try_parse_dates=True, 

494 low_memory=False, # Use more memory for faster processing 

495 ) 

496 if read_columns: 

497 df = df.select(read_columns) 

498 df = df.collect() 

499 else: 

500 # Eager loading with optimizations 

501 df = pl.read_csv( 

502 file_path, 

503 infer_schema_length=1000, # Reduced from 10000 

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

505 try_parse_dates=True, 

506 low_memory=False, # Use more memory for faster processing 

507 rechunk=True, # Rechunk for better memory layout 

508 columns=read_columns, 

509 ) 

510 except Exception: 

511 # Fallback to more robust parsing if initial attempt fails 

512 console.print( 

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

514 ) 

515 if use_lazy: 

516 df = pl.scan_csv( 

517 file_path, 

518 infer_schema_length=5000, # Increased for better type detection 

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

520 try_parse_dates=True, 

521 ignore_errors=True, # More lenient parsing 

522 low_memory=False, 

523 ).collect() 

524 else: 

525 df = pl.read_csv( 

526 file_path, 

527 infer_schema_length=5000, # Increased for better type detection 

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

529 try_parse_dates=True, 

530 ignore_errors=True, # More lenient parsing 

531 low_memory=False, 

532 rechunk=True, 

533 ) 

534 finally: 

535 # Stop progress thread and complete 

536 stop_progress.set() 

537 progress.update(task, completed=100) 

538 

539 console.print( 

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

541 ) 

542 return df 

543 except Exception as e: 

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

545 sys.exit(1) 

546 

547 

548def show_column_info(df: pl.DataFrame) -> None: 

549 """Display column names and basic info.""" 

550 table = Table(title="📊 Column Information", box=box.ROUNDED) 

551 table.add_column("#", style="cyan", no_wrap=True) 

552 table.add_column("Column Name", style="magenta") 

553 table.add_column("Data Type", style="green") 

554 table.add_column("Null Count", style="yellow", justify="right") 

555 

556 for i, col in enumerate(df.columns, 1): 

557 dtype = str(df[col].dtype) 

558 null_count = df[col].null_count() 

559 null_display = f"{null_count:,}" if null_count > 0 else "0" 

560 null_style = "red" if null_count > 0 else "green" 

561 

562 table.add_row( 

563 str(i), col, dtype, f"[{null_style}]{null_display}[/{null_style}]" 

564 ) 

565 

566 console.print(table) 

567 

568 

569def get_column_choice(df: pl.DataFrame) -> str: 

570 """Interactive column selection.""" 

571 while True: 

572 try: 

573 choice = console.input( 

574 f"\n💡 [cyan]Enter column number (1-{len(df.columns)}) or column name:[/cyan] " 

575 ).strip() 

576 

577 # Try as number first 

578 if choice.isdigit(): 

579 idx = int(choice) - 1 

580 if 0 <= idx < len(df.columns): 

581 return df.columns[idx] 

582 

583 # Try as column name 

584 if choice in df.columns: 

585 return choice 

586 

587 console.print("❌ [red]Invalid choice. Please try again.[/red]") 

588 

589 except KeyboardInterrupt: 

590 console.print("\n👋 [yellow]Exiting...[/yellow]") 

591 sys.exit(0) 

592 

593 

594def analyze_column(df: pl.DataFrame, column: str) -> None: 

595 """Analyze distinct values in the selected column.""" 

596 try: 

597 # Get distinct values and their counts 

598 value_counts = ( 

599 df.group_by(column) 

600 .agg(pl.len().alias("count")) 

601 .sort("count", descending=True) 

602 ) 

603 

604 total_distinct = len(value_counts) 

605 total_rows = len(df) 

606 

607 # Create summary panel 

608 summary_text = f"""📈 [cyan]Total distinct values:[/cyan] [yellow]{total_distinct:,}[/yellow] 

609📊 [cyan]Total rows:[/cyan] [yellow]{total_rows:,}[/yellow] 

610🎯 [cyan]Uniqueness:[/cyan] [yellow]{total_distinct / total_rows:.2%}[/yellow]""" 

611 

612 console.print( 

613 Panel( 

614 summary_text, 

615 title=f"🔍 Analyzing Column: [magenta]{column}[/magenta]", 

616 box=box.ROUNDED, 

617 border_style="blue", 

618 ) 

619 ) 

620 

621 # Create table for top values 

622 table = Table(title="🏆 Top 20 Most Frequent Values", box=box.ROUNDED) 

623 table.add_column("Value", style="cyan", no_wrap=True, max_width=25) 

624 table.add_column("Count", style="green", justify="right") 

625 table.add_column("Percentage", style="yellow", justify="right") 

626 

627 top_values = value_counts.head(20) 

628 for row in top_values.iter_rows(): 

629 value, count = row 

630 percentage = count / total_rows * 100 

631 

632 # Color code percentages 

633 if percentage >= 50: 

634 pct_style = "red" 

635 elif percentage >= 10: 

636 pct_style = "yellow" 

637 else: 

638 pct_style = "green" 

639 

640 table.add_row( 

641 str(value), 

642 f"{count:,}", 

643 f"[{pct_style}]{percentage:.1f}%[/{pct_style}]", 

644 ) 

645 

646 console.print(table) 

647 

648 if total_distinct > 20: 

649 console.print( 

650 f"\n💡 [dim]... and {total_distinct - 20:,} more distinct values[/dim]" 

651 ) 

652 

653 except Exception as e: 

654 console.print(f"❌ [red]Error analyzing column {column}: {e}[/red]") 

655 

656 

657def analyze_timing_patterns(df: pl.DataFrame) -> None: 

658 """Analyze timing patterns by runtime model.""" 

659 print("\n" + "=" * 80) 

660 print("TIMING PATTERNS BY RUNTIME MODEL") 

661 print("=" * 80) 

662 

663 timing_analysis = ( 

664 df.filter(pl.col("RUNTIME_MODEL_NAME").is_not_null()) 

665 .group_by("RUNTIME_MODEL_NAME") 

666 .agg( 

667 [ 

668 pl.col("TOTAL_TIME").mean().alias("avg_time"), 

669 pl.col("TOTAL_DECODING_LATENCY").mean().alias("avg_latency"), 

670 pl.len().alias("count"), 

671 ] 

672 ) 

673 .sort("count", descending=True) 

674 ) 

675 

676 print(f"{'Model':<35} {'Avg Time (s)':<12} {'Avg Latency':<12} {'Count':<10}") 

677 print("-" * 75) 

678 

679 for row in timing_analysis.iter_rows(): 

680 model, avg_time, avg_latency, count = row 

681 print(f"{model:<35} {avg_time:<12.2f} {avg_latency:<12.6f} {count:<10,}") 

682 

683 

684def analyze_pricing_patterns(df: pl.DataFrame) -> None: 

685 """Analyze pricing patterns by runtime model.""" 

686 print("\n" + "=" * 80) 

687 print("PRICING PATTERNS BY RUNTIME MODEL") 

688 print("=" * 80) 

689 

690 pricing_analysis = ( 

691 df.filter(pl.col("RUNTIME_MODEL_NAME").is_not_null()) 

692 .group_by("RUNTIME_MODEL_NAME") 

693 .agg( 

694 [ 

695 pl.col("PRICE_PER_INPUT_TOKEN").mean().alias("avg_input_price"), 

696 pl.col("PRICE_PER_OUTPUT_TOKEN").mean().alias("avg_output_price"), 

697 pl.col("TOTAL_COST").mean().alias("avg_total_cost"), 

698 pl.len().alias("count"), 

699 ] 

700 ) 

701 .sort("count", descending=True) 

702 ) 

703 

704 print( 

705 f"{'Model':<35} {'Input Price':<12} {'Output Price':<12} {'Avg Cost':<12} {'Count':<10}" 

706 ) 

707 print("-" * 90) 

708 

709 for row in pricing_analysis.iter_rows(): 

710 model, input_price, output_price, avg_cost, count = row 

711 print( 

712 f"{model:<35} ${input_price:<11.6f} ${output_price:<11.6f} ${avg_cost:<11.6f} {count:<10,}" 

713 ) 

714 

715 

716def analyze_customer_patterns(df: pl.DataFrame) -> None: 

717 """Analyze customer usage patterns by API key.""" 

718 print("\n" + "=" * 80) 

719 print("TOP CUSTOMER USAGE PATTERNS") 

720 print("=" * 80) 

721 

722 customer_analysis = ( 

723 df.group_by("API_KEY_SUFFIX") 

724 .agg( 

725 [ 

726 pl.len().alias("total_requests"), 

727 pl.col("TOTAL_COST").sum().alias("total_spend"), 

728 pl.col("MODEL").n_unique().alias("unique_models"), 

729 pl.col("RUNTIME_MODEL_NAME").n_unique().alias("unique_runtime_models"), 

730 ] 

731 ) 

732 .sort("total_requests", descending=True) 

733 .head(15) 

734 ) 

735 

736 print( 

737 f"{'API Key':<10} {'Requests':<12} {'Total Spend':<15} {'Models':<8} {'Runtime Models':<15}" 

738 ) 

739 print("-" * 70) 

740 

741 for row in customer_analysis.iter_rows(): 

742 api_key, requests, spend, models, runtime_models = row 

743 print( 

744 f"{api_key:<10} {requests:<12,} ${spend:<14.2f} {models:<8} {runtime_models:<15}" 

745 ) 

746 

747 

748def analyze_search_patterns(df: pl.DataFrame) -> None: 

749 """Analyze search vs non-search usage patterns.""" 

750 print("\n" + "=" * 80) 

751 print("SEARCH QUERY USAGE PATTERNS") 

752 print("=" * 80) 

753 

754 search_analysis = ( 

755 df.group_by("NUM_SEARCH_QUERIES") 

756 .agg( 

757 [ 

758 pl.len().alias("request_count"), 

759 pl.col("TOTAL_COST").mean().alias("avg_cost"), 

760 pl.col("TOTAL_TIME").mean().alias("avg_time"), 

761 ] 

762 ) 

763 .sort("NUM_SEARCH_QUERIES") 

764 .head(10) # Show first 10 search query levels 

765 ) 

766 

767 print( 

768 f"{'Search Queries':<15} {'Request Count':<15} {'Avg Cost':<12} {'Avg Time (s)':<12}" 

769 ) 

770 print("-" * 60) 

771 

772 for row in search_analysis.iter_rows(): 

773 queries, count, avg_cost, avg_time = row 

774 print(f"{queries:<15} {count:<15,} ${avg_cost:<11.6f} {avg_time:<12.2f}") 

775 

776 

777def analyze_internal_external(df: pl.DataFrame) -> None: 

778 """Analyze internal vs external usage.""" 

779 print("\n" + "=" * 80) 

780 print("INTERNAL VS EXTERNAL USAGE") 

781 print("=" * 80) 

782 

783 internal_analysis = df.group_by("IS_PERPLEXITY_INTERNAL").agg( 

784 [ 

785 pl.len().alias("request_count"), 

786 pl.col("TOTAL_COST").sum().alias("total_revenue"), 

787 pl.col("TOTAL_COST").mean().alias("avg_cost_per_request"), 

788 ] 

789 ) 

790 

791 print( 

792 f"{'Type':<15} {'Request Count':<15} {'Total Revenue':<15} {'Avg Cost/Request':<20}" 

793 ) 

794 print("-" * 70) 

795 

796 for row in internal_analysis.iter_rows(): 

797 is_internal, count, revenue, avg_cost = row 

798 usage_type = "Internal" if is_internal else "External" 

799 print(f"{usage_type:<15} {count:<15,} ${revenue:<14.2f} ${avg_cost:<19.6f}") 

800 

801 

802def analyze_business_summary(df: pl.DataFrame) -> None: 

803 """Provide high-level business summary.""" 

804 print("\n" + "=" * 80) 

805 print("BUSINESS SUMMARY") 

806 print("=" * 80) 

807 

808 total_requests = len(df) 

809 total_revenue = df["TOTAL_COST"].sum() 

810 avg_revenue_per_request = df["TOTAL_COST"].mean() 

811 unique_customers = df["API_KEY_SUFFIX"].n_unique() 

812 unique_models = df["MODEL"].n_unique() 

813 unique_runtime_models = df.filter(pl.col("RUNTIME_MODEL_NAME").is_not_null())[ 

814 "RUNTIME_MODEL_NAME" 

815 ].n_unique() 

816 

817 print(f"Total Requests: {total_requests:,}") 

818 print(f"Total Revenue: ${total_revenue:,.2f}") 

819 print(f"Average Revenue per Request: ${avg_revenue_per_request:.6f}") 

820 print(f"Unique Customers (API Keys): {unique_customers:,}") 

821 print(f"Unique Models: {unique_models}") 

822 print(f"Unique Runtime Models: {unique_runtime_models}") 

823 

824 # Search usage breakdown 

825 search_requests = df.filter(pl.col("NUM_SEARCH_QUERIES") > 0) 

826 search_percentage = len(search_requests) / total_requests * 100 

827 print("\nSearch Usage:") 

828 print( 

829 f" Requests with search: {len(search_requests):,} ({search_percentage:.1f}%)" 

830 ) 

831 print( 

832 f" Requests without search: {total_requests - len(search_requests):,} ({100 - search_percentage:.1f}%)" 

833 ) 

834 

835 

836def parse_rollup_spec(rollup_spec: str) -> tuple[list[str], dict[str, str]]: 

837 """Parse rollup specification string into group columns and aggregation columns. 

838 

839 Format: "group_col1,group_col2:count(col1),sum(col2),count(col3)" 

840 Returns: (group_columns, {column: operation}) 

841 """ 

842 if ":" not in rollup_spec: 

843 # Old format: just group columns, use default aggregations 

844 group_cols = [col.strip() for col in rollup_spec.split(",")] 

845 return group_cols, {} 

846 

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

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

849 

850 # Parse aggregation specifications 

851 agg_specs = {} 

852 for agg_spec in agg_part.split(","): 

853 agg_spec = agg_spec.strip() 

854 if "(" in agg_spec and ")" in agg_spec: 

855 # Format: operation(column) 

856 op, rest = agg_spec.split("(", 1) 

857 col = rest.rstrip(")") 

858 agg_specs[col.strip()] = op.strip().lower() 

859 else: 

860 # Format: column (default to count) 

861 agg_specs[agg_spec] = "count" 

862 

863 return group_cols, agg_specs 

864 

865 

866def create_rollup( 

867 df: pl.DataFrame, 

868 rollup_spec: str, 

869 display_limit: int = None, 

870 group_by_time: str = None, 

871 datetime_col: str = None, 

872 ignore_none: bool = False, 

873 output_file: str = None, 

874) -> None: 

875 """Create a rollup/pivot table for the specified columns with custom aggregations. 

876 

877 Args: 

878 df: DataFrame to analyze 

879 rollup_spec: Specification string in format "group_col1,group_col2:count(col1),sum(col2)" 

880 display_limit: Maximum rows to display (only applies to screen output) 

881 group_by_time: Time grouping ('hour' or 'day'), or None for no time grouping 

882 datetime_col: Specific datetime column to use, or None for auto-detection 

883 ignore_none: If True, filter out rows with None/null values in any non-aggregate columns 

884 output_file: If provided, export results to CSV file instead of displaying on screen 

885 """ 

886 try: 

887 group_cols, agg_specs = parse_rollup_spec(rollup_spec) 

888 

889 # Handle datetime grouping 

890 datetime_group_col = None 

891 if group_by_time: 

892 if group_by_time not in ["hour", "day"]: 

893 console.print( 

894 f"ERROR [red]Error: group_by_time must be 'hour' or 'day', got '{group_by_time}'[/red]" 

895 ) 

896 return 

897 

898 detected_datetime_col = detect_datetime_column(df, datetime_col) 

899 if not detected_datetime_col: 

900 console.print( 

901 "ERROR [red]Error: No suitable datetime column found for time grouping[/red]" 

902 ) 

903 return 

904 

905 # Create datetime grouping column 

906 datetime_group_col = f"{detected_datetime_col}_{group_by_time}" 

907 

908 # Add datetime parsing and grouping to DataFrame 

909 try: 

910 # Check if column is already datetime type or needs parsing 

911 col_dtype = df[detected_datetime_col].dtype 

912 

913 if ( 

914 str(col_dtype).startswith("Datetime") 

915 or "datetime" in str(col_dtype).lower() 

916 ): 

917 # Column is already datetime, just truncate and format 

918 if group_by_time == "hour": 

919 df = df.with_columns( 

920 [ 

921 pl.col(detected_datetime_col) 

922 .dt.truncate("1h") 

923 .dt.strftime("%Y-%m-%d %H:00") 

924 .alias(datetime_group_col) 

925 ] 

926 ) 

927 else: # day 

928 df = df.with_columns( 

929 [ 

930 pl.col(detected_datetime_col) 

931 .dt.truncate("1d") 

932 .dt.strftime("%Y-%m-%d") 

933 .alias(datetime_group_col) 

934 ] 

935 ) 

936 else: 

937 # Column is string, need to parse first 

938 if group_by_time == "hour": 

939 df = df.with_columns( 

940 [ 

941 pl.col(detected_datetime_col) 

942 .str.to_datetime(format=None, strict=False) 

943 .dt.truncate("1h") 

944 .dt.strftime("%Y-%m-%d %H:00") 

945 .alias(datetime_group_col) 

946 ] 

947 ) 

948 else: # day 

949 df = df.with_columns( 

950 [ 

951 pl.col(detected_datetime_col) 

952 .str.to_datetime(format=None, strict=False) 

953 .dt.truncate("1d") 

954 .dt.strftime("%Y-%m-%d") 

955 .alias(datetime_group_col) 

956 ] 

957 ) 

958 

959 # Add datetime grouping column to group_cols 

960 group_cols = [datetime_group_col] + group_cols 

961 console.print( 

962 f"TIME [green]Added {group_by_time}ly grouping using column: {detected_datetime_col}[/green]" 

963 ) 

964 

965 except Exception as e: 

966 console.print( 

967 f"ERROR [red]Error processing datetime column '{detected_datetime_col}': {e}[/red]" 

968 ) 

969 console.print( 

970 "HINT [cyan]Hint: Try specifying a different datetime column with --datetime-col[/cyan]" 

971 ) 

972 return 

973 

974 # Validate group columns exist 

975 missing_cols = [col for col in group_cols if col not in df.columns] 

976 if missing_cols: 

977 console.print( 

978 f"❌ [red]Error: Group columns not found: {', '.join(missing_cols)}[/red]" 

979 ) 

980 console.print(f"📋 [cyan]Available columns: {', '.join(df.columns)}[/cyan]") 

981 return 

982 

983 # Validate aggregation columns exist 

984 if agg_specs: 

985 missing_agg_cols = [ 

986 col for col in agg_specs.keys() if col not in df.columns 

987 ] 

988 if missing_agg_cols: 

989 console.print( 

990 f"❌ [red]Error: Aggregation columns not found: {', '.join(missing_agg_cols)}[/red]" 

991 ) 

992 console.print( 

993 f"📋 [cyan]Available columns: {', '.join(df.columns)}[/cyan]" 

994 ) 

995 return 

996 

997 # Filter out rows with None/null values in non-aggregate columns if requested 

998 if ignore_none: 

999 original_rows = len(df) 

1000 # Get all non-aggregate columns (group columns + aggregation columns) 

1001 non_agg_cols = group_cols.copy() 

1002 if agg_specs: 

1003 # Add any columns that are being aggregated to the filter list 

1004 non_agg_cols.extend(agg_specs.keys()) 

1005 

1006 # Remove duplicates and filter 

1007 non_agg_cols = list(set(non_agg_cols)) 

1008 filter_conditions = [pl.col(col).is_not_null() for col in non_agg_cols] 

1009 df = df.filter(pl.all_horizontal(filter_conditions)) 

1010 filtered_rows = len(df) 

1011 excluded_rows = original_rows - filtered_rows 

1012 if excluded_rows > 0: 

1013 console.print( 

1014 f"FILTER [yellow]Filtered out {excluded_rows:,} rows with null values in rollup columns ({excluded_rows / original_rows:.1%})[/yellow]" 

1015 ) 

1016 console.print( 

1017 f"FILTER [dim]Filtered columns: {', '.join(non_agg_cols)}[/dim]" 

1018 ) 

1019 

1020 with Progress( 

1021 SpinnerColumn(), 

1022 TextColumn("[progress.description]{task.description}"), 

1023 console=console, 

1024 ) as progress: 

1025 task = progress.add_task("Creating rollup analysis...", total=None) 

1026 

1027 # Build aggregation expressions 

1028 agg_exprs = [] 

1029 agg_display_names = [] 

1030 

1031 if agg_specs: 

1032 # Use custom aggregations 

1033 for col, operation in agg_specs.items(): 

1034 if operation == "count": 

1035 agg_exprs.append(pl.col(col).count().alias(f"COUNT_{col}")) 

1036 agg_display_names.append(f"COUNT of {col}") 

1037 elif operation == "sum": 

1038 agg_exprs.append(pl.col(col).sum().alias(f"SUM_{col}")) 

1039 agg_display_names.append(f"SUM of {col}") 

1040 elif operation == "mean" or operation == "avg": 

1041 agg_exprs.append(pl.col(col).mean().alias(f"AVG_{col}")) 

1042 agg_display_names.append(f"AVG of {col}") 

1043 elif operation == "max": 

1044 agg_exprs.append(pl.col(col).max().alias(f"MAX_{col}")) 

1045 agg_display_names.append(f"MAX of {col}") 

1046 elif operation == "min": 

1047 agg_exprs.append(pl.col(col).min().alias(f"MIN_{col}")) 

1048 agg_display_names.append(f"MIN of {col}") 

1049 else: 

1050 console.print( 

1051 f"⚠️ [yellow]Warning: Unknown operation '{operation}' for column '{col}', using count[/yellow]" 

1052 ) 

1053 agg_exprs.append(pl.col(col).count().alias(f"COUNT_{col}")) 

1054 agg_display_names.append(f"COUNT of {col}") 

1055 else: 

1056 # Default aggregations for compatibility 

1057 default_agg_cols = [ 

1058 "TOTAL_COST", 

1059 "CT_PROMPT_TOKENS", 

1060 "CT_TOTAL_COMPLETION_TOKENS", 

1061 "NUM_SEARCH_QUERIES", 

1062 ] 

1063 available_agg_cols = [ 

1064 col for col in default_agg_cols if col in df.columns 

1065 ] 

1066 

1067 agg_exprs = [pl.len().alias("count")] 

1068 agg_display_names = ["Count"] 

1069 

1070 for col in available_agg_cols: 

1071 agg_exprs.extend( 

1072 [ 

1073 pl.col(col).sum().alias(f"total_{col.lower()}"), 

1074 pl.col(col).mean().alias(f"avg_{col.lower()}"), 

1075 ] 

1076 ) 

1077 agg_display_names.extend([f"Total {col}", f"Avg {col}"]) 

1078 

1079 # Select only necessary columns and use optimized lazy operations 

1080 required_cols = set( 

1081 group_cols + list(agg_specs.keys()) 

1082 if agg_specs 

1083 else group_cols 

1084 + [ 

1085 "TOTAL_COST", 

1086 "CT_PROMPT_TOKENS", 

1087 "CT_TOTAL_COMPLETION_TOKENS", 

1088 "NUM_SEARCH_QUERIES", 

1089 ] 

1090 ) 

1091 available_cols = [col for col in required_cols if col in df.columns] 

1092 

1093 # Use lazy evaluation for better performance on large datasets 

1094 if len(df) > 100000: # Use lazy for large datasets 

1095 rollup_data = ( 

1096 df.lazy() 

1097 .select(available_cols) 

1098 .group_by(group_cols) 

1099 .agg(agg_exprs) 

1100 .sort(agg_exprs[0].meta.output_name(), descending=True) 

1101 .collect() 

1102 ) 

1103 else: 

1104 # Use eager evaluation for smaller datasets 

1105 rollup_data = ( 

1106 df.select(available_cols) 

1107 .group_by(group_cols) 

1108 .agg(agg_exprs) 

1109 .sort(agg_exprs[0].meta.output_name(), descending=True) 

1110 ) 

1111 

1112 progress.update(task, completed=True) 

1113 

1114 # Export to CSV if output file is specified 

1115 if output_file: 

1116 try: 

1117 rollup_data.write_csv(output_file) 

1118 total_combinations = len(rollup_data) 

1119 console.print( 

1120 f"✅ [green]Exported {total_combinations:,} rollup combinations to {output_file}[/green]" 

1121 ) 

1122 return 

1123 except Exception as e: 

1124 console.print(f"❌ [red]Error exporting to CSV: {e}[/red]") 

1125 return 

1126 

1127 # Limit display for performance and readability 

1128 if display_limit is not None: 

1129 display_data = rollup_data.head(display_limit) 

1130 else: 

1131 display_data = rollup_data 

1132 

1133 # Create rich table 

1134 title = f"📊 Rollup Analysis: [magenta]{', '.join(group_cols)}[/magenta]" 

1135 if agg_specs: 

1136 agg_summary = ", ".join([f"{op}({col})" for col, op in agg_specs.items()]) 

1137 title += f" → [yellow]{agg_summary}[/yellow]" 

1138 

1139 table = Table(title=title, box=box.ROUNDED) 

1140 

1141 # Add group columns 

1142 for col in group_cols: 

1143 table.add_column(col[:20], style="cyan", no_wrap=True) 

1144 

1145 # Add aggregation columns 

1146 colors = ["green", "yellow", "blue", "magenta", "red", "white"] 

1147 for i, display_name in enumerate(agg_display_names): 

1148 color = colors[i % len(colors)] 

1149 table.add_column(display_name[:20], style=color, justify="right") 

1150 

1151 # Add data rows 

1152 for row in display_data.iter_rows(): 

1153 group_values = row[: len(group_cols)] 

1154 agg_values = row[len(group_cols) :] 

1155 

1156 row_data = [str(val)[:20] for val in group_values] 

1157 

1158 # Format aggregation values 

1159 for i, val in enumerate(agg_values): 

1160 if agg_display_names[i].startswith("COUNT"): 

1161 row_data.append(f"{val:,}") 

1162 elif agg_display_names[i].startswith("SUM") or agg_display_names[ 

1163 i 

1164 ].startswith("Total"): 

1165 if "COST" in agg_display_names[i].upper(): 

1166 row_data.append(f"${val:.2f}") 

1167 else: 

1168 row_data.append(f"{val:,}") 

1169 elif agg_display_names[i].startswith("AVG"): 

1170 if "COST" in agg_display_names[i].upper(): 

1171 row_data.append(f"${val:.6f}") 

1172 else: 

1173 row_data.append(f"{val:.2f}") 

1174 else: 

1175 row_data.append(f"{val:,}") 

1176 

1177 table.add_row(*row_data) 

1178 

1179 # Show if there are more rows 

1180 if display_limit is not None and len(rollup_data) > display_limit: 

1181 table.add_row( 

1182 *["[dim]...[/dim]"] * (len(group_cols) + len(agg_display_names)) 

1183 ) 

1184 

1185 # Add grand total row if using custom aggregations 

1186 if agg_specs: 

1187 total_exprs = [] 

1188 for col, operation in agg_specs.items(): 

1189 if operation == "count": 

1190 total_exprs.append(pl.col(col).count()) 

1191 elif operation == "sum": 

1192 total_exprs.append(pl.col(col).sum()) 

1193 elif operation in ["mean", "avg"]: 

1194 total_exprs.append(pl.col(col).mean()) 

1195 elif operation == "max": 

1196 total_exprs.append(pl.col(col).max()) 

1197 elif operation == "min": 

1198 total_exprs.append(pl.col(col).min()) 

1199 else: 

1200 total_exprs.append(pl.col(col).count()) 

1201 

1202 grand_totals = df.select(total_exprs).row(0) 

1203 

1204 summary_data = ["[bold]Grand Total[/bold]"] + [""] * (len(group_cols) - 1) 

1205 for i, val in enumerate(grand_totals): 

1206 if agg_display_names[i].startswith("COUNT"): 

1207 summary_data.append(f"[bold]{val:,}[/bold]") 

1208 elif agg_display_names[i].startswith("SUM"): 

1209 if "COST" in agg_display_names[i].upper(): 

1210 summary_data.append(f"[bold]${val:.2f}[/bold]") 

1211 else: 

1212 summary_data.append(f"[bold]{val:,}[/bold]") 

1213 elif agg_display_names[i].startswith("AVG"): 

1214 if "COST" in agg_display_names[i].upper(): 

1215 summary_data.append(f"[bold]${val:.6f}[/bold]") 

1216 else: 

1217 summary_data.append(f"[bold]{val:.2f}[/bold]") 

1218 else: 

1219 summary_data.append(f"[bold]{val:,}[/bold]") 

1220 

1221 table.add_section() 

1222 table.add_row(*summary_data) 

1223 

1224 console.print(table) 

1225 

1226 # Summary message 

1227 total_combinations = len(rollup_data) 

1228 if display_limit is not None and total_combinations > display_limit: 

1229 console.print( 

1230 f"\n📈 [green]Rollup shows top {display_limit} of {total_combinations:,} unique combinations[/green]" 

1231 ) 

1232 else: 

1233 console.print( 

1234 f"\n📈 [green]Rollup shows {total_combinations:,} unique combinations[/green]" 

1235 ) 

1236 

1237 except Exception as e: 

1238 console.print(f"❌ [red]Error creating rollup: {e}[/red]") 

1239 

1240 

1241def run_comprehensive_analysis(df: pl.DataFrame) -> None: 

1242 """Run all analysis functions for comprehensive insights.""" 

1243 print("\n" + "#" * 80) 

1244 print("COMPREHENSIVE PERPLEXITY DATA ANALYSIS") 

1245 print("#" * 80) 

1246 

1247 analyze_business_summary(df) 

1248 analyze_timing_patterns(df) 

1249 analyze_pricing_patterns(df) 

1250 analyze_customer_patterns(df) 

1251 analyze_search_patterns(df) 

1252 analyze_internal_external(df) 

1253 

1254 print("\n" + "#" * 80) 

1255 print("ANALYSIS COMPLETE") 

1256 print("#" * 80) 

1257 

1258 

1259def main(): 

1260 parser = argparse.ArgumentParser( 

1261 description="Analyze Perplexity CSV data with interactive column exploration" 

1262 ) 

1263 parser.add_argument( 

1264 "--in", 

1265 type=Path, 

1266 dest="input_file", 

1267 help="Path to CSV file to analyze (required for analysis operations)", 

1268 ) 

1269 parser.add_argument( 

1270 "--column", type=str, help="Specific column to analyze (skips interactive mode)" 

1271 ) 

1272 parser.add_argument( 

1273 "--full-analysis", 

1274 action="store_true", 

1275 help="Run comprehensive analysis of the entire dataset", 

1276 ) 

1277 parser.add_argument( 

1278 "--rollup", 

1279 type=str, 

1280 help="Create rollup/pivot table. Format: 'group_cols:agg_specs' (e.g., 'MODEL_SKU,RUNTIME_MODEL_NAME:count(LLM_INPUT_TOKENS),sum(CT_TOTAL_COMPLETION_TOKENS)'). Use with --group-by for time-based analysis.", 

1281 ) 

1282 parser.add_argument( 

1283 "--rollup-limit", 

1284 type=int, 

1285 default=None, 

1286 help="Maximum number of rollup rows to display (default: show all)", 

1287 ) 

1288 parser.add_argument( 

1289 "--out", 

1290 type=str, 

1291 help="Export rollup results to CSV file instead of displaying on screen", 

1292 ) 

1293 parser.add_argument( 

1294 "--group-by", 

1295 choices=["hour", "day"], 

1296 help="Group rollup data by time period (hour or day). Requires --rollup option.", 

1297 ) 

1298 parser.add_argument( 

1299 "--datetime-col", 

1300 type=str, 

1301 help="Specify datetime column for time grouping. If not provided, will auto-detect.", 

1302 ) 

1303 parser.add_argument( 

1304 "--ignore-none", 

1305 action="store_true", 

1306 help="Ignore rows with None/null values in any non-aggregate columns when creating rollup. Use with --rollup option.", 

1307 ) 

1308 parser.add_argument( 

1309 "--lazy", 

1310 action="store_true", 

1311 help="Use lazy loading for better memory efficiency with very large datasets", 

1312 ) 

1313 parser.add_argument( 

1314 "--fast", 

1315 action="store_true", 

1316 help="Enable all performance optimizations (lazy loading, column selection, etc.)", 

1317 ) 

1318 parser.add_argument( 

1319 "--sample", 

1320 type=int, 

1321 help="Process only a sample of N rows for faster analysis of large files", 

1322 ) 

1323 parser.add_argument( 

1324 "--generate-schema", 

1325 type=Path, 

1326 help="Generate schema manifest for all CSV files in the specified folder", 

1327 ) 

1328 parser.add_argument( 

1329 "--interactive", 

1330 action="store_true", 

1331 help="Start interactive mode for column exploration", 

1332 ) 

1333 

1334 args = parser.parse_args() 

1335 

1336 # Validate group-by requires rollup 

1337 if args.group_by and not args.rollup: 

1338 console.print( 

1339 "ERROR [red]Error: --group-by option requires --rollup to be specified[/red]" 

1340 ) 

1341 sys.exit(1) 

1342 

1343 # Validate ignore-none requires rollup 

1344 if args.ignore_none and not args.rollup: 

1345 console.print( 

1346 "ERROR [red]Error: --ignore-none option requires --rollup to be specified[/red]" 

1347 ) 

1348 sys.exit(1) 

1349 

1350 # Handle schema generation first (doesn't require --in argument) 

1351 if args.generate_schema: 

1352 generate_schema_manifest(args.generate_schema) 

1353 return 

1354 

1355 # Check if no arguments provided, show help and exit 

1356 if not any([args.full_analysis, args.rollup, args.column, args.interactive]): 

1357 parser.print_help() 

1358 sys.exit(0) 

1359 

1360 # Validate --in argument is provided for analysis operations 

1361 if not args.input_file: 

1362 console.print( 

1363 "ERROR [red]Error: --in argument is required for analysis operations[/red]" 

1364 ) 

1365 sys.exit(1) 

1366 

1367 if not args.input_file.exists(): 

1368 console.print(f"❌ [red]Error: File {args.input_file} not found[/red]") 

1369 sys.exit(1) 

1370 

1371 # Determine required columns for optimization 

1372 required_cols = get_required_columns( 

1373 rollup_spec=args.rollup, 

1374 column=args.column, 

1375 full_analysis=args.full_analysis, 

1376 group_by_time=args.group_by, 

1377 ) 

1378 

1379 # Apply performance optimizations 

1380 use_lazy_loading = args.lazy or args.fast 

1381 use_optimized_loading = args.fast or (required_cols and len(required_cols) < 20) 

1382 

1383 # Load data with optimizations 

1384 if use_optimized_loading: 

1385 console.print( 

1386 f"🚀 [cyan]Optimizing for {len(required_cols) if required_cols else 'all'} columns: {', '.join((required_cols or [])[:5])}{'...' if required_cols and len(required_cols) > 5 else ''}[/cyan]" 

1387 ) 

1388 df = load_data_optimized( 

1389 args.input_file, required_cols, use_lazy=use_lazy_loading 

1390 ) 

1391 else: 

1392 df = load_data(args.input_file, use_lazy=use_lazy_loading) 

1393 

1394 # Apply sampling if requested 

1395 if args.sample and args.sample < len(df): 

1396 original_size = len(df) 

1397 df = df.sample(n=args.sample, seed=42) 

1398 console.print( 

1399 f"📊 [yellow]Sampling {args.sample:,} rows from {original_size:,} total rows[/yellow]" 

1400 ) 

1401 

1402 if args.full_analysis: 

1403 # Run comprehensive analysis 

1404 run_comprehensive_analysis(df) 

1405 elif args.rollup: 

1406 # Create rollup/pivot table with new format 

1407 create_rollup( 

1408 df, 

1409 args.rollup, 

1410 args.rollup_limit, 

1411 args.group_by, 

1412 args.datetime_col, 

1413 args.ignore_none, 

1414 args.out, 

1415 ) 

1416 elif args.column: 

1417 # Direct column analysis 

1418 if args.column not in df.columns: 

1419 console.print(f"❌ [red]Error: Column '{args.column}' not found[/red]") 

1420 console.print(f"📋 [cyan]Available columns: {', '.join(df.columns)}[/cyan]") 

1421 sys.exit(1) 

1422 analyze_column(df, args.column) 

1423 elif args.interactive: 

1424 # Interactive mode 

1425 while True: 

1426 show_column_info(df) 

1427 column = get_column_choice(df) 

1428 analyze_column(df, column) 

1429 

1430 continue_choice = ( 

1431 console.input("\n🔄 [cyan]Analyze another column? (y/N):[/cyan] ") 

1432 .strip() 

1433 .lower() 

1434 ) 

1435 if continue_choice not in ["y", "yes"]: 

1436 break 

1437 

1438 console.print("\n✨ [green]Analysis complete![/green]") 

1439 

1440 

1441if __name__ == "__main__": 

1442 main()