snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/worker_util.py

@@ -0,0 +1,46 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Util functions for workers.
+"""
+import sys
+from typing import IO
+
+from pyspark.errors import PySparkRuntimeError
+from pyspark.serializers import (
+    UTF8Deserializer,
+    CPickleSerializer,
+)
+
+pickleSer = CPickleSerializer()
+utf8_deserializer = UTF8Deserializer()
+
+
+def check_python_version(infile: IO) -> None:
+    """
+    Check the Python version between the running process and the one used to serialize the command.
+    """
+    version = utf8_deserializer.loads(infile)
+    if version != "%d.%d" % sys.version_info[:2]:
+        raise PySparkRuntimeError(
+            error_class="PYTHON_VERSION_MISMATCH",
+            message_parameters={
+                "worker_version": str(sys.version_info[:2]),
+                "driver_version": str(version),
+            },
+        )

snowflake/snowpark_connect/proto/__init__.py

@@ -0,0 +1,10 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+# N.B. The generated protobuf code uses a local import, which breaks without this.
+
+import os
+import sys
+
+sys.path.append(os.path.abspath(os.path.dirname(__file__)))

snowflake/snowpark_connect/proto/control_pb2.py

@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: control.proto
+"""Generated protocol buffer code."""
+from google.protobuf.internal import builder as _builder
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\rcontrol.proto\x12\rsnowflake.ses\"*\n\x06\x43onfig\x12\x14\n\x07log_ast\x18\x01 \x01(\x08H\x00\x88\x01\x01\x42\n\n\x08_log_ast\"\x1e\n\x0bPingRequest\x12\x0f\n\x07payload\x18\x01 \x01(\t\"\x1f\n\x0cPingResponse\x12\x0f\n\x07payload\x18\x01 \x01(\t\"+\n\x14GetRequestAstRequest\x12\x13\n\x0b\x66orce_flush\x18\x01 \x01(\x08\"M\n\x15GetRequestAstResponse\x12\x16\n\x0espark_requests\x18\x01 \x03(\x0c\x12\x1c\n\x14snowpark_ast_batches\x18\x02 \x03(\t2\xe8\x01\n\x0e\x43ontrolService\x12\x39\n\tConfigure\x12\x15.snowflake.ses.Config\x1a\x15.snowflake.ses.Config\x12?\n\x04Ping\x12\x1a.snowflake.ses.PingRequest\x1a\x1b.snowflake.ses.PingResponse\x12Z\n\rGetRequestAst\x12#.snowflake.ses.GetRequestAstRequest\x1a$.snowflake.ses.GetRequestAstResponseb\x06proto3')
+
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'control_pb2', globals())
+if _descriptor._USE_C_DESCRIPTORS == False:
+
+  DESCRIPTOR._options = None
+  _CONFIG._serialized_start=32
+  _CONFIG._serialized_end=74
+  _PINGREQUEST._serialized_start=76
+  _PINGREQUEST._serialized_end=106
+  _PINGRESPONSE._serialized_start=108
+  _PINGRESPONSE._serialized_end=139
+  _GETREQUESTASTREQUEST._serialized_start=141
+  _GETREQUESTASTREQUEST._serialized_end=184
+  _GETREQUESTASTRESPONSE._serialized_start=186
+  _GETREQUESTASTRESPONSE._serialized_end=263
+  _CONTROLSERVICE._serialized_start=266
+  _CONTROLSERVICE._serialized_end=498
+# @@protoc_insertion_point(module_scope)

snowflake/snowpark_connect/proto/control_pb2.pyi

@@ -0,0 +1,38 @@
+from google.protobuf.internal import containers as _containers
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Iterable as _Iterable, Optional as _Optional
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class Config(_message.Message):
+    __slots__ = ["log_ast"]
+    LOG_AST_FIELD_NUMBER: _ClassVar[int]
+    log_ast: bool
+    def __init__(self, log_ast: bool = ...) -> None: ...
+
+class GetRequestAstRequest(_message.Message):
+    __slots__ = ["force_flush"]
+    FORCE_FLUSH_FIELD_NUMBER: _ClassVar[int]
+    force_flush: bool
+    def __init__(self, force_flush: bool = ...) -> None: ...
+
+class GetRequestAstResponse(_message.Message):
+    __slots__ = ["snowpark_ast_batches", "spark_requests"]
+    SNOWPARK_AST_BATCHES_FIELD_NUMBER: _ClassVar[int]
+    SPARK_REQUESTS_FIELD_NUMBER: _ClassVar[int]
+    snowpark_ast_batches: _containers.RepeatedScalarFieldContainer[str]
+    spark_requests: _containers.RepeatedScalarFieldContainer[bytes]
+    def __init__(self, spark_requests: _Optional[_Iterable[bytes]] = ..., snowpark_ast_batches: _Optional[_Iterable[str]] = ...) -> None: ...
+
+class PingRequest(_message.Message):
+    __slots__ = ["payload"]
+    PAYLOAD_FIELD_NUMBER: _ClassVar[int]
+    payload: str
+    def __init__(self, payload: _Optional[str] = ...) -> None: ...
+
+class PingResponse(_message.Message):
+    __slots__ = ["payload"]
+    PAYLOAD_FIELD_NUMBER: _ClassVar[int]
+    payload: str
+    def __init__(self, payload: _Optional[str] = ...) -> None: ...

snowflake/snowpark_connect/proto/control_pb2_grpc.py

@@ -0,0 +1,183 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+import warnings
+
+import control_pb2 as control__pb2
+
+GRPC_GENERATED_VERSION = '1.66.1'
+GRPC_VERSION = grpc.__version__
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    raise RuntimeError(
+        f'The grpc package installed is at version {GRPC_VERSION},'
+        + f' but the generated code in control_pb2_grpc.py depends on'
+        + f' grpcio>={GRPC_GENERATED_VERSION}.'
+        + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
+        + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
+    )
+
+
+class ControlServiceStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.Configure = channel.unary_unary(
+                '/snowflake.ses.ControlService/Configure',
+                request_serializer=control__pb2.Config.SerializeToString,
+                response_deserializer=control__pb2.Config.FromString,
+                _registered_method=True)
+        self.Ping = channel.unary_unary(
+                '/snowflake.ses.ControlService/Ping',
+                request_serializer=control__pb2.PingRequest.SerializeToString,
+                response_deserializer=control__pb2.PingResponse.FromString,
+                _registered_method=True)
+        self.GetRequestAst = channel.unary_unary(
+                '/snowflake.ses.ControlService/GetRequestAst',
+                request_serializer=control__pb2.GetRequestAstRequest.SerializeToString,
+                response_deserializer=control__pb2.GetRequestAstResponse.FromString,
+                _registered_method=True)
+
+
+class ControlServiceServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def Configure(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def Ping(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def GetRequestAst(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_ControlServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'Configure': grpc.unary_unary_rpc_method_handler(
+                    servicer.Configure,
+                    request_deserializer=control__pb2.Config.FromString,
+                    response_serializer=control__pb2.Config.SerializeToString,
+            ),
+            'Ping': grpc.unary_unary_rpc_method_handler(
+                    servicer.Ping,
+                    request_deserializer=control__pb2.PingRequest.FromString,
+                    response_serializer=control__pb2.PingResponse.SerializeToString,
+            ),
+            'GetRequestAst': grpc.unary_unary_rpc_method_handler(
+                    servicer.GetRequestAst,
+                    request_deserializer=control__pb2.GetRequestAstRequest.FromString,
+                    response_serializer=control__pb2.GetRequestAstResponse.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'snowflake.ses.ControlService', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+    server.add_registered_method_handlers('snowflake.ses.ControlService', rpc_method_handlers)
+
+
+ # This class is part of an EXPERIMENTAL API.
+class ControlService(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def Configure(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/snowflake.ses.ControlService/Configure',
+            control__pb2.Config.SerializeToString,
+            control__pb2.Config.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def Ping(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/snowflake.ses.ControlService/Ping',
+            control__pb2.PingRequest.SerializeToString,
+            control__pb2.PingResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def GetRequestAst(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/snowflake.ses.ControlService/GetRequestAst',
+            control__pb2.GetRequestAstRequest.SerializeToString,
+            control__pb2.GetRequestAstResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)

snowflake/snowpark_connect/proto/snowflake_expression_ext_pb2.py

@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: snowflake_expression_ext.proto
+"""Generated protocol buffer code."""
+from google.protobuf.internal import builder as _builder
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from pyspark.sql.connect.proto import expressions_pb2 as spark_dot_connect_dot_expressions__pb2
+from pyspark.sql.connect.proto import relations_pb2 as spark_dot_connect_dot_relations__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1esnowflake_expression_ext.proto\x12\rsnowflake.ext\x1a\x1fspark/connect/expressions.proto\x1a\x1dspark/connect/relations.proto\"\x98\x01\n\x0c\x45xpExtension\x12@\n\x0enamed_argument\x18\x01 \x01(\x0b\x32&.snowflake.ext.NamedArgumentExpressionH\x00\x12@\n\x13subquery_expression\x18\x02 \x01(\x0b\x32!.snowflake.ext.SubqueryExpressionH\x00\x42\x04\n\x02op\"P\n\x17NamedArgumentExpression\x12\x0b\n\x03key\x18\x01 \x01(\t\x12(\n\x05value\x18\x02 \x01(\x0b\x32\x19.spark.connect.Expression\"\xf4\x04\n\x12SubqueryExpression\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x45\n\rsubquery_type\x18\x02 \x01(\x0e\x32..snowflake.ext.SubqueryExpression.SubqueryType\x12Q\n\x11table_arg_options\x18\x03 \x01(\x0b\x32\x31.snowflake.ext.SubqueryExpression.TableArgOptionsH\x00\x88\x01\x01\x12\x35\n\x12in_subquery_values\x18\x04 \x03(\x0b\x32\x19.spark.connect.Expression\x1a\xbb\x01\n\x0fTableArgOptions\x12\x31\n\x0epartition_spec\x18\x01 \x03(\x0b\x32\x19.spark.connect.Expression\x12\x37\n\norder_spec\x18\x02 \x03(\x0b\x32#.spark.connect.Expression.SortOrder\x12\"\n\x15with_single_partition\x18\x03 \x01(\x08H\x00\x88\x01\x01\x42\x18\n\x16_with_single_partition\"\x90\x01\n\x0cSubqueryType\x12\x19\n\x15SUBQUERY_TYPE_UNKNOWN\x10\x00\x12\x18\n\x14SUBQUERY_TYPE_SCALAR\x10\x01\x12\x18\n\x14SUBQUERY_TYPE_EXISTS\x10\x02\x12\x1b\n\x17SUBQUERY_TYPE_TABLE_ARG\x10\x03\x12\x14\n\x10SUBQUERY_TYPE_IN\x10\x04\x42\x14\n\x12_table_arg_optionsb\x06proto3')
+
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'snowflake_expression_ext_pb2', globals())
+if _descriptor._USE_C_DESCRIPTORS == False:
+
+  DESCRIPTOR._options = None
+  _EXPEXTENSION._serialized_start=114
+  _EXPEXTENSION._serialized_end=266
+  _NAMEDARGUMENTEXPRESSION._serialized_start=268
+  _NAMEDARGUMENTEXPRESSION._serialized_end=348
+  _SUBQUERYEXPRESSION._serialized_start=351
+  _SUBQUERYEXPRESSION._serialized_end=979
+  _SUBQUERYEXPRESSION_TABLEARGOPTIONS._serialized_start=623
+  _SUBQUERYEXPRESSION_TABLEARGOPTIONS._serialized_end=810
+  _SUBQUERYEXPRESSION_SUBQUERYTYPE._serialized_start=813
+  _SUBQUERYEXPRESSION_SUBQUERYTYPE._serialized_end=957
+# @@protoc_insertion_point(module_scope)

snowflake/snowpark_connect/proto/snowflake_expression_ext_pb2.pyi

@@ -0,0 +1,53 @@
+from pyspark.sql.connect.proto import expressions_pb2 as _expressions_pb2
+from pyspark.sql.connect.proto import relations_pb2 as _relations_pb2
+from google.protobuf.internal import containers as _containers
+from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class ExpExtension(_message.Message):
+    __slots__ = ["named_argument", "subquery_expression"]
+    NAMED_ARGUMENT_FIELD_NUMBER: _ClassVar[int]
+    SUBQUERY_EXPRESSION_FIELD_NUMBER: _ClassVar[int]
+    named_argument: NamedArgumentExpression
+    subquery_expression: SubqueryExpression
+    def __init__(self, named_argument: _Optional[_Union[NamedArgumentExpression, _Mapping]] = ..., subquery_expression: _Optional[_Union[SubqueryExpression, _Mapping]] = ...) -> None: ...
+
+class NamedArgumentExpression(_message.Message):
+    __slots__ = ["key", "value"]
+    KEY_FIELD_NUMBER: _ClassVar[int]
+    VALUE_FIELD_NUMBER: _ClassVar[int]
+    key: str
+    value: _expressions_pb2.Expression
+    def __init__(self, key: _Optional[str] = ..., value: _Optional[_Union[_expressions_pb2.Expression, _Mapping]] = ...) -> None: ...
+
+class SubqueryExpression(_message.Message):
+    __slots__ = ["in_subquery_values", "input", "subquery_type", "table_arg_options"]
+    class SubqueryType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+        __slots__ = []
+    class TableArgOptions(_message.Message):
+        __slots__ = ["order_spec", "partition_spec", "with_single_partition"]
+        ORDER_SPEC_FIELD_NUMBER: _ClassVar[int]
+        PARTITION_SPEC_FIELD_NUMBER: _ClassVar[int]
+        WITH_SINGLE_PARTITION_FIELD_NUMBER: _ClassVar[int]
+        order_spec: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression.SortOrder]
+        partition_spec: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
+        with_single_partition: bool
+        def __init__(self, partition_spec: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ..., order_spec: _Optional[_Iterable[_Union[_expressions_pb2.Expression.SortOrder, _Mapping]]] = ..., with_single_partition: bool = ...) -> None: ...
+    INPUT_FIELD_NUMBER: _ClassVar[int]
+    IN_SUBQUERY_VALUES_FIELD_NUMBER: _ClassVar[int]
+    SUBQUERY_TYPE_EXISTS: SubqueryExpression.SubqueryType
+    SUBQUERY_TYPE_FIELD_NUMBER: _ClassVar[int]
+    SUBQUERY_TYPE_IN: SubqueryExpression.SubqueryType
+    SUBQUERY_TYPE_SCALAR: SubqueryExpression.SubqueryType
+    SUBQUERY_TYPE_TABLE_ARG: SubqueryExpression.SubqueryType
+    SUBQUERY_TYPE_UNKNOWN: SubqueryExpression.SubqueryType
+    TABLE_ARG_OPTIONS_FIELD_NUMBER: _ClassVar[int]
+    in_subquery_values: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
+    input: _relations_pb2.Relation
+    subquery_type: SubqueryExpression.SubqueryType
+    table_arg_options: SubqueryExpression.TableArgOptions
+    def __init__(self, input: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., subquery_type: _Optional[_Union[SubqueryExpression.SubqueryType, str]] = ..., table_arg_options: _Optional[_Union[SubqueryExpression.TableArgOptions, _Mapping]] = ..., in_subquery_values: _Optional[_Iterable[_Union[_expressions_pb2.Expression, _Mapping]]] = ...) -> None: ...

snowflake/snowpark_connect/proto/snowflake_rdd_pb2.pyi

@@ -0,0 +1,39 @@
+from pyspark.sql.connect.proto import relations_pb2 as _relations_pb2
+from google.protobuf.internal import containers as _containers
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class Rdd(_message.Message):
+    __slots__ = ["map", "reduce"]
+    MAP_FIELD_NUMBER: _ClassVar[int]
+    REDUCE_FIELD_NUMBER: _ClassVar[int]
+    map: RddMap
+    reduce: RddReduce
+    def __init__(self, map: _Optional[_Union[RddMap, _Mapping]] = ..., reduce: _Optional[_Union[RddReduce, _Mapping]] = ...) -> None: ...
+
+class RddMap(_message.Message):
+    __slots__ = ["func", "input"]
+    FUNC_FIELD_NUMBER: _ClassVar[int]
+    INPUT_FIELD_NUMBER: _ClassVar[int]
+    func: bytes
+    input: _relations_pb2.Relation
+    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]
+    INPUT_FIELD_NUMBER: _ClassVar[int]
+    func: bytes
+    input: _relations_pb2.Relation
+    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]
+    INPUT_FIELD_NUMBER: _ClassVar[int]
+    aliases: _containers.RepeatedScalarFieldContainer[str]
+    input: _relations_pb2.Relation
+    def __init__(self, input: _Optional[_Union[_relations_pb2.Relation, _Mapping]] = ..., aliases: _Optional[_Iterable[str]] = ...) -> None: ...

snowflake/snowpark_connect/proto/snowflake_relation_ext_pb2.py

@@ -0,0 +1,47 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: snowflake_relation_ext.proto
+"""Generated protocol buffer code."""
+from google.protobuf.internal import builder as _builder
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from pyspark.sql.connect.proto import relations_pb2 as spark_dot_connect_dot_relations__pb2
+from pyspark.sql.connect.proto import expressions_pb2 as spark_dot_connect_dot_expressions__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1csnowflake_relation_ext.proto\x12\rsnowflake.ext\x1a\x1dspark/connect/relations.proto\x1a\x1fspark/connect/expressions.proto\"\xe3\x02\n\tExtension\x12(\n\x07rdd_map\x18\x01 \x01(\x0b\x32\x15.snowflake.ext.RddMapH\x00\x12.\n\nrdd_reduce\x18\x02 \x01(\x0b\x32\x18.snowflake.ext.RddReduceH\x00\x12G\n\x17subquery_column_aliases\x18\x03 \x01(\x0b\x32$.snowflake.ext.SubqueryColumnAliasesH\x00\x12\x32\n\x0clateral_join\x18\x04 \x01(\x0b\x32\x1a.snowflake.ext.LateralJoinH\x00\x12J\n\x19udtf_with_table_arguments\x18\x05 \x01(\x0b\x32%.snowflake.ext.UDTFWithTableArgumentsH\x00\x12-\n\taggregate\x18\x06 \x01(\x0b\x32\x18.snowflake.ext.AggregateH\x00\x42\x04\n\x02op\">\n\x06RddMap\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x0c\n\x04\x66unc\x18\x02 \x01(\x0c\"A\n\tRddReduce\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x0c\n\x04\x66unc\x18\x02 \x01(\x0c\"P\n\x15SubqueryColumnAliases\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x0f\n\x07\x61liases\x18\x02 \x03(\t\"\\\n\x0bLateralJoin\x12%\n\x04left\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12&\n\x05right\x18\x02 \x01(\x0b\x32\x17.spark.connect.Relation\"\x98\x01\n\x16UDTFWithTableArguments\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12,\n\targuments\x18\x02 \x03(\x0b\x32\x19.spark.connect.Expression\x12\x39\n\x0ftable_arguments\x18\x03 \x03(\x0b\x32 .snowflake.ext.TableArgumentInfo\"`\n\x11TableArgumentInfo\x12/\n\x0etable_argument\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x1a\n\x12table_argument_idx\x18\x02 \x01(\x05\"\x92\x05\n\tAggregate\x12&\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.Relation\x12\x36\n\ngroup_type\x18\x02 \x01(\x0e\x32\".snowflake.ext.Aggregate.GroupType\x12\x37\n\x14grouping_expressions\x18\x03 \x03(\x0b\x32\x19.spark.connect.Expression\x12\x38\n\x15\x61ggregate_expressions\x18\x04 \x03(\x0b\x32\x19.spark.connect.Expression\x12-\n\x05pivot\x18\x05 \x01(\x0b\x32\x1e.snowflake.ext.Aggregate.Pivot\x12<\n\rgrouping_sets\x18\x06 \x03(\x0b\x32%.snowflake.ext.Aggregate.GroupingSets\x1a\x62\n\x05Pivot\x12&\n\x03\x63ol\x18\x01 \x01(\x0b\x32\x19.spark.connect.Expression\x12\x31\n\x06values\x18\x02 \x03(\x0b\x32!.spark.connect.Expression.Literal\x1a?\n\x0cGroupingSets\x12/\n\x0cgrouping_set\x18\x01 \x03(\x0b\x32\x19.spark.connect.Expression\"\x9f\x01\n\tGroupType\x12\x1a\n\x16GROUP_TYPE_UNSPECIFIED\x10\x00\x12\x16\n\x12GROUP_TYPE_GROUPBY\x10\x01\x12\x15\n\x11GROUP_TYPE_ROLLUP\x10\x02\x12\x13\n\x0fGROUP_TYPE_CUBE\x10\x03\x12\x14\n\x10GROUP_TYPE_PIVOT\x10\x04\x12\x1c\n\x18GROUP_TYPE_GROUPING_SETS\x10\x05\x62\x06proto3')
+
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'snowflake_relation_ext_pb2', globals())
+if _descriptor._USE_C_DESCRIPTORS == False:
+
+  DESCRIPTOR._options = None
+  _EXTENSION._serialized_start=112
+  _EXTENSION._serialized_end=467
+  _RDDMAP._serialized_start=469
+  _RDDMAP._serialized_end=531
+  _RDDREDUCE._serialized_start=533
+  _RDDREDUCE._serialized_end=598
+  _SUBQUERYCOLUMNALIASES._serialized_start=600
+  _SUBQUERYCOLUMNALIASES._serialized_end=680
+  _LATERALJOIN._serialized_start=682
+  _LATERALJOIN._serialized_end=774
+  _UDTFWITHTABLEARGUMENTS._serialized_start=777
+  _UDTFWITHTABLEARGUMENTS._serialized_end=929
+  _TABLEARGUMENTINFO._serialized_start=931
+  _TABLEARGUMENTINFO._serialized_end=1027
+  _AGGREGATE._serialized_start=1030
+  _AGGREGATE._serialized_end=1688
+  _AGGREGATE_PIVOT._serialized_start=1363
+  _AGGREGATE_PIVOT._serialized_end=1461
+  _AGGREGATE_GROUPINGSETS._serialized_start=1463
+  _AGGREGATE_GROUPINGSETS._serialized_end=1526
+  _AGGREGATE_GROUPTYPE._serialized_start=1529
+  _AGGREGATE_GROUPTYPE._serialized_end=1688
+# @@protoc_insertion_point(module_scope)

snowflake/snowpark_connect/proto/snowflake_relation_ext_pb2.pyi

@@ -0,0 +1,111 @@
+from pyspark.sql.connect.proto import relations_pb2 as _relations_pb2
+from pyspark.sql.connect.proto import expressions_pb2 as _expressions_pb2
+from google.protobuf.internal import containers as _containers
+from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union
+
+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]
+    RDD_MAP_FIELD_NUMBER: _ClassVar[int]
+    RDD_REDUCE_FIELD_NUMBER: _ClassVar[int]
+    SUBQUERY_COLUMN_ALIASES_FIELD_NUMBER: _ClassVar[int]
+    UDTF_WITH_TABLE_ARGUMENTS_FIELD_NUMBER: _ClassVar[int]
+    aggregate: Aggregate
+    lateral_join: LateralJoin
+    rdd_map: RddMap
+    rdd_reduce: RddReduce
+    subquery_column_aliases: SubqueryColumnAliases
+    udtf_with_table_arguments: UDTFWithTableArguments
+    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]
+    INPUT_FIELD_NUMBER: _ClassVar[int]
+    func: bytes
+    input: _relations_pb2.Relation
+    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]
+    INPUT_FIELD_NUMBER: _ClassVar[int]
+    func: bytes
+    input: _relations_pb2.Relation
+    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]
+    INPUT_FIELD_NUMBER: _ClassVar[int]
+    aliases: _containers.RepeatedScalarFieldContainer[str]
+    input: _relations_pb2.Relation
+    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 UDTFWithTableArguments(_message.Message):
+    __slots__ = ["arguments", "function_name", "table_arguments"]
+    ARGUMENTS_FIELD_NUMBER: _ClassVar[int]
+    FUNCTION_NAME_FIELD_NUMBER: _ClassVar[int]
+    TABLE_ARGUMENTS_FIELD_NUMBER: _ClassVar[int]
+    arguments: _containers.RepeatedCompositeFieldContainer[_expressions_pb2.Expression]
+    function_name: str
+    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: ...

snowflake/snowpark_connect/relation/__init__.py

@@ -0,0 +1,3 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#

snowflake/snowpark_connect/relation/catalogs/__init__.py

@@ -0,0 +1,12 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+from .snowflake_catalog import SnowflakeCatalog
+
+SNOWFLAKE_CATALOG = SnowflakeCatalog()
+
+CATALOGS = {
+    "snowflake_catalog": SNOWFLAKE_CATALOG,
+    "spark_catalog": SNOWFLAKE_CATALOG,
+}