Preserve query as Parquet file-level metadata

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)
 ==============
 

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,

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})

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,13 +1870,15 @@ def run(
         results: Iterator[SQLResult],
         path: str,
         *,
+        original_query: str,
         image_protocol: str,
         plot_scale_factor: float,
         plot_ppi: int,
         plot_theme: str,
     ) -> 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,

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)", '<test>', '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)