snowpark-connect 0.20.2__py3-none-any.whl → 0.22.1__py3-none-any.whl

snowflake/snowpark_connect/relation/map_sample_by.py

@@ -6,7 +6,7 @@ import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 from snowflake import snowpark
-from snowflake.snowpark_connect.column_name_handler import set_schema_getter
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.expression.literal import get_literal_field_and_name
 from snowflake.snowpark_connect.expression.map_expression import (
     map_single_column_expression,
@@ -15,21 +15,28 @@ from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 
 
-def map_sample_by(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_sample_by(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Sample by an expression on the input DataFrame.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.sample_by.input)
+    input_container = map_relation(rel.sample_by.input)
+    input_df = input_container.dataframe
+
     exp: expressions_proto.Expression = rel.sample_by.col
     _, col_expr = map_single_column_expression(
-        exp, input_df._column_map, ExpressionTyper(input_df)
+        exp, input_container.column_map, ExpressionTyper(input_df)
     )
     fractions = {
         get_literal_field_and_name(frac.stratum)[0]: frac.fraction
         for frac in rel.sample_by.fractions
     }
     result: snowpark.DataFrame = input_df.sampleBy(col_expr.col, fractions)
-    result._column_map = input_df._column_map
-    result._table_name = input_df._table_name
-    set_schema_getter(result, lambda: input_df.schema)
-    return result
+    return DataFrameContainer(
+        result,
+        column_map=input_container.column_map,
+        table_name=input_container.table_name,
+        alias=input_container.alias,
+        cached_schema_getter=lambda: input_df.schema,
+    )

snowflake/snowpark_connect/relation/map_show_string.py

@@ -2,11 +2,17 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
+import copy
+
 import pandas
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 from snowflake import snowpark
 from snowflake.snowpark._internal.analyzer import analyzer_utils
+from snowflake.snowpark.functions import col
+from snowflake.snowpark.types import DateType, StringType, StructField, StructType
+from snowflake.snowpark_connect.column_name_handler import set_schema_getter
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 
 
@@ -18,12 +24,15 @@ def map_show_string(rel: relation_proto.Relation) -> pandas.DataFrame:
     message creates a string. The client expects this string to be packed into an Arrow
     Buffer object as a single cell.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.show_string.input)
+    input_df_container: DataFrameContainer = map_relation(rel.show_string.input)
+    raw_input_df = input_df_container.dataframe
+    input_df = _handle_datetype_columns(raw_input_df)
+
     show_string = input_df._show_string_spark(
         num_rows=rel.show_string.num_rows,
         truncate=rel.show_string.truncate,
         vertical=rel.show_string.vertical,
-        _spark_column_names=input_df._column_map.get_spark_columns(),
+        _spark_column_names=input_df_container.column_map.get_spark_columns(),
     )
     return pandas.DataFrame({"show_string": [show_string]})
 
@@ -32,13 +41,15 @@ def map_repr_html(rel: relation_proto.Relation) -> pandas.DataFrame:
     """
     Generate the html string representation of the input dataframe.
     """
-    input_df: snowpark.DataFrame = map_relation(rel.html_string.input)
+    input_df_container: DataFrameContainer = map_relation(rel.html_string.input)
+    input_df = input_df_container.dataframe
+
     input_panda = input_df.toPandas()
     input_panda.rename(
         columns={
             analyzer_utils.unquote_if_quoted(
-                input_df._column_map.get_snowpark_columns()[i]
-            ): input_df._column_map.get_spark_columns()[i]
+                input_df_container.column_map.get_snowpark_columns()[i]
+            ): input_df_container.column_map.get_spark_columns()[i]
             for i in range(len(input_panda.columns))
         },
         inplace=True,
@@ -48,3 +59,29 @@ def map_repr_html(rel: relation_proto.Relation) -> pandas.DataFrame:
         max_rows=rel.html_string.num_rows,
     )
     return pandas.DataFrame({"html_string": [html_string]})
+
+
+def _handle_datetype_columns(input_df: snowpark.DataFrame) -> snowpark.DataFrame:
+    """
+    Maps DateType columns to strings it aims to allow showing the dates which are out of range of datetime.datetime.
+    """
+    new_column_mapping = []
+    new_fields = []
+    transformation_required = False
+    for field in input_df.schema:
+        if isinstance(field.datatype, DateType):
+            transformation_required = True
+            new_column_mapping.append(col(field.name).cast(StringType()))
+            new_fields.append(StructField(field.name, StringType()))
+        else:
+            new_column_mapping.append(col(field.name))
+            new_fields.append(field)
+
+    if not transformation_required:
+        return input_df
+
+    transformed_df = input_df.select(new_column_mapping)
+    set_schema_getter(transformed_df, lambda: StructType(new_fields))
+    transformed_df._column_map = copy.deepcopy(input_df._column_map)
+
+    return transformed_df