ml-analytics-tools 0.4.3__tar.gz → 0.4.4__tar.gz

{ml_analytics_tools-0.4.3/ml_analytics_tools.egg-info → ml_analytics_tools-0.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ml-analytics-tools
-Version: 0.4.3
+Version: 0.4.4
 Summary: Tools for ML projects and data management
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
@@ -220,6 +220,14 @@ sf.sql("queries/experiment.sql", save_table=True, schema="analytics", table="exp
 
 # or save any Spark DataFrame to Unity Catalog
 sf.save_to_uc(df, table="exp", schema="analytics", catalog="prod")
+
+# save a YAML-ordered folder of SQL queries as Unity Catalog tables
+df = sf.save_pipeline_to_uc(
+    "queries/churn_pipeline",
+    pipeline="daily",
+    catalog="prod",
+    schema="analytics",
+)
 ```
 
 Credentials resolve per field as: explicit argument → `SNOWFLAKE_*` environment

{ml_analytics_tools-0.4.3 → ml_analytics_tools-0.4.4}/README.md

@@ -185,6 +185,14 @@ sf.sql("queries/experiment.sql", save_table=True, schema="analytics", table="exp
 
 # or save any Spark DataFrame to Unity Catalog
 sf.save_to_uc(df, table="exp", schema="analytics", catalog="prod")
+
+# save a YAML-ordered folder of SQL queries as Unity Catalog tables
+df = sf.save_pipeline_to_uc(
+    "queries/churn_pipeline",
+    pipeline="daily",
+    catalog="prod",
+    schema="analytics",
+)
 ```
 
 Credentials resolve per field as: explicit argument → `SNOWFLAKE_*` environment

{ml_analytics_tools-0.4.3 → ml_analytics_tools-0.4.4}/ml_analytics/__init__.py

@@ -19,6 +19,7 @@ from .utils import (
     get_sql_files,
     load_sql_query,
     log_and_raise_error,
+    resolve_sql_query_paths,
 )
 
 # Automatically load .env file when the package is imported
@@ -50,6 +51,7 @@ __all__ = [
     "load_sql_query",
     "log_and_raise_error",
     "ModelManager",
+    "resolve_sql_query_paths",
     "S3Connector",
     "SFConnector",
     "SlackConnector",

{ml_analytics_tools-0.4.3 → ml_analytics_tools-0.4.4}/ml_analytics/data_connector.py

@@ -232,6 +232,7 @@ class DataConnector:
         schema=None,
         role=None,
         authenticator=None,
+        token=None,
         private_key=None,
         private_key_path=None,
         private_key_passphrase=None,
@@ -262,6 +263,7 @@ class DataConnector:
                 schema=schema,
                 role=role,
                 authenticator=authenticator,
+                token=token,
                 private_key=private_key,
                 private_key_path=private_key_path,
                 private_key_passphrase=private_key_passphrase,
@@ -335,6 +337,7 @@ class DataConnector:
         schema=None,
         role=None,
         authenticator=None,
+        token=None,
         private_key=None,
         private_key_path=None,
         private_key_passphrase=None,
@@ -361,6 +364,7 @@ class DataConnector:
         )
         token = _get_snowflake_config_value(
             "SNOWFLAKE_TOKEN",
+            explicit=token,
             secret_scope=secret_scope,
             aliases=("SNOWFLAKE_OAUTH_TOKEN", "SNOWFLAKE_ACCESS_TOKEN"),
         )

{ml_analytics_tools-0.4.3 → ml_analytics_tools-0.4.4}/ml_analytics/sf_connector.py

@@ -18,7 +18,7 @@ from .data_connector import (
     _load_private_key_pem_for_spark,
     _snowflake_secret_scope,
 )
-from .utils import get_logger, load_sql_query, log_and_raise_error
+from .utils import get_logger, load_sql_query, log_and_raise_error, resolve_sql_query_paths
 
 # Cached Spark session shared across SFConnector instances. Populated lazily by
 # get_spark(); never created at import time so the package stays importable
@@ -254,6 +254,10 @@ class SFConnector:
         schema: str = None,
         catalog: str = None,
         mode: str = "overwrite",
+        optimize: bool = True,
+        zorder_by=None,
+        merge_schema: bool = True,
+        comment: str = None,
         **kwargs,
     ):
         """
@@ -284,6 +288,15 @@ class SFConnector:
         mode : str, optional
             Spark write mode for the saved table ('overwrite', 'append',
             'ignore', 'error'). Defaults to 'overwrite'.
+        optimize : bool, optional
+            If saving to Unity Catalog, run ``OPTIMIZE`` after the write.
+            Defaults to True.
+        zorder_by : str or list[str], optional
+            Optional columns for Delta ``ZORDER BY`` during optimize.
+        merge_schema : bool, optional
+            If saving to Unity Catalog, set Delta ``mergeSchema=true``. Defaults to True.
+        comment : str, optional
+            Optional table comment stored as a Unity Catalog table property.
         **kwargs
             Template variables substituted into the SQL file using ``str.format()``.
         """
@@ -295,12 +308,120 @@ class SFConnector:
             log_and_raise_error(self._logger, f"Error reading from Snowflake: {e}")
 
         if save_table:
-            self.save_to_uc(df, table=table, schema=schema, catalog=catalog, mode=mode)
+            self.save_to_uc(
+                df,
+                table=table,
+                schema=schema,
+                catalog=catalog,
+                mode=mode,
+                optimize=optimize,
+                zorder_by=zorder_by,
+                merge_schema=merge_schema,
+                comment=comment,
+            )
 
         if return_pandas:
             return df.toPandas()
         return df
 
+    def save_pipeline_to_uc(
+        self,
+        query_paths,
+        *,
+        pipeline: str | None = None,
+        catalog: str = None,
+        schema: str = None,
+        tables: dict[str, str] = None,
+        table_prefix: str = "",
+        table_suffix: str = "",
+        mode: str = "overwrite",
+        modes: dict[str, str] = None,
+        optimize: bool = True,
+        zorder_by=None,
+        merge_schema: bool = True,
+        comment: str = None,
+        comments: dict[str, str] = None,
+        return_all: bool = False,
+        **kwargs,
+    ):
+        """
+        Run YAML-ordered Snowflake queries and save each result as a Unity Catalog table.
+
+        This is a convenience wrapper around ``sql(..., save_table=True)``. It
+        uses the same folder/YAML resolution as ``execute_sql_scripts``:
+        ``steps`` define the SQL files to run and their order.
+
+        Parameters
+        ----------
+        query_paths
+            Folder, file, list, or ordered dict of SQL files.
+        pipeline
+            Optional YAML pipeline name.
+        catalog, schema
+            Default Unity Catalog destination for unqualified table names.
+        tables
+            Optional mapping of step name to destination table. Values may be
+            unqualified (using ``catalog`` / ``schema``) or fully qualified.
+        table_prefix, table_suffix
+            Applied to step names when ``tables`` does not define a destination.
+        mode
+            Default Spark write mode for every table.
+        modes
+            Optional mapping of step name to Spark write mode.
+        optimize
+            If True, run ``OPTIMIZE`` after saving each Unity Catalog table.
+        zorder_by
+            Optional columns for Delta ``ZORDER BY``. Pass a dict to configure
+            columns per step, or a string/list to use the same columns for every
+            saved table.
+        merge_schema
+            If True, set Delta ``mergeSchema=true`` for every saved table.
+        comment
+            Optional table comment applied to every saved table.
+        comments
+            Optional mapping of step name to table comment.
+        return_all
+            If True, return a dict of step name to Spark DataFrame. Otherwise
+            return the last step's Spark DataFrame.
+        **kwargs
+            Template variables substituted into SQL files via ``str.format()``.
+        """
+        resolved_paths = resolve_sql_query_paths(query_paths, pipeline=pipeline)
+        if not resolved_paths:
+            log_and_raise_error(self._logger, "No SQL files found for pipeline.")
+
+        tables = tables or {}
+        modes = modes or {}
+        comments = comments or {}
+        results = {}
+        last_df = None
+
+        for name, query_path in resolved_paths.items():
+            destination = tables.get(name) or f"{table_prefix}{name}{table_suffix}"
+            if not destination:
+                log_and_raise_error(self._logger, f"No Unity Catalog table configured for step '{name}'.")
+
+            step_mode = modes.get(name, mode)
+            step_zorder_by = zorder_by.get(name) if isinstance(zorder_by, dict) else zorder_by
+            step_comment = comments.get(name, comment)
+            self._logger.info(f"[{name}] saving to Unity Catalog table {destination} (mode={step_mode}) ...")
+            last_df = self.sql(
+                str(query_path),
+                save_table=True,
+                table=destination,
+                schema=schema,
+                catalog=catalog,
+                mode=step_mode,
+                optimize=optimize,
+                zorder_by=step_zorder_by,
+                merge_schema=merge_schema,
+                comment=step_comment,
+                **kwargs,
+            )
+            results[name] = last_df
+
+        return results if return_all else last_df
+
     @staticmethod
     def _qualified_uc_name(table: str, schema: str = None, catalog: str = None) -> str:
         """Build a Unity Catalog table identifier from its parts.
@@ -313,12 +434,91 @@ class SFConnector:
         parts = [part for part in (catalog, schema, table) if part]
         return ".".join(parts)
 
-    def save_to_uc(self, df, table: str, schema: str = None, catalog: str = None, mode: str = "overwrite"):
+    @staticmethod
+    def _zorder_clause(zorder_by=None) -> str:
+        """Build the optional Delta ZORDER BY clause."""
+        if not zorder_by:
+            return ""
+        if isinstance(zorder_by, str):
+            columns = [column.strip() for column in zorder_by.split(",")]
+        else:
+            columns = [str(column).strip() for column in zorder_by]
+        columns = [column for column in columns if column]
+        if not columns:
+            return ""
+        return f" ZORDER BY ({', '.join(columns)})"
+
+    @staticmethod
+    def _sql_string_literal(value: str) -> str:
+        """Escape a value for use inside a single-quoted SQL string literal."""
+        return str(value).replace("'", "''")
+
+    def set_uc_table_comment(self, table: str, comment: str, schema: str = None, catalog: str = None, spark=None):
+        """
+        Set a Unity Catalog table comment using Databricks table properties.
+
+        Parameters
+        ----------
+        table
+            Table name. May be fully qualified.
+        comment
+            Comment text to store.
+        schema, catalog
+            Optional qualifiers when ``table`` is not fully qualified.
+        spark
+            Optional SparkSession to use. Defaults to this connector's Spark session.
+        """
+        full_name = self._qualified_uc_name(table, schema=schema, catalog=catalog)
+        spark = spark or self._get_spark()
+        escaped_comment = self._sql_string_literal(comment)
+        try:
+            spark.sql(f"ALTER TABLE {full_name} SET TBLPROPERTIES ('comment' = '{escaped_comment}')")
+        except Exception as e:
+            log_and_raise_error(self._logger, f"Error setting comment for Unity Catalog table '{full_name}': {e}")
+        self._logger.info(f"Comment set for Unity Catalog table '{full_name}'.")
+
+    def optimize_uc_table(self, table: str, schema: str = None, catalog: str = None, zorder_by=None, spark=None):
+        """
+        Run Databricks Delta ``OPTIMIZE`` on a Unity Catalog table.
+
+        Parameters
+        ----------
+        table
+            Table name. May be fully qualified.
+        schema, catalog
+            Optional qualifiers when ``table`` is not fully qualified.
+        zorder_by
+            Optional column or columns for ``ZORDER BY``.
+        spark
+            Optional SparkSession to use. Defaults to this connector's Spark session.
+        """
+        full_name = self._qualified_uc_name(table, schema=schema, catalog=catalog)
+        spark = spark or self._get_spark()
+        optimize_sql = f"OPTIMIZE {full_name}{self._zorder_clause(zorder_by)}"
+        try:
+            spark.sql(optimize_sql)
+        except Exception as e:
+            log_and_raise_error(self._logger, f"Error optimizing Unity Catalog table '{full_name}': {e}")
+        self._logger.info(f"Table '{full_name}' optimized.")
+
+    def save_to_uc(
+        self,
+        df,
+        table: str,
+        schema: str = None,
+        catalog: str = None,
+        mode: str = "overwrite",
+        optimize: bool = True,
+        zorder_by=None,
+        merge_schema: bool = True,
+        comment: str = None,
+    ):
         """
         Write a Spark DataFrame to a Databricks Unity Catalog table.
 
         Uses Spark's native ``df.write.saveAsTable(...)`` (a managed UC table),
-        not the Snowflake connector.
+        not the Snowflake connector. By default, runs Delta ``OPTIMIZE`` after
+        the write.
 
         Parameters
         ----------
@@ -333,13 +533,31 @@ class SFConnector:
         mode : str, optional
             Spark write mode: 'overwrite', 'append', 'ignore', or 'error'.
             Defaults to 'overwrite'.
+        optimize : bool, optional
+            If True, run ``OPTIMIZE`` after saving. Defaults to True.
+        zorder_by : str or list[str], optional
+            Optional columns for Delta ``ZORDER BY`` during optimize.
+        merge_schema : bool, optional
+            If True, writes as Delta with ``mergeSchema=true``. Defaults to True.
+        comment : str, optional
+            Optional table comment stored as a Unity Catalog table property.
         """
         if not table:
             log_and_raise_error(self._logger, "A destination table name is required.")
 
         full_name = self._qualified_uc_name(table, schema=schema, catalog=catalog)
+        spark = getattr(df, "sparkSession", None) or self._spark
         try:
-            df.write.mode(mode).saveAsTable(full_name)
+            writer = df.write.format("delta")
+            if merge_schema:
+                writer = writer.option("mergeSchema", "true")
+            writer.mode(mode).saveAsTable(full_name)
         except Exception as e:
             log_and_raise_error(self._logger, f"Error writing to Unity Catalog table '{full_name}': {e}")
         self._logger.info(f"Table '{full_name}' written to Unity Catalog (mode={mode}).")
+
+        if comment is not None:
+            self.set_uc_table_comment(full_name, comment, spark=spark)
+
+        if optimize:
+            self.optimize_uc_table(full_name, zorder_by=zorder_by, spark=spark)

{ml_analytics_tools-0.4.3 → ml_analytics_tools-0.4.4}/ml_analytics/utils.py

@@ -624,6 +624,61 @@ def _is_select_statement(statement: str) -> bool:
     return True
 
 
+def resolve_sql_query_paths(query_paths, pipeline: str | None = None) -> dict[str, Path]:
+    """
+    Normalize SQL pipeline input into an ordered mapping of query name to file path.
+
+    Args:
+        query_paths: one of:
+            - str: relative folder path from project root; SQL files are discovered
+              via get_sql_files() (respects pipeline.yaml if present).
+            - Path pointing to a directory: same as str, resolved relative to project root.
+            - Path pointing to a single .sql file: executes that file only.
+            - list[str | Path]: ordered list of individual SQL file paths.
+            - dict[str, str | Path]: explicit ordered mapping of name -> path; preserves insertion order.
+        pipeline: Optional pipeline name passed to get_sql_files() when query_paths is a folder.
+
+    Returns:
+        Ordered dict[str, Path].
+    """
+    logger = get_logger("ml_analytics.utils.resolve_sql_query_paths")
+
+    if isinstance(query_paths, str):
+        try:
+            project_root = find_project_root()
+            candidate = project_root / query_paths
+        except FileNotFoundError:
+            candidate = Path(query_paths)
+        if candidate.is_file():
+            return {candidate.stem: candidate}
+        resolved = get_sql_files(query_paths, pipeline=pipeline)
+        if not resolved:
+            log_and_raise_error(logger, f"No SQL files found for folder '{query_paths}'.")
+        return resolved
+
+    if isinstance(query_paths, Path):
+        if query_paths.is_dir():
+            try:
+                project_root = find_project_root()
+                relative = query_paths.relative_to(project_root)
+            except ValueError:
+                relative = query_paths
+            resolved = get_sql_files(str(relative), pipeline=pipeline)
+            if not resolved:
+                log_and_raise_error(logger, f"No SQL files found in directory '{query_paths}'.")
+            return resolved
+        return {query_paths.stem: query_paths}
+
+    if isinstance(query_paths, list):
+        return {Path(p).stem: Path(p) for p in query_paths}
+
+    if isinstance(query_paths, dict):
+        return {k: Path(v) if isinstance(v, str) else v for k, v in query_paths.items()}
+
+    log_and_raise_error(logger, f"Expected a folder path, list, or dict, got: {type(query_paths)}")
+    return {}
+
+
 def execute_sql_scripts(
     query_paths,
     data_connector=None,
@@ -662,42 +717,7 @@ def execute_sql_scripts(
 
     from ml_analytics.data_connector import DataConnector
 
-    # Normalize input to an ordered dict[str, Path]
-    if isinstance(query_paths, str):
-        try:
-            project_root = find_project_root()
-            candidate = project_root / query_paths
-        except FileNotFoundError:
-            candidate = Path(query_paths)
-        if candidate.is_file():
-            query_paths = {candidate.stem: candidate}
-        else:
-            resolved = get_sql_files(query_paths, pipeline=pipeline)
-            if not resolved:
-                log_and_raise_error(logger, f"No SQL files found for folder '{query_paths}'.")
-                return
-            query_paths = resolved
-    elif isinstance(query_paths, Path):
-        if query_paths.is_dir():
-            try:
-                project_root = find_project_root()
-                relative = query_paths.relative_to(project_root)
-            except ValueError:
-                relative = query_paths
-            resolved = get_sql_files(str(relative), pipeline=pipeline)
-            if not resolved:
-                log_and_raise_error(logger, f"No SQL files found in directory '{query_paths}'.")
-                return
-            query_paths = resolved
-        else:
-            query_paths = {query_paths.stem: query_paths}
-    elif isinstance(query_paths, list):
-        query_paths = {Path(p).stem: Path(p) for p in query_paths}
-    elif isinstance(query_paths, dict):
-        query_paths = {k: Path(v) if isinstance(v, str) else v for k, v in query_paths.items()}
-    else:
-        log_and_raise_error(logger, f"Expected a folder path, list, or dict, got: {type(query_paths)}")
-        return
+    query_paths = resolve_sql_query_paths(query_paths, pipeline=pipeline)
 
     def _run_scripts(dc):
         """Execute all scripts on the given DataConnector instance.

{ml_analytics_tools-0.4.3 → ml_analytics_tools-0.4.4/ml_analytics_tools.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ml-analytics-tools
-Version: 0.4.3
+Version: 0.4.4
 Summary: Tools for ML projects and data management
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
@@ -220,6 +220,14 @@ sf.sql("queries/experiment.sql", save_table=True, schema="analytics", table="exp
 
 # or save any Spark DataFrame to Unity Catalog
 sf.save_to_uc(df, table="exp", schema="analytics", catalog="prod")
+
+# save a YAML-ordered folder of SQL queries as Unity Catalog tables
+df = sf.save_pipeline_to_uc(
+    "queries/churn_pipeline",
+    pipeline="daily",
+    catalog="prod",
+    schema="analytics",
+)
 ```
 
 Credentials resolve per field as: explicit argument → `SNOWFLAKE_*` environment

{ml_analytics_tools-0.4.3 → ml_analytics_tools-0.4.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ml-analytics-tools"
-version = "0.4.3"
+version = "0.4.4"
 description = "Tools for ML projects and data management"
 readme = "README.md"
 requires-python = ">=3.11"

{ml_analytics_tools-0.4.3 → ml_analytics_tools-0.4.4}/tests/test_sf_connector.py

@@ -58,6 +58,7 @@ def _mock_spark():
     """Spark double whose read chain returns a DataFrame mock."""
     spark = MagicMock()
     df = MagicMock()
+    df.sparkSession = spark
     reader = spark.read.format.return_value
     reader.options.return_value.option.return_value.load.return_value = df
     reader.options.return_value.load.return_value = df
@@ -220,11 +221,20 @@ def test_qualified_uc_name_already_qualified():
 
 def test_save_to_uc_uses_saveastable(monkeypatch):
     _clear_snowflake_env(monkeypatch)
-    sf = SFConnector(account="acct", user="u")
+    spark = MagicMock()
+    sf = SFConnector(account="acct", user="u", spark=spark)
 
     calls = {}
 
     class _Writer:
+        def format(self, fmt):
+            calls["format"] = fmt
+            return self
+
+        def option(self, key, value):
+            calls["option"] = (key, value)
+            return self
+
         def mode(self, m):
             calls["mode"] = m
             return self
@@ -234,9 +244,42 @@ def test_save_to_uc_uses_saveastable(monkeypatch):
 
     class _DF:
         write = _Writer()
+        sparkSession = spark
 
     sf.save_to_uc(_DF(), table="tbl", schema="sch", catalog="cat", mode="append")
-    assert calls == {"mode": "append", "name": "cat.sch.tbl"}
+    assert calls == {
+        "format": "delta",
+        "option": ("mergeSchema", "true"),
+        "mode": "append",
+        "name": "cat.sch.tbl",
+    }
+    spark.sql.assert_called_once_with("OPTIMIZE cat.sch.tbl")
+
+
+def test_save_to_uc_can_zorder_comment_or_skip_optimize(monkeypatch):
+    _clear_snowflake_env(monkeypatch)
+    spark = MagicMock()
+    sf = SFConnector(account="acct", user="u", spark=spark)
+
+    df = MagicMock()
+    df.sparkSession = spark
+
+    sf.save_to_uc(
+        df,
+        table="tbl",
+        schema="sch",
+        catalog="cat",
+        zorder_by=["customer_id", "event_date"],
+        comment="Tutor's metrics",
+    )
+    assert [call.args[0] for call in spark.sql.call_args_list] == [
+        "ALTER TABLE cat.sch.tbl SET TBLPROPERTIES ('comment' = 'Tutor''s metrics')",
+        "OPTIMIZE cat.sch.tbl ZORDER BY (customer_id, event_date)",
+    ]
+
+    spark.reset_mock()
+    sf.save_to_uc(df, table="tbl", schema="sch", catalog="cat", optimize=False)
+    spark.sql.assert_not_called()
 
 
 def test_save_to_uc_requires_table(monkeypatch):
@@ -317,3 +360,82 @@ def test_sql_return_pandas(monkeypatch):
     sf.sql("select 1", return_pandas=True)
 
     df.toPandas.assert_called_once()
+
+
+def test_save_pipeline_to_uc_uses_yaml_order_and_file_stem_tables(monkeypatch, tmp_path):
+    _clear_snowflake_env(monkeypatch)
+    folder = tmp_path / "queries"
+    folder.mkdir()
+    (folder / "base.sql").write_text("SELECT '{run_date}' AS run_date;")
+    (folder / "features.sql").write_text("SELECT 1 AS feature;")
+    (folder / "daily.yaml").write_text(
+        """
+steps:
+  - features
+  - base
+"""
+    )
+    monkeypatch.setattr("ml_analytics.utils.find_project_root", lambda *args, **kwargs: tmp_path)
+
+    spark, df = _mock_spark()
+    sf = SFConnector(account="acct", user="u", password="p", spark=spark)
+
+    result = sf.save_pipeline_to_uc(
+        "queries",
+        pipeline="daily",
+        catalog="prod",
+        schema="analytics",
+        run_date="2026-06-17",
+    )
+
+    assert result is df
+    query_calls = spark.read.format.return_value.options.return_value.option.call_args_list
+    assert [call.args for call in query_calls] == [
+        ("query", "SELECT 1 AS feature;"),
+        ("query", "SELECT '2026-06-17' AS run_date;"),
+    ]
+    save_calls = df.write.format.return_value.option.return_value.mode.return_value.saveAsTable.call_args_list
+    assert [call.args[0] for call in save_calls] == [
+        "prod.analytics.features",
+        "prod.analytics.base",
+    ]
+    assert [call.args[0] for call in spark.sql.call_args_list] == [
+        "OPTIMIZE prod.analytics.features",
+        "OPTIMIZE prod.analytics.base",
+    ]
+
+
+def test_save_pipeline_to_uc_allows_table_and_mode_overrides(monkeypatch, tmp_path):
+    _clear_snowflake_env(monkeypatch)
+    folder = tmp_path / "queries"
+    folder.mkdir()
+    (folder / "base.sql").write_text("SELECT 1 AS col_1;")
+    (folder / "final.sql").write_text("SELECT 2 AS col_2;")
+    monkeypatch.setattr("ml_analytics.utils.find_project_root", lambda *args, **kwargs: tmp_path)
+
+    spark, df = _mock_spark()
+    sf = SFConnector(account="acct", user="u", password="p", spark=spark)
+
+    result = sf.save_pipeline_to_uc(
+        "queries",
+        schema="analytics",
+        catalog="prod",
+        tables={"final": "churn_daily"},
+        table_prefix="stg_",
+        modes={"final": "append"},
+        zorder_by={"final": "customer_id"},
+        return_all=True,
+    )
+
+    assert result == {"base": df, "final": df}
+    mode_calls = df.write.format.return_value.option.return_value.mode.call_args_list
+    assert [call.args[0] for call in mode_calls] == ["overwrite", "append"]
+    save_calls = df.write.format.return_value.option.return_value.mode.return_value.saveAsTable.call_args_list
+    assert [call.args[0] for call in save_calls] == [
+        "prod.analytics.stg_base",
+        "prod.analytics.churn_daily",
+    ]
+    assert [call.args[0] for call in spark.sql.call_args_list] == [
+        "OPTIMIZE prod.analytics.stg_base",
+        "OPTIMIZE prod.analytics.churn_daily ZORDER BY (customer_id)",
+    ]