snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/expression/typer.py

@@ -0,0 +1,125 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+from functools import cached_property
+
+from snowflake import snowpark
+from snowflake.snowpark import Column, DataFrame
+from snowflake.snowpark._internal.analyzer.expression import (
+    Attribute,
+    Literal,
+    UnresolvedAttribute,
+)
+from snowflake.snowpark.exceptions import SnowparkClientException
+from snowflake.snowpark.types import DataType, LongType, StructField, StructType
+from snowflake.snowpark_connect.config import global_config
+from snowflake.snowpark_connect.utils.context import (
+    get_df_before_projection,
+    get_outer_dataframes,
+)
+from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
+
+
+class ExpressionTyper:
+    def __init__(self, df: DataFrame) -> None:
+        self.df = df
+
+    def type(self, column: Column) -> list[DataType]:
+        types = self._try_to_type_attribute_or_literal(self.df, column)
+        if not types and get_df_before_projection():
+            types = self._try_to_type_attribute_or_literal(
+                get_df_before_projection(), column
+            )
+        if not types:
+            # df.select().schema results in DESCRIBE call to Snowflake, so avoid it if possible
+            types = self._type_using_select(self.df, column)
+        return types
+
+    def _type_using_select(self, df: DataFrame, column: Column) -> list[DataType]:
+        df = self._join_df_with_outer_dataframes(df)
+
+        try:
+            return self._get_df_datatypes(df, column)
+        except SnowparkClientException:  # Fallback to the df before projection
+            df = get_df_before_projection()
+            if df is None:
+                raise
+
+            df = self._join_df_with_outer_dataframes(df)
+            return self._get_df_datatypes(df, column)
+
+    @staticmethod
+    def _join_df_with_outer_dataframes(df: DataFrame) -> DataFrame:
+        for outer_df in get_outer_dataframes():
+            df = df.join(outer_df)
+
+        return df
+
+    @staticmethod
+    def _get_df_datatypes(df: DataFrame, column: Column) -> list[DataType]:
+        return [f.datatype for f in df.select(column).schema.fields]
+
+    def _try_to_type_attribute_or_literal(
+        self, df: DataFrame, column: Column
+    ) -> list[DataType] | None:
+        types = None
+        expr = column._expression if hasattr(column, "_expression") else None
+        match expr:
+            case UnresolvedAttribute() | Attribute():
+                # there is a chance that df.schema is already evaluated
+                types = [
+                    f.datatype
+                    for f in df.schema.fields
+                    if (
+                        f.name == expr.name
+                        or (
+                            (not global_config.spark_sql_caseSensitive)
+                            and f.name.lower() == expr.name.lower()
+                        )
+                    )
+                ]  # doesn't work for nested attributes e.g. `"properties-3":"salary"`
+            case Literal():
+                types = [expr.datatype]
+        return types
+
+    @staticmethod
+    def dummy_typer(session: snowpark.Session = None):
+        """
+        Get a dummy typer, which can be used to get expression types. Since typer requires a dataframe,
+        the dummy typer is mainly used when there is no existing handy dataframe.
+
+        Example:
+            (_, typed_column) = map_single_column_expression(
+                    expression, column_map, ExpressionTyper.dummy_typer(session)
+                )
+        """
+        if session is None:
+            session = get_or_create_snowpark_session()
+        empty_df = session.create_dataframe(
+            [], schema=StructType([StructField("id", LongType(), True)])
+        )
+        typer = ExpressionTyper(empty_df)
+        return typer
+
+
+class JoinExpressionTyper(ExpressionTyper):
+    def __init__(self, left: DataFrame, right: DataFrame) -> None:
+        self.left = left
+        self.right = right
+
+    def type(self, column: Column) -> list[DataType]:
+        types = self._try_to_type_attribute_or_literal(self.left, column)
+        if not types:
+            types = self._try_to_type_attribute_or_literal(self.right, column)
+        if not types:
+            types = self._type_using_select(self.df, column)
+        return types
+
+    @cached_property
+    def df(self) -> DataFrame:
+        return self.left.join(
+            self.right,
+            lsuffix="_left",
+            rsuffix="_right",
+        )

snowflake/snowpark_connect/includes/python/__init__.py

@@ -0,0 +1,21 @@
+from pyspark.errors.exceptions.base import (
+    AnalysisException,
+    ArithmeticException,
+    ArrayIndexOutOfBoundsException,
+    DateTimeException,
+    IllegalArgumentException,
+    NumberFormatException,
+    ParseException,
+    PySparkException,
+    PySparkTypeError,
+    PySparkValueError,
+    PythonException,
+    SparkRuntimeException,
+)
+
+# We change the module so that these exceptions are shown as pyspark exceptions in the client
+AnalysisException.__module__ = ArithmeticException.__module__ = ArrayIndexOutOfBoundsException.__module__ = \
+    DateTimeException.__module__ = IllegalArgumentException.__module__ = NumberFormatException.__module__ = \
+    ParseException.__module__ = PySparkException.__module__ = PySparkTypeError.__module__ = \
+    PySparkValueError.__module__ = PythonException.__module__ = \
+    SparkRuntimeException.__module__ = "pyspark.errors.exceptions.base"

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

@@ -0,0 +1,173 @@
+#
+# 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.
+#
+
+"""
+PySpark is the Python API for Spark.
+
+Public classes:
+
+  - :class:`SparkContext`:
+      Main entry point for Spark functionality.
+  - :class:`RDD`:
+      A Resilient Distributed Dataset (RDD), the basic abstraction in Spark.
+  - :class:`Broadcast`:
+      A broadcast variable that gets reused across tasks.
+  - :class:`Accumulator`:
+      An "add-only" shared variable that tasks can only add values to.
+  - :class:`SparkConf`:
+      For configuring Spark.
+  - :class:`SparkFiles`:
+      Access files shipped with jobs.
+  - :class:`StorageLevel`:
+      Finer-grained cache persistence levels.
+  - :class:`TaskContext`:
+      Information about the current running task, available on the workers and experimental.
+  - :class:`RDDBarrier`:
+      Wraps an RDD under a barrier stage for barrier execution.
+  - :class:`BarrierTaskContext`:
+      A :class:`TaskContext` that provides extra info and tooling for barrier execution.
+  - :class:`BarrierTaskInfo`:
+      Information about a barrier task.
+  - :class:`InheritableThread`:
+      A inheritable thread to use in Spark when the pinned thread mode is on.
+"""
+
+from functools import wraps
+import types
+from typing import cast, Any, Callable, Optional, TypeVar, Union
+from warnings import filterwarnings
+
+filterwarnings(
+    "ignore", message="distutils Version classes are deprecated. Use packaging.version instead."
+)
+
+from pyspark.conf import SparkConf
+from pyspark.rdd import RDD, RDDBarrier
+from pyspark.files import SparkFiles
+from pyspark.status import StatusTracker, SparkJobInfo, SparkStageInfo
+from pyspark.util import InheritableThread, inheritable_thread_target
+from pyspark.storagelevel import StorageLevel
+from pyspark.accumulators import Accumulator, AccumulatorParam
+from pyspark.broadcast import Broadcast
+from pyspark.serializers import MarshalSerializer, CPickleSerializer
+from pyspark.taskcontext import TaskContext, BarrierTaskContext, BarrierTaskInfo
+from pyspark.profiler import Profiler, BasicProfiler
+from pyspark.version import __version__
+from pyspark._globals import _NoValue  # noqa: F401
+
+_F = TypeVar("_F", bound=Callable)
+
+
+def since(version: Union[str, float]) -> Callable[[_F], _F]:
+    """
+    A decorator that annotates a function to append the version of Spark the function was added.
+    """
+    import re
+
+    indent_p = re.compile(r"\n( +)")
+
+    def deco(f: _F) -> _F:
+        assert f.__doc__ is not None
+
+        indents = indent_p.findall(f.__doc__)
+        indent = " " * (min(len(m) for m in indents) if indents else 0)
+        f.__doc__ = f.__doc__.rstrip() + "\n\n%s.. versionadded:: %s" % (indent, version)
+        return f
+
+    return deco
+
+
+def copy_func(
+    f: _F,
+    name: Optional[str] = None,
+    sinceversion: Optional[Union[str, float]] = None,
+    doc: Optional[str] = None,
+) -> _F:
+    """
+    Returns a function with same code, globals, defaults, closure, and
+    name (or provide a new name).
+    """
+    # See
+    # http://stackoverflow.com/questions/6527633/how-can-i-make-a-deepcopy-of-a-function-in-python
+    assert isinstance(f, types.FunctionType)
+
+    fn = types.FunctionType(
+        f.__code__,
+        f.__globals__,
+        name or f.__name__,
+        f.__defaults__,
+        f.__closure__,
+    )
+    # in case f was given attrs (note this dict is a shallow copy):
+    fn.__dict__.update(f.__dict__)
+    if doc is not None:
+        fn.__doc__ = doc
+    if sinceversion is not None:
+        fn = since(sinceversion)(fn)
+    return cast(_F, fn)
+
+
+def keyword_only(func: _F) -> _F:
+    """
+    A decorator that forces keyword arguments in the wrapped method
+    and saves actual input keyword arguments in `_input_kwargs`.
+
+    Notes
+    -----
+    Should only be used to wrap a method where first arg is `self`
+    """
+
+    @wraps(func)
+    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+        if len(args) > 0:
+            raise TypeError("Method %s forces keyword arguments." % func.__name__)
+        self._input_kwargs = kwargs
+        return func(self, **kwargs)
+
+    return cast(_F, wrapper)
+
+
+# To avoid circular dependencies
+from pyspark.context import SparkContext
+
+# for back compatibility
+from pyspark.sql import SQLContext, HiveContext, Row  # noqa: F401
+
+__all__ = [
+    "SparkConf",
+    "SparkContext",
+    "SparkFiles",
+    "RDD",
+    "StorageLevel",
+    "Broadcast",
+    "Accumulator",
+    "AccumulatorParam",
+    "MarshalSerializer",
+    "CPickleSerializer",
+    "StatusTracker",
+    "SparkJobInfo",
+    "SparkStageInfo",
+    "Profiler",
+    "BasicProfiler",
+    "TaskContext",
+    "RDDBarrier",
+    "BarrierTaskContext",
+    "BarrierTaskInfo",
+    "InheritableThread",
+    "inheritable_thread_target",
+    "__version__",
+]

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

@@ -0,0 +1,71 @@
+#
+# 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.
+#
+
+"""
+Module defining global singleton classes.
+
+This module raises a RuntimeError if an attempt to reload it is made. In that
+way the identities of the classes defined here are fixed and will remain so
+even if pyspark itself is reloaded. In particular, a function like the following
+will still work correctly after pyspark is reloaded:
+
+    def foo(arg=pyspark._NoValue):
+        if arg is pyspark._NoValue:
+            ...
+
+See gh-7844 for a discussion of the reload problem that motivated this module.
+
+Note that this approach is taken after from NumPy.
+"""
+
+__ALL__ = ["_NoValue"]
+
+
+# Disallow reloading this module so as to preserve the identities of the
+# classes defined here.
+if "_is_loaded" in globals():
+    raise RuntimeError("Reloading pyspark._globals is not allowed")
+_is_loaded = True
+
+
+class _NoValueType:
+    """Special keyword value.
+
+    The instance of this class may be used as the default value assigned to a
+    deprecated keyword in order to check if it has been given a user defined
+    value.
+
+    This class was copied from NumPy.
+    """
+
+    __instance = None
+
+    def __new__(cls):
+        # ensure that only one instance exists
+        if not cls.__instance:
+            cls.__instance = super(_NoValueType, cls).__new__(cls)
+        return cls.__instance
+
+    # needed for python 2 to preserve identity through a pickle
+    def __reduce__(self):
+        return (self.__class__, ())
+
+    def __repr__(self):
+        return "<no value>"
+
+
+_NoValue = _NoValueType()

snowflake/snowpark_connect/includes/python/pyspark/_typing.pyi

@@ -0,0 +1,43 @@
+#
+# 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.
+
+from typing import Callable, Iterable, Sized, TypeVar, Union
+from typing_extensions import Literal, Protocol
+
+from numpy import int32, int64, float32, float64, ndarray
+
+F = TypeVar("F", bound=Callable)
+T_co = TypeVar("T_co", covariant=True)
+
+PrimitiveType = Union[bool, float, int, str]
+
+NonUDFType = Literal[0]
+
+class SupportsIAdd(Protocol):
+    def __iadd__(self, other: SupportsIAdd) -> SupportsIAdd: ...
+
+class SupportsOrdering(Protocol):
+    def __lt__(self, other: SupportsOrdering) -> bool: ...
+
+class SizedIterable(Protocol, Sized, Iterable[T_co]): ...
+
+S = TypeVar("S", bound=SupportsOrdering)
+
+NumberOrArray = TypeVar(
+    "NumberOrArray", float, int, complex, int32, int64, float32, float64, ndarray
+)