Coverage for src/common/rollup.py: 73%
232 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
5"""
6Rollup and aggregation functions for CloudZero Data Tool.
7Handles pivot table creation, time-based grouping, and custom aggregations.
8"""
10import re
11from typing import Tuple
13import polars as pl
14from rich.progress import Progress, SpinnerColumn, TextColumn
16try:
17 from ..utils.console_helpers import print_error, print_success, print_warning, console
18 from ..utils.table_helpers import create_rollup_results_table
19except ImportError:
20 # Fallback for direct execution or test imports
21 import sys
22 from pathlib import Path
23 sys.path.append(str(Path(__file__).parent.parent))
24 from utils.console_helpers import print_error, print_success, print_warning, console
25 from utils.table_helpers import create_rollup_results_table
28def _abbreviate_column_name(column_name: str) -> str:
29 """Create a readable abbreviation for column names in rollup headers."""
30 # Common abbreviations for known patterns
31 abbreviations = {
32 "CT_TOTAL_COMPLETION_TOKENS": "COMPLETION_TOKENS",
33 "CT_PROMPT_TOKENS": "PROMPT_TOKENS",
34 "LLM_INPUT_TOKENS": "INPUT_TOKENS",
35 "TOTAL_COST": "COST",
36 "MODEL_SKU": "MODEL",
37 "RUNTIME_MODEL_NAME": "RUNTIME",
38 "API_KEY_SUFFIX": "API_KEY",
39 "NUM_SEARCH_QUERIES": "SEARCH_QUERIES",
40 "IS_PERPLEXITY_INTERNAL": "INTERNAL"
41 }
43 # Return known abbreviation if available
44 if column_name in abbreviations:
45 return abbreviations[column_name]
47 # For unknown columns, create intelligent abbreviation
48 # Remove common prefixes and use title case
49 name = column_name.replace("CT_", "").replace("LLM_", "").replace("_", " ")
51 # If still too long, use first letters of words
52 if len(name) > 12:
53 words = name.split()
54 if len(words) > 1:
55 return "".join(word[0] for word in words).upper()
56 else:
57 return name[:12].upper()
59 return name.replace(" ", "_")
62def detect_datetime_column(df: pl.DataFrame, datetime_col: str = None) -> str:
63 """Detect the most appropriate datetime column in the DataFrame.
65 Args:
66 df: DataFrame to analyze
67 datetime_col: Specific column name to use, or None for auto-detection
69 Returns:
70 Column name to use for datetime grouping
71 """
72 if datetime_col:
73 if datetime_col in df.columns:
74 return datetime_col
75 else:
76 print_error(f"Specified datetime column '{datetime_col}' not found")
77 console.print(f"📋 [cyan]Available columns: {', '.join(df.columns)}[/cyan]")
78 return None
80 # Auto-detect datetime column by analyzing column types and names
81 datetime_candidates = []
83 # First, prioritize columns with Datetime type that have time information
84 for col in df.columns:
85 col_dtype = df[col].dtype
86 if (
87 str(col_dtype).startswith("Datetime")
88 or "datetime" in str(col_dtype).lower()
89 ):
90 # Check if it actually contains time information (not just date)
91 sample_values = df[col].drop_nulls().head(5)
92 has_time_info = False
93 for value in sample_values:
94 value_str = str(value)
95 if " " in value_str and ":" in value_str: # Contains time component
96 has_time_info = True
97 break
99 if has_time_info:
100 datetime_candidates.append(
101 (col, "datetime_type", 100)
102 ) # Highest priority
103 else:
104 datetime_candidates.append(
105 (col, "datetime_type", 50)
106 ) # Medium priority
108 # Then look for columns with datetime-like names
109 datetime_name_patterns = [
110 (r".*created.*at.*utc.*", 90), # created_at_utc gets high priority
111 (r".*created.*at.*", 80), # created_at gets high priority
112 (r".*timestamp.*", 70), # timestamp columns
113 (r".*time.*", 60), # time columns
114 (r".*date.*", 40), # date columns (lower priority)
115 ]
117 for col in df.columns:
118 col_lower = col.lower()
119 for pattern, priority in datetime_name_patterns:
120 if re.match(pattern, col_lower):
121 datetime_candidates.append((col, "name_match", priority))
122 break
124 if not datetime_candidates:
125 print_warning("No datetime columns detected")
126 return None
128 # Sort by priority (highest first) and select the best candidate
129 datetime_candidates.sort(key=lambda x: x[2], reverse=True)
130 selected_col = datetime_candidates[0][0]
132 print_success(f"Auto-detected datetime column: {selected_col}")
133 return selected_col
136def parse_rollup_spec(rollup_spec: str) -> Tuple[list[str], dict[str, str]]:
137 """Parse rollup specification string into group columns and aggregation specs.
139 Args:
140 rollup_spec: String in format "group_col1,group_col2:agg1(col1),agg2(col2)"
142 Returns:
143 Tuple of (group_columns, aggregation_specs)
144 """
145 if ":" not in rollup_spec:
146 # Only group columns specified, no aggregations
147 group_cols = [col.strip() for col in rollup_spec.split(",")]
148 return group_cols, {}
150 group_part, agg_part = rollup_spec.split(":", 1)
151 group_cols = [col.strip() for col in group_part.split(",")]
153 # Parse aggregations
154 agg_specs = {}
156 # Pattern to match aggregation functions: operation(column)
157 agg_pattern = r"(count|sum|avg|min|max)\(([^)]+)\)"
158 matches = re.findall(agg_pattern, agg_part)
160 for operation, column in matches:
161 column = column.strip()
162 agg_specs[column] = operation
164 return group_cols, agg_specs
167def create_rollup(
168 df: pl.DataFrame,
169 rollup_spec: str,
170 display_limit: int = None,
171 group_by_time: str = None,
172 datetime_col: str = None,
173 ignore_none: bool = False,
174 output_file: str = None,
175) -> None:
176 """Create a rollup/pivot table for the specified columns with custom aggregations.
178 Args:
179 df: DataFrame to analyze
180 rollup_spec: Specification string in format "group_col1,group_col2:count(col1),sum(col2)"
181 display_limit: Maximum rows to display (only applies to screen output)
182 group_by_time: Time grouping ('hour' or 'day'), or None for no time grouping
183 datetime_col: Specific datetime column to use, or None for auto-detection
184 ignore_none: If True, filter out rows with None/null values in any non-aggregate columns
185 output_file: If provided, export results to CSV file instead of displaying on screen
186 """
187 try:
188 group_cols, agg_specs = parse_rollup_spec(rollup_spec)
190 # Handle datetime grouping
191 if group_by_time:
192 detected_datetime_col = detect_datetime_column(df, datetime_col)
193 if detected_datetime_col:
194 # Create time-based grouping column
195 time_col_name = f"{detected_datetime_col}_{group_by_time}"
197 if group_by_time == "hour":
198 df = df.with_columns(
199 [
200 pl.col(detected_datetime_col)
201 .dt.truncate("1h")
202 .alias(time_col_name)
203 ]
204 )
205 elif group_by_time == "day":
206 df = df.with_columns(
207 [pl.col(detected_datetime_col).dt.date().alias(time_col_name)]
208 )
210 # Add time column to group columns
211 group_cols.insert(0, time_col_name)
212 print_success(f"Added {group_by_time}ly grouping using column: {detected_datetime_col}")
213 else:
214 print_error("Could not detect datetime column for time grouping")
215 return
217 # Validate group columns exist
218 missing_group_cols = [col for col in group_cols if col not in df.columns]
219 if missing_group_cols:
220 print_error(f"Group columns not found: {', '.join(missing_group_cols)}")
221 console.print(f"📋 [cyan]Available columns: {', '.join(df.columns)}[/cyan]")
222 try:
223 from ..utils.validation import ValidationError
224 except ImportError:
225 import sys
226 from pathlib import Path
227 sys.path.append(str(Path(__file__).parent.parent))
228 from utils.validation import ValidationError
229 raise ValidationError(f"Group columns not found: {', '.join(missing_group_cols)}")
231 # Validate aggregation columns exist (if specified)
232 if agg_specs:
233 missing_agg_cols = [
234 col for col in agg_specs.keys() if col not in df.columns and col != "*"
235 ]
236 if missing_agg_cols:
237 print_error(f"Aggregation columns not found: {', '.join(missing_agg_cols)}")
238 console.print(f"📋 [cyan]Available columns: {', '.join(df.columns)}[/cyan]")
239 try:
240 from ..utils.validation import ValidationError
241 except ImportError:
242 import sys
243 from pathlib import Path
244 sys.path.append(str(Path(__file__).parent.parent))
245 from utils.validation import ValidationError
246 raise ValidationError(f"Aggregation columns not found: {', '.join(missing_agg_cols)}")
248 # Filter out null values if requested
249 if ignore_none:
250 # Get all columns that will be used (group + aggregation, excluding count(*))
251 filter_cols = group_cols.copy()
252 if agg_specs:
253 filter_cols.extend([col for col in agg_specs.keys() if col != "*"])
255 # Remove duplicates and filter
256 filter_cols = list(set(filter_cols))
257 original_size = len(df)
259 # Create filter expression for non-null values
260 filter_expr = pl.col(filter_cols[0]).is_not_null()
261 for col in filter_cols[1:]:
262 filter_expr = filter_expr & pl.col(col).is_not_null()
264 df = df.filter(filter_expr)
265 filtered_size = len(df)
267 console.print(
268 f"🔍 [yellow]Filtered out {original_size - filtered_size:,} rows with null values in rollup columns ({(original_size - filtered_size) / original_size * 100:.1f}%)[/yellow]"
269 )
270 console.print(
271 f"🔍 [yellow]Filtered columns: {', '.join(filter_cols)}[/yellow]"
272 )
274 # Create rollup with progress tracking
275 with Progress(
276 SpinnerColumn(),
277 TextColumn("[green]Creating rollup analysis..."),
278 console=console,
279 ) as progress:
280 task = progress.add_task("Processing...", total=100)
282 if agg_specs:
283 # Build aggregation expressions
284 agg_exprs = []
285 for col, operation in agg_specs.items():
286 if operation == "count":
287 if col == "*":
288 agg_exprs.append(pl.len().alias("COUNT"))
289 else:
290 agg_exprs.append(
291 pl.col(col).count().alias(f"COUNT_{col.upper()}")
292 )
293 elif operation == "sum":
294 agg_exprs.append(pl.col(col).sum().alias(f"SUM_{col.upper()}"))
295 elif operation == "avg":
296 agg_exprs.append(pl.col(col).mean().alias(f"AVG_{col.upper()}"))
297 elif operation == "min":
298 agg_exprs.append(pl.col(col).min().alias(f"MIN_{col.upper()}"))
299 elif operation == "max":
300 agg_exprs.append(pl.col(col).max().alias(f"MAX_{col.upper()}"))
302 rollup_data = (
303 df.group_by(group_cols)
304 .agg(agg_exprs)
305 .sort(agg_exprs[0].meta.output_name(), descending=True)
306 )
307 else:
308 # Simple group by with count
309 rollup_data = (
310 df.group_by(group_cols)
311 .agg(pl.len().alias("COUNT"))
312 .sort("COUNT", descending=True)
313 )
315 progress.update(task, completed=True)
317 # Export to CSV if output file is specified
318 if output_file:
319 try:
320 rollup_data.write_csv(output_file)
321 total_combinations = len(rollup_data)
322 print_success(f"Exported {total_combinations:,} rollup combinations to {output_file}")
323 return
324 except Exception as e:
325 print_error(f"Error exporting to CSV: {e}")
326 return
328 # Limit display for performance and readability
329 if display_limit is not None:
330 display_data = rollup_data.head(display_limit)
331 else:
332 display_data = rollup_data
334 # Create rich table
335 title = f"📊 Rollup Analysis: [magenta]{', '.join(group_cols)}[/magenta]"
337 # Prepare aggregation column names
338 agg_display_names = []
339 if agg_specs:
340 agg_summary = ", ".join([f"{op}({col})" for col, op in agg_specs.items()])
341 title += f" → [yellow]{agg_summary}[/yellow]"
343 for col, operation in agg_specs.items():
344 # Create more concise but readable column names
345 if operation == "count":
346 if col == "*":
347 display_name = "COUNT"
348 else:
349 # Use abbreviated column name for count
350 short_col = _abbreviate_column_name(col)
351 display_name = f"COUNT_{short_col}"
352 elif operation == "sum":
353 short_col = _abbreviate_column_name(col)
354 display_name = f"SUM_{short_col}"
355 elif operation == "avg":
356 short_col = _abbreviate_column_name(col)
357 display_name = f"AVG_{short_col}"
358 elif operation == "min":
359 short_col = _abbreviate_column_name(col)
360 display_name = f"MIN_{short_col}"
361 elif operation == "max":
362 short_col = _abbreviate_column_name(col)
363 display_name = f"MAX_{short_col}"
364 agg_display_names.append(display_name)
365 else:
366 agg_display_names = ["COUNT"]
368 # Create table with all columns at once
369 table = create_rollup_results_table(group_cols, agg_display_names)
370 table.title = title
372 # Add data rows
373 for row in display_data.iter_rows():
374 row_data = []
376 # Group column values
377 for i, val in enumerate(row[: len(group_cols)]):
378 row_data.append(str(val))
380 # Aggregation values
381 for val in row[len(group_cols) :]:
382 if isinstance(val, (int, float)):
383 row_data.append(f"{val:,}")
384 else:
385 row_data.append(f"{val:,}")
387 table.add_row(*row_data)
389 # Show if there are more rows
390 if display_limit is not None and len(rollup_data) > display_limit:
391 table.add_row(
392 *["[dim]...[/dim]"] * (len(group_cols) + len(agg_display_names))
393 )
395 # Add grand total row if using custom aggregations
396 if agg_specs:
397 total_exprs = []
398 for col, operation in agg_specs.items():
399 if operation == "count":
400 total_exprs.append(pl.col(col).count())
401 elif operation == "sum":
402 total_exprs.append(pl.col(col).sum())
403 elif operation == "avg":
404 total_exprs.append(pl.col(col).mean())
405 elif operation == "min":
406 total_exprs.append(pl.col(col).min())
407 elif operation == "max":
408 total_exprs.append(pl.col(col).max())
410 if total_exprs:
411 totals = df.select(total_exprs).row(0)
412 summary_data = ["Grand Total"] + [""] * (len(group_cols) - 1)
413 for val in totals:
414 if isinstance(val, (int, float)):
415 summary_data.append(f"{val:,}")
416 else:
417 summary_data.append(str(val))
419 table.add_row(*summary_data)
421 console.print(table)
423 # Summary message
424 total_combinations = len(rollup_data)
425 if display_limit is not None and total_combinations > display_limit:
426 console.print(
427 f"\n📈 [green]Rollup shows top {display_limit} of {total_combinations:,} unique combinations[/green]"
428 )
429 else:
430 console.print(
431 f"\n📈 [green]Rollup shows {total_combinations:,} unique combinations[/green]"
432 )
434 except Exception as e:
435 print_error(f"Error creating rollup: {e}")
436 # Re-raise ValidationError so main.py can handle it properly
437 try:
438 from ..utils.validation import ValidationError
439 except ImportError:
440 import sys
441 from pathlib import Path
442 sys.path.append(str(Path(__file__).parent.parent))
443 from utils.validation import ValidationError
444 if isinstance(e, ValidationError):
445 raise