From 87c6e4445f76fe198b50c47ef7745c50b7bce4dc Mon Sep 17 00:00:00 2001 From: Roland Walker Date: Thu, 10 Sep 2026 14:59:54 -0400 Subject: [PATCH] preserve query as Parquet file-level metadata Since the Parquet file format supports file-level metadata, when saving a Parquet with .>, we can preserve the originating query in the file. The .> operator and target filename are preserved in the query, which might be debatable. The Polars documentation also warns that the metadata interface is experimental, but our tests should catch any issues on a dependency upgrade. --- changelog.md | 8 ++++++ mycli/main_modes/repl.py | 1 + mycli/packages/polars_transform.py | 12 ++++++-- test/pytests/test_main_modes_repl.py | 22 +++++++++------ test/pytests/test_polars_transform.py | 40 +++++++++++++++++++++++++-- 5 files changed, 69 insertions(+), 14 deletions(-) diff --git a/changelog.md b/changelog.md index 24d591caa..0b2f1e32f 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 c5d97d914..866839398 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 e6d3bb57e..65ae67977 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 e6141a3fb..badf448ab 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 a25cd2ab5..42a54fb99 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)