snowflake-ml-python 1.4.1__py3-none-any.whl → 1.5.1__py3-none-any.whl

snowflake/ml/model/_client/sql/model.py

@@ -1,14 +1,11 @@
 from typing import Any, Dict, List, Optional
 
-from snowflake.ml._internal.utils import (
-    identifier,
-    query_result_checker,
-    sql_identifier,
-)
-from snowflake.snowpark import row, session
+from snowflake.ml._internal.utils import query_result_checker, sql_identifier
+from snowflake.ml.model._client.sql import _base
+from snowflake.snowpark import row
 
 
-class ModelSQLClient:
+class ModelSQLClient(_base._BaseSQLClient):
     MODEL_NAME_COL_NAME = "name"
     MODEL_COMMENT_COL_NAME = "comment"
     MODEL_DEFAULT_VERSION_NAME_COL_NAME = "default_version_name"
@@ -18,35 +15,18 @@ class ModelSQLClient:
     MODEL_VERSION_METADATA_COL_NAME = "metadata"
     MODEL_VERSION_MODEL_SPEC_COL_NAME = "model_spec"
 
-    def __init__(
-        self,
-        session: session.Session,
-        *,
-        database_name: sql_identifier.SqlIdentifier,
-        schema_name: sql_identifier.SqlIdentifier,
-    ) -> None:
-        self._session = session
-        self._database_name = database_name
-        self._schema_name = schema_name
-
-    def __eq__(self, __value: object) -> bool:
-        if not isinstance(__value, ModelSQLClient):
-            return False
-        return self._database_name == __value._database_name and self._schema_name == __value._schema_name
-
-    def fully_qualified_model_name(self, model_name: sql_identifier.SqlIdentifier) -> str:
-        return identifier.get_schema_level_object_identifier(
-            self._database_name.identifier(), self._schema_name.identifier(), model_name.identifier()
-        )
-
     def show_models(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: Optional[sql_identifier.SqlIdentifier] = None,
         validate_result: bool = True,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> List[row.Row]:
-        fully_qualified_schema_name = ".".join([self._database_name.identifier(), self._schema_name.identifier()])
+        actual_database_name = database_name or self._database_name
+        actual_schema_name = schema_name or self._schema_name
+        fully_qualified_schema_name = ".".join([actual_database_name.identifier(), actual_schema_name.identifier()])
         like_sql = ""
         if model_name:
             like_sql = f" LIKE '{model_name.resolved()}'"
@@ -69,6 +49,8 @@ class ModelSQLClient:
     def show_versions(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: Optional[sql_identifier.SqlIdentifier] = None,
         validate_result: bool = True,
@@ -82,7 +64,10 @@ class ModelSQLClient:
         res = (
             query_result_checker.SqlResultValidator(
                 self._session,
-                f"SHOW VERSIONS{like_sql} IN MODEL {self.fully_qualified_model_name(model_name)}",
+                (
+                    f"SHOW VERSIONS{like_sql} IN "
+                    f"MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+                ),
                 statement_params=statement_params,
             )
             .has_column(ModelSQLClient.MODEL_VERSION_NAME_COL_NAME, allow_empty=True)
@@ -99,43 +84,53 @@ class ModelSQLClient:
     def set_comment(
         self,
         *,
-        comment: str,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
+        comment: str,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
-            f"COMMENT ON MODEL {self.fully_qualified_model_name(model_name)} IS $${comment}$$",
+            (
+                f"COMMENT ON MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+                f" IS $${comment}$$"
+            ),
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
     def drop_model(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
-            f"DROP MODEL {self.fully_qualified_model_name(model_name)}",
+            f"DROP MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}",
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
-    def config_model_details(
+    def rename(
         self,
         *,
-        enable: bool,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        new_model_db: Optional[sql_identifier.SqlIdentifier],
+        new_model_schema: Optional[sql_identifier.SqlIdentifier],
+        new_model_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
-        if enable:
-            query_result_checker.SqlResultValidator(
-                self._session,
-                "ALTER SESSION SET SHOW_MODEL_DETAILS_IN_SHOW_VERSIONS_IN_MODEL=true",
-                statement_params=statement_params,
-            ).has_dimensions(expected_rows=1, expected_cols=1).validate()
-        else:
-            query_result_checker.SqlResultValidator(
-                self._session,
-                "ALTER SESSION UNSET SHOW_MODEL_DETAILS_IN_SHOW_VERSIONS_IN_MODEL",
-                statement_params=statement_params,
-            ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+        # Use registry's database and schema if a non fully qualified new model name is provided.
+        new_fully_qualified_name = self.fully_qualified_object_name(new_model_db, new_model_schema, new_model_name)
+        query_result_checker.SqlResultValidator(
+            self._session,
+            (
+                f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+                f" RENAME TO {new_fully_qualified_name}"
+            ),
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()

snowflake/ml/model/_client/sql/model_version.py

@@ -9,7 +9,8 @@ from snowflake.ml._internal.utils import (
     query_result_checker,
     sql_identifier,
 )
-from snowflake.snowpark import dataframe, functions as F, row, session, types as spt
+from snowflake.ml.model._client.sql import _base
+from snowflake.snowpark import dataframe, functions as F, row, types as spt
 from snowflake.snowpark._internal import utils as snowpark_utils
 
 
@@ -20,34 +21,15 @@ def _normalize_url_for_sql(url: str) -> str:
     return f"'{url}'"
 
 
-class ModelVersionSQLClient:
+class ModelVersionSQLClient(_base._BaseSQLClient):
     FUNCTION_NAME_COL_NAME = "name"
     FUNCTION_RETURN_TYPE_COL_NAME = "return_type"
 
-    def __init__(
-        self,
-        session: session.Session,
-        *,
-        database_name: sql_identifier.SqlIdentifier,
-        schema_name: sql_identifier.SqlIdentifier,
-    ) -> None:
-        self._session = session
-        self._database_name = database_name
-        self._schema_name = schema_name
-
-    def __eq__(self, __value: object) -> bool:
-        if not isinstance(__value, ModelVersionSQLClient):
-            return False
-        return self._database_name == __value._database_name and self._schema_name == __value._schema_name
-
-    def fully_qualified_model_name(self, model_name: sql_identifier.SqlIdentifier) -> str:
-        return identifier.get_schema_level_object_identifier(
-            self._database_name.identifier(), self._schema_name.identifier(), model_name.identifier()
-        )
-
     def create_from_stage(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         stage_path: str,
@@ -56,8 +38,8 @@ class ModelVersionSQLClient:
         query_result_checker.SqlResultValidator(
             self._session,
             (
-                f"CREATE MODEL {self.fully_qualified_model_name(model_name)} WITH VERSION {version_name.identifier()}"
-                f" FROM {stage_path}"
+                f"CREATE MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+                f" WITH VERSION {version_name.identifier()} FROM {stage_path}"
             ),
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
@@ -66,6 +48,8 @@ class ModelVersionSQLClient:
     def add_version_from_stage(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         stage_path: str,
@@ -74,8 +58,8 @@ class ModelVersionSQLClient:
         query_result_checker.SqlResultValidator(
             self._session,
             (
-                f"ALTER MODEL {self.fully_qualified_model_name(model_name)} ADD VERSION {version_name.identifier()}"
-                f" FROM {stage_path}"
+                f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+                f" ADD VERSION {version_name.identifier()} FROM {stage_path}"
             ),
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
@@ -83,6 +67,8 @@ class ModelVersionSQLClient:
     def set_default_version(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
@@ -90,15 +76,54 @@ class ModelVersionSQLClient:
         query_result_checker.SqlResultValidator(
             self._session,
             (
-                f"ALTER MODEL {self.fully_qualified_model_name(model_name)} "
+                f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)} "
                 f"SET DEFAULT_VERSION = {version_name.identifier()}"
             ),
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    def list_file(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        file_path: pathlib.PurePosixPath,
+        is_dir: bool = False,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> List[row.Row]:
+        # Workaround for snowURL bug.
+        trailing_slash = "/" if is_dir else ""
+
+        stage_location = (
+            pathlib.PurePosixPath(
+                self.fully_qualified_object_name(database_name, schema_name, model_name),
+                "versions",
+                version_name.resolved(),
+                file_path,
+            ).as_posix()
+            + trailing_slash
+        )
+        stage_location_url = ParseResult(
+            scheme="snow", netloc="model", path=stage_location, params="", query="", fragment=""
+        ).geturl()
+
+        return (
+            query_result_checker.SqlResultValidator(
+                self._session,
+                f"List {_normalize_url_for_sql(stage_location_url)}",
+                statement_params=statement_params,
+            )
+            .has_column("name", allow_empty=True)
+            .validate()
+        )
+
     def get_file(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         file_path: pathlib.PurePosixPath,
@@ -106,7 +131,10 @@ class ModelVersionSQLClient:
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> pathlib.Path:
         stage_location = pathlib.PurePosixPath(
-            self.fully_qualified_model_name(model_name), "versions", version_name.resolved(), file_path
+            self.fully_qualified_object_name(database_name, schema_name, model_name),
+            "versions",
+            version_name.resolved(),
+            file_path,
         ).as_posix()
         stage_location_url = ParseResult(
             scheme="snow", netloc="model", path=stage_location, params="", query="", fragment=""
@@ -130,6 +158,8 @@ class ModelVersionSQLClient:
     def show_functions(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
@@ -137,7 +167,7 @@ class ModelVersionSQLClient:
         res = query_result_checker.SqlResultValidator(
             self._session,
             (
-                f"SHOW FUNCTIONS IN MODEL {self.fully_qualified_model_name(model_name)}"
+                f"SHOW FUNCTIONS IN MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
                 f" VERSION {version_name.identifier()}"
             ),
             statement_params=statement_params,
@@ -148,23 +178,27 @@ class ModelVersionSQLClient:
     def set_comment(
         self,
         *,
-        comment: str,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
+        comment: str,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
             (
-                f"ALTER MODEL {self.fully_qualified_model_name(model_name)} "
+                f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)} "
                 f"MODIFY VERSION {version_name.identifier()} SET COMMENT=$${comment}$$"
             ),
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
-    def invoke_method(
+    def invoke_function_method(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         method_name: sql_identifier.SqlIdentifier,
@@ -178,10 +212,12 @@ class ModelVersionSQLClient:
             INTERMEDIATE_TABLE_NAME = "SNOWPARK_ML_MODEL_INFERENCE_INPUT"
             with_statements.append(f"{INTERMEDIATE_TABLE_NAME} AS ({input_df.queries['queries'][0]})")
         else:
+            actual_database_name = database_name or self._database_name
+            actual_schema_name = schema_name or self._schema_name
             tmp_table_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.TABLE)
             INTERMEDIATE_TABLE_NAME = identifier.get_schema_level_object_identifier(
-                self._database_name.identifier(),
-                self._schema_name.identifier(),
+                actual_database_name.identifier(),
+                actual_schema_name.identifier(),
                 tmp_table_name,
             )
             input_df.write.save_as_table(  # type: ignore[call-overload]
@@ -196,7 +232,8 @@ class ModelVersionSQLClient:
         module_version_alias = "MODEL_VERSION_ALIAS"
         with_statements.append(
             f"{module_version_alias} AS "
-            f"MODEL {self.fully_qualified_model_name(model_name)} VERSION {version_name.identifier()}"
+            f"MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+            f" VERSION {version_name.identifier()}"
         )
 
         args_sql_list = []
@@ -232,10 +269,93 @@ class ModelVersionSQLClient:
 
         return output_df
 
+    def invoke_table_function_method(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        method_name: sql_identifier.SqlIdentifier,
+        input_df: dataframe.DataFrame,
+        input_args: List[sql_identifier.SqlIdentifier],
+        returns: List[Tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
+        partition_column: Optional[sql_identifier.SqlIdentifier],
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> dataframe.DataFrame:
+        with_statements = []
+        if len(input_df.queries["queries"]) == 1 and len(input_df.queries["post_actions"]) == 0:
+            INTERMEDIATE_TABLE_NAME = "SNOWPARK_ML_MODEL_INFERENCE_INPUT"
+            with_statements.append(f"{INTERMEDIATE_TABLE_NAME} AS ({input_df.queries['queries'][0]})")
+        else:
+            actual_database_name = database_name or self._database_name
+            actual_schema_name = schema_name or self._schema_name
+            tmp_table_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.TABLE)
+            INTERMEDIATE_TABLE_NAME = identifier.get_schema_level_object_identifier(
+                actual_database_name.identifier(),
+                actual_schema_name.identifier(),
+                tmp_table_name,
+            )
+            input_df.write.save_as_table(  # type: ignore[call-overload]
+                table_name=INTERMEDIATE_TABLE_NAME,
+                mode="errorifexists",
+                table_type="temporary",
+                statement_params=statement_params,
+            )
+
+        module_version_alias = "MODEL_VERSION_ALIAS"
+        with_statements.append(
+            f"{module_version_alias} AS "
+            f"MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+            f" VERSION {version_name.identifier()}"
+        )
+
+        partition_by = partition_column.identifier() if partition_column is not None else "1"
+
+        args_sql_list = []
+        for input_arg_value in input_args:
+            args_sql_list.append(input_arg_value)
+
+        args_sql = ", ".join(args_sql_list)
+
+        sql = textwrap.dedent(
+            f"""WITH {','.join(with_statements)}
+                 SELECT *,
+                 FROM {INTERMEDIATE_TABLE_NAME},
+                     TABLE({module_version_alias}!{method_name.identifier()}({args_sql})
+                     OVER (PARTITION BY {partition_by}))"""
+        )
+
+        output_df = self._session.sql(sql)
+
+        # Prepare the output
+        output_cols = []
+        output_names = []
+
+        for output_name, output_type, output_col_name in returns:
+            output_cols.append(F.col(output_name).astype(output_type))
+            output_names.append(output_col_name)
+
+        if partition_column is not None:
+            output_cols.append(F.col(partition_column.identifier()))
+            output_names.append(partition_column)
+
+        output_df = output_df.with_columns(
+            col_names=output_names,
+            values=output_cols,
+        )
+
+        if statement_params:
+            output_df._statement_params = statement_params  # type: ignore[assignment]
+
+        return output_df
+
     def set_metadata(
         self,
         metadata_dict: Dict[str, Any],
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
@@ -244,8 +364,8 @@ class ModelVersionSQLClient:
         query_result_checker.SqlResultValidator(
             self._session,
             (
-                f"ALTER MODEL {self.fully_qualified_model_name(model_name)} MODIFY VERSION {version_name.identifier()}"
-                f" SET METADATA=$${json_metadata}$$"
+                f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+                f" MODIFY VERSION {version_name.identifier()} SET METADATA=$${json_metadata}$$"
             ),
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
@@ -253,12 +373,17 @@ class ModelVersionSQLClient:
     def drop_version(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
-            f"ALTER MODEL {self.fully_qualified_model_name(model_name)} DROP VERSION {version_name.identifier()}",
+            (
+                f"ALTER MODEL {self.fully_qualified_object_name(database_name, schema_name, model_name)}"
+                f" DROP VERSION {version_name.identifier()}"
+            ),
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()

snowflake/ml/model/_client/sql/stage.py

@@ -1,46 +1,20 @@
 from typing import Any, Dict, Optional
 
-from snowflake.ml._internal.utils import (
-    identifier,
-    query_result_checker,
-    sql_identifier,
-)
-from snowflake.snowpark import session
+from snowflake.ml._internal.utils import query_result_checker, sql_identifier
+from snowflake.ml.model._client.sql import _base
 
 
-class StageSQLClient:
-    def __init__(
-        self,
-        session: session.Session,
-        *,
-        database_name: sql_identifier.SqlIdentifier,
-        schema_name: sql_identifier.SqlIdentifier,
-    ) -> None:
-        self._session = session
-        self._database_name = database_name
-        self._schema_name = schema_name
-
-    def __eq__(self, __value: object) -> bool:
-        if not isinstance(__value, StageSQLClient):
-            return False
-        return self._database_name == __value._database_name and self._schema_name == __value._schema_name
-
-    def fully_qualified_stage_name(
-        self,
-        stage_name: sql_identifier.SqlIdentifier,
-    ) -> str:
-        return identifier.get_schema_level_object_identifier(
-            self._database_name.identifier(), self._schema_name.identifier(), stage_name.identifier()
-        )
-
+class StageSQLClient(_base._BaseSQLClient):
     def create_tmp_stage(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         stage_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
-            f"CREATE TEMPORARY STAGE {self.fully_qualified_stage_name(stage_name)}",
+            f"CREATE TEMPORARY STAGE {self.fully_qualified_object_name(database_name, schema_name, stage_name)}",
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()