diff --git a/changelog.md b/changelog.md index 24d591ca..0b2f1e32 100644 --- a/changelog.md +++ b/changelog.md @@ -1,3 +1,11 @@ +Upcoming (TBD) +============== + +Features +-------- +* Preserve the query as metadata when saving to Parquet with `.>`. + + 2.23.0 (2026/09/09) ============== diff --git a/mycli/main_modes/repl.py b/mycli/main_modes/repl.py index c5d97d91..86683939 100644 --- a/mycli/main_modes/repl.py +++ b/mycli/main_modes/repl.py @@ -940,6 +940,7 @@ def _one_iteration( polars_transform, results, polars_pipeline.output_path, + original_query=original_text, image_protocol=mycli.image_protocol, plot_scale_factor=mycli.plot_scale_factor, plot_ppi=mycli.plot_ppi, diff --git a/mycli/packages/polars_transform.py b/mycli/packages/polars_transform.py index e6d3bb57..65ae6797 100644 --- a/mycli/packages/polars_transform.py +++ b/mycli/packages/polars_transform.py @@ -17,6 +17,7 @@ delimiter_command = DelimiterCommand() PLOT_FORMATS = ('png', 'pdf', 'svg', 'html') +PARQUET_QUERY_METADATA_KEY = 'mycli_query' class PolarsTransformError(RuntimeError): @@ -151,7 +152,7 @@ def _pipeline_operator_indexes( def _parse_output_path(path: str) -> str: - if path[0] in ('\'', '"'): + if path[0] in ("'", '"'): if len(path) < 2 or path[-1] != path[0]: raise PolarsTransformError('File save paths must use matching quotes.') path = path[1:-1] @@ -223,6 +224,7 @@ def run_polars_transform( results: Iterable[SQLResult], output_path: str | None = None, *, + original_query: str | None = None, image_protocol: ImageProtocol = 'none', plot_scale_factor: float = 1.0, plot_ppi: int = 200, @@ -278,7 +280,7 @@ def run_polars_transform( if not output_path.lower().endswith('.parquet'): raise PolarsTransformError('Polars DataFrame results can only be written to ".parquet" files.') try: - value.write_parquet(output_path) + _write_parquet(value, output_path, original_query or transform.sql) except Exception as exc: raise PolarsTransformError(f'Unable to write Parquet file "{output_path}": {type(exc).__name__}: {exc}') from exc return SQLResult(status=f'Wrote {len(value)} rows to {output_path}.') @@ -290,7 +292,7 @@ def run_polars_transform( raise PolarsTransformError('Polars Series results can only be written to ".parquet" files.') try: series_dataframe = value.rename(column_name).to_frame() - series_dataframe.write_parquet(output_path) + _write_parquet(series_dataframe, output_path, original_query or transform.sql) except Exception as exc: raise PolarsTransformError(f'Unable to write Parquet file "{output_path}": {type(exc).__name__}: {exc}') from exc return SQLResult(status=f'Wrote {len(series_dataframe)} rows to {output_path}.') @@ -339,3 +341,7 @@ def run_polars_transform( return SQLResult() return SQLResult(status=f'Nothing could be displayed for return type: {type(value)}') + + +def _write_parquet(dataframe: Any, parquet_path: str, original_query: str) -> None: + dataframe.write_parquet(parquet_path, metadata={PARQUET_QUERY_METADATA_KEY: original_query}) diff --git a/test/pytests/test_main_modes_repl.py b/test/pytests/test_main_modes_repl.py index e6141a3f..badf448a 100644 --- a/test/pytests/test_main_modes_repl.py +++ b/test/pytests/test_main_modes_repl.py @@ -1731,7 +1731,7 @@ def run(self, text: str) -> Iterator[SQLResult]: cli.post_redirect_command = 'post {}' transform = object() prepare_calls: list[tuple[str, str | None]] = [] - run_calls: list[tuple[object, str]] = [] + run_calls: list[tuple[object, str, str]] = [] def prepare(sql: str, expression: str | None) -> object: prepare_calls.append((sql, expression)) @@ -1742,6 +1742,7 @@ def run( results: Iterator[SQLResult], path: str, *, + original_query: str, image_protocol: str, plot_scale_factor: float, plot_ppi: int, @@ -1752,7 +1753,7 @@ def run( assert plot_ppi == 200 assert plot_theme == 'carbong90' assert list(results) == [SQLResult(header=['id'], rows=[(1,)])] - run_calls.append((received_transform, path)) + run_calls.append((received_transform, path, original_query)) return SQLResult(status=f'Wrote 1 rows to {path}.') monkeypatch.setattr(repl_mode, 'prepare_polars_transform', prepare) @@ -1770,7 +1771,7 @@ def run_hook(command: str, filename: str) -> None: assert sqlexecute.calls == ['SELECT * FROM orders'] assert prepare_calls == [('SELECT * FROM orders', None)] - assert run_calls == [(transform, 'orders.parquet')] + assert run_calls == [(transform, 'orders.parquet', command)] assert cli.output_calls[-1][1] == SQLResult(status='Wrote 1 rows to orders.parquet.') assert cli.output_calls[-1][1].rows is None assert hook_calls == [('post {}', 'orders.parquet')] @@ -1794,7 +1795,7 @@ def run(self, text: str) -> Iterator[SQLResult]: cli.post_redirect_command = 'post {}' transform = object() prepare_calls: list[tuple[str, str | None]] = [] - run_calls: list[str] = [] + run_calls: list[tuple[str, str]] = [] def prepare(sql: str, expression: str | None) -> object: prepare_calls.append((sql, expression)) @@ -1805,6 +1806,7 @@ def run( results: Iterator[SQLResult], path: str, *, + original_query: str, image_protocol: str, plot_scale_factor: float, plot_ppi: int, @@ -1816,7 +1818,7 @@ def run( assert plot_theme == 'carbong90' assert received_transform is transform assert list(results) == [SQLResult(header=['id'], rows=[(1,)])] - run_calls.append(path) + run_calls.append((path, original_query)) return SQLResult(status=f'Wrote 1 rows to {path}.') monkeypatch.setattr(repl_mode, 'prepare_polars_transform', prepare) @@ -1828,14 +1830,15 @@ def run( lambda command, filename: hook_calls.append((command, filename)), ) + command = 'SELECT * FROM orders .| df.filter(pl.col(\'id\') > 0) .> orders.parquet' repl_mode._one_iteration( cli, repl_mode.ReplState(), - 'SELECT * FROM orders .| df.filter(pl.col(\'id\') > 0) .> orders.parquet', + command, ) assert prepare_calls == [('SELECT * FROM orders', "df.filter(pl.col('id') > 0)")] - assert run_calls == ['orders.parquet'] + assert run_calls == [('orders.parquet', command)] assert hook_calls == [('post {}', 'orders.parquet')] @@ -1867,6 +1870,7 @@ def run( results: Iterator[SQLResult], path: str, *, + original_query: str, image_protocol: str, plot_scale_factor: float, plot_ppi: int, @@ -1874,6 +1878,7 @@ def run( ) -> SQLResult: assert received_transform is transform assert list(results) == [SQLResult(header=['id'], rows=[(1,)])] + assert original_query == f'SELECT * FROM orders .| alt.Plot(df) .> {path}' assert image_protocol == 'none' assert plot_scale_factor == 1.0 assert plot_ppi == 200 @@ -1917,7 +1922,7 @@ def run(self, text: str) -> Iterator[SQLResult]: monkeypatch.setattr( repl_mode, 'run_polars_transform', - lambda received_transform, results, path, *, image_protocol, plot_scale_factor, plot_ppi, plot_theme: SQLResult( + lambda received_transform, results, path, *, original_query, image_protocol, plot_scale_factor, plot_ppi, plot_theme: SQLResult( status=f'Wrote 1 rows to {path}.' ), ) @@ -1957,6 +1962,7 @@ def fail_write( results: Iterator[SQLResult], path: str, *, + original_query: str, image_protocol: str, plot_scale_factor: float, plot_ppi: int, diff --git a/test/pytests/test_polars_transform.py b/test/pytests/test_polars_transform.py index a25cd2ab..42a54fb9 100644 --- a/test/pytests/test_polars_transform.py +++ b/test/pytests/test_polars_transform.py @@ -21,6 +21,7 @@ class FakeDataFrame: written_paths: list[str] = [] written_dataframes: list['FakeDataFrame'] = [] + written_metadata: list[dict[str, str]] = [] def __init__( self, @@ -42,10 +43,11 @@ def iter_rows(self) -> Iterator[tuple[Any, ...]]: def __len__(self) -> int: return len(self.rows) - def write_parquet(self, path: str) -> None: + def write_parquet(self, path: str, *, metadata: dict[str, str]) -> None: self.parquet_paths.append(path) self.written_paths.append(path) self.written_dataframes.append(self) + self.written_metadata.append(metadata) class FakeSeries: @@ -80,7 +82,7 @@ class FakePolars: class FailingDataFrame(FakeDataFrame): - def write_parquet(self, path: str) -> None: + def write_parquet(self, path: str, *, metadata: dict[str, str]) -> None: raise OSError('disk full') @@ -572,15 +574,42 @@ def test_run_polars_transform_writes_raw_dataframe_to_parquet() -> None: transform = make_transform('df') FakeDataFrame.written_paths = [] FakeDataFrame.written_dataframes = [] + FakeDataFrame.written_metadata = [] result = run_polars_transform( transform, iter([SQLResult(header=['id'], rows=[(1,), (2,)])]), 'orders.parquet', + original_query='SELECT * FROM orders .> orders.parquet', ) assert result == SQLResult(status='Wrote 2 rows to orders.parquet.') assert FakeDataFrame.written_paths == ['orders.parquet'] + assert FakeDataFrame.written_metadata == [{polars_transform.PARQUET_QUERY_METADATA_KEY: 'SELECT * FROM orders .> orders.parquet'}] + + +def test_run_polars_transform_writes_query_to_parquet_file_metadata(tmp_path: Path) -> None: + import polars as pl + + command = 'SELECT id FROM orders .| df.filter(pl.col(\'id\') > 0) .> orders.parquet' + path = tmp_path / 'orders.parquet' + transform = PolarsTransform( + sql='SELECT id FROM orders', + expression="df.filter(pl.col('id') > 0)", + code=compile("df.filter(pl.col('id') > 0)", '', 'eval'), + polars=pl, + altair=None, + ) + + result = run_polars_transform( + transform, + iter([SQLResult(header=['id'], rows=[(1,), (2,)])]), + str(path), + original_query=command, + ) + + assert result == SQLResult(status=f'Wrote 2 rows to {path}.') + assert pl.read_parquet_metadata(path)[polars_transform.PARQUET_QUERY_METADATA_KEY] == command @pytest.mark.parametrize( @@ -597,21 +626,26 @@ def test_run_polars_transform_writes_series_to_parquet( ) -> None: FakeDataFrame.written_paths = [] FakeDataFrame.written_dataframes = [] + FakeDataFrame.written_metadata = [] result = run_polars_transform( make_transform(expression), iter([SQLResult(header=['id'], rows=[(1,)])]), 'series.parquet', + original_query=f'SELECT id FROM orders .| {expression} .> series.parquet', ) assert result == SQLResult(status=f'Wrote {len(rows)} rows to series.parquet.') assert FakeDataFrame.written_paths == ['series.parquet'] assert FakeDataFrame.written_dataframes[-1].columns == [column_name] assert FakeDataFrame.written_dataframes[-1].rows == rows + assert FakeDataFrame.written_metadata == [ + {polars_transform.PARQUET_QUERY_METADATA_KEY: (f'SELECT id FROM orders .| {expression} .> series.parquet')} + ] def test_run_polars_transform_reports_series_parquet_write_error(monkeypatch: pytest.MonkeyPatch) -> None: - def fail_write(self: FakeDataFrame, path: str) -> None: + def fail_write(self: FakeDataFrame, path: str, *, metadata: dict[str, str]) -> None: raise OSError('disk full') monkeypatch.setattr(FakeDataFrame, 'write_parquet', fail_write)