snowpark-connect 0.23.0__py3-none-any.whl → 0.25.0__py3-none-any.whl

snowflake/snowpark_connect/type_mapping.py

@@ -19,8 +19,9 @@ from pyspark.sql.connect.proto import expressions_pb2
 
 from snowflake import snowpark
 from snowflake.snowpark._internal.utils import quote_name
-from snowflake.snowpark.types import TimestampTimeZone
+from snowflake.snowpark.types import TimestampTimeZone, TimestampType
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+from snowflake.snowpark_connect.config import get_timestamp_type
 from snowflake.snowpark_connect.constants import (
     COLUMN_METADATA_COLLISION_KEY,
     STRUCTURED_TYPES_ENABLED,
@@ -29,6 +30,7 @@ from snowflake.snowpark_connect.date_time_format_mapping import (
     convert_spark_format_to_snowflake,
 )
 from snowflake.snowpark_connect.expression.literal import get_literal_field_and_name
+from snowflake.snowpark_connect.utils.context import get_is_evaluating_sql
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
@@ -63,7 +65,7 @@ def get_python_sql_utils_class():
     return jpype.JClass("org.apache.spark.sql.api.python.PythonSQLUtils")
 
 
-def parse_ddl_with_spark_scala(ddl_string: str) -> pyspark.sql.types.DataType:
+def _parse_ddl_with_spark_scala(ddl_string: str) -> pyspark.sql.types.DataType:
     """
     Parse DDL string using PySpark's Scala StructType.fromDDL() method.
 
@@ -869,51 +871,53 @@ def map_json_schema_to_snowpark(
             return map_simple_types(schema["type"])
 
 
-def map_type_string_to_snowpark_type(type_string: str) -> snowpark.types.DataType:
+def _map_type_string(
+    type_string: str,
+) -> tuple[pyspark.sql.types.DataType, snowpark.types.DataType]:
     """
-    Converts a pyspark type string like x: int or struct<x: int, y: string> etc. to a snowpark type.
+    Map a ddl-string type to a tuple of PySpark data type and a Snowpark data type.
     """
-    pyspark_type = parse_ddl_with_spark_scala(type_string)
-    return map_pyspark_types_to_snowpark_types(pyspark_type)
+
+    pyspark_type = _parse_ddl_with_spark_scala(type_string)
+    match type_string:
+        case "timestamp" if not get_is_evaluating_sql():
+            return pyspark_type, get_timestamp_type()
+        case "timestamp" | "timestamp_ltz":
+            return pyspark_type, TimestampType(TimestampTimeZone.LTZ)
+        case _:
+            return pyspark_type, map_pyspark_types_to_snowpark_types(pyspark_type)
 
 
-def preprocess_tinyint_to_byte(ddl_string):
-    if ddl_string.strip().lower() == "tinyint":
-        # If the ddl_string is just "tinyint", we can directly return "byte"
-        return "byte"
-    # Replace any occurrence of ": tinyint" (with optional spaces) with ": byte"
-    return re.sub(r"(:\s*)tinyint\b", r"\1byte", ddl_string, flags=re.IGNORECASE)
+def map_type_string_to_snowpark_type(type_string: str) -> snowpark.types.DataType:
+    _, snowpark_type = _map_type_string(type_string)
+    return snowpark_type
 
 
-def parse_ddl_string(
-    ddl_string: str, return_raw_type: bool = False
+def map_type_string_to_proto(
+    ddl_string: str,
 ) -> types_proto.DataType | snowpark.types.DataType:
     """
     Parse a DDL string and return the Proto data type.
-
-    This function is used to parse a DDL string and return the Proto data type.
     """
-    pyspark_type = parse_ddl_with_spark_scala(ddl_string)
-    schema = map_pyspark_types_to_snowpark_types(pyspark_type)
-
-    if return_raw_type:
-        return schema
+    pyspark_type, snowpark_type = _map_type_string(ddl_string)
 
     column_name_map = None
     if isinstance(pyspark_type, pyspark.sql.types.StructType):
         spark_column_names = [field.name for field in pyspark_type.fields]
-        snowpark_column_names = [field.name for field in schema.fields]
+        snowpark_column_names = [field.name for field in snowpark_type.fields]
 
         if len(spark_column_names) == len(snowpark_column_names):
             column_name_map = ColumnNameMap(spark_column_names, snowpark_column_names)
 
-    struct_proto = list(snowpark_to_proto_type(schema, column_name_map).values())[0]
+    struct_proto = list(
+        snowpark_to_proto_type(snowpark_type, column_name_map).values()
+    )[0]
 
-    if isinstance(schema, snowpark.types.StructType):
+    if isinstance(snowpark_type, snowpark.types.StructType):
         return types_proto.DataType(struct=struct_proto)
     else:
         # Handle non-struct types (arrays, maps, primitives, etc.)
-        proto_type_mapping = snowpark_to_proto_type(schema, column_name_map)
+        proto_type_mapping = snowpark_to_proto_type(snowpark_type, column_name_map)
         proto_field_name = list(proto_type_mapping.keys())[0]
         return types_proto.DataType(**{proto_field_name: struct_proto})
 

snowflake/snowpark_connect/typed_column.py

@@ -21,6 +21,7 @@ class TypedColumn:
         self.col = col
         self._regex_matched_columns: list[str] = list()
         self._type_resolver = type_resolver
+        self._catalog_database_info: dict[str, str] = {}
 
     def __iter__(self):
         return iter((self.col, self._type_resolver))
@@ -49,6 +50,19 @@ class TypedColumn:
     def get_qualifiers(self) -> list[str]:
         return getattr(self, "qualifiers", [])
 
+    def set_catalog_database_info(self, catalog_database_info: dict[str, str]) -> None:
+        self._catalog_database_info = catalog_database_info
+
+    def get_catalog_database_info(self) -> dict[str, str]:
+        """Get catalog and database information for this column."""
+        return self._catalog_database_info
+
+    def get_catalog(self) -> str | None:
+        return self._catalog_database_info.get("catalog")
+
+    def get_database(self) -> str | None:
+        return self._catalog_database_info.get("database")
+
     def set_multi_col_qualifiers(self, qualifiers: list[list[str]]):
         self.multi_col_qualifiers = qualifiers
 

snowflake/snowpark_connect/utils/artifacts.py

@@ -3,9 +3,11 @@
 #
 
 import gzip
+import hashlib
 import os
 import pathlib
 import tempfile
+import zipfile
 import zlib
 
 from snowflake import snowpark
@@ -46,3 +48,24 @@ def write_temporary_artifact(
     with open_file(filepath, write_mode) as in_memory_file:
         in_memory_file.write(data)
     return filepath
+
+
+def write_class_files_to_stage(
+    session: snowpark.Session, files: dict[str, str]
+) -> None:
+    if os.name != "nt":
+        filepath = f"/tmp/sas-{session.session_id}"
+    else:
+        filepath = f"{tempfile.gettempdir()}/sas-{session.session_id}"
+    jar_name = (
+        f'{filepath}/{hashlib.sha256(str(files).encode("utf-8")).hexdigest()[:10]}.jar'
+    )
+    with zipfile.ZipFile(jar_name, "w", zipfile.ZIP_DEFLATED) as jar:
+        for name, path in files.items():
+            jar.write(path, name)
+    session.file.put(
+        jar_name,
+        session.get_session_stage(),
+        auto_compress=False,
+        overwrite=True,
+    )

snowflake/snowpark_connect/utils/context.py

@@ -94,7 +94,12 @@ def resolve_lca_alias(name: str) -> Optional[TypedColumn]:
 
 
 def set_current_grouping_columns(columns: list[str]) -> None:
-    """Set the current grouping columns for grouping_id() function."""
+    """Set the current grouping columns for grouping_id() function.
+
+    TODO: This should use a push/pop mechanism to properly handle nested queries
+    with GROUP BY. Currently, nested queries may interfere with parent queries.
+    Also need to add clearing in clear_context_data() to prevent cross-request contamination.
+    """
     _current_grouping_columns.set(columns)