snowflake-ml-python 1.6.2__py3-none-any.whl → 1.6.4__py3-none-any.whl

snowflake/cortex/__init__.py

@@ -1,5 +1,7 @@
 from snowflake.cortex._classify_text import ClassifyText
 from snowflake.cortex._complete import Complete, CompleteOptions
+from snowflake.cortex._embed_text_768 import EmbedText768
+from snowflake.cortex._embed_text_1024 import EmbedText1024
 from snowflake.cortex._extract_answer import ExtractAnswer
 from snowflake.cortex._sentiment import Sentiment
 from snowflake.cortex._summarize import Summarize
@@ -9,6 +11,8 @@ __all__ = [
     "ClassifyText",
     "Complete",
     "CompleteOptions",
+    "EmbedText768",
+    "EmbedText1024",
     "ExtractAnswer",
     "Sentiment",
     "Summarize",

snowflake/cortex/_classify_text.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import List, Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -33,4 +33,4 @@ def _classify_text_impl(
     categories: Union[List[str], snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, str_input, categories)
+    return cast(Union[str, snowpark.Column], call_sql_function(function, session, str_input, categories))

snowflake/cortex/_embed_text_1024.py

@@ -0,0 +1,37 @@
+from typing import List, Optional, Union, cast
+
+from snowflake import snowpark
+from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.ml._internal import telemetry
+
+
+@telemetry.send_api_usage_telemetry(
+    project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+)
+def EmbedText1024(
+    model: Union[str, snowpark.Column],
+    text: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[List[float], snowpark.Column]:
+    """TextEmbed calls into the LLM inference service to embed the text.
+
+    Args:
+        model: A Column of strings representing the model to use for embedding. The value
+               of the strings must be within the SUPPORTED_MODELS list.
+        text: A Column of strings representing input text.
+        session: The snowpark session to use. Will be inferred by context if not specified.
+
+    Returns:
+        A column of vectors containing embeddings.
+    """
+
+    return _embed_text_1024_impl("snowflake.cortex.embed_text_1024", model, text, session=session)
+
+
+def _embed_text_1024_impl(
+    function: str,
+    model: Union[str, snowpark.Column],
+    text: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[List[float], snowpark.Column]:
+    return cast(Union[List[float], snowpark.Column], call_sql_function(function, session, model, text))

snowflake/cortex/_embed_text_768.py

@@ -0,0 +1,37 @@
+from typing import List, Optional, Union, cast
+
+from snowflake import snowpark
+from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.ml._internal import telemetry
+
+
+@telemetry.send_api_usage_telemetry(
+    project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+)
+def EmbedText768(
+    model: Union[str, snowpark.Column],
+    text: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[List[float], snowpark.Column]:
+    """TextEmbed calls into the LLM inference service to embed the text.
+
+    Args:
+        model: A Column of strings representing the model to use for embedding. The value
+               of the strings must be within the SUPPORTED_MODELS list.
+        text: A Column of strings representing input text.
+        session: The snowpark session to use. Will be inferred by context if not specified.
+
+    Returns:
+        A column of vectors containing embeddings.
+    """
+
+    return _embed_text_768_impl("snowflake.cortex.embed_text_768", model, text, session=session)
+
+
+def _embed_text_768_impl(
+    function: str,
+    model: Union[str, snowpark.Column],
+    text: Union[str, snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[List[float], snowpark.Column]:
+    return cast(Union[List[float], snowpark.Column], call_sql_function(function, session, model, text))

snowflake/cortex/_extract_answer.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -33,4 +33,4 @@ def _extract_answer_impl(
     question: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, from_text, question)
+    return cast(Union[str, snowpark.Column], call_sql_function(function, session, from_text, question))

snowflake/cortex/_sentiment.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -30,4 +30,4 @@ def _sentiment_impl(
     output = call_sql_function(function, session, text)
     if isinstance(output, snowpark.Column):
         return output
-    return float(output)
+    return float(cast(str, output))

snowflake/cortex/_summarize.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -30,4 +30,4 @@ def _summarize_impl(
     text: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, text)
+    return cast(Union[str, snowpark.Column], call_sql_function(function, session, text))

snowflake/cortex/_translate.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
@@ -36,4 +36,4 @@ def _translate_impl(
     to_language: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    return call_sql_function(function, session, text, from_language, to_language)
+    return cast(Union[str, snowpark.Column], call_sql_function(function, session, text, from_language, to_language))

snowflake/cortex/_util.py

@@ -24,7 +24,7 @@ def call_sql_function(
     function: str,
     session: Optional[snowpark.Session],
     *args: Union[str, List[str], snowpark.Column, Dict[str, Union[int, float]]],
-) -> Union[str, snowpark.Column]:
+) -> Union[str, List[float], snowpark.Column]:
     handle_as_column = False
 
     for arg in args:
@@ -32,9 +32,9 @@ def call_sql_function(
             handle_as_column = True
 
     if handle_as_column:
-        return cast(Union[str, snowpark.Column], _call_sql_function_column(function, *args))
+        return cast(Union[str, List[float], snowpark.Column], _call_sql_function_column(function, *args))
     return cast(
-        Union[str, snowpark.Column],
+        Union[str, List[float], snowpark.Column],
         _call_sql_function_immediate(function, session, *args),
     )
 
@@ -49,7 +49,7 @@ def _call_sql_function_immediate(
     function: str,
     session: Optional[snowpark.Session],
     *args: Union[str, List[str], snowpark.Column, Dict[str, Union[int, float]]],
-) -> str:
+) -> Union[str, List[float]]:
     session = session or context.get_active_session()
     if session is None:
         raise SnowflakeAuthenticationException(

snowflake/ml/_internal/env_utils.py

@@ -9,7 +9,7 @@ from importlib import metadata as importlib_metadata
 from typing import Any, DefaultDict, Dict, List, Optional, Tuple
 
 import yaml
-from packaging import requirements, specifiers, utils as packaging_utils, version
+from packaging import requirements, specifiers, version
 
 import snowflake.connector
 from snowflake.ml._internal import env as snowml_env
@@ -54,15 +54,12 @@ def _validate_pip_requirement_string(req_str: str) -> requirements.Requirement:
     """
     try:
         r = requirements.Requirement(req_str)
-        r.name = packaging_utils.canonicalize_name(r.name)
 
         if r.name == "python":
             raise ValueError("Don't specify python as a dependency, use python version argument instead.")
     except requirements.InvalidRequirement:
         raise ValueError(f"Invalid package requirement {req_str} found.")
 
-    if r.marker:
-        raise ValueError("Markers is not supported in conda dependency.")
     return r
 
 
@@ -84,6 +81,8 @@ def _validate_conda_dependency_string(dep_str: str) -> Tuple[str, requirements.R
     channel_str, _, requirement_str = dep_str.rpartition("::")
     r = _validate_pip_requirement_string(requirement_str)
     if channel_str != "pip":
+        if r.marker:
+            raise ValueError("Markers is not supported in conda dependency.")
         if r.extras:
             raise ValueError("Extras is not supported in conda dependency.")
         if r.url:
@@ -221,7 +220,7 @@ def get_local_installed_version_of_pip_package(pip_req: requirements.Requirement
         else:
             return pip_req
     new_pip_req = copy.deepcopy(pip_req)
-    new_pip_req.specifier = specifiers.SpecifierSet(specifiers=f"=={local_dist_version}")
+    new_pip_req.specifier = specifiers.SpecifierSet(specifiers=f"=={version.parse(local_dist_version).base_version}")
     if not pip_req.specifier.contains(local_dist_version):
         warnings.warn(
             f"Package requirement {str(pip_req)} specified, while version {local_dist_version} is installed. "
@@ -513,6 +512,7 @@ def save_conda_env_file(
         )
 
     with open(path, "w", encoding="utf-8") as f:
+        yaml.SafeDumper.ignore_aliases = lambda *args: True  # type: ignore[method-assign]
         yaml.safe_dump(env, stream=f, default_flow_style=False)
 
 

snowflake/ml/_internal/exceptions/error_codes.py

@@ -58,6 +58,8 @@ NOT_FOUND = "2101"
 METHOD_NOT_ALLOWED = "2102"
 # Not implemented.
 NOT_IMPLEMENTED = "2103"
+# User needs to opt in to use a feature.
+OPT_IN_REQUIRED = "2104"
 
 # Calling an API with unsupported keywords/values.
 INVALID_ARGUMENT = "2110"

snowflake/ml/_internal/utils/db_utils.py

@@ -0,0 +1,50 @@
+from enum import Enum
+from typing import Any, Dict, Optional
+
+from snowflake.ml._internal.utils import query_result_checker, sql_identifier
+from snowflake.snowpark import session
+
+MAX_IDENTIFIER_LENGTH = 255
+
+
+class SnowflakeDbObjectType(Enum):
+    TABLE = "TABLE"
+    WAREHOUSE = "WAREHOUSE"
+
+
+def db_object_exists(
+    session: session.Session,
+    object_type: SnowflakeDbObjectType,
+    object_name: sql_identifier.SqlIdentifier,
+    *,
+    database_name: Optional[sql_identifier.SqlIdentifier] = None,
+    schema_name: Optional[sql_identifier.SqlIdentifier] = None,
+    statement_params: Optional[Dict[str, Any]] = None,
+) -> bool:
+    """Check if object exists in database.
+
+    Args:
+        session: Active Snowpark Session.
+        object_type: Type of object to search for.
+        object_name: Name of object to search for.
+        database_name: Optional database name to search in. Only used if both schema is also provided.
+        schema_name: Optional schema to search in.
+        statement_params: Optional set of statement_params to include with queries.
+
+    Returns:
+        boolean indicating whether object exists.
+    """
+    optional_in_clause = ""
+    if database_name and schema_name:
+        optional_in_clause = f" IN {database_name}.{schema_name}"
+
+    result = (
+        query_result_checker.SqlResultValidator(
+            session,
+            f"""SHOW {object_type.value}S LIKE '{object_name}'{optional_in_clause}""",
+            statement_params=statement_params,
+        )
+        .has_column("name", allow_empty=True)  # TODO: Check this is actually what is returned from server
+        .validate()
+    )
+    return len(result) == 1

snowflake/ml/_internal/utils/service_logger.py

@@ -0,0 +1,63 @@
+import enum
+import logging
+import sys
+
+
+class LogColor(enum.Enum):
+    GREY = "\x1b[38;20m"
+    RED = "\x1b[31;20m"
+    BOLD_RED = "\x1b[31;1m"
+    YELLOW = "\x1b[33;20m"
+    BLUE = "\x1b[34;20m"
+    GREEN = "\x1b[32;20m"
+
+
+class CustomFormatter(logging.Formatter):
+
+    reset = "\x1b[0m"
+    log_format = "%(name)s [%(asctime)s] [%(levelname)s] %(message)s"
+
+    def __init__(self, info_color: LogColor) -> None:
+        super().__init__()
+        self.level_colors = {
+            logging.DEBUG: LogColor.GREY.value,
+            logging.INFO: info_color.value,
+            logging.WARNING: LogColor.YELLOW.value,
+            logging.ERROR: LogColor.RED.value,
+            logging.CRITICAL: LogColor.BOLD_RED.value,
+        }
+
+    def format(self, record: logging.LogRecord) -> str:
+        # default to DEBUG color
+        fmt = self.level_colors.get(record.levelno, self.level_colors[logging.DEBUG]) + self.log_format + self.reset
+        formatter = logging.Formatter(fmt)
+
+        # split the log message by lines and format each line individually
+        original_message = record.getMessage()
+        message_lines = original_message.splitlines()
+        formatted_lines = [
+            formatter.format(
+                logging.LogRecord(
+                    name=record.name,
+                    level=record.levelno,
+                    pathname=record.pathname,
+                    lineno=record.lineno,
+                    msg=line,
+                    args=None,
+                    exc_info=None,
+                )
+            )
+            for line in message_lines
+        ]
+
+        return "\n".join(formatted_lines)
+
+
+def get_logger(logger_name: str, info_color: LogColor) -> logging.Logger:
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(logging.INFO)
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setLevel(logging.INFO)
+    handler.setFormatter(CustomFormatter(info_color))
+    logger.addHandler(handler)
+    return logger

snowflake/ml/_internal/utils/sql_identifier.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Union
 
 from snowflake.ml._internal.utils import identifier
 
@@ -92,3 +92,27 @@ def parse_fully_qualified_name(
         SqlIdentifier(schema) if schema else None,
         SqlIdentifier(object),
     )
+
+
+def get_fully_qualified_name(
+    db: Union[SqlIdentifier, str, None],
+    schema: Union[SqlIdentifier, str, None],
+    object: Union[SqlIdentifier, str],
+    session_db: Optional[str] = None,
+    session_schema: Optional[str] = None,
+) -> str:
+    db_name: Optional[SqlIdentifier] = None
+    schema_name: Optional[SqlIdentifier] = None
+    if not db and session_db:
+        db_name = SqlIdentifier(session_db)
+    elif isinstance(db, str):
+        db_name = SqlIdentifier(db)
+    if not schema and session_schema:
+        schema_name = SqlIdentifier(session_schema)
+    elif isinstance(schema, str):
+        schema_name = SqlIdentifier(schema)
+    return identifier.get_schema_level_object_identifier(
+        db=db_name.identifier() if db_name else None,
+        schema=schema_name.identifier() if schema_name else None,
+        object_name=object.identifier() if isinstance(object, SqlIdentifier) else SqlIdentifier(object).identifier(),
+    )

snowflake/ml/data/_internal/arrow_ingestor.py

@@ -11,7 +11,6 @@ import pyarrow as pa
 import pyarrow.dataset as pds
 
 from snowflake import snowpark
-from snowflake.connector import result_batch
 from snowflake.ml.data import data_ingestor, data_source, ingestor_utils
 
 _EMPTY_RECORD_BATCH = pa.RecordBatch.from_arrays([], [])
@@ -140,16 +139,7 @@ class ArrowIngestor(data_ingestor.DataIngestor):
                 #        We may be able to optimize this by splitting the result batches into
                 #        in-memory (first batch) and file URLs (subsequent batches) and creating a
                 #        union dataset.
-                result_batches = ingestor_utils.get_dataframe_result_batches(self._session, source)
-                sources.extend(
-                    b.to_arrow(self._session.connection)
-                    if isinstance(b, result_batch.ArrowResultBatch)
-                    else b.to_arrow()
-                    for b in result_batches
-                )
-                # HACK: Mitigate typing inconsistencies in Snowpark results
-                if len(sources) > 0:
-                    sources = [_cast_if_needed(s, sources[-1].schema) for s in sources]
+                sources.append(_cast_if_needed(ingestor_utils.get_dataframe_arrow_table(self._session, source)))
                 source_format = None  # Arrow Dataset expects "None" for in-memory datasets
             else:
                 raise RuntimeError(f"Unsupported data source type: {type(source)}")

snowflake/ml/data/ingestor_utils.py

@@ -1,19 +1,17 @@
 from typing import List, Optional
 
 import fsspec
+import pyarrow as pa
 
 from snowflake import snowpark
-from snowflake.connector import result_batch
+from snowflake.connector import cursor as sf_cursor, result_batch
 from snowflake.ml.data import data_source
 from snowflake.ml.fileset import snowfs
 
 _TARGET_FILE_SIZE = 32 * 2**20  # The max file size for data loading.
 
 
-def get_dataframe_result_batches(
-    session: snowpark.Session, df_info: data_source.DataFrameInfo
-) -> List[result_batch.ResultBatch]:
-    """Retrieve the ResultBatches for a given query"""
+def _get_dataframe_cursor(session: snowpark.Session, df_info: data_source.DataFrameInfo) -> sf_cursor.SnowflakeCursor:
     cursor = session._conn._cursor
 
     if df_info.query_id:
@@ -29,12 +27,24 @@ def get_dataframe_result_batches(
     if cursor._prefetch_hook is None:
         raise RuntimeError("Loading data from result query failed unexpectedly. Please contact Snowflake support.")
     cursor._prefetch_hook()
+
+    return cursor
+
+
+def get_dataframe_result_batches(
+    session: snowpark.Session, df_info: data_source.DataFrameInfo
+) -> List[result_batch.ResultBatch]:
+    """Retrieve the ResultBatches for a given query"""
+    cursor = _get_dataframe_cursor(session, df_info)
     batches = cursor.get_result_batches()
-    if batches is None:
-        raise ValueError(
-            "Failed to retrieve training data. Query status:" f" {session._conn._conn.get_query_status(query_id)}"
-        )
-    return batches
+    return batches or []
+
+
+def get_dataframe_arrow_table(session: snowpark.Session, df_info: data_source.DataFrameInfo) -> pa.Table:
+    """Retrieve the full in-memory result for a given query"""
+    cursor = _get_dataframe_cursor(session, df_info)
+    table = cursor.fetch_arrow_all()  # type: ignore[call-overload]
+    return table
 
 
 def get_dataset_filesystem(

snowflake/ml/feature_store/access_manager.py

@@ -30,6 +30,7 @@ class _Privilege:
     object_name: str
     privileges: List[str]
     scope: Optional[str] = None
+    optional: bool = False
 
 
 @dataclass(frozen=True)
@@ -72,8 +73,7 @@ _PRE_INIT_PRIVILEGES: Dict[_FeatureStoreRole, List[_Privilege]] = {
         _Privilege("VIEW", _ALL_OBJECTS, ["SELECT", "REFERENCES"], "SCHEMA {database}.{schema}"),
         _Privilege("TABLE", _ALL_OBJECTS, ["SELECT", "REFERENCES"], "SCHEMA {database}.{schema}"),
         _Privilege("DATASET", _ALL_OBJECTS, ["USAGE"], "SCHEMA {database}.{schema}"),
-        # User should decide whether they want to grant warehouse usage to CONSUMER
-        # _Privilege("WAREHOUSE", "{warehouse}", ["USAGE"]),
+        _Privilege("WAREHOUSE", "{warehouse}", ["USAGE"], optional=True),
     ],
     _FeatureStoreRole.NONE: [],
 }
@@ -109,7 +109,7 @@ def _grant_privileges(
                 query += f" TO ROLE {role_name}"
                 session.sql(query).collect()
         except exceptions.SnowparkSQLException as e:
-            if any(
+            if p.optional or any(
                 s in e.message
                 for s in (
                     "Ask your account admin",

snowflake/ml/feature_store/feature_store.py

@@ -122,6 +122,14 @@ _DT_OR_VIEW_QUERY_PATTERN = re.compile(
     flags=re.DOTALL | re.IGNORECASE | re.X,
 )
 
+_DT_INITIALIZE_PATTERN = re.compile(
+    r"""CREATE\ DYNAMIC\ TABLE\ .*
+        initialize\ =\ '(?P<initialize>.*)'\ .*?
+        AS\ .*
+    """,
+    flags=re.DOTALL | re.IGNORECASE | re.X,
+)
+
 _LIST_FEATURE_VIEW_SCHEMA = StructType(
     [
         StructField("name", StringType()),
@@ -565,11 +573,15 @@ class FeatureStore:
         tagging_clause_str = ",\n".join(tagging_clause)
 
         def create_col_desc(col: StructField) -> str:
-            desc = feature_view.feature_descs.get(SqlIdentifier(col.name), None)
+            desc = feature_view.feature_descs.get(SqlIdentifier(col.name), None)  # type: ignore[union-attr]
             desc = "" if desc is None else f"COMMENT '{desc}'"
             return f"{col.name} {desc}"
 
-        column_descs = ", ".join([f"{create_col_desc(col)}" for col in feature_view.output_schema.fields])
+        column_descs = (
+            ", ".join([f"{create_col_desc(col)}" for col in feature_view.output_schema.fields])
+            if feature_view.feature_descs is not None
+            else ""
+        )
 
         if refresh_freq is not None:
             schedule_task = refresh_freq != "DOWNSTREAM" and timeparse(refresh_freq) is None
@@ -1819,6 +1831,7 @@ class FeatureStore:
                 )
                 WAREHOUSE = {warehouse}
                 REFRESH_MODE = {feature_view.refresh_mode}
+                INITIALIZE = {feature_view.initialize}
                 AS {feature_view.query}
             """
             self._session.sql(query).collect(block=block, statement_params=self._telemetry_stmp)
@@ -2293,6 +2306,8 @@ class FeatureStore:
             entities = [find_and_compose_entity(n) for n in fv_metadata.entities]
             ts_col = fv_metadata.timestamp_col
             timestamp_col = ts_col if ts_col not in _LEGACY_TIMESTAMP_COL_PLACEHOLDER_VALS else None
+            re_initialize = re.match(_DT_INITIALIZE_PATTERN, row["text"])
+            initialize = re_initialize.group("initialize") if re_initialize is not None else "ON_CREATE"
 
             fv = FeatureView._construct_feature_view(
                 name=name,
@@ -2317,6 +2332,7 @@ class FeatureStore:
                 ),
                 refresh_mode=row["refresh_mode"],
                 refresh_mode_reason=row["refresh_mode_reason"],
+                initialize=initialize,
                 owner=row["owner"],
                 infer_schema_df=infer_schema_df,
                 session=self._session,
@@ -2343,6 +2359,7 @@ class FeatureStore:
                 warehouse=None,
                 refresh_mode=None,
                 refresh_mode_reason=None,
+                initialize="ON_CREATE",
                 owner=row["owner"],
                 infer_schema_df=infer_schema_df,
                 session=self._session,