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

snowflake/snowpark_connect/relation/read/map_read_table.py

@@ -11,12 +11,13 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
     unquote_if_quoted,
 )
 from snowflake.snowpark.exceptions import SnowparkSQLException
-from snowflake.snowpark_connect.column_name_handler import with_column_map
-from snowflake.snowpark_connect.config import auto_uppercase_dml
+from snowflake.snowpark_connect.column_name_handler import ALREADY_QUOTED
+from snowflake.snowpark_connect.config import auto_uppercase_non_column_identifiers
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read.utils import (
     rename_columns_as_snowflake_standard,
 )
-from snowflake.snowpark_connect.utils.attribute_handling import (
+from snowflake.snowpark_connect.utils.identifiers import (
     split_fully_qualified_spark_name,
 )
 from snowflake.snowpark_connect.utils.session import _get_current_snowpark_session
@@ -27,9 +28,16 @@ from snowflake.snowpark_connect.utils.telemetry import (
 
 def post_process_df(
     df: snowpark.DataFrame, plan_id: int, source_table_name: str = None
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
+    def _lower_or_unquote(string):
+        return (
+            string[1:-1].replace('""', '"')
+            if ALREADY_QUOTED.match(string)
+            else string.lower()
+        )
+
     try:
-        true_names = list(map(lambda x: unquote_if_quoted(x).lower(), df.columns))
+        true_names = list(map(lambda x: _lower_or_unquote(x), df.columns))
         renamed_df, snowpark_column_names = rename_columns_as_snowflake_standard(
             df, plan_id
         )
@@ -44,11 +52,11 @@ def post_process_df(
             if current_schema:
                 name_parts = [unquote_if_quoted(current_schema)] + name_parts
 
-        return with_column_map(
-            renamed_df,
-            true_names,
-            snowpark_column_names,
-            [f.datatype for f in df.schema.fields],
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=renamed_df,
+            spark_column_names=true_names,
+            snowpark_column_names=snowpark_column_names,
+            snowpark_column_types=[f.datatype for f in df.schema.fields],
             column_qualifiers=[name_parts] * len(true_names)
             if source_table_name
             else None,
@@ -66,19 +74,18 @@ def post_process_df(
 
 def get_table_from_name(
     table_name: str, session: snowpark.Session, plan_id: int
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
+    """Get table from name returning a container."""
     snowpark_name = ".".join(
         quote_name_without_upper_casing(part)
         for part in split_fully_qualified_spark_name(table_name)
     )
 
-    if auto_uppercase_dml():
+    if auto_uppercase_non_column_identifiers():
         snowpark_name = snowpark_name.upper()
 
     df = session.read.table(snowpark_name)
-    post_processed_df = post_process_df(df, plan_id, table_name)
-    post_processed_df._table_name = table_name
-    return post_processed_df
+    return post_process_df(df, plan_id, table_name)
 
 
 def get_table_from_query(
@@ -88,7 +95,9 @@ def get_table_from_query(
     return post_process_df(df, plan_id)
 
 
-def map_read_table(rel: relation_proto.Relation) -> snowpark.DataFrame:
+def map_read_table(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
     """
     Read a table into a Snowpark DataFrame.
     """

snowflake/snowpark_connect/relation/read/map_read_text.py

@@ -7,7 +7,7 @@ import typing
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 from snowflake import snowpark
-from snowflake.snowpark_connect.column_name_handler import with_column_map
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read.utils import (
     get_spark_column_names_from_snowpark_columns,
     rename_columns_as_snowflake_standard,
@@ -71,7 +71,7 @@ def map_read_text(
     schema: snowpark.types.StructType | None,
     session: snowpark.Session,
     paths: list[str],
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Read a TEXT file into a Snowpark DataFrame.
     """
@@ -98,9 +98,9 @@ def map_read_text(
     renamed_df, snowpark_column_names = rename_columns_as_snowflake_standard(
         df, rel.common.plan_id
     )
-    return with_column_map(
-        renamed_df,
-        spark_column_names,
-        snowpark_column_names,
-        [f.datatype for f in df.schema.fields],
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=renamed_df,
+        spark_column_names=spark_column_names,
+        snowpark_column_names=snowpark_column_names,
+        snowpark_column_types=[f.datatype for f in df.schema.fields],
     )

snowflake/snowpark_connect/relation/read/reader_config.py

@@ -346,6 +346,7 @@ class JsonReaderConfig(ReaderWriterConfig):
                     "compression",
                     # "ignoreNullFields",
                     "rowsToInferSchema",
+                    # "inferTimestamp",
                 },
                 boolean_config_list=[
                     "multiLine",

snowflake/snowpark_connect/relation/utils.py

@@ -32,6 +32,7 @@ from snowflake.snowpark_connect.column_name_handler import (
     ColumnNameMap,
     make_column_names_snowpark_compatible,
 )
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 
 TYPE_MAP_FOR_TO_SCHEMA = {
@@ -91,7 +92,9 @@ TYPE_MAP_FOR_TO_SCHEMA = {
 
 
 def get_df_with_partition_row_number(
-    df: snowpark.DataFrame, plan_id: int | None, row_number_column_name: str
+    container: DataFrameContainer,
+    plan_id: int | None,
+    row_number_column_name: str,
 ) -> snowpark.DataFrame:
     """
     Add a row number for each row in each partition for the given df, where
@@ -106,21 +109,24 @@ def get_df_with_partition_row_number(
     |  c|  4|                               |  c|  4|  0         |
     +---+---+                               +---+---+------------+
     """
+    df = container.dataframe
+    column_map = container.column_map
+
     row_number_snowpark_column_name = make_column_names_snowpark_compatible(
-        [row_number_column_name], plan_id, len(df._column_map.get_spark_columns())
+        [row_number_column_name], plan_id, len(column_map.get_spark_columns())
     )[0]
     row_number_snowpark_column = (
         snowpark_fn.row_number()
         .over(
             snowpark.window.Window.partition_by(
-                *df._column_map.get_snowpark_columns()
+                *column_map.get_snowpark_columns()
             ).order_by(snowpark_fn.lit(1))
         )
         .alias(row_number_snowpark_column_name)
     )
 
     df_with_partition_number = df.select(
-        *df._column_map.get_snowpark_columns(), row_number_snowpark_column
+        *column_map.get_snowpark_columns(), row_number_snowpark_column
     )
     return df_with_partition_number
 
@@ -197,7 +203,7 @@ def get_semantic_string(rel: relation_proto.Relation) -> str:
     """
     queries = [
         query
-        for query_list in map_relation(rel)._plan.execution_queries.values()
+        for query_list in map_relation(rel).dataframe._plan.execution_queries.values()
         for query in query_list
     ]
 

snowflake/snowpark_connect/relation/write/jdbc_write_dbapi.py

@@ -10,6 +10,7 @@ import snowflake.snowpark
 from snowflake import snowpark
 from snowflake.snowpark import DataFrameWriter
 from snowflake.snowpark.dataframe import DataFrame
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.read import jdbc_read_dbapi
 from snowflake.snowpark_connect.relation.read.jdbc_read_dbapi import JdbcDialect
 from snowflake.snowpark_connect.relation.read.utils import Connection
@@ -36,7 +37,7 @@ class JdbcDataFrameWriter(DataFrameWriter):
 
     def jdbc_write_dbapi(
         self,
-        input_df: DataFrame,
+        container: DataFrameContainer,
         create_connection: Callable[[dict[str, str]], "Connection"],
         close_connection: Callable[[Connection], None],
         table: str,
@@ -46,6 +47,7 @@ class JdbcDataFrameWriter(DataFrameWriter):
         Write a Snowpark Dataframe data into table of a JDBC datasource.
         """
 
+        input_df = container.dataframe
         conn = create_connection(self.jdbc_options)
         try:
             url = self.jdbc_options.get("url", None)
@@ -53,32 +55,32 @@ class JdbcDataFrameWriter(DataFrameWriter):
 
             table_exist = self._does_table_exist(conn, table)
             insert_query = self._generate_insert_query(
-                input_df,
+                container,
                 table,
             )
 
             match write_mode:
                 case "append":
                     if not table_exist:
-                        self._create_table(conn, table, input_df, jdbc_dialect)
+                        self._create_table(conn, table, container, jdbc_dialect)
                 case "errorifexists":
                     if table_exist:
                         raise ValueError(
                             "table is already exist and write mode is ERROR_IF_EXISTS"
                         )
                     else:
-                        self._create_table(conn, table, input_df, jdbc_dialect)
+                        self._create_table(conn, table, container, jdbc_dialect)
                 case "overwrite":
                     if table_exist:
                         self._drop_table(conn, table)
-                    self._create_table(conn, table, input_df, jdbc_dialect)
+                    self._create_table(conn, table, container, jdbc_dialect)
                 case "ignore":
                     if table_exist:
                         # With Ignore write mode, if table already exists, the save operation is expected
                         # to not save the contents of the DataFrame and to not change the existing data.
                         return
                     else:
-                        self._create_table(conn, table, input_df, jdbc_dialect)
+                        self._create_table(conn, table, container, jdbc_dialect)
                 case _:
                     raise ValueError(f"Invalid write mode value{write_mode}")
 
@@ -92,14 +94,14 @@ class JdbcDataFrameWriter(DataFrameWriter):
         finally:
             close_connection(conn)
 
-    def _generate_insert_query(self, input_df: DataFrame, table: str) -> str:
+    def _generate_insert_query(self, container: DataFrameContainer, table: str) -> str:
         """
         Generates INSERT statement with placeholders.
-        :param input_df: Snowpark dataframe to save
+        :param container: Snowpark dataframe container
         :param table: JDBC datasource table name
         :return: INSERT SQL statement
         """
-        true_names = input_df._column_map.get_spark_columns()
+        true_names = container.column_map.get_spark_columns()
         # quote each column name to match PySpark's case-sensitive column naming behavior.
         quoted_column_names = ",".join([f'"{col}"' for col in true_names])
         place_holders = ",".join(["?"] * len(true_names))
@@ -145,7 +147,7 @@ class JdbcDataFrameWriter(DataFrameWriter):
         self,
         conn: Connection,
         table: str,
-        input_df: DataFrame,
+        container,
         jdbc_dialect: JdbcDialect,
     ) -> None:
         """
@@ -154,14 +156,15 @@ class JdbcDataFrameWriter(DataFrameWriter):
 
         :param conn: A Python DBAPI connection over JDBC connection
         :param table: DBC datasource table name
-        :param input_df: Snowpark dataframe to save
+        :param container: Snowpark dataframe container
         :param jdbc_dialect: JDBC specific dialect
         :return: None
         """
+        input_df = container.dataframe
         columns_str = ""
         fields = input_df.schema.fields
         total_columns = len(fields)
-        column_map = input_df._column_map
+        column_map = container.column_map
 
         column_index = 0
         for field in fields: