snowpark-connect 0.27.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

snowflake/snowpark_connect/utils/cache.py

@@ -26,7 +26,6 @@ def df_cache_map_get(key: Tuple[str, any]) -> DataFrameContainer | None:
 def df_cache_map_put_if_absent(
     key: Tuple[str, any],
     compute_fn: Callable[[], DataFrameContainer | pandas.DataFrame],
-    materialize: bool,
 ) -> DataFrameContainer | pandas.DataFrame:
     """
     Put a DataFrame container into the cache map if the key is absent. Optionally, as side effect, materialize
@@ -35,7 +34,6 @@ def df_cache_map_put_if_absent(
     Args:
         key (Tuple[str, int]): The key to insert into the cache map (session_id, plan_id).
         compute_fn (Callable[[], DataFrameContainer | pandas.DataFrame]): A function to compute the DataFrame container if the key is absent.
-        materialize (bool): Whether to materialize the DataFrame.
 
     Returns:
         DataFrameContainer | pandas.DataFrame: The cached or newly computed DataFrame container.
@@ -45,7 +43,7 @@ def df_cache_map_put_if_absent(
         container: DataFrameContainer,
     ) -> DataFrameContainer:
 
-        if materialize:
+        if container.can_be_materialized:
             df = container.dataframe
             cached_result = df.cache_result()
             return DataFrameContainer(
@@ -58,30 +56,54 @@ def df_cache_map_put_if_absent(
         return container
 
     with _cache_map_lock:
-        if key not in df_cache_map:
-            result = compute_fn()
-
-            # check cache again, since recursive call in compute_fn could've already cached the result.
-            # we want return it, instead of saving it again. This is important if materialize = True
-            # because materialization is expensive operation that we don't want to do twice.
-            if key in df_cache_map:
-                return df_cache_map[key]
-
-            # only cache DataFrameContainer, but not pandas result.
-            # Pandas result is only returned when df.show() is called, where we convert
-            # a dataframe to a string representation.
-            # We don't expect map_relation would return pandas df here because that would
-            # be equivalent to calling df.show().cache(), which is not allowed.
-            if isinstance(result, DataFrameContainer):
-                df_cache_map[key] = _object_to_cache(result)
-            else:
-                # This is not expected, but we will just log a warning
-                logger.warning(
-                    "Unexpected pandas dataframe returned for caching. Ignoring the cache call."
-                )
-                return result
-
-        return df_cache_map[key]
+        if key in df_cache_map:
+            return df_cache_map[key]
+
+    # the compute_fn is not guaranteed to be called only once, but it's acceptable based upon the following analysis:
+    # there are in total 5 occurrences of passing compute_fn callback falling into two categories:
+    # 2 occurrences as lambda that needs to be computed:
+    #     1) server::AnalyzePlan case "persist"
+    #     2) server::AddArtifacts case "read"
+    # 3 occurrences as lambda that simply returns pre-computed dataframe without any computation:
+    #     1) map_relation case "local_relation"
+    #     2) map_relation case "sample"
+    #     3) map_read case "data_source"
+    # based upon the analysis of the code, the chance of concurrently calling compute_fn for the same key is very low and if it happens
+    # repeating the computation will not affect the result.
+    # This is a trade-off between implementation simplicity and fine-grained locking.
+    result = compute_fn()
+
+    if isinstance(result, DataFrameContainer) and not result.can_be_cached:
+        return result
+
+    # check cache again, since recursive call in compute_fn could've already cached the result.
+    # we want return it, instead of saving it again. This is important if materialize = True
+    # because materialization is expensive operation that we don't want to do twice.
+    with _cache_map_lock:
+        if key in df_cache_map:
+            return df_cache_map[key]
+
+    # only cache DataFrameContainer, but not pandas result.
+    # Pandas result is only returned when df.show() is called, where we convert
+    # a dataframe to a string representation.
+    # We don't expect map_relation would return pandas df here because that would
+    # be equivalent to calling df.show().cache(), which is not allowed.
+    if isinstance(result, DataFrameContainer):
+        # The _object_to_cache function is not guaranteed to be called only once.
+        # In rare multithreading cases, this may result in duplicate temporary table
+        # creation because df.cache_result() materializes the DataFrame into a temp table each time.
+        # This is acceptable because correctness is not affected, the likelihood is very low, and
+        # it simplifies the implementation by avoiding fine-grained locking.
+        cached_result = _object_to_cache(result)
+        with _cache_map_lock:
+            df_cache_map[key] = cached_result
+            return df_cache_map[key]
+    else:
+        # This is not expected, but we will just log a warning
+        logger.warning(
+            "Unexpected pandas dataframe returned for caching. Ignoring the cache call."
+        )
+        return result
 
 
 def df_cache_map_pop(key: Tuple[str, any]) -> None:

snowflake/snowpark_connect/utils/concurrent.py

@@ -5,10 +5,11 @@
 import threading
 from collections.abc import Mapping
 from copy import copy
-from typing import TypeVar
+from typing import Callable, TypeVar
 
 K = TypeVar("K")
 V = TypeVar("V")
+T = TypeVar("T")
 
 
 class SynchronizedDict(Mapping[K, V]):
@@ -52,6 +53,10 @@ class SynchronizedDict(Mapping[K, V]):
         with self._lock.writer():
             self._dict[key] = value
 
+    def __delitem__(self, key: K) -> None:
+        with self._lock.writer():
+            del self._dict[key]
+
     def __contains__(self, key: K) -> bool:
         with self._lock.reader():
             return key in self._dict
@@ -69,6 +74,36 @@ class SynchronizedDict(Mapping[K, V]):
             self._dict.clear()
 
 
+class SynchronizedList:
+    def __init__(self, in_list: list[T] | None = None) -> None:
+        self._lock = ReadWriteLock()
+        self._list = in_list if in_list is not None else []
+
+    def append(self, item: T) -> None:
+        with self._lock.writer():
+            self._list.append(item)
+
+    def clear(self) -> None:
+        with self._lock.writer():
+            self._list.clear()
+
+    def copy(self) -> list[T]:
+        with self._lock.reader():
+            return self._list.copy()
+
+    def filter(self, predicate: Callable[[T], bool]) -> None:
+        with self._lock.writer():
+            self._list = [item for item in self._list if predicate(item)]
+
+    def __len__(self) -> int:
+        with self._lock.reader():
+            return len(self._list)
+
+    def __iter__(self):
+        with self._lock.reader():
+            return iter(self._list.copy())
+
+
 class ReadWriteLock:
     class _Reader:
         def __init__(self, lock) -> None:

snowflake/snowpark_connect/utils/context.py

@@ -2,10 +2,12 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
+import os
 import re
+import threading
 from contextlib import contextmanager
 from contextvars import ContextVar
-from typing import Mapping, Optional
+from typing import Iterator, Mapping, Optional
 
 import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 
@@ -14,25 +16,29 @@ 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")
+_spark_session_id = ContextVar[str]("_spark_session_id")
 _plan_id_map = ContextVar[Mapping[int, DataFrameContainer]]("_plan_id_map")
 _alias_map = ContextVar[Mapping[str, DataFrameContainer | None]]("_alias_map")
 _spark_version = ContextVar[str]("_spark_version")
 _is_aggregate_function = ContextVar(
     "_is_aggregate_function", default=("default", False)
 )
+_grouping_by_scala_udf_key = ContextVar[bool](
+    "_grouping_by_scala_udf_key", default=False
+)
 _is_evaluating_sql = ContextVar[bool]("_is_evaluating_sql", default=False)
 _is_evaluating_join_condition = ContextVar(
     "_is_evaluating_join_condition", default=("default", False, [], [])
 )
+_is_processing_order_by = ContextVar[bool]("_is_processing_order_by", default=False)
+_is_processing_aliased_relation = ContextVar[bool](
+    "_is_processing_aliased_relation", default=False
+)
 
 _sql_aggregate_function_count = ContextVar[int](
     "_contains_aggregate_function", default=0
 )
 
-# Context for parsing map_partitions
-_map_partitions_stack = ContextVar[int]("_map_partitions_stack", 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
@@ -55,10 +61,26 @@ _resolving_lambda_fun = ContextVar[bool]("_resolving_lambdas", default=False)
 _current_lambda_params = ContextVar[list[str]]("_current_lambda_params", default=[])
 
 _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)
 
+# Thread-safe lock for JPype JClass creation to prevent access violations
+_jpype_jclass_lock = threading.Lock()
+
+
+@contextmanager
+def get_jpype_jclass_lock() -> Iterator[None]:
+    """
+    Context manager that acquires the JPype JClass lock on Windows platforms.
+    On non-Windows (os.name != 'nt'), it yields without acquiring the lock.
+    """
+    if os.name == "nt":
+        with _jpype_jclass_lock:
+            yield
+    else:
+        yield
+
+
 # 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.
@@ -70,12 +92,60 @@ _lca_alias_map: ContextVar[dict[str, TypedColumn]] = ContextVar(
     default={},
 )
 
+_view_process_context = ContextVar("_view_process_context", default=[])
+
+
+@contextmanager
+def push_processed_view(name: str):
+    _view_process_context.set(_view_process_context.get() + [name])
+    yield
+    _view_process_context.set(_view_process_context.get()[:-1])
+
+
+def get_processed_views() -> list[str]:
+    return _view_process_context.get()
+
+
+def register_processed_view(name: str) -> None:
+    context = _view_process_context.get()
+    context.append(name)
+    _view_process_context.set(context)
+
+
+_request_external_tables = ContextVar[list[str]]("_used_external_tables", default=[])
+
+
+def register_request_external_table(name: str) -> None:
+    _request_external_tables.set(_request_external_tables.get() + [name])
+
+
+def get_request_external_tables() -> list[str]:
+    return _request_external_tables.get()
+
+
+def clean_request_external_tables() -> None:
+    _request_external_tables.set([])
+
+
 # Context variable to track current grouping columns for grouping_id() function
 _current_grouping_columns: ContextVar[list[str]] = ContextVar(
     "_current_grouping_columns",
     default=[],
 )
 
+# Context variable to capture all original_attr_name values during subquery resolution
+# This is a stack of lists to handle nested subqueries correctly
+_captured_attribute_names: ContextVar[list[list[str]]] = ContextVar(
+    "_captured_attribute_names",
+    default=[],
+)
+
+# Context variable to track if we're resolving a subquery expression
+_is_resolving_subquery_exp: ContextVar[bool] = ContextVar(
+    "_is_resolving_subquery_exp",
+    default=False,
+)
+
 
 def clear_lca_alias_map() -> None:
     _lca_alias_map.set({})
@@ -112,14 +182,56 @@ def get_current_grouping_columns() -> list[str]:
     return _current_grouping_columns.get()
 
 
-def set_session_id(value: str) -> None:
-    """Set the session ID for the current context"""
-    _session_id.set(value)
+def capture_attribute_name(attr_name: str) -> None:
+    """Capture an original_attr_name during expression resolution."""
+    stack = _captured_attribute_names.get()
+    if stack:
+        stack[-1].append(attr_name)
+        _captured_attribute_names.set(stack)
+
+
+def get_captured_attribute_names() -> list[str]:
+    """Get the list of captured attribute names from the current top of the stack."""
+    stack = _captured_attribute_names.get()
+    return stack[-1] if stack else []
+
+
+def is_resolving_subquery_exp() -> bool:
+    """
+    Returns True if currently resolving a subquery expression.
+    """
+    return _is_resolving_subquery_exp.get()
+
+
+@contextmanager
+def resolving_subquery_exp():
+    """
+    Context manager that captures all original_attr_name values during subquery expression resolution.
+    Sets a flag to indicate we're in a subquery context and pushes a new list onto the stack.
+    When the context exits, pops the list from the stack.
+    """
+    stack = _captured_attribute_names.get()
+    stack.append([])
+    _captured_attribute_names.set(stack)
+    token = _is_resolving_subquery_exp.set(True)
+    try:
+        yield
+    finally:
+        stack = _captured_attribute_names.get()
+        if stack:
+            stack.pop()
+            _captured_attribute_names.set(stack)
+        _is_resolving_subquery_exp.reset(token)
+
+
+def set_spark_session_id(value: str) -> None:
+    """Set the Spark session ID for the current context"""
+    _spark_session_id.set(value)
 
 
-def get_session_id() -> str:
-    """Get the session ID for the current context."""
-    return _session_id.get(None)
+def get_spark_session_id() -> str:
+    """Get the Spark session ID for the current context."""
+    return _spark_session_id.get(None)
 
 
 def set_plan_id_map(plan_id: int, container: DataFrameContainer) -> None:
@@ -190,6 +302,66 @@ def push_evaluating_sql_scope():
         _is_evaluating_sql.set(prev)
 
 
+def get_grouping_by_scala_udf_key() -> bool:
+    """
+    Gets the value of _grouping_by_scala_udf_key for the current context, defaults to False.
+    """
+    return _grouping_by_scala_udf_key.get()
+
+
+@contextmanager
+def grouping_by_scala_udf_key(value: bool):
+    """
+    Context manager that conditionally sets a flag indicating grouping by scala_udf key.
+    Only activates the flag when value=True, otherwise leaves the current context unchanged
+    """
+    prev = _grouping_by_scala_udf_key.get()
+    try:
+        if value:
+            _grouping_by_scala_udf_key.set(True)
+        yield
+    finally:
+        _grouping_by_scala_udf_key.set(prev)
+
+
+def get_is_processing_order_by() -> bool:
+    """
+    Gets the value of _is_processing_order_by for the current context, defaults to False.
+    """
+    return _is_processing_order_by.get()
+
+
+@contextmanager
+def push_processing_order_by_scope():
+    """
+    Context manager that sets a flag indicating if ORDER BY expressions are being evaluated.
+    This enables optimizations like reusing already-computed UDF columns.
+    """
+    prev = _is_processing_order_by.get()
+    try:
+        _is_processing_order_by.set(True)
+        yield
+    finally:
+        _is_processing_order_by.set(prev)
+
+
+def get_is_processing_aliased_relation() -> bool:
+    return _is_processing_aliased_relation.get()
+
+
+@contextmanager
+def push_processing_aliased_relation_scope(process_aliased_relation: bool):
+    """
+    Context manager that sets a flag indicating if an aliased relation is being resolved.
+    """
+    prev = _is_processing_aliased_relation.get()
+    try:
+        _is_processing_aliased_relation.set(process_aliased_relation)
+        yield
+    finally:
+        _is_processing_aliased_relation.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.
@@ -210,16 +382,6 @@ def push_evaluating_join_condition(join_type, left_keys, right_keys):
         _is_evaluating_join_condition.set(prev)
 
 
-@contextmanager
-def push_map_partitions():
-    _map_partitions_stack.set(_map_partitions_stack.get() + 1)
-    yield
-
-
-def map_partitions_depth() -> int:
-    return _map_partitions_stack.get()
-
-
 @contextmanager
 def push_sql_scope():
     """
@@ -383,13 +545,14 @@ def get_outer_dataframes() -> list[DataFrameContainer]:
 
 
 def clear_context_data() -> None:
-    _session_id.set(None)
+    _spark_session_id.set(None)
     _plan_id_map.set({})
     _alias_map.set({})
 
+    _request_external_tables.set([])
+    _view_process_context.set([])
     _next_sql_plan_id.set(_STARTING_SQL_PLAN_ID)
     _sql_plan_name_map.set({})
-    _map_partitions_stack.set(0)
     _sql_aggregate_function_count.set(0)
     _sql_named_args.set({})
     _sql_pos_args.set({})
@@ -419,19 +582,6 @@ 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.

snowflake/snowpark_connect/utils/describe_query_cache.py

@@ -6,20 +6,28 @@ 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.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
+from snowflake.snowpark_connect.utils.concurrent import (
+    SynchronizedDict,
+    SynchronizedList,
+)
 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)
+DDL_DETECTION_PATTERN = re.compile(r"\s*(CREATE|ALTER|DROP)\b", re.IGNORECASE)
+PLAIN_CREATE_PATTERN = re.compile(
+    r"\s*CREATE\s+((LOCAL|GLOBAL)\s+)?(TRANSIENT\s+)?TABLE\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}
@@ -32,8 +40,7 @@ SIMPLE_CONSTANT_PATTERN = re.compile(
 
 class DescribeQueryCache:
     def __init__(self) -> None:
-        self._cache = {}
-        self._lock = threading.Lock()
+        self._cache = SynchronizedDict()
 
     @staticmethod
     def _hash_query(sql_query: str) -> str:
@@ -48,49 +55,53 @@ class DescribeQueryCache:
         return sql_query
 
     def get(self, sql_query: str) -> list[ResultMetadataV2] | None:
+        from snowflake.snowpark_connect.config import get_describe_cache_ttl_seconds
+
+        telemetry.report_describe_query_cache_lookup()
+
         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]
+        if key in self._cache:
+            result, timestamp = self._cache[key]
+
+            expired_by = current_time - (timestamp + get_describe_cache_ttl_seconds())
+            if expired_by < 0:
+                logger.debug(
+                    f"Returning query result from cache for query: {sql_query[:20]}"
+                )
+                self._cache[key] = (result, current_time)
+
+                # 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,
+                        )
+
+                        telemetry.report_describe_query_cache_hit()
+                        return [new_metadata]
+
+                telemetry.report_describe_query_cache_hit()
+                return result
+            else:
+                telemetry.report_describe_query_cache_expired(expired_by)
+                del self._cache[key]
         return None
 
     def put(self, sql_query: str, result: list[ResultMetadataV2] | None) -> None:
@@ -102,12 +113,18 @@ class DescribeQueryCache:
 
         logger.debug(f"Putting query into cache: {sql_query[:50]}...")
 
-        with self._lock:
-            self._cache[key] = (result, time.monotonic())
+        self._cache[key] = (result, time.monotonic())
 
     def clear(self) -> None:
-        with self._lock:
-            self._cache.clear()
+        self._cache.clear()
+
+    def update_cache_for_query(self, query: str) -> None:
+        # Clear cache for DDL operations that modify existing objects (exclude CREATE TABLE)
+        if DDL_DETECTION_PATTERN.search(query) and not PLAIN_CREATE_PATTERN.search(
+            query
+        ):
+            self.clear()
+            telemetry.report_describe_query_cache_clear()
 
 
 def instrument_session_for_describe_cache(session: snowpark.Session):
@@ -118,7 +135,7 @@ def instrument_session_for_describe_cache(session: snowpark.Session):
         return
 
     session._describe_query_cache = DescribeQueryCache()
-    session._snowpark_api_describe_calls = []
+    session._snowpark_api_describe_calls = SynchronizedList()
 
     def update_cache_for_query(query: str):
         cache = None
@@ -126,10 +143,7 @@ def instrument_session_for_describe_cache(session: snowpark.Session):
         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()
+        cache.update_cache_for_query(query)
 
     def wrap_execute(wrapped_fn):
         def fn(query: str, **kwargs):
@@ -139,6 +153,7 @@ def instrument_session_for_describe_cache(session: snowpark.Session):
                 telemetry.report_query(result, **kwargs)
             except Exception as e:
                 telemetry.report_query(e, **kwargs)
+                attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
                 raise e
             return result