Merge pull request #2045 from dbcli/RW/save-plots-like-parquets

Ability to save plots with the `.>` operator

Author: rolandwalker

changelog.md

@@ -7,6 +7,7 @@ Features
 * Allow file target of `$>` redirection to be quoted.
 * Display of inline plots returned from `.|` operations.
 * Don't attempt inline plots in the Windows console.
+* Save plots, like parquets, with the `.>` operator.
 
 
 Bug Fixes

doc/transforms.md

@@ -22,9 +22,8 @@ functionality may still change.
 Here are some known limitations:
 
  * transforms can't be composed with `$|` shell redirection
- * multiple transform steps are not permitted
+ * multiple transform steps are not supported
  * results from `UNION`s may be unable to be transformed
- * images cannot yet be saved
  * PNG images are static and do not support all Altair features
 
 And there are inherent limitations to the post-processing model: the entire
@@ -84,30 +83,31 @@ the `[dataframe]` section of `~/.myclirc`.
 ## Saving
 
 A query result, transformed `DataFrame`, or transformed `Series` can be
-written directly to a Parquet file with the `.>` operator.
+written directly to a Parquet file with the `.>` operator. An Altair plot can
+also be written to a PNG file with the same operator.
 
 Save example:
 
 ```sql
 SELECT * FROM orders .> orders.parquet;
 ```
 
-The `.>` operator must be last, requires a `.parquet` destination, and
+The `.>` operator must be last, requires a `.parquet` or `.png` destination, and
 overwrites any existing file.  Spaces may be required around the operator.
 Destination paths containing whitespace must be quoted.  A successful write
-reports its destination and row count.
+reports its destination and row count if appropriate.
 
 When `post_redirect_command` is set in `~/.myclirc`, the given command runs
-after a successful Parquet save.
+after a successful file save.
 
 `.>` cannot be combined with the `\x` or `\G` special display terminators.
 
 ## Combining
 
-Parquet saves may be combined with dataframe transforms.  Again, `.>` must
-be the last operator.
+Parquet saves may be combined with dataframe transforms.  Again, the save
+operator `.>` must be the last operator.
 
-Combined transform and save example:
+Combined transform and save examples:
 
 ```sql
 SELECT * FROM orders .| df.group_by('customer_id').len() .> customer_counts.parquet;
@@ -116,3 +116,7 @@ SELECT * FROM orders .| df.group_by('customer_id').len() .> customer_counts.parq
 ```sql
 SELECT * FROM orders .| df['order_id'] .> order_ids.parquet;
 ```
+
+```sql
+SELECT * FROM orders .| df['total'].plot.hist() .> total_histogram.png;
+```

mycli/main_modes/repl.py

@@ -750,7 +750,7 @@ def _one_iteration(
         mycli.explorer_formatter.query = text
         if polars_transform is not None:
             assert polars_pipeline is not None
-            if polars_pipeline.parquet_path is None:
+            if polars_pipeline.output_path is None:
                 polars_result = run_polars_transform(
                     polars_transform,
                     results,
@@ -763,22 +763,22 @@ def _one_iteration(
                 polars_result = run_polars_transform(
                     polars_transform,
                     results,
-                    polars_pipeline.parquet_path,
+                    polars_pipeline.output_path,
                     image_protocol=mycli.image_protocol,
                     plot_scale_factor=mycli.plot_scale_factor,
                     plot_ppi=mycli.plot_ppi,
                     plot_theme=mycli.plot_theme,
                 )
-            if polars_pipeline.parquet_path is None:
+            if polars_pipeline.output_path is None:
                 if polars_pipeline.output_mode == 'explorer':
                     special.set_explorer_output(True)
                 elif polars_pipeline.output_mode == 'expanded':
                     special.set_expanded_output(True)
             _output_results(mycli, state, iter([polars_result]), start)
-            if polars_pipeline.parquet_path is not None:
+            if polars_pipeline.output_path is not None:
                 special.run_post_redirect_hook(
                     mycli.post_redirect_command,
-                    polars_pipeline.parquet_path,
+                    polars_pipeline.output_path,
                 )
         else:
             _output_results(mycli, state, results, start)

mycli/packages/polars_transform.py

@@ -23,7 +23,7 @@ class PolarsTransformError(RuntimeError):
 class PolarsPipeline:
     sql: str
     expression: str | None
-    parquet_path: str | None
+    output_path: str | None
     output_mode: OutputMode
 
 
@@ -37,7 +37,7 @@ class PolarsTransform:
 
 
 def parse_polars_transform(command: str) -> PolarsPipeline | None:
-    """Parse a SQL statement with optional Polars transform and Parquet output."""
+    """Parse a SQL statement with optional Polars transform and file output."""
     try:
         tokens = sqlglot.tokenize(command)
     except sqlglot.errors.TokenError as exc:
@@ -56,7 +56,7 @@ def parse_polars_transform(command: str) -> PolarsPipeline | None:
         if following.end + 1 >= len(command):
             if following.token_type == sqlglot.TokenType.PIPE:
                 raise PolarsTransformError('Polars transforms require a Python expression.')
-            raise PolarsTransformError('Parquet saves require a destination path.')
+            raise PolarsTransformError('File saves require a destination path.')
         if not command[following.end + 1].isspace():
             continue
         if following.token_type == sqlglot.TokenType.PIPE:
@@ -65,7 +65,7 @@ def parse_polars_transform(command: str) -> PolarsPipeline | None:
             pipe_index = index
         else:
             if parquet_index is not None:
-                raise PolarsTransformError('Parquet saves support only one ".>" operator.')
+                raise PolarsTransformError('File saves support only one ".>" operator.')
             parquet_index = index
 
     if pipe_index is None and parquet_index is None:
@@ -77,20 +77,20 @@ def parse_polars_transform(command: str) -> PolarsPipeline | None:
     assert first_index is not None
     sql = command[: tokens[first_index].start].strip()
     expression: str | None = None
-    parquet_path: str | None = None
+    output_path: str | None = None
     if pipe_index is not None:
         pipe_operator = tokens[pipe_index + 1]
         expression_end = tokens[parquet_index].start if parquet_index is not None else len(command)
         expression = command[pipe_operator.end + 1 : expression_end].strip()
     if parquet_index is not None:
         parquet_operator = tokens[parquet_index + 1]
-        parquet_path = command[parquet_operator.end + 1 :].strip()
-        parquet_path = parquet_path.removesuffix(delimiter_command.current).rstrip()
-        parquet_path = parquet_path.removesuffix(r'\g').rstrip()
+        output_path = command[parquet_operator.end + 1 :].strip()
+        output_path = output_path.removesuffix(delimiter_command.current).rstrip()
+        output_path = output_path.removesuffix(r'\g').rstrip()
 
-    has_display_terminator = any(value is not None and value.endswith((r'\x', r'\G')) for value in (sql, expression, parquet_path))
-    if parquet_path is not None and has_display_terminator:
-        raise PolarsTransformError('Parquet saves cannot use special display terminators.')
+    has_display_terminator = any(value is not None and value.endswith((r'\x', r'\G')) for value in (sql, expression, output_path))
+    if output_path is not None and has_display_terminator:
+        raise PolarsTransformError('File saves cannot use special display terminators.')
     if sql.endswith(r'\x') or expression is not None and expression.endswith(r'\x'):
         output_mode: OutputMode = 'explorer'
     elif sql.endswith(r'\G') or expression is not None and expression.endswith(r'\G'):
@@ -107,23 +107,23 @@ def parse_polars_transform(command: str) -> PolarsPipeline | None:
         raise PolarsTransformError('Polars transforms require a SQL statement.')
     if expression is not None and not expression:
         raise PolarsTransformError('Polars transforms require a Python expression.')
-    if parquet_path is not None:
-        if not parquet_path:
-            raise PolarsTransformError('Parquet saves require a destination path.')
-        parquet_path = _parse_parquet_path(parquet_path)
+    if output_path is not None:
+        if not output_path:
+            raise PolarsTransformError('File saves require a destination path.')
+        output_path = _parse_output_path(output_path)
     _validate_sql(sql)
-    return PolarsPipeline(sql=sql, expression=expression, parquet_path=parquet_path, output_mode=output_mode)
+    return PolarsPipeline(sql=sql, expression=expression, output_path=output_path, output_mode=output_mode)
 
 
-def _parse_parquet_path(path: str) -> str:
+def _parse_output_path(path: str) -> str:
     if path[0] in ('\'', '"'):
         if len(path) < 2 or path[-1] != path[0]:
-            raise PolarsTransformError('Parquet save paths must use matching quotes.')
+            raise PolarsTransformError('File save paths must use matching quotes.')
         path = path[1:-1]
     elif any(character.isspace() for character in path):
-        raise PolarsTransformError('Parquet save paths containing spaces must be quoted.')
-    if not path.lower().endswith('.parquet'):
-        raise PolarsTransformError('Parquet save paths must end in ".parquet".')
+        raise PolarsTransformError('File save paths containing spaces must be quoted.')
+    if not path.lower().endswith(('.parquet', '.png')):
+        raise PolarsTransformError('File save paths must end in ".parquet" or ".png".')
     return path
 
 
@@ -181,7 +181,7 @@ def _load_vl_convert() -> None:
 def run_polars_transform(
     transform: PolarsTransform,
     results: Iterable[SQLResult],
-    parquet_path: str | None = None,
+    output_path: str | None = None,
     *,
     image_protocol: ImageProtocol = 'none',
     plot_scale_factor: float = 1.0,
@@ -212,34 +212,49 @@ def run_polars_transform(
     except Exception as exc:
         raise PolarsTransformError(f'Polars expression failed: {type(exc).__name__}: {exc}') from exc
     if isinstance(value, transform.polars.DataFrame):
-        if parquet_path is not None:
+        if output_path is not None:
+            if not output_path.lower().endswith('.parquet'):
+                raise PolarsTransformError('Polars DataFrame results can only be written to ".parquet" files.')
             try:
-                value.write_parquet(parquet_path)
+                value.write_parquet(output_path)
             except Exception as exc:
-                raise PolarsTransformError(f'Unable to write Parquet file "{parquet_path}": {type(exc).__name__}: {exc}') from exc
-            return SQLResult(status=f'Wrote {len(value)} rows to {parquet_path}.')
+                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}.')
         return SQLResult(header=list(value.columns), rows=list(value.iter_rows()))
     if isinstance(value, transform.polars.Series):
         column_name = value.name or 'value'
-        if parquet_path is not None:
+        if output_path is not None:
+            if not output_path.lower().endswith('.parquet'):
+                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(parquet_path)
+                series_dataframe.write_parquet(output_path)
             except Exception as exc:
-                raise PolarsTransformError(f'Unable to write Parquet file "{parquet_path}": {type(exc).__name__}: {exc}') from exc
-            return SQLResult(status=f'Wrote {len(series_dataframe)} rows to {parquet_path}.')
+                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}.')
         return SQLResult(header=[column_name], rows=[(item,) for item in value])
     if transform.altair is not None and isinstance(value, transform.altair.TopLevelMixin):
-        if parquet_path is not None:
-            raise PolarsTransformError('Polars transforms must return a DataFrame or Series before writing Parquet output.')
-        if image_protocol == 'none':
+        if output_path is not None and not output_path.lower().endswith('.png'):
+            raise PolarsTransformError('Altair charts can only be written to ".png" files.')
+        if output_path is None and image_protocol == 'none':
             return SQLResult(status='image_protocol is unset in ~/.myclirc. Inline plotting is disabled.')
         _load_vl_convert()
-        png = BytesIO()
         try:
             transform.altair.theme.enable(plot_theme)
         except Exception as exc:
             raise PolarsTransformError(f'Unable to enable Altair plot theme "{plot_theme}": {type(exc).__name__}: {exc}') from exc
+        if output_path is not None:
+            try:
+                value.save(
+                    output_path,
+                    format='png',
+                    scale_factor=plot_scale_factor,
+                    ppi=plot_ppi,
+                )
+            except Exception as exc:
+                raise PolarsTransformError(f'Unable to write PNG file "{output_path}": {type(exc).__name__}: {exc}') from exc
+            return SQLResult(status=f'Wrote PNG image to {output_path}.')
+        png = BytesIO()
         try:
             value.save(
                 png,
@@ -250,6 +265,8 @@ def run_polars_transform(
         except Exception as exc:
             raise PolarsTransformError(f'Unable to render Altair chart: {type(exc).__name__}: {exc}') from exc
         return SQLResult(image=png.getvalue(), image_protocol=image_protocol)
-    if parquet_path is not None:
-        raise PolarsTransformError('Polars transforms must return a DataFrame or Series before writing Parquet output.')
+    if output_path is not None:
+        raise PolarsTransformError(
+            'Polars transforms must return a DataFrame or Series for Parquet output, or an Altair chart for PNG output.'
+        )
     return SQLResult(status=f'Nothing could be displayed for return type: {type(value)}')

test/pytests/test_main_modes_repl.py

@@ -1374,6 +1374,61 @@ def run(
     assert hook_calls == [('post {}', 'orders.parquet')]
 
 
+def test_one_iteration_writes_polars_png_and_runs_post_redirect_hook(monkeypatch: pytest.MonkeyPatch) -> None:
+    patch_repl_runtime_defaults(monkeypatch)
+
+    class FakeSQLExecute:
+        dbname = 'db'
+        connection_id = 0
+
+        def run(self, text: str) -> Iterator[SQLResult]:
+            assert text == 'SELECT * FROM orders'
+            return iter([SQLResult(header=['id'], rows=[(1,)])])
+
+    cli = make_repl_cli(FakeSQLExecute())
+    cli.post_redirect_command = 'post {}'
+    transform = object()
+    run_calls: list[str] = []
+    hook_calls: list[tuple[str, str]] = []
+
+    monkeypatch.setattr(repl_mode, 'prepare_polars_transform', lambda sql, expression: transform)
+
+    def run(
+        received_transform: object,
+        results: Iterator[SQLResult],
+        path: 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 image_protocol == 'none'
+        assert plot_scale_factor == 1.0
+        assert plot_ppi == 200
+        assert plot_theme == 'carbong90'
+        run_calls.append(path)
+        return SQLResult(status=f'Wrote PNG image to {path}.')
+
+    monkeypatch.setattr(repl_mode, 'run_polars_transform', run)
+    monkeypatch.setattr(
+        repl_mode.special,
+        'run_post_redirect_hook',
+        lambda command, filename: hook_calls.append((command, filename)),
+    )
+
+    repl_mode._one_iteration(
+        cli,
+        repl_mode.ReplState(),
+        'SELECT * FROM orders .| alt.Chart(df) .> orders.png',
+    )
+
+    assert run_calls == ['orders.png']
+    assert hook_calls == [('post {}', 'orders.png')]
+
+
 def test_one_iteration_reports_polars_post_redirect_hook_error(monkeypatch: pytest.MonkeyPatch) -> None:
     patch_repl_runtime_defaults(monkeypatch)