snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/utils/io_utils.py

@@ -0,0 +1,35 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import functools
+
+from snowflake.snowpark import Session
+
+
+@functools.cache
+def file_format(
+    session: Session, compression: str, record_delimiter: str = None
+) -> str:
+    """
+    Create a temporary file format for reading text files in Snowpark Connect.
+    """
+    if record_delimiter is None:
+        record_delimiter = "NONE"
+        identifier_delimiter = "NONE"
+    else:
+        record_delimiter = record_delimiter
+        # Encode delimiter to ensure that it is a valid identifier
+        identifier_delimiter = record_delimiter.encode("utf-8").hex()
+
+    file_format_name = f"IDENTIFIER('__SNOWPARK_CONNECT_TEXT_FILE_FORMAT__{compression}_{identifier_delimiter}')"
+    session.sql(
+        f"""
+    CREATE TEMPORARY FILE FORMAT IF NOT EXISTS  {file_format_name}
+    RECORD_DELIMITER = '{record_delimiter}'
+    FIELD_DELIMITER = 'NONE'
+    EMPTY_FIELD_AS_NULL = FALSE
+    COMPRESSION = '{compression}'"""
+    ).collect()
+
+    return file_format_name

snowflake/snowpark_connect/utils/pandas_udtf_utils.py

@@ -0,0 +1,117 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+from typing import Any, Callable, Iterator
+
+import cloudpickle
+import pandas as pd
+import pyarrow as pa
+from pyspark.sql.connect.proto.expressions_pb2 import CommonInlineUserDefinedFunction
+
+import snowflake.snowpark.functions as snowpark_fn
+from snowflake import snowpark
+from snowflake.snowpark.types import IntegerType, PandasDataFrameType, StructType
+
+
+def get_map_in_arrow_udtf(
+    user_function: Callable,
+    spark_column_names: list[str],
+    output_column_names: list[str],
+) -> Any:
+    """
+    Create and return a MapInArrowUDTF class with the given parameters.
+
+    Args:
+        user_function: Arrow function that processes RecordBatch iterators
+        spark_column_names: List of spark column names of the given dataframe.
+        output_column_names: List of expected output column names
+
+    Returns:
+        MapInArrowUDTF class that can be used with pandas_udtf
+    """
+
+    class MapInArrowUDTF:
+        def __init__(self) -> None:
+            self.user_function = user_function
+            self.output_column_names = output_column_names
+            self.spark_column_names = spark_column_names
+
+        def end_partition(self, df: pd.DataFrame):
+            if df.empty:
+                empty_df = pd.DataFrame(columns=self.output_column_names)
+                yield empty_df
+                return
+
+            df_without_dummy = df.drop(
+                columns=["_DUMMY_PARTITION_KEY"], errors="ignore"
+            )
+            df_without_dummy.columns = self.spark_column_names
+
+            # Convert pandas DataFrame to Arrow format
+            table = pa.Table.from_pandas(df_without_dummy, preserve_index=False)
+            batch_iterator = table.to_batches()
+
+            result_iterator = self.user_function(batch_iterator)
+
+            result_batches = []
+
+            if not isinstance(result_iterator, Iterator) and not hasattr(
+                result_iterator, "__iter__"
+            ):
+                raise RuntimeError(
+                    f"snowpark_connect::UDF_RETURN_TYPE Return type of the user-defined function should be "
+                    f"iterator of pyarrow.RecordBatch, but is {type(result_iterator).__name__}"
+                )
+
+            for batch in result_iterator:
+                if not isinstance(batch, pa.RecordBatch):
+                    raise RuntimeError(
+                        f"snowpark_connect::UDF_RETURN_TYPE Return type of the user-defined function should "
+                        f"be iterator of pyarrow.RecordBatch, but is iterator of {type(batch).__name__}"
+                    )
+                if batch.num_rows > 0:
+                    result_batches.append(batch)
+
+            if result_batches:
+                combined_table = pa.Table.from_batches(result_batches)
+                result_df = combined_table.to_pandas()
+                yield result_df
+            else:
+                empty_df = pd.DataFrame(columns=self.output_column_names)
+                yield empty_df
+
+    return MapInArrowUDTF
+
+
+def create_pandas_udtf(
+    udtf_proto: CommonInlineUserDefinedFunction,
+    spark_column_names: list[str],
+    input_schema: StructType | None = None,
+    return_schema: StructType | None = None,
+) -> str | snowpark.udtf.UserDefinedTableFunction:
+    user_function, _ = cloudpickle.loads(udtf_proto.python_udf.command)
+    output_column_names = [field.name for field in return_schema.fields]
+
+    MapInArrowUDTF = get_map_in_arrow_udtf(
+        user_function, spark_column_names, output_column_names
+    )
+
+    return snowpark_fn.pandas_udtf(
+        MapInArrowUDTF,
+        output_schema=PandasDataFrameType(
+            [field.datatype for field in return_schema.fields],
+            [field.name for field in return_schema.fields],
+        ),
+        input_types=[
+            PandasDataFrameType(
+                [field.datatype for field in input_schema.fields] + [IntegerType()]
+            )
+        ],
+        input_names=[field.name for field in input_schema.fields]
+        + ["_DUMMY_PARTITION_KEY"],
+        name="mapinarrow_udtf",
+        replace=True,
+        packages=["pyarrow", "pandas"],
+        is_permanent=False,
+    )

snowflake/snowpark_connect/utils/profiling.py

@@ -0,0 +1,47 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+# Some basic profiling that is enabled with env variable SNOWPARK_CONNECT_ENABLE_PROFILING = 1
+# By default, profiles are written for each annotated method into /tmp/snowpark_connect_profiles, this is
+# controlled by SNOWPARK_CONNECT_PROFILE_OUTPUT_DIR
+#
+# By shipping this, we can ask customers/partners to create a profile for us to investigate further if needed.
+
+import cProfile
+import functools
+import os
+from datetime import datetime
+from typing import Any, Callable
+
+PROFILING_ENABLED = os.environ.get("SNOWPARK_CONNECT_ENABLE_PROFILING", "0") == "1"
+PROFILE_OUTPUT_DIR = os.environ.get(
+    "SNOWPARK_CONNECT_PROFILE_OUTPUT_DIR", "/tmp/snowpark_connect_profiles"
+)
+
+
+def profile_method(method: Callable) -> Callable:
+    """Decorator to profile a specific method."""
+
+    @functools.wraps(method)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        if not PROFILING_ENABLED:
+            return method(*args, **kwargs)
+
+        os.makedirs(PROFILE_OUTPUT_DIR, exist_ok=True)
+
+        method_name = method.__name__
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
+        profile_filename = f"{PROFILE_OUTPUT_DIR}/{method_name}_{timestamp}.prof"
+
+        profiler = cProfile.Profile()
+        profiler.enable()
+
+        try:
+            result = method(*args, **kwargs)
+            return result
+        finally:
+            profiler.disable()
+            profiler.dump_stats(profile_filename)
+
+    return wrapper

snowflake/snowpark_connect/utils/session.py

@@ -0,0 +1,180 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import logging
+import os
+from collections.abc import Sequence
+from contextlib import contextmanager
+from typing import Any
+
+from snowflake import snowpark
+from snowflake.snowpark.exceptions import SnowparkClientException
+from snowflake.snowpark.session import _get_active_session
+from snowflake.snowpark_connect.constants import DEFAULT_CONNECTION_NAME
+from snowflake.snowpark_connect.utils.describe_query_cache import (
+    instrument_session_for_describe_cache,
+)
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
+from snowflake.snowpark_connect.utils.telemetry import telemetry
+from snowflake.snowpark_connect.utils.udf_cache import init_builtin_udf_cache
+
+
+def _get_current_snowpark_session() -> snowpark.Session | None:
+    # TODO: this is a temporary solution to get the current session, it would be better to add a function in snowpark
+    try:
+        session = _get_active_session()
+        # if session._conn._conn.expired:
+        #     _remove_session(session)
+        #     return self.create()
+        return session
+    except SnowparkClientException as ex:
+        if ex.error_code == "1403":  # No session
+            return None
+        raise
+
+
+@contextmanager
+def suppress_experimental_warnings():
+    """
+    Suppress experimental parameter warnings from snowpark logging.
+
+    This context manager filters out logging messages containing
+    "is experimental since" and "Do not use it in production"
+    from the snowpark logger, while preserving other important warnings.
+    """
+    snowpark_logger = logging.getLogger("snowflake.snowpark")
+
+    def filter_experimental_warnings(record):
+        """Filter function to suppress experimental parameter warnings."""
+        message = record.getMessage()
+        return not (
+            "is experimental since" in message
+            and "Do not use it in production" in message
+        )
+
+    snowpark_logger.addFilter(filter_experimental_warnings)
+
+    try:
+        yield
+    finally:
+        snowpark_logger.removeFilter(filter_experimental_warnings)
+
+
+def configure_snowpark_session(session: snowpark.Session):
+    """Configure a snowpark session with required parameters and settings."""
+    from snowflake.snowpark_connect.config import global_config
+
+    logger.info(f"Configuring session {session}")
+
+    telemetry.initialize(session)
+    # custom udfs
+    session._udfs = {}
+    session._udtfs = {}
+    session._sprocs = set()
+
+    # custom udf imports
+    session._python_files = set()
+    session._import_files = set()
+
+    # built-in udf cache
+    init_builtin_udf_cache(session)
+
+    # Set experimental parameters with warning suppression
+    with suppress_experimental_warnings():
+        session.ast_enabled = False
+        session.eliminate_numeric_sql_value_cast_enabled = False
+        session.reduce_describe_query_enabled = True
+
+    session._join_alias_fix = True
+    session.connection.arrow_number_to_decimal_setter = True
+    session.custom_package_usage_config["enabled"] = True
+
+    session_params = {
+        "TIMESTAMP_TYPE_MAPPING": "TIMESTAMP_LTZ",
+        "TIMEZONE": f"'{global_config.spark_sql_session_timeZone}'",
+        "QUOTED_IDENTIFIERS_IGNORE_CASE": "false",
+        "PYTHON_SNOWPARK_ENABLE_THREAD_SAFE_SESSION": "true",
+        "PYTHON_SNOWPARK_USE_SCOPED_TEMP_OBJECTS": "false",  # this is required for creating udfs from sproc
+    }
+
+    session.sql(
+        f"ALTER SESSION SET {', '.join([f'{k} = {v}' for k, v in session_params.items()])}"
+    ).collect()
+    # Instrument the snowpark session to use a cache for describe queries.
+    instrument_session_for_describe_cache(session)
+
+
+def _is_running_in_SPCS():
+    return (
+        os.path.exists("/snowflake/session/token")
+        and os.getenv("SNOWFLAKE_ACCOUNT") is not None
+        and os.getenv("SNOWFLAKE_HOST") is not None
+    )
+
+
+def _get_session_configs_from_ENV() -> dict[str, Any]:
+    session_configs = {
+        "account": os.getenv("SNOWFLAKE_ACCOUNT"),
+        "protocol": "https",
+        "host": os.getenv("SNOWFLAKE_HOST"),
+        "port": os.getenv("SNOWFLAKE_PORT", 443),
+        "authenticator": "oauth",
+        "token_file_path": "/snowflake/session/token",
+        "client_session_keep_alive": True,
+    }
+    return session_configs
+
+
+def get_or_create_snowpark_session(
+    custom_configs: dict | None = None,
+) -> snowpark.Session:
+    """
+    snowpark connect code should use this function to create or get snowpark session
+    """
+    session_configs = {}
+    if _is_running_in_SPCS():
+        # Running in SPCS, use environment variables injected by SPCS run time
+        # We don't use connections.toml file created by SPCS because of the 0600 permissions issue
+        session_configs = _get_session_configs_from_ENV()
+    else:
+        session_configs["connection_name"] = DEFAULT_CONNECTION_NAME
+
+    if os.getenv("SNOWFLAKE_DATABASE") is not None:
+        session_configs["database"] = os.getenv("SNOWFLAKE_DATABASE")
+
+    if os.getenv("SNOWFLAKE_SCHEMA") is not None:
+        session_configs["schema"] = os.getenv("SNOWFLAKE_SCHEMA")
+
+    if os.getenv("SNOWFLAKE_WAREHOUSE") is not None:
+        session_configs["warehouse"] = os.getenv("SNOWFLAKE_WAREHOUSE")
+
+    # add custom session configs
+    if custom_configs:
+        session_configs.update(custom_configs)
+
+    old_session = _get_current_snowpark_session()
+    new_session = snowpark.Session.builder.configs(session_configs).getOrCreate()
+    if old_session is None or old_session.session_id != new_session.session_id:
+        # every new session needs to be configured
+        configure_snowpark_session(new_session)
+
+    return new_session
+
+
+def set_query_tags(spark_tags: Sequence[str]) -> None:
+    """Sets Snowpark session query_tag value to the tag from the Spark request."""
+
+    if any("," in tag for tag in spark_tags):
+        raise ValueError("Tags cannot contain ','.")
+
+    # TODO: Tags might not be set correctly in parallel workloads or multi-threaded code.
+    snowpark_session = get_or_create_snowpark_session()
+    spark_tags_str = ",".join(sorted(spark_tags)) if spark_tags else None
+
+    if spark_tags_str != snowpark_session.query_tag:
+        snowpark_session.query_tag = spark_tags_str
+
+
+def get_python_udxf_import_files(session: snowpark.Session) -> str:
+    return ",".join([file for file in [*session._python_files, *session._import_files]])

snowflake/snowpark_connect/utils/snowpark_connect_logging.py

@@ -0,0 +1,38 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import logging
+
+from pyspark import StorageLevel
+
+logger = logging.getLogger("snowflake_connect_server")
+logger.setLevel(logging.WARN)
+
+console_handler = logging.StreamHandler()
+console_handler.setLevel(logging.INFO)
+formatter = logging.Formatter(
+    "%(asctime)s - %(name)s - %(levelname)s - [Thread %(thread)d] - %(message)s"
+)
+console_handler.setFormatter(formatter)
+# Display the logs to the console
+logger.addHandler(console_handler)
+
+
+def run_once_decorator(func):
+    def wrapper(*args, **kwargs):
+        if not hasattr(wrapper, "has_run"):
+            wrapper.has_run = True
+            return func(*args, **kwargs)
+
+    return wrapper
+
+
+@run_once_decorator
+def log_waring_once_storage_level(storage_level: StorageLevel):
+    logger.warning(
+        f"Ignored unsupported Spark storage level:\n{storage_level}"
+        "Snowflake will always create materialized temp table from the dataframe "
+        "when dataframe.cache or dataframe.persist is called.\n"
+        "The behavior is similar with Spark's StorageLevel.DISK_ONLY."
+    )