diff --git a/owilix/cli/query.py b/owilix/cli/query.py index d7579d1..74977bc 100644 --- a/owilix/cli/query.py +++ b/owilix/cli/query.py @@ -190,8 +190,6 @@ def aggregate( ), limit: Optional[int] = typer.Option(None, "--limit", "-l", help="Per-file row limit"), files: str = typer.Option("**/*.parquet", "--files", "-f", help="File glob pattern"), - as_json: bool = typer.Option(False, "--json", help="Output as JSON"), - json_file: Optional[str] = typer.Option(None, "--output", "-o", help="Output to JSON file"), batch_size: int = typer.Option(100, "--batch-size", help="Rows per batch"), pq_batch_size: int = typer.Option(1, "--pq-batch", help="Parquet files per batch"), verbose: bool = typer.Option(False, "--verbose", help="Enable verbose output"), @@ -213,8 +211,8 @@ def aggregate( aggregate_query=aggregate_query, limit=limit, files=files, - as_json=as_json, - json_file=json_file, + format=cli_ctx.output_format, + output_file=cli_ctx.output_file, batch_size=batch_size, pq_batch_size=pq_batch_size, console=cli_ctx.console, diff --git a/owilix/core/db/models.py b/owilix/core/db/models.py index 974c07d..5a92ee2 100644 --- a/owilix/core/db/models.py +++ b/owilix/core/db/models.py @@ -13,7 +13,7 @@ sql_templates = { SELECT 'Rows imported' AS message, ${owi_len_files} as num_files, COUNT(*) AS count FROM owi_slice; """, "pq_select": """ - SELECT ${select} FROM read_parquet(${owi_remote_files}) as p ${where} ${groupby} ${postfix}; + SELECT ${select} FROM read_parquet(${owi_remote_files}) as p ${where} ${groupby} ${postfix} ${limit}; """, "pq_slice": """ BEGIN; @@ -134,31 +134,35 @@ class OWIlixSQLQuery: """ return self.format(owi_remote_files=str(files)) - def where(self, where: str = "") -> 'OWIlixSQLQuery': + def where(self, where = "") -> 'OWIlixSQLQuery': """ Add a WHERE clause to the SQL query. Args: - where (str): The WHERE clause to add. If empty, no WHERE clause is added. + where: The WHERE clause to add. If empty/None, no WHERE clause is added. Returns: OWIlixSQLQuery: The updated OWIlixSQLQuery instance. """ - if where and not where.strip().lower().startswith("where"): + if where is None or where == "": + return self.format(where="") + if not where.strip().lower().startswith("where"): where = "WHERE " + where return self.format(where=where) - def limit(self, limit: str = "") -> 'OWIlixSQLQuery': + def limit(self, limit = "") -> 'OWIlixSQLQuery': """ Add a LIMIT clause to the SQL query. Args: - limit (str): The LIMIT clause to add. If empty, no LIMIT clause is added. + limit: The LIMIT value. If empty/None, no LIMIT clause is added. Returns: OWIlixSQLQuery: The updated OWIlixSQLQuery instance. """ - if isinstance(limit, int) or limit and not limit.strip().lower().startswith("limit"): + if limit is None or limit == "": + return self.format(limit="") + if isinstance(limit, int) or (isinstance(limit, str) and not limit.strip().lower().startswith("limit")): limit = "LIMIT " + str(limit) return self.format(limit=limit) @@ -203,7 +207,10 @@ class OWIlixSQLQuery: groupby = f"GROUP BY ({groupby})" return self.format(groupby=groupby) - def postfix(self, postfix:str = "") -> 'OWIlixSQLQuery': + def postfix(self, postfix = "") -> 'OWIlixSQLQuery': + """Add a postfix to the SQL query.""" + if postfix is None: + postfix = "" return self.format(postfix=postfix) @property diff --git a/owilix/core/tasks/query.py b/owilix/core/tasks/query.py index f7d36b7..dd00a24 100644 --- a/owilix/core/tasks/query.py +++ b/owilix/core/tasks/query.py @@ -607,19 +607,22 @@ def query_aggregate( aggregate_query: str = "Select * from aggregates", files: str = "**/metadata*.parquet", explain: bool = False, - as_json: bool = False, + format: str = "table", interactive: bool = False, pq_batch_size: int = 1, batch_size: int = 100, prefetch: int = 1, page_size: int = 10, - json_file: Optional[str] = None, + output_file: Optional[str] = None, console: Optional[Console] = None ) -> CommandResult: """Run SQL aggregation queries across datasets.""" if console is None: console = Console() + # Determine format flags + is_json_mode = format in ["json", "jsonl", "json+", "jsonl+"] + try: files_pattern = eval(files) if files is not None else None except: @@ -627,11 +630,11 @@ def query_aggregate( all_files_by_dataset = get_all_files( manager, files_pattern, local_specifier, remote_specifier, - console=console, print_it=not as_json + console=console, print_it=(format == "table" or output_file) ) all_files = group_all_files_by_fs(all_files_by_dataset) - if not as_json: + if format == "table" or output_file: total_files = sum([len(v) for v in all_files.values()]) console.print(f"Found '{total_files}' parquet files in {len(all_files)} filesystems. Running queries against them.") @@ -668,6 +671,10 @@ def query_aggregate( progress.update(spinner_task, records=processed_count) _aggregates.extend(results.rows) + if not _aggregates: + console.print("[yellow]No records aggregated. Check your query parameters.[/yellow]") + return CommandResult(success=True, msg="No records found") + aggregates = pd.DataFrame(_aggregates) try: @@ -677,20 +684,15 @@ def query_aggregate( "You can run duckdb.query('...from aggregates...')") embed() else: - # Recreate context for duckdb query - # Note: duckdb.query() uses `aggregates` dataframe from local scope automagically in some versions, - # but explicit registration is safer or ensuring variable name matches. - # `duckdb.query(..., connection=...)` - # `aggregates` DF is local variable. - df = duckdb.query(aggregate_query).to_df() - if as_json: - if json_file: - with open(json_file, "w", encoding="utf-8") as f: - f.write(df.to_json(orient="records", lines=True)) + if is_json_mode: + json_output = df.to_json(orient="records", lines=True) + if output_file: + with open(output_file, "w", encoding="utf-8") as f: + f.write(json_output) else: - console.print(df.to_json(orient="records", lines=True)) + console.print(json_output) else: console.print(df) @@ -698,6 +700,5 @@ def query_aggregate( logger.exception(f"Error when running query: {e}") console.print(f"[red]Error when running query (note that aggregate query requires 'FROM aggregates' clause): [/red]" + str(e)) return CommandResult(success=False, msg=str(e)) - - return CommandResult(success=True, object=all_files, msg=f"Aggregated {len(all_files)} filesystems") - + + return CommandResult(success=True, object=df, msg=f"Aggregated {len(_aggregates)} records")