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

snowflake/snowpark_connect/proto/snowflake_relation_ext_pb2.pyi

@@ -8,104 +8,110 @@ from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Map
 
 DESCRIPTOR: _descriptor.FileDescriptor
 
-class Aggregate(_message.Message):
-    __slots__ = ["aggregate_expressions", "group_type", "grouping_expressions", "grouping_sets", "input", "pivot"]
-    class GroupType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
-        __slots__ = []
-    class GroupingSets(_message.Message):
-        __slots__ = ["grouping_set"]
-        GROUPING_SET_FIELD_NUMBER: _ClassVar[int]
-        grouping_set: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
-        def __init__(self, grouping_set: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ...) -> None: ...
-    class Pivot(_message.Message):
-        __slots__ = ["col", "values"]
-        COL_FIELD_NUMBER: _ClassVar[int]
-        VALUES_FIELD_NUMBER: _ClassVar[int]
-        col: _expressions_pb2.Expression
-        values: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression.Literal]
-        def __init__(self, col: _Optional[_Union[_expressions_pb2.Expression, _Mapping]] = ..., values: _Optional[_Iterable[_Union[_expressions_pb2.Expression.Literal, _Mapping]]] = ...) -> None: ...
-    AGGREGATE_EXPRESSIONS_FIELD_NUMBER: _ClassVar[int]
-    GROUPING_EXPRESSIONS_FIELD_NUMBER: _ClassVar[int]
-    GROUPING_SETS_FIELD_NUMBER: _ClassVar[int]
-    GROUP_TYPE_CUBE: Aggregate.GroupType
-    GROUP_TYPE_FIELD_NUMBER: _ClassVar[int]
-    GROUP_TYPE_GROUPBY: Aggregate.GroupType
-    GROUP_TYPE_GROUPING_SETS: Aggregate.GroupType
-    GROUP_TYPE_PIVOT: Aggregate.GroupType
-    GROUP_TYPE_ROLLUP: Aggregate.GroupType
-    GROUP_TYPE_UNSPECIFIED: Aggregate.GroupType
-    INPUT_FIELD_NUMBER: _ClassVar[int]
-    PIVOT_FIELD_NUMBER: _ClassVar[int]
-    aggregate_expressions: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
-    group_type: Aggregate.GroupType
-    grouping_expressions: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
-    grouping_sets: _containers.RepeatedCompositeFieldContainer[Aggregate.GroupingSets]
-    input: _relations_pb2.Relation
-    pivot: Aggregate.Pivot
-    def __init__(self, input: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., group_type: _Optional[_Union[Aggregate.GroupType, str]] = ..., grouping_expressions: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ..., aggregate_expressions: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ..., pivot: _Optional[_Union[Aggregate.Pivot, _Mapping]] = ..., grouping_sets: _Optional[_Iterable[_Union[Aggregate.GroupingSets, _Mapping]]] = ...) -> None: ...
-
 class Extension(_message.Message):
-    __slots__ = ["aggregate", "lateral_join", "rdd_map", "rdd_reduce", "subquery_column_aliases", "udtf_with_table_arguments"]
-    AGGREGATE_FIELD_NUMBER: _ClassVar[int]
-    LATERAL_JOIN_FIELD_NUMBER: _ClassVar[int]
+    __slots__ = ("rdd_map", "rdd_reduce", "subquery_column_aliases", "lateral_join", "udtf_with_table_arguments", "aggregate")
     RDD_MAP_FIELD_NUMBER: _ClassVar[int]
     RDD_REDUCE_FIELD_NUMBER: _ClassVar[int]
     SUBQUERY_COLUMN_ALIASES_FIELD_NUMBER: _ClassVar[int]
+    LATERAL_JOIN_FIELD_NUMBER: _ClassVar[int]
     UDTF_WITH_TABLE_ARGUMENTS_FIELD_NUMBER: _ClassVar[int]
-    aggregate: Aggregate
-    lateral_join: LateralJoin
+    AGGREGATE_FIELD_NUMBER: _ClassVar[int]
     rdd_map: RddMap
     rdd_reduce: RddReduce
     subquery_column_aliases: SubqueryColumnAliases
+    lateral_join: LateralJoin
     udtf_with_table_arguments: UDTFWithTableArguments
+    aggregate: Aggregate
     def __init__(self, rdd_map: _Optional[_Union[RddMap, _Mapping]] = ..., rdd_reduce: _Optional[_Union[RddReduce, _Mapping]] = ..., subquery_column_aliases: _Optional[_Union[SubqueryColumnAliases, _Mapping]] = ..., lateral_join: _Optional[_Union[LateralJoin, _Mapping]] = ..., udtf_with_table_arguments: _Optional[_Union[UDTFWithTableArguments, _Mapping]] = ..., aggregate: _Optional[_Union[Aggregate, _Mapping]] = ...) -> None: ...
 
-class LateralJoin(_message.Message):
-    __slots__ = ["left", "right"]
-    LEFT_FIELD_NUMBER: _ClassVar[int]
-    RIGHT_FIELD_NUMBER: _ClassVar[int]
-    left: _relations_pb2.Relation
-    right: _relations_pb2.Relation
-    def __init__(self, left: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., right: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ...) -> None: ...
-
 class RddMap(_message.Message):
-    __slots__ = ["func", "input"]
-    FUNC_FIELD_NUMBER: _ClassVar[int]
+    __slots__ = ("input", "func")
     INPUT_FIELD_NUMBER: _ClassVar[int]
-    func: bytes
+    FUNC_FIELD_NUMBER: _ClassVar[int]
     input: _relations_pb2.Relation
+    func: bytes
     def __init__(self, input: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., func: _Optional[bytes] = ...) -> None: ...
 
 class RddReduce(_message.Message):
-    __slots__ = ["func", "input"]
-    FUNC_FIELD_NUMBER: _ClassVar[int]
+    __slots__ = ("input", "func")
     INPUT_FIELD_NUMBER: _ClassVar[int]
-    func: bytes
+    FUNC_FIELD_NUMBER: _ClassVar[int]
     input: _relations_pb2.Relation
+    func: bytes
     def __init__(self, input: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., func: _Optional[bytes] = ...) -> None: ...
 
 class SubqueryColumnAliases(_message.Message):
-    __slots__ = ["aliases", "input"]
-    ALIASES_FIELD_NUMBER: _ClassVar[int]
+    __slots__ = ("input", "aliases")
     INPUT_FIELD_NUMBER: _ClassVar[int]
-    aliases: _containers.RepeatedScalarFieldContainer[str]
+    ALIASES_FIELD_NUMBER: _ClassVar[int]
     input: _relations_pb2.Relation
+    aliases: _containers.RepeatedScalarFieldContainer[str]
     def __init__(self, input: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., aliases: _Optional[_Iterable[str]] = ...) -> None: ...
 
-class TableArgumentInfo(_message.Message):
-    __slots__ = ["table_argument", "table_argument_idx"]
-    TABLE_ARGUMENT_FIELD_NUMBER: _ClassVar[int]
-    TABLE_ARGUMENT_IDX_FIELD_NUMBER: _ClassVar[int]
-    table_argument: _relations_pb2.Relation
-    table_argument_idx: int
-    def __init__(self, table_argument: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., table_argument_idx: _Optional[int] = ...) -> None: ...
+class LateralJoin(_message.Message):
+    __slots__ = ("left", "right")
+    LEFT_FIELD_NUMBER: _ClassVar[int]
+    RIGHT_FIELD_NUMBER: _ClassVar[int]
+    left: _relations_pb2.Relation
+    right: _relations_pb2.Relation
+    def __init__(self, left: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., right: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ...) -> None: ...
 
 class UDTFWithTableArguments(_message.Message):
-    __slots__ = ["arguments", "function_name", "table_arguments"]
-    ARGUMENTS_FIELD_NUMBER: _ClassVar[int]
+    __slots__ = ("function_name", "arguments", "table_arguments")
     FUNCTION_NAME_FIELD_NUMBER: _ClassVar[int]
+    ARGUMENTS_FIELD_NUMBER: _ClassVar[int]
     TABLE_ARGUMENTS_FIELD_NUMBER: _ClassVar[int]
-    arguments: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
     function_name: str
+    arguments: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
     table_arguments: _containers.RepeatedCompositeFieldContainer[TableArgumentInfo]
     def __init__(self, function_name: _Optional[str] = ..., arguments: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ..., table_arguments: _Optional[_Iterable[_Union[TableArgumentInfo, _Mapping]]] = ...) -> None: ...
+
+class TableArgumentInfo(_message.Message):
+    __slots__ = ("table_argument", "table_argument_idx")
+    TABLE_ARGUMENT_FIELD_NUMBER: _ClassVar[int]
+    TABLE_ARGUMENT_IDX_FIELD_NUMBER: _ClassVar[int]
+    table_argument: _relations_pb2.Relation
+    table_argument_idx: int
+    def __init__(self, table_argument: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., table_argument_idx: _Optional[int] = ...) -> None: ...
+
+class Aggregate(_message.Message):
+    __slots__ = ("input", "group_type", "grouping_expressions", "aggregate_expressions", "pivot", "grouping_sets")
+    class GroupType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+        __slots__ = ()
+        GROUP_TYPE_UNSPECIFIED: _ClassVar[Aggregate.GroupType]
+        GROUP_TYPE_GROUPBY: _ClassVar[Aggregate.GroupType]
+        GROUP_TYPE_ROLLUP: _ClassVar[Aggregate.GroupType]
+        GROUP_TYPE_CUBE: _ClassVar[Aggregate.GroupType]
+        GROUP_TYPE_PIVOT: _ClassVar[Aggregate.GroupType]
+        GROUP_TYPE_GROUPING_SETS: _ClassVar[Aggregate.GroupType]
+    GROUP_TYPE_UNSPECIFIED: Aggregate.GroupType
+    GROUP_TYPE_GROUPBY: Aggregate.GroupType
+    GROUP_TYPE_ROLLUP: Aggregate.GroupType
+    GROUP_TYPE_CUBE: Aggregate.GroupType
+    GROUP_TYPE_PIVOT: Aggregate.GroupType
+    GROUP_TYPE_GROUPING_SETS: Aggregate.GroupType
+    class Pivot(_message.Message):
+        __slots__ = ("col", "values")
+        COL_FIELD_NUMBER: _ClassVar[int]
+        VALUES_FIELD_NUMBER: _ClassVar[int]
+        col: _expressions_pb2.Expression
+        values: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression.Literal]
+        def __init__(self, col: _Optional[_Union[_expressions_pb2.Expression, _Mapping]] = ..., values: _Optional[_Iterable[_Union[_expressions_pb2.Expression.Literal, _Mapping]]] = ...) -> None: ...
+    class GroupingSets(_message.Message):
+        __slots__ = ("grouping_set",)
+        GROUPING_SET_FIELD_NUMBER: _ClassVar[int]
+        grouping_set: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
+        def __init__(self, grouping_set: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ...) -> None: ...
+    INPUT_FIELD_NUMBER: _ClassVar[int]
+    GROUP_TYPE_FIELD_NUMBER: _ClassVar[int]
+    GROUPING_EXPRESSIONS_FIELD_NUMBER: _ClassVar[int]
+    AGGREGATE_EXPRESSIONS_FIELD_NUMBER: _ClassVar[int]
+    PIVOT_FIELD_NUMBER: _ClassVar[int]
+    GROUPING_SETS_FIELD_NUMBER: _ClassVar[int]
+    input: _relations_pb2.Relation
+    group_type: Aggregate.GroupType
+    grouping_expressions: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
+    aggregate_expressions: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
+    pivot: Aggregate.Pivot
+    grouping_sets: _containers.RepeatedCompositeFieldContainer[Aggregate.GroupingSets]
+    def __init__(self, input: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., group_type: _Optional[_Union[Aggregate.GroupType, str]] = ..., grouping_expressions: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ..., aggregate_expressions: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ..., pivot: _Optional[_Union[Aggregate.Pivot, _Mapping]] = ..., grouping_sets: _Optional[_Iterable[_Union[Aggregate.GroupingSets, _Mapping]]] = ...) -> None: ...

snowflake/snowpark_connect/proto/snowflake_relation_ext_pb2_grpc.py

@@ -0,0 +1,4 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+

snowflake/snowpark_connect/relation/catalogs/abstract_spark_catalog.py

@@ -9,10 +9,10 @@ import pandas
 import pyspark.sql.connect.proto.common_pb2 as common_proto
 import pyspark.sql.connect.proto.types_pb2 as types_proto
 
-from snowflake import snowpark
 from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.error.exceptions import MissingDatabase, MissingSchema
-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_or_create_snowpark_session
@@ -38,7 +38,7 @@ class AbstractSparkCatalog(ABC):
         schema: types_proto.DataType,
         description: str,
         **options: typing.Any,
-    ) -> snowpark.DataFrame:
+    ) -> DataFrameContainer:
         raise SnowparkConnectNotImplementedError("createTable is not implemented")
 
     @abstractmethod
@@ -110,7 +110,7 @@ class AbstractSparkCatalog(ABC):
     def dropGlobalTempView(
         self,
         spark_view_name: str,
-    ) -> snowpark.DataFrame:
+    ) -> DataFrameContainer:
         raise SnowparkConnectNotImplementedError(
             "dropGlobalTempView is not implemented"
         )
@@ -119,7 +119,7 @@ class AbstractSparkCatalog(ABC):
     def dropTempView(
         self,
         spark_view_name: str,
-    ) -> snowpark.DataFrame:
+    ) -> DataFrameContainer:
         raise SnowparkConnectNotImplementedError("dropTempView is not implemented")
 
     def cacheTable(

snowflake/snowpark_connect/relation/catalogs/snowflake_catalog.py

@@ -10,7 +10,6 @@ import pyspark.sql.connect.proto.common_pb2 as common_proto
 import pyspark.sql.connect.proto.types_pb2 as types_proto
 from snowflake.core.exceptions import NotFoundError
 
-from snowflake import snowpark
 from snowflake.snowpark import functions
 from snowflake.snowpark._internal.analyzer.analyzer_utils import (
     quote_name_without_upper_casing,
@@ -18,8 +17,11 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
 )
 from snowflake.snowpark.functions import lit
 from snowflake.snowpark.types import BooleanType, StringType
-from snowflake.snowpark_connect import column_name_handler
-from snowflake.snowpark_connect.config import auto_uppercase_ddl, global_config
+from snowflake.snowpark_connect.config import (
+    auto_uppercase_non_column_identifiers,
+    global_config,
+)
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.catalogs.abstract_spark_catalog import (
     AbstractSparkCatalog,
     _get_current_snowflake_schema,
@@ -27,7 +29,7 @@ from snowflake.snowpark_connect.relation.catalogs.abstract_spark_catalog import
     _process_multi_layer_identifier,
 )
 from snowflake.snowpark_connect.type_mapping import proto_to_snowpark_type
-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_or_create_snowpark_session
@@ -40,7 +42,7 @@ from snowflake.snowpark_connect.utils.udf_cache import cached_udf
 def _normalize_identifier(identifier: str | None) -> str | None:
     if identifier is None:
         return None
-    return identifier.upper() if auto_uppercase_ddl() else identifier
+    return identifier.upper() if auto_uppercase_non_column_identifiers() else identifier
 
 
 def sf_quote(name: str | None) -> str | None:
@@ -345,7 +347,7 @@ class SnowflakeCatalog(AbstractSparkCatalog):
     def dropGlobalTempView(
         self,
         spark_view_name: str,
-    ) -> snowpark.DataFrame:
+    ) -> DataFrameContainer:
         session = get_or_create_snowpark_session()
         schema = global_config.spark_sql_globalTempDatabase
         result_df = session.sql(
@@ -358,14 +360,17 @@ class SnowflakeCatalog(AbstractSparkCatalog):
             )
         )
         columns = ["value"]
-        return column_name_handler.with_column_map(
-            result_df, columns, columns, [BooleanType()]
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=result_df,
+            spark_column_names=columns,
+            snowpark_column_names=columns,
+            snowpark_column_types=[BooleanType()],
         )
 
     def dropTempView(
         self,
         spark_view_name: str,
-    ) -> snowpark.DataFrame:
+    ) -> DataFrameContainer:
         """Drop the current temporary view."""
         session = get_or_create_snowpark_session()
         result = session.sql(
@@ -377,8 +382,11 @@ class SnowflakeCatalog(AbstractSparkCatalog):
         )
         result_df = session.createDataFrame([(view_was_dropped,)], schema=["value"])
         columns = ["value"]
-        return column_name_handler.with_column_map(
-            result_df, columns, columns, [BooleanType()]
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=result_df,
+            spark_column_names=columns,
+            snowpark_column_names=columns,
+            snowpark_column_types=[BooleanType()],
         )
 
     def createTable(
@@ -389,7 +397,7 @@ class SnowflakeCatalog(AbstractSparkCatalog):
         schema: types_proto.DataType,
         description: str,
         **options: typing.Any,
-    ) -> snowpark.DataFrame:
+    ) -> DataFrameContainer:
         """Create either an external, or a managed table.
 
         If path is supplied in which the data for this table exists. When path is specified, an external table is
@@ -422,11 +430,11 @@ class SnowflakeCatalog(AbstractSparkCatalog):
         table_name_parts = split_fully_qualified_spark_name(tableName)
         qualifiers = [table_name_parts for _ in columns]
         column_types = [f.datatype for f in sp_schema.fields]
-        return column_name_handler.with_column_map(
-            session.createDataFrame([], sp_schema),
-            columns,
-            columns,
-            column_types,
+        return DataFrameContainer.create_with_column_mapping(
+            dataframe=session.createDataFrame([], sp_schema),
+            spark_column_names=columns,
+            snowpark_column_names=columns,
+            snowpark_column_types=column_types,
             column_qualifiers=qualifiers,
         )