snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/utils/context.py

@@ -0,0 +1,390 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import re
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import Mapping, Optional
+
+import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
+
+from snowflake import snowpark
+from snowflake.snowpark_connect.typed_column import TypedColumn
+
+# TODO: remove session id from context when we host SAS in Snowflake server
+
+_session_id = ContextVar[str]("_session_id")
+_plan_id_map = ContextVar[Mapping[int, snowpark.DataFrame]]("_plan_id_map")
+_alias_map = ContextVar[Mapping[str, snowpark.DataFrame | None]]("_alias_map")
+_spark_version = ContextVar[str]("_spark_version")
+_is_aggregate_function = ContextVar(
+    "_is_aggregate_function", default=("default", False)
+)
+_is_evaluating_sql = ContextVar[bool]("_is_evaluating_sql", default=False)
+_is_evaluating_join_condition = ContextVar(
+    "_is_evaluating_join_condition", default=("default", False, [], [])
+)
+
+_sql_aggregate_function_count = ContextVar[int](
+    "_contains_aggregate_function", default=0
+)
+
+# We have to generate our own plan IDs that are different from Spark's.
+# Spark plan IDs start at 0, so pick a "big enough" number to avoid overlaps.
+_STARTING_SQL_PLAN_ID = 0x80000000
+
+_next_sql_plan_id = ContextVar[int]("_next_sql_plan_id")
+_sql_plan_name_map = ContextVar[dict[str, int]]("_sql_plan_name_map")
+_sql_named_args = ContextVar[dict[str, expressions_proto.Expression]]("_sql_named_args")
+_sql_pos_args = ContextVar[dict[int, expressions_proto.Expression]]("_sql_pos_args")
+
+# Used to store the df before the last projection operation
+_df_before_projection = ContextVar[snowpark.DataFrame | None](
+    "_df_before_projection", default=None
+)
+_outer_dataframes = ContextVar[list[snowpark.DataFrame]]("_parent_dataframes")
+
+_spark_client_type_regex = re.compile(r"spark/(?P<spark_version>\d+\.\d+\.\d+)")
+_current_operation = ContextVar[str]("_current_operation", default="default")
+_resolving_fun_args = ContextVar[bool]("_resolving_fun_args", default=False)
+_resolving_lambda_fun = ContextVar[bool]("_resolving_lambdas", default=False)
+
+_is_window_enabled = ContextVar[bool]("_is_window_enabled", default=False)
+_is_in_pivot = ContextVar[bool]("_is_in_pivot", default=False)
+_is_in_udtf_context = ContextVar[bool]("_is_in_udtf_context", default=False)
+_accessing_temp_object = ContextVar[bool]("_accessing_temp_object", default=False)
+
+# Lateral Column Alias helpers
+# We keep a thread-local mapping from alias name -> TypedColumn that is
+# populated incrementally while the projection list is being processed.
+# The map is cleared at the beginning of every `map_project` call and read
+# by `map_unresolved_attribute` as a last-resort resolver.
+
+_lca_alias_map: ContextVar[dict[str, TypedColumn]] = ContextVar(
+    "_lca_alias_map",
+    default={},
+)
+
+
+def clear_lca_alias_map() -> None:
+    _lca_alias_map.set({})
+
+
+def _normalize(name: str) -> str:
+    from snowflake.snowpark_connect.config import global_config
+
+    return name if global_config.spark_sql_caseSensitive else name.lower()
+
+
+def register_lca_alias(name: str, typed_col: TypedColumn) -> None:
+    alias_map = _lca_alias_map.get()
+    alias_map[_normalize(name)] = typed_col
+    _lca_alias_map.set(alias_map)
+
+
+def resolve_lca_alias(name: str) -> Optional[TypedColumn]:
+    return _lca_alias_map.get().get(_normalize(name))
+
+
+def set_session_id(value: str) -> None:
+    """Set the session ID for the current context"""
+    _session_id.set(value)
+
+
+def get_session_id() -> str:
+    """Get the session ID for the current context."""
+    return _session_id.get(None)
+
+
+def set_plan_id_map(plan_id: int, df: snowpark.DataFrame) -> None:
+    """Set the plan id map for the current context."""
+    _plan_id_map.get()[plan_id] = df
+
+
+def get_plan_id_map(plan_id: int) -> snowpark.DataFrame | None:
+    """Set the plan id map for the current context."""
+    return _plan_id_map.get().get(plan_id)
+
+
+def gen_sql_plan_id() -> int:
+    next = _next_sql_plan_id.get()
+    _next_sql_plan_id.set(next + 1)
+    return next
+
+
+def get_spark_version() -> str:
+    """
+    Get the spark version for the current context, as sent by the client in `client_type`.
+    """
+    return _spark_version.get()
+
+
+def set_spark_version(client_type: str) -> None:
+    """
+    Set the client spark version for the current context.
+    Takes a single `client_type: str` argument in format as sent by pyspark, e.g.
+    "_SPARK_CONNECT_PYTHON spark/3.5.3 os/darwin python/3.11.11" and extracts the spark version - 3.5.3 in this case.
+    """
+    match = _spark_client_type_regex.search(client_type)
+    version = match.group("spark_version") if match else ""
+    _spark_version.set(version)
+
+
+def get_is_aggregate_function() -> tuple[str, bool]:
+    """
+    Gets the value of _is_aggregate_function for the current context, defaults to False.
+    """
+    return _is_aggregate_function.get()
+
+
+def set_is_aggregate_function(is_agg: tuple[str, bool]) -> None:
+    """
+    Sets the value of _is_aggregate_function for the current context.
+    """
+    _is_aggregate_function.set(is_agg)
+
+
+def get_is_evaluating_sql() -> bool:
+    """
+    Gets the value of _is_evaluating_sql for the current context, defaults to False.
+    """
+    return _is_evaluating_sql.get()
+
+
+@contextmanager
+def push_evaluating_sql_scope():
+    """
+    Context manager that sets a flag indicating if a sql statement is being resolved.
+    """
+    prev = _is_evaluating_sql.get()
+    try:
+        _is_evaluating_sql.set(True)
+        yield
+    finally:
+        _is_evaluating_sql.set(prev)
+
+
+def get_is_evaluating_join_condition() -> tuple[str, bool, list, list]:
+    """
+    Gets the value of _is_evaluating_join_condition for the current context, defaults to False.
+    """
+    return _is_evaluating_join_condition.get()
+
+
+@contextmanager
+def push_evaluating_join_condition(join_type, left_keys, right_keys):
+    """
+    Context manager that sets a flag indicating if a join statement is being resolved.
+    """
+    prev = _is_evaluating_join_condition.get()
+    try:
+        _is_evaluating_join_condition.set((join_type, True, left_keys, right_keys))
+        yield
+    finally:
+        _is_evaluating_join_condition.set(prev)
+
+
+@contextmanager
+def push_sql_scope():
+    """
+    Creates a new variable scope when evaluating nested SQL expressions.
+    E.g., in `SELECT x, (SELECT 1 AS x)`, the two `x`s are different variables.
+    """
+    cur = _sql_plan_name_map.get()
+    map_token = _sql_plan_name_map.set(cur.copy())
+    agg_token = _sql_aggregate_function_count.set(0)
+    try:
+        yield
+    finally:
+        _sql_aggregate_function_count.reset(agg_token)
+        _sql_plan_name_map.reset(map_token)
+
+
+@contextmanager
+def push_operation_scope(operation: str):
+    """
+    Context manager that sets the current operation scope for column name resolution.
+    Example:
+    with push_operation_scope('filter'):
+        df.filter(col('historical_name') > 0)  # Historical names allowed in filter
+    """
+    token = _current_operation.set(operation)
+    try:
+        yield
+    finally:
+        _current_operation.reset(token)
+
+
+@contextmanager
+def resolving_lambda_function():
+    """
+    Context manager that sets a flag indicating lambda function is being resolved.
+    """
+    prev = _resolving_lambda_fun.get()
+    try:
+        _resolving_lambda_fun.set(True)
+        yield
+    finally:
+        _resolving_lambda_fun.set(prev)
+
+
+def is_lambda_being_resolved() -> bool:
+    """
+    Returns True if lambda function is being resolved.
+    """
+    return _resolving_lambda_fun.get()
+
+
+@contextmanager
+def resolving_fun_args():
+    """
+    Context manager that sets a flag indicating function arguments are being resolved.
+    """
+    prev = _resolving_fun_args.get()
+    try:
+        _resolving_fun_args.set(True)
+        yield
+    finally:
+        _resolving_fun_args.set(prev)
+
+
+def is_function_argument_being_resolved() -> bool:
+    """
+    Returns True if function arguments are being resolved.
+    """
+    return _resolving_fun_args.get()
+
+
+def get_current_operation_scope() -> str:
+    """
+    Returns the current operation scope for column name resolution.
+    """
+    return _current_operation.get()
+
+
+def set_sql_plan_name(name: str, plan_id: int) -> None:
+    _sql_plan_name_map.get()[name] = plan_id
+
+
+def get_sql_plan(name: str) -> int:
+    return _sql_plan_name_map.get().get(name)
+
+
+def add_sql_aggregate_function() -> None:
+    _sql_aggregate_function_count.set(_sql_aggregate_function_count.get() + 1)
+
+
+def get_sql_aggregate_function_count() -> int:
+    return _sql_aggregate_function_count.get()
+
+
+def set_sql_args(
+    named: dict[str, expressions_proto.Expression],
+    pos: dict[int, expressions_proto.Expression],
+) -> None:
+    _sql_named_args.set(named)
+    _sql_pos_args.set(pos)
+
+
+def get_sql_named_arg(name: str) -> expressions_proto.Expression:
+    return _sql_named_args.get()[name]
+
+
+def get_sql_pos_arg(pos: int) -> expressions_proto.Expression:
+    return _sql_pos_args.get()[pos]
+
+
+def set_df_before_projection(df: snowpark.DataFrame | None) -> None:
+    """
+    Sets the current DataFrame in the context.
+    This is used to track the DataFrame in the current context.
+    """
+    _df_before_projection.set(df)
+
+
+def get_df_before_projection() -> snowpark.DataFrame | None:
+    """
+    Returns the current DataFrame if set, otherwise None.
+    This is used to track the DataFrame in the current context.
+    """
+    return _df_before_projection.get()
+
+
+@contextmanager
+def push_outer_dataframe(df: snowpark.DataFrame):
+    _outer_dataframes.get().append(df)
+    yield
+    _outer_dataframes.get().pop()
+
+
+def get_outer_dataframes() -> list[snowpark.DataFrame]:
+    return _outer_dataframes.get()
+
+
+def clear_context_data() -> None:
+    _session_id.set(None)
+    _plan_id_map.set({})
+    _alias_map.set({})
+
+    _next_sql_plan_id.set(_STARTING_SQL_PLAN_ID)
+    _sql_plan_name_map.set({})
+    _sql_aggregate_function_count.set(0)
+    _sql_named_args.set({})
+    _sql_pos_args.set({})
+    _is_aggregate_function.set(("default", False))
+    _df_before_projection.set(None)
+    _outer_dataframes.set([])
+
+
+@contextmanager
+def temporary_window_expression():
+    """
+    Temporarily sets the 'is_window_enabled' attribute of the given typer to True.
+
+    Yields:
+        None: The context manager only yields control; it does not return a value.
+    -------
+
+    """
+    _is_window_enabled.set(True)
+    try:
+        yield
+    finally:
+        _is_window_enabled.set(False)
+
+
+def is_window_enabled():
+    return _is_window_enabled.get()
+
+
+@contextmanager
+def temporary_pivot_expression(value: bool):
+    token = _is_in_pivot.set(value)
+    try:
+        yield
+    finally:
+        _is_in_pivot.reset(token)
+
+
+def is_in_pivot() -> bool:
+    return _is_in_pivot.get()
+
+
+def get_is_in_udtf_context() -> bool:
+    """
+    Gets the value of _is_in_udtf_context for the current context, defaults to False.
+    """
+    return _is_in_udtf_context.get()
+
+
+@contextmanager
+def push_udtf_context():
+    """
+    Context manager that sets a flag indicating UDTF arguments are being processed.
+    """
+    token = _is_in_udtf_context.set(True)
+    try:
+        yield
+    finally:
+        _is_in_udtf_context.reset(token)

snowflake/snowpark_connect/utils/describe_query_cache.py

@@ -0,0 +1,231 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import hashlib
+import inspect
+import random
+import re
+import threading
+import time
+from typing import Any
+
+from snowflake import snowpark
+from snowflake.connector.cursor import ResultMetadataV2
+from snowflake.snowpark._internal.server_connection import ServerConnection
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
+from snowflake.snowpark_connect.utils.telemetry import telemetry
+
+DESCRIBE_CACHE_TTL_SECONDS = 15
+USE_DESCRIBE_QUERY_CACHE = True
+
+DDL_DETECTION_PATTERN = re.compile(r"^\s*(CREATE|ALTER|DROP|RENAME)\b", re.IGNORECASE)
+# Pattern for simple constant queries like: SELECT 3 :: INT AS "3-80000030-0" FROM ( SELECT $1 AS "__DUMMY" FROM  VALUES (NULL :: STRING))
+# Using exact spacing pattern from generated SQL for deterministic matching
+# Column ID format: {original_name}-{8_digit_hex_plan_id}-{column_index}
+# Examples: "i-8000002a-0", "1-8000002c-0"
+SIMPLE_CONSTANT_PATTERN = re.compile(
+    r'^\s*SELECT (\d+) :: INT AS "([^"]+)" FROM \( SELECT \$1 AS "__DUMMY" FROM  VALUES \(NULL :: STRING\)\)\s*$',
+    re.IGNORECASE,
+)
+
+
+class DescribeQueryCache:
+    def __init__(self) -> None:
+        self._cache = {}
+        self._lock = threading.Lock()
+
+    @staticmethod
+    def _hash_query(sql_query: str) -> str:
+        return hashlib.md5(sql_query.encode("utf-8")).hexdigest()
+
+    @staticmethod
+    def _get_cache_key(sql_query: str) -> str:
+        """Get cache key for a query, normalizing simple constants."""
+        if SIMPLE_CONSTANT_PATTERN.match(sql_query):
+            # Normalize simple constants to share cache entries
+            return 'SELECT PLACEHOLDER :: INT AS "PLACEHOLDER" FROM ( SELECT $1 AS "__DUMMY" FROM  VALUES (NULL :: STRING))'
+        return sql_query
+
+    def get(self, sql_query: str) -> list[ResultMetadataV2] | None:
+        cache_key = self._get_cache_key(sql_query)
+        key = self._hash_query(cache_key)
+        current_time = time.monotonic()
+
+        # TODO: maybe too much locking, we could use read-write lock also. Or a thread safe dictionary.
+        with self._lock:
+            if key in self._cache:
+                result, timestamp = self._cache[key]
+                if current_time < timestamp + DESCRIBE_CACHE_TTL_SECONDS:
+                    logger.debug(
+                        f"Returning query result from cache for query: {sql_query[:20]}"
+                    )
+
+                    # If this is a constant query, we need to transform the result metadata
+                    # to match the actual query's column name
+                    if (
+                        cache_key != sql_query
+                    ):  # Only transform if we normalized the key
+                        match = SIMPLE_CONSTANT_PATTERN.match(sql_query)
+                        if match:
+                            number, column_id = match.groups()
+                            expected_column_name = column_id
+
+                            # Transform the cached result to match this query's column name
+                            # There should only be one column in these constant queries
+                            metadata = result[0]
+                            new_metadata = ResultMetadataV2(
+                                name=expected_column_name,
+                                type_code=metadata.type_code,
+                                display_size=metadata.display_size,
+                                internal_size=metadata.internal_size,
+                                precision=metadata.precision,
+                                scale=metadata.scale,
+                                is_nullable=metadata.is_nullable,
+                            )
+                            return [new_metadata]
+
+                    return result
+                else:
+                    logger.debug(
+                        f"Had a cached entry, but it expired for query: {sql_query[:20]}"
+                    )
+                    del self._cache[key]
+        return None
+
+    def put(self, sql_query: str, result: list[ResultMetadataV2] | None) -> None:
+        if result is None:
+            return
+
+        cache_key = self._get_cache_key(sql_query)
+        key = self._hash_query(cache_key)
+
+        logger.debug(f"Putting query into cache: {sql_query[:50]}...")
+
+        with self._lock:
+            self._cache[key] = (result, time.monotonic())
+
+    def clear(self) -> None:
+        with self._lock:
+            self._cache.clear()
+
+
+def instrument_session_for_describe_cache(session: snowpark.Session):
+    if hasattr(session, "_describe_query_cache"):
+        logger.debug(
+            "Session already instrumented for describe query cache, skipping..."
+        )
+        return
+
+    session._describe_query_cache = DescribeQueryCache()
+    session._snowpark_api_describe_calls = []
+
+    def update_cache_for_query(query: str):
+        cache = None
+        cache_instance = getattr(session, "_describe_query_cache", None)
+        if isinstance(cache_instance, DescribeQueryCache):
+            cache = cache_instance
+
+        # TODO: This is very broad right now. We should be able to reduce the scope of clearing.
+        if DDL_DETECTION_PATTERN.search(query):
+            logger.debug(f"DDL detected, clearing describe query cache: '{query}'")
+            cache.clear()
+
+    def report_query(qid: str, is_internal: bool) -> None:
+        if is_internal:
+            telemetry.report_internal_query()
+        elif qid:
+            telemetry.report_query_id(qid)
+
+    def wrap_execute(wrapped_fn):
+        def fn(query: str, **kwargs):
+            update_cache_for_query(query)
+            is_internal = kwargs.get("_is_internal", False)
+            try:
+                result = wrapped_fn(query, **kwargs)
+                report_query(result.sfqid, is_internal)
+            except Exception as e:
+                report_query(e.sfqid, is_internal)
+                raise e
+            return result
+
+        return fn
+
+    # This is a wrapper to cache describe queries.
+    def cached_describe_internal(session, wrapped_fn):
+        def fn(*args: Any, **kwargs: Any) -> list[ResultMetadataV2] | None:
+            def _should_skip_query():
+                # We skip queries that are not relevant, but can introduce non-determinism since many tests share the same session
+                f = inspect.currentframe().f_back
+                while f is not None:
+                    if "resources_initializer" in f.f_code.co_filename:
+                        # async resource initialization
+                        return True
+                    f = f.f_back
+                return False
+
+            # In tests, we want to record describe queries before reaching to cache (`_test_name` is only present in expectation tests).
+            if hasattr(session, "_test_name") and not _should_skip_query():
+                session._snowpark_api_describe_calls.append(
+                    {
+                        "test_name": session._test_name,
+                        "query_hash": DescribeQueryCache._hash_query(args[0])
+                        if isinstance(args[0], str)
+                        else f"unknown{random.randint(0, 2**32)}",
+                    }
+                )
+
+            if not USE_DESCRIBE_QUERY_CACHE:
+                return wrapped_fn(*args, **kwargs)
+
+            # Query should be first arg
+            if args and isinstance(args[0], str):
+                query = args[0]
+            else:
+                logger.debug(
+                    "Warning: Could not extract query string from args in cached_describe_internal. Skipping cache."
+                )
+                return wrapped_fn(*args, **kwargs)  # Call original
+
+            # TODO cache on cursor - does it make sense?
+            cache = None
+            cache_instance = getattr(session, "_describe_query_cache", None)
+            if isinstance(cache_instance, DescribeQueryCache):
+                cache = cache_instance
+
+            if cache:
+                cached_result = cache.get(query)
+                if cached_result is not None:
+                    return cached_result
+
+            # No cache hit, so we call the orig describe function and make a snowflake query
+            result = wrapped_fn(*args, **kwargs)
+
+            if cache and result is not None:
+                cache.put(query, result)
+
+            return result
+
+        return fn
+
+    # TODO: is there a better way to intercept queries?
+    orig_cursor_getter = ServerConnection._cursor.fget
+
+    def cursor_wrapper(conn):
+        cursor = orig_cursor_getter(conn)
+        if hasattr(cursor, "_instrumented_describe") and cursor._instrumented_describe:
+            return cursor
+
+        cursor.execute = wrap_execute(cursor.execute)
+
+        # Override describe so that we can stop queries from being sent to Snowflake if they are
+        # already in the cache.
+        # TODO: Test with multiple sessions to confirm we don't cross them
+        cursor._describe_internal = cached_describe_internal(
+            session, cursor._describe_internal
+        )
+
+        cursor._instrumented_describe = True
+        return cursor
+
+    setattr(ServerConnection, "_cursor", property(cursor_wrapper))  # noqa: B010

snowflake/snowpark_connect/utils/interrupt.py

@@ -0,0 +1,85 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import concurrent.futures
+
+from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
+
+
+def interrupt_all_queries() -> list[str]:
+    snowpark_session = get_or_create_snowpark_session()
+
+    sql = _build_sql_for_select_running_queries()
+    running_queries = snowpark_session.sql(sql).collect()
+
+    snowpark_session.cancel_all()
+
+    return [row[0] for row in running_queries]
+
+
+def interrupt_queries_with_tag(tag: str) -> list[str]:
+    snowpark_session = get_or_create_snowpark_session()
+
+    sql = _build_sql_for_select_running_queries(tag=tag)
+    running_queries_with_tag_result = [
+        row[0] for row in snowpark_session.sql(sql).collect()
+    ]
+
+    # Final list of canceled queries might be slightly smaller than running_queries_with_tag_result, because
+    # some jobs can finish in the meantime
+    canceled_query_ids = []
+
+    max_workers = max(1, min(32, len(running_queries_with_tag_result)))
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as exc:
+        futures = []
+
+        for query_id in running_queries_with_tag_result:
+            future = exc.submit(
+                lambda query_id_: (
+                    snowpark_session.sql(
+                        f"SELECT SYSTEM$CANCEL_QUERY('{query_id_}')"
+                    ).collect(),
+                    query_id_,
+                ),
+                query_id_=query_id,
+            )
+            futures.append(future)
+
+        for future in futures:
+            cancel_query_result, query_id = future.result()
+
+            if _is_cancel_query_successful(cancel_query_result[0][0]):
+                canceled_query_ids.append(query_id)
+
+    return canceled_query_ids
+
+
+def interrupt_query(query_id: str) -> list[str]:
+    snowpark_session = get_or_create_snowpark_session()
+
+    cancel_result = snowpark_session.sql(
+        f"SELECT SYSTEM$CANCEL_QUERY('{query_id}')"
+    ).collect()
+
+    return [query_id] if _is_cancel_query_successful(cancel_result[0][0]) else []
+
+
+def _build_sql_for_select_running_queries(tag: str | None = None) -> str:
+    sql = "select query_id"
+    sql += " from table(information_schema.query_history_by_session(include_client_generated_statement => true, result_limit => 10000))"
+    sql += " where execution_status = 'RUNNING'"
+
+    if tag:
+        sql += f" and array_contains('{tag}', split(query_tag, ',')::array(string))"
+
+    sql_escaped_as_str = sql.replace("'", "\\'")
+
+    # Filter out the currently running query_history_by_session query from the result
+    sql += f" and query_text not like '{sql_escaped_as_str}%'"
+
+    return sql
+
+
+def _is_cancel_query_successful(response_message: str) -> bool:
+    return response_message.endswith("terminated.")