snowflake-ml-python 1.7.2__py3-none-any.whl → 1.7.3__py3-none-any.whl

snowflake/cortex/__init__.py

@@ -1,24 +1,32 @@
-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._classify_text import ClassifyText, classify_text
+from snowflake.cortex._complete import Complete, CompleteOptions, complete
+from snowflake.cortex._embed_text_768 import EmbedText768, embed_text_768
+from snowflake.cortex._embed_text_1024 import EmbedText1024, embed_text_1024
+from snowflake.cortex._extract_answer import ExtractAnswer, extract_answer
 from snowflake.cortex._finetune import Finetune, FinetuneJob, FinetuneStatus
-from snowflake.cortex._sentiment import Sentiment
-from snowflake.cortex._summarize import Summarize
-from snowflake.cortex._translate import Translate
+from snowflake.cortex._sentiment import Sentiment, sentiment
+from snowflake.cortex._summarize import Summarize, summarize
+from snowflake.cortex._translate import Translate, translate
 
 __all__ = [
     "ClassifyText",
+    "classify_text",
     "Complete",
+    "complete",
     "CompleteOptions",
     "EmbedText768",
+    "embed_text_768",
     "EmbedText1024",
+    "embed_text_1024",
     "ExtractAnswer",
+    "extract_answer",
     "Finetune",
     "FinetuneJob",
     "FinetuneStatus",
     "Sentiment",
+    "sentiment",
     "Summarize",
+    "summarize",
     "Translate",
+    "translate",
 ]

snowflake/cortex/_classify_text.py

@@ -1,5 +1,7 @@
 from typing import List, Optional, Union, cast
 
+from typing_extensions import deprecated
+
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
 from snowflake.ml._internal import telemetry
@@ -8,7 +10,7 @@ from snowflake.ml._internal import telemetry
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )
-def ClassifyText(
+def classify_text(
     str_input: Union[str, snowpark.Column],
     categories: Union[List[str], snowpark.Column],
     session: Optional[snowpark.Session] = None,
@@ -34,3 +36,12 @@ def _classify_text_impl(
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
     return cast(Union[str, snowpark.Column], call_sql_function(function, session, str_input, categories))
+
+
+ClassifyText = deprecated(
+    "ClassifyText() is deprecated and will be removed in a future release. Please use classify_text() instead."
+)(
+    telemetry.send_api_usage_telemetry(
+        project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+    )(classify_text)
+)

snowflake/cortex/_complete.py

@@ -6,7 +6,7 @@ from typing import Any, Callable, Dict, Iterator, List, Optional, TypedDict, Uni
 from urllib.parse import urlunparse
 
 import requests
-from typing_extensions import NotRequired
+from typing_extensions import NotRequired, deprecated
 
 from snowflake import snowpark
 from snowflake.cortex._sse_client import SSEClient
@@ -127,8 +127,26 @@ def _xp_dict_to_response(raw_resp: Dict[str, Any]) -> requests.Response:
     response.status_code = int(raw_resp["status"])
     response.headers = raw_resp["headers"]
 
+    request_id = None
+    for key, value in raw_resp["headers"].items():
+        # Note: there is some whitespace in the headers making it not possible
+        # to directly index the header reliably.
+        if key.strip().lower() == "x-snowflake-request-id":
+            request_id = value
+            break
+
     data = raw_resp["content"]
-    data = json.loads(data)
+    try:
+        data = json.loads(data)
+    except json.JSONDecodeError:
+        raise ValueError(f"Request failed (request id: {request_id})")
+
+    if response.status_code < 200 or response.status_code >= 300:
+        if "message" not in data:
+            raise ValueError(f"Request failed (request id: {request_id})")
+        message = data["message"]
+        raise ValueError(f"Request failed: {message} (request id: {request_id})")
+
     # Convert the dictionary to a string format that resembles the SSE event format
     # For example, if the dict is {'event': 'message', 'data': 'your data'}, it should be formatted like this:
     sse_format_data = ""
@@ -144,6 +162,7 @@ def _xp_dict_to_response(raw_resp: Dict[str, Any]) -> requests.Response:
 
 @retry
 def _call_complete_xp(
+    snow_api_xp_request_handler: Optional[Callable[..., Dict[str, Any]]],
     model: str,
     prompt: Union[str, List[ConversationMessage]],
     options: Optional[CompleteOptions] = None,
@@ -151,9 +170,8 @@ def _call_complete_xp(
 ) -> requests.Response:
     headers = _make_common_request_headers()
     body = _make_request_body(model, prompt, options)
-    import _snowflake
-
-    raw_resp = _snowflake.send_snow_api_request("POST", _REST_COMPLETE_URL, {}, headers, body, {}, deadline)
+    assert snow_api_xp_request_handler is not None
+    raw_resp = snow_api_xp_request_handler("POST", _REST_COMPLETE_URL, {}, headers, body, {}, deadline)
     return _xp_dict_to_response(raw_resp)
 
 
@@ -218,17 +236,26 @@ def _complete_call_sql_function_snowpark(
 
 
 def _complete_non_streaming_immediate(
+    snow_api_xp_request_handler: Optional[Callable[..., Dict[str, Any]]],
     model: str,
     prompt: Union[str, List[ConversationMessage]],
     options: Optional[CompleteOptions],
     session: Optional[snowpark.Session] = None,
     deadline: Optional[float] = None,
 ) -> str:
-    response = _complete_rest(model=model, prompt=prompt, options=options, session=session, deadline=deadline)
+    response = _complete_rest(
+        snow_api_xp_request_handler=snow_api_xp_request_handler,
+        model=model,
+        prompt=prompt,
+        options=options,
+        session=session,
+        deadline=deadline,
+    )
     return "".join(response)
 
 
 def _complete_non_streaming_impl(
+    snow_api_xp_request_handler: Optional[Callable[..., Dict[str, Any]]],
     function: str,
     model: Union[str, snowpark.Column],
     prompt: Union[str, List[ConversationMessage], snowpark.Column],
@@ -246,19 +273,31 @@ def _complete_non_streaming_impl(
     if isinstance(options, snowpark.Column):
         raise ValueError("'options' cannot be a snowpark.Column when 'prompt' is a string.")
     return _complete_non_streaming_immediate(
-        model=model, prompt=prompt, options=options, session=session, deadline=deadline
+        snow_api_xp_request_handler=snow_api_xp_request_handler,
+        model=model,
+        prompt=prompt,
+        options=options,
+        session=session,
+        deadline=deadline,
     )
 
 
 def _complete_rest(
+    snow_api_xp_request_handler: Optional[Callable[..., Dict[str, Any]]],
     model: str,
     prompt: Union[str, List[ConversationMessage]],
     options: Optional[CompleteOptions] = None,
     session: Optional[snowpark.Session] = None,
     deadline: Optional[float] = None,
 ) -> Iterator[str]:
-    if is_in_stored_procedure():  # type: ignore[no-untyped-call]
-        response = _call_complete_xp(model=model, prompt=prompt, options=options, deadline=deadline)
+    if snow_api_xp_request_handler is not None:
+        response = _call_complete_xp(
+            snow_api_xp_request_handler=snow_api_xp_request_handler,
+            model=model,
+            prompt=prompt,
+            options=options,
+            deadline=deadline,
+        )
     else:
         response = _call_complete_rest(model=model, prompt=prompt, options=options, session=session, deadline=deadline)
     assert response.status_code >= 200 and response.status_code < 300
@@ -268,10 +307,11 @@ def _complete_rest(
 def _complete_impl(
     model: Union[str, snowpark.Column],
     prompt: Union[str, List[ConversationMessage], snowpark.Column],
+    snow_api_xp_request_handler: Optional[Callable[..., Dict[str, Any]]] = None,
+    function: str = "snowflake.cortex.complete",
     options: Optional[CompleteOptions] = None,
     session: Optional[snowpark.Session] = None,
     stream: bool = False,
-    function: str = "snowflake.cortex.complete",
     timeout: Optional[float] = None,
     deadline: Optional[float] = None,
 ) -> Union[str, Iterator[str], snowpark.Column]:
@@ -284,14 +324,29 @@ def _complete_impl(
             raise ValueError("in REST mode, 'model' must be a string")
         if not isinstance(prompt, str) and not isinstance(prompt, List):
             raise ValueError("in REST mode, 'prompt' must be a string or a list of ConversationMessage")
-        return _complete_rest(model=model, prompt=prompt, options=options, session=session, deadline=deadline)
-    return _complete_non_streaming_impl(function, model, prompt, options, session, deadline)
+        return _complete_rest(
+            snow_api_xp_request_handler=snow_api_xp_request_handler,
+            model=model,
+            prompt=prompt,
+            options=options,
+            session=session,
+            deadline=deadline,
+        )
+    return _complete_non_streaming_impl(
+        snow_api_xp_request_handler=snow_api_xp_request_handler,
+        function=function,
+        model=model,
+        prompt=prompt,
+        options=options,
+        session=session,
+        deadline=deadline,
+    )
 
 
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )
-def Complete(
+def complete(
     model: Union[str, snowpark.Column],
     prompt: Union[str, List[ConversationMessage], snowpark.Column],
     *,
@@ -319,10 +374,19 @@ def Complete(
     Returns:
         A column of string responses.
     """
+
+    # Set the XP snow api function, if available.
+    snow_api_xp_request_handler = None
+    if is_in_stored_procedure():  # type: ignore[no-untyped-call]
+        import _snowflake
+
+        snow_api_xp_request_handler = _snowflake.send_snow_api_request
+
     try:
         return _complete_impl(
             model,
             prompt,
+            snow_api_xp_request_handler=snow_api_xp_request_handler,
             options=options,
             session=session,
             stream=stream,
@@ -331,3 +395,8 @@ def Complete(
         )
     except ValueError as err:
         raise err
+
+
+Complete = deprecated("Complete() is deprecated and will be removed in a future release. Use complete() instead")(
+    telemetry.send_api_usage_telemetry(project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT)(complete)
+)

snowflake/cortex/_embed_text_1024.py

@@ -1,5 +1,7 @@
 from typing import List, Optional, Union, cast
 
+from typing_extensions import deprecated
+
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
 from snowflake.ml._internal import telemetry
@@ -8,12 +10,12 @@ from snowflake.ml._internal import telemetry
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )
-def EmbedText1024(
+def embed_text_1024(
     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.
+    """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
@@ -35,3 +37,8 @@ def _embed_text_1024_impl(
     session: Optional[snowpark.Session] = None,
 ) -> Union[List[float], snowpark.Column]:
     return cast(Union[List[float], snowpark.Column], call_sql_function(function, session, model, text))
+
+
+EmbedText1024 = deprecated(
+    "EmbedText1024() is deprecated and will be removed in a future release. Use embed_text_1024() instead"
+)(telemetry.send_api_usage_telemetry(project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT)(embed_text_1024))

snowflake/cortex/_embed_text_768.py

@@ -1,5 +1,7 @@
 from typing import List, Optional, Union, cast
 
+from typing_extensions import deprecated
+
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
 from snowflake.ml._internal import telemetry
@@ -8,12 +10,12 @@ from snowflake.ml._internal import telemetry
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )
-def EmbedText768(
+def embed_text_768(
     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.
+    """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
@@ -35,3 +37,8 @@ def _embed_text_768_impl(
     session: Optional[snowpark.Session] = None,
 ) -> Union[List[float], snowpark.Column]:
     return cast(Union[List[float], snowpark.Column], call_sql_function(function, session, model, text))
+
+
+EmbedText768 = deprecated(
+    "EmbedText768() is deprecated and will be removed in a future release. Use embed_text_768() instead"
+)(telemetry.send_api_usage_telemetry(project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT)(embed_text_768))

snowflake/cortex/_extract_answer.py

@@ -1,5 +1,7 @@
 from typing import Optional, Union, cast
 
+from typing_extensions import deprecated
+
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
 from snowflake.ml._internal import telemetry
@@ -8,12 +10,12 @@ from snowflake.ml._internal import telemetry
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )
-def ExtractAnswer(
+def extract_answer(
     from_text: Union[str, snowpark.Column],
     question: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    """ExtractAnswer calls into the LLM inference service to extract an answer from within specified text.
+    """Calls into the LLM inference service to extract an answer from within specified text.
 
     Args:
         from_text: A Column of strings representing input text.
@@ -34,3 +36,8 @@ def _extract_answer_impl(
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
     return cast(Union[str, snowpark.Column], call_sql_function(function, session, from_text, question))
+
+
+ExtractAnswer = deprecated(
+    "ExtractAnswer() is deprecated and will be removed in a future release. Use extract_answer() instead"
+)(telemetry.send_api_usage_telemetry(project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT)(extract_answer))

snowflake/cortex/_sentiment.py

@@ -1,5 +1,7 @@
 from typing import Optional, Union, cast
 
+from typing_extensions import deprecated
+
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
 from snowflake.ml._internal import telemetry
@@ -8,10 +10,10 @@ from snowflake.ml._internal import telemetry
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )
-def Sentiment(
+def sentiment(
     text: Union[str, snowpark.Column], session: Optional[snowpark.Session] = None
 ) -> Union[float, snowpark.Column]:
-    """Sentiment calls into the LLM inference service to perform sentiment analysis on the input text.
+    """Calls into the LLM inference service to perform sentiment analysis on the input text.
 
     Args:
         text: A Column of text strings to send to the LLM.
@@ -31,3 +33,8 @@ def _sentiment_impl(
     if isinstance(output, snowpark.Column):
         return output
     return float(cast(str, output))
+
+
+Sentiment = deprecated("Sentiment() is deprecated and will be removed in a future release. Use sentiment() instead")(
+    telemetry.send_api_usage_telemetry(project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT)(sentiment)
+)

snowflake/cortex/_summarize.py

@@ -1,5 +1,7 @@
 from typing import Optional, Union, cast
 
+from typing_extensions import deprecated
+
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
 from snowflake.ml._internal import telemetry
@@ -8,11 +10,11 @@ from snowflake.ml._internal import telemetry
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )
-def Summarize(
+def summarize(
     text: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    """Summarize calls into the LLM inference service to summarize the input text.
+    """Calls into the LLM inference service to summarize the input text.
 
     Args:
         text: A Column of strings to summarize.
@@ -31,3 +33,8 @@ def _summarize_impl(
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
     return cast(Union[str, snowpark.Column], call_sql_function(function, session, text))
+
+
+Summarize = deprecated("Summarize() is deprecated and will be removed in a future release. Use summarize() instead")(
+    telemetry.send_api_usage_telemetry(project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT)(summarize)
+)

snowflake/cortex/_translate.py

@@ -1,5 +1,7 @@
 from typing import Optional, Union, cast
 
+from typing_extensions import deprecated
+
 from snowflake import snowpark
 from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
 from snowflake.ml._internal import telemetry
@@ -8,13 +10,13 @@ from snowflake.ml._internal import telemetry
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )
-def Translate(
+def translate(
     text: Union[str, snowpark.Column],
     from_language: Union[str, snowpark.Column],
     to_language: Union[str, snowpark.Column],
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
-    """Translate calls into the LLM inference service to perform translation.
+    """Calls into the LLM inference service to perform translation.
 
     Args:
         text: A Column of strings to translate.
@@ -37,3 +39,8 @@ def _translate_impl(
     session: Optional[snowpark.Session] = None,
 ) -> Union[str, snowpark.Column]:
     return cast(Union[str, snowpark.Column], call_sql_function(function, session, text, from_language, to_language))
+
+
+Translate = deprecated("Translate() is deprecated and will be removed in a future release. Use translate() instead")(
+    telemetry.send_api_usage_telemetry(project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT)(translate)
+)

snowflake/ml/_internal/env_utils.py

@@ -15,7 +15,6 @@ import snowflake.connector
 from snowflake.ml._internal import env as snowml_env
 from snowflake.ml._internal.utils import query_result_checker
 from snowflake.snowpark import context, exceptions, session
-from snowflake.snowpark._internal import utils as snowpark_utils
 
 
 class CONDA_OS(Enum):
@@ -344,55 +343,6 @@ def relax_requirement_version(req: requirements.Requirement) -> requirements.Req
     return new_req
 
 
-def get_matched_package_versions_in_snowflake_conda_channel(
-    req: requirements.Requirement,
-    python_version: str = snowml_env.PYTHON_VERSION,
-    conda_os: CONDA_OS = CONDA_OS.LINUX_64,
-) -> List[version.Version]:
-    """Search the snowflake anaconda channel for packages that matches the specifier. Note that this will be the
-    source of truth for checking whether a package indeed exists in Snowflake conda channel.
-
-    Given that a package comes in different architectures, we only check for the Linux x86_64 architecture and assume
-    the package exists in other architectures. If such an assumption does not hold true for a certain package, the
-    caller should specify the architecture to search for.
-
-    Args:
-        req: Requirement specifier.
-        python_version: A string of python version where model is run.
-        conda_os: Specified platform to search availability of the package.
-
-    Returns:
-        List of package versions that meet the requirement specifier.
-    """
-    # Move the retryable_http import here as when UDF import this file, it won't have the "requests" dependency.
-    from snowflake.ml._internal.utils import retryable_http
-
-    assert not snowpark_utils.is_in_stored_procedure()  # type: ignore[no-untyped-call]
-
-    url = f"{SNOWFLAKE_CONDA_CHANNEL_URL}/{conda_os.value}/repodata.json"
-
-    if req.name not in _SNOWFLAKE_CONDA_PACKAGE_CACHE:
-        try:
-            http_client = retryable_http.get_http_client()
-            parsed_python_version = version.Version(python_version)
-            python_version_build_str = f"py{parsed_python_version.major}{parsed_python_version.minor}"
-            repodata = http_client.get(url).json()
-            assert isinstance(repodata, dict)
-            packages_info = repodata["packages"]
-            assert isinstance(packages_info, dict)
-            version_list = [
-                version.parse(package_info["version"])
-                for package_info in packages_info.values()
-                if package_info["name"] == req.name and python_version_build_str in package_info["build"]
-            ]
-            _SNOWFLAKE_CONDA_PACKAGE_CACHE[req.name] = version_list
-        except Exception:
-            pass
-
-    matched_versions = list(req.specifier.filter(set(_SNOWFLAKE_CONDA_PACKAGE_CACHE.get(req.name, []))))
-    return matched_versions
-
-
 def get_matched_package_versions_in_information_schema_with_active_session(
     reqs: List[requirements.Requirement], python_version: str
 ) -> Dict[str, List[version.Version]]:
@@ -404,7 +354,10 @@ def get_matched_package_versions_in_information_schema_with_active_session(
 
 
 def get_matched_package_versions_in_information_schema(
-    session: session.Session, reqs: List[requirements.Requirement], python_version: str
+    session: session.Session,
+    reqs: List[requirements.Requirement],
+    python_version: str,
+    statement_params: Optional[Dict[str, Any]] = None,
 ) -> Dict[str, List[version.Version]]:
     """Look up the information_schema table to check if a package with the specified specifier exists in the Snowflake
     Conda channel. Note that this is not the source of truth due to the potential delay caused by a package that might
@@ -414,6 +367,7 @@ def get_matched_package_versions_in_information_schema(
         session: Snowflake connection session.
         reqs: List of requirement specifiers.
         python_version: A string of python version where model is run.
+        statement_params: Optional statement parameters.
 
     Returns:
         A Dict, whose key is the package name, and value is a list of versions match the requirements.
@@ -451,8 +405,9 @@ def get_matched_package_versions_in_information_schema(
                 query_result_checker.SqlResultValidator(
                     session=session,
                     query=sql,
+                    statement_params=statement_params,
                 )
-                .has_column("VERSION")
+                .has_column("VERSION", allow_empty=True)
                 .has_dimensions(expected_rows=None, expected_cols=2)
                 .validate()
             )

snowflake/ml/_internal/utils/identifier.py

@@ -158,8 +158,10 @@ def parse_schema_level_object_identifier(
     res = _SF_SCHEMA_LEVEL_OBJECT_RE.fullmatch(object_name)
     if not res:
         raise ValueError(
-            "Invalid identifier because it does not follow the pattern. "
-            f"It should start with [[database.]schema.]object. Getting {object_name}"
+            f"Invalid object name `{object_name}` cannot be parsed as a SQL identifier. "
+            "Alphanumeric characters and underscores are permitted. "
+            "See https://docs.snowflake.com/en/sql-reference/identifiers-syntax for "
+            "more information."
         )
     return (
         res.group("db"),

snowflake/ml/data/__init__.py

@@ -1,5 +1,8 @@
+from pkgutil import extend_path
+
 from .data_connector import DataConnector
 from .data_ingestor import DataIngestor, DataIngestorType
 from .data_source import DataFrameInfo, DatasetInfo, DataSource
 
 __all__ = ["DataConnector", "DataSource", "DataFrameInfo", "DatasetInfo", "DataIngestor", "DataIngestorType"]
+__path__ = extend_path(__path__, __name__)

snowflake/ml/data/_internal/arrow_ingestor.py

@@ -2,7 +2,7 @@ import collections
 import logging
 import os
 import time
-from typing import Any, Deque, Dict, Iterator, List, Optional, Union
+from typing import Any, Deque, Dict, Iterator, List, Optional, Sequence, Union
 
 import numpy as np
 import numpy.typing as npt
@@ -47,7 +47,7 @@ class ArrowIngestor(data_ingestor.DataIngestor):
     def __init__(
         self,
         session: snowpark.Session,
-        data_sources: List[data_source.DataSource],
+        data_sources: Sequence[data_source.DataSource],
         format: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
@@ -60,14 +60,14 @@ class ArrowIngestor(data_ingestor.DataIngestor):
             kwargs: Miscellaneous arguments passed to underlying PyArrow Dataset initializer.
         """
         self._session = session
-        self._data_sources = data_sources
+        self._data_sources = list(data_sources)
         self._format = format
         self._kwargs = kwargs
 
         self._schema: Optional[pa.Schema] = None
 
     @classmethod
-    def from_sources(cls, session: snowpark.Session, sources: List[data_source.DataSource]) -> "ArrowIngestor":
+    def from_sources(cls, session: snowpark.Session, sources: Sequence[data_source.DataSource]) -> "ArrowIngestor":
         return cls(session, sources)
 
     @property