sqlframe 2.0.0__tar.gz → 2.2.0__tar.gz

{sqlframe-2.0.0 → sqlframe-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 2.0.0
+Version: 2.2.0
 Summary: Turning PySpark Into a Universal DataFrame API
 Home-page: https://github.com/eakmanrq/sqlframe
 Author: Ryan Eakman
@@ -90,7 +90,7 @@ from sqlframe.bigquery import functions as F
 from sqlframe.bigquery import Window
 
 session = BigQuerySession()
-table_path = "bigquery-public-data.samples.natality"
+table_path = '"bigquery-public-data".samples.natality'
 # Top 5 years with the greatest year-over-year % change in new families with single child
 df = (
     session.table(table_path)

{sqlframe-2.0.0 → sqlframe-2.2.0}/README.md

@@ -60,7 +60,7 @@ from sqlframe.bigquery import functions as F
 from sqlframe.bigquery import Window
 
 session = BigQuerySession()
-table_path = "bigquery-public-data.samples.natality"
+table_path = '"bigquery-public-data".samples.natality'
 # Top 5 years with the greatest year-over-year % change in new families with single child
 df = (
     session.table(table_path)

{sqlframe-2.0.0 → sqlframe-2.2.0}/blogs/add_chatgpt_support.md

@@ -47,7 +47,7 @@ from sqlframe.bigquery import functions as F
 from sqlframe.bigquery import Window
 
 session = BigQuerySession()
-table_path = "bigquery-public-data.samples.natality"
+table_path = '"bigquery-public-data".samples.natality'
 # Top 5 years with the greatest year-over-year % change in new families with single child
 df = (
     session.table(table_path)

{sqlframe-2.0.0 → sqlframe-2.2.0}/blogs/sqlframe_universal_dataframe_api.md

@@ -39,7 +39,7 @@ from sqlframe.bigquery import Window
 
 # Unique to SQLFrame: Ability to connect directly to BigQuery
 session = BigQuerySession()
-table_path = "bigquery-public-data.samples.natality"
+table_path = '"bigquery-public-data".samples.natality'
 # Get the top 5 years with the greatest year-over-year % change in new families with a single child
 df = (
     session.table(table_path)

{sqlframe-2.0.0 → sqlframe-2.2.0}/docs/bigquery.md

@@ -72,7 +72,7 @@ from sqlframe.bigquery import functions as F
 
 session = BigQuerySession(default_dataset="sqlframe.db1")
 (
-    session.table("bigquery-public-data.samples.natality")
+    session.table('"bigquery-public-data".samples.natality')
     .select(F.call_function("FARM_FINGERPRINT", F.col("source")).alias("source_hash"))
     .show()
 )
@@ -86,7 +86,7 @@ from sqlframe.bigquery import functions as F
 from sqlframe.bigquery import Window
 
 session = BigQuerySession(default_dataset="sqlframe.db1")
-table_path = "bigquery-public-data.samples.natality"
+table_path = '"bigquery-public-data".samples.natality'
 # Get columns in the table
 print(session.catalog.listColumns(table_path))
 # Get the top 5 years with the greatest year-over-year % change in new families with a single child
@@ -214,6 +214,8 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [intersectAll](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.intersectAll.html)
 * [join](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.join.html)
 * [limit](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.limit.html)
+* lineage
+   * Get lineage for a specific column. [Returns a SQLGlot Node](https://sqlglot.com/sqlglot/lineage.html#Node). Can be used to get lineage SQL or HTML representation.
 * [na](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.na.html)
 * [orderBy](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.orderBy.html)
 * [persist](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.persist.html)

{sqlframe-2.0.0 → sqlframe-2.2.0}/docs/configuration.md

@@ -62,7 +62,7 @@ from sqlframe.bigquery import functions as F
 from sqlframe.bigquery import Window
 
 session = BigQuerySession()
-table_path = "bigquery-public-data.samples.natality"
+table_path = '"bigquery-public-data".samples.natality'
 # Top 5 years with the greatest year-over-year % change in new families with single child
 df = (
     session.table(table_path)

{sqlframe-2.0.0 → sqlframe-2.2.0}/docs/duckdb.md

@@ -187,6 +187,8 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [intersectAll](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.intersectAll.html)
 * [join](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.join.html)
 * [limit](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.limit.html)
+* lineage
+   * Get lineage for a specific column. [Returns a SQLGlot Node](https://sqlglot.com/sqlglot/lineage.html#Node). Can be used to get lineage SQL or HTML representation.
 * [na](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.na.html)
 * [orderBy](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.orderBy.html)
 * [persist](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.persist.html)
@@ -258,6 +260,7 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [concat](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.concat.html)                               
     * Only works on strings (does not work on arrays)
 * [concat_ws](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.concat_ws.html)
+* [convert_timezone](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.convert_timezone.html)
 * [corr](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.corr.html)
 * [cos](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.cos.html)
 * [cot](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.cot.html)

{sqlframe-2.0.0 → sqlframe-2.2.0}/docs/postgres.md

@@ -198,6 +198,8 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [intersectAll](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.intersectAll.html)
 * [join](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.join.html)
 * [limit](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.limit.html)
+* lineage
+   * Get lineage for a specific column. [Returns a SQLGlot Node](https://sqlglot.com/sqlglot/lineage.html#Node). Can be used to get lineage SQL or HTML representation.
 * [na](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.na.html)
 * [orderBy](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.orderBy.html)
 * [persist](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.persist.html)

{sqlframe-2.0.0 → sqlframe-2.2.0}/docs/snowflake.md

@@ -209,6 +209,8 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [intersectAll](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.intersectAll.html)
 * [join](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.join.html)
 * [limit](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.limit.html)
+* lineage
+   * Get lineage for a specific column. [Returns a SQLGlot Node](https://sqlglot.com/sqlglot/lineage.html#Node). Can be used to get lineage SQL or HTML representation.
 * [na](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.na.html)
 * [orderBy](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.orderBy.html)
 * [persist](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.persist.html)
@@ -286,6 +288,7 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [concat](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.concat.html)
   * Can only concat strings not arrays
 * [concat_ws](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.concat_ws.html)
+* [convert_timezone](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.convert_timezone.html)
 * [corr](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.corr.html)
 * [cos](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.cos.html)
 * [cosh](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.cosh.html)

{sqlframe-2.0.0 → sqlframe-2.2.0}/docs/spark.md

@@ -156,6 +156,8 @@ df.show(5)
 * [intersectAll](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.intersectAll.html)
 * [join](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.join.html)
 * [limit](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.limit.html)
+* lineage
+   * Get lineage for a specific column. [Returns a SQLGlot Node](https://sqlglot.com/sqlglot/lineage.html#Node). Can be used to get lineage SQL or HTML representation.
 * [na](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.na.html)
 * [orderBy](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.orderBy.html)
 * [persist](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.persist.html)

{sqlframe-2.0.0 → sqlframe-2.2.0}/docs/standalone.md

@@ -133,6 +133,8 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [intersectAll](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.intersectAll.html)
 * [join](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.join.html)
 * [limit](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.limit.html)
+* lineage
+   * Get lineage for a specific column. [Returns a SQLGlot Node](https://sqlglot.com/sqlglot/lineage.html#Node). Can be used to get lineage SQL or HTML representation.
 * [na](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.na.html)
 * [orderBy](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.orderBy.html)
 * [persist](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.persist.html)

{sqlframe-2.0.0 → sqlframe-2.2.0}/setup.py

@@ -20,7 +20,7 @@ setup(
     python_requires=">=3.8",
     install_requires=[
         "prettytable<3.11.0",
-        "sqlglot>=24.0.0,<25.9",
+        "sqlglot>=24.0.0,<25.11",
         "typing_extensions>=4.8,<5",
     ],
     extras_require={
@@ -31,7 +31,7 @@ setup(
         "dev": [
             "duckdb>=0.9,<1.1",
             "mypy>=1.10.0,<1.12",
-            "openai>=1.30,<1.38",
+            "openai>=1.30,<1.41",
             "pandas>=2,<3",
             "pandas-stubs>=2,<3",
             "psycopg>=3.1,<4",
@@ -57,7 +57,7 @@ setup(
             "pandas>=2,<3",
         ],
         "openai": [
-            "openai>=1.30,<1.38",
+            "openai>=1.30,<1.41",
         ],
         "pandas": [
             "pandas>=2,<3",

{sqlframe-2.0.0 → sqlframe-2.2.0}/sqlframe/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.0.0'
-__version_tuple__ = version_tuple = (2, 0, 0)
+__version__ = version = '2.2.0'
+__version_tuple__ = version_tuple = (2, 2, 0)

{sqlframe-2.0.0 → sqlframe-2.2.0}/sqlframe/base/_typing.py

@@ -24,6 +24,7 @@ OutputExpressionContainer = t.Union[exp.Select, exp.Create, exp.Insert]
 StorageLevel = str
 PathOrPaths = t.Union[str, t.List[str]]
 OptionalPrimitiveType = t.Optional[PrimitiveType]
+DataTypeOrString = t.Union[DataType, str]
 
 
 class UserDefinedFunctionLike(t.Protocol):

{sqlframe-2.0.0 → sqlframe-2.2.0}/sqlframe/base/dataframe.py

@@ -17,6 +17,7 @@ import sqlglot
 from prettytable import PrettyTable
 from sqlglot import Dialect
 from sqlglot import expressions as exp
+from sqlglot import lineage as sqlglot_lineage
 from sqlglot.helper import ensure_list, flatten, object_to_dict, seq_get
 from sqlglot.optimizer.pushdown_projections import pushdown_projections
 from sqlglot.optimizer.qualify import qualify
@@ -1613,6 +1614,13 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
                 0,
             )
 
+    def lineage(self, col: ColumnOrName, optimize: bool = True) -> sqlglot_lineage.Node:
+        return sqlglot_lineage.lineage(
+            column=self._ensure_and_normalize_col(col).alias_or_name,
+            sql=self._get_expressions(optimize=optimize)[0],
+            schema=self.session.catalog._schema,
+        )
+
     def toPandas(self) -> pd.DataFrame:
         return self.session._fetchdf(self._get_expressions(optimize=False))
 

{sqlframe-2.0.0 → sqlframe-2.2.0}/sqlframe/base/functions.py

@@ -2070,14 +2070,19 @@ def contains(left: ColumnOrName, right: ColumnOrName) -> Column:
     return Column.invoke_anonymous_function(left, "contains", right)
 
 
-@meta(unsupported_engines="*")
+@meta(unsupported_engines=["bigquery", "postgres"])
 def convert_timezone(
     sourceTz: t.Optional[Column], targetTz: Column, sourceTs: ColumnOrName
 ) -> Column:
-    if sourceTz is None:
-        return Column.invoke_anonymous_function(targetTz, "convert_timezone", sourceTs)
-    else:
-        return Column.invoke_anonymous_function(sourceTz, "convert_timezone", targetTz, sourceTs)
+    to_timestamp = get_func_from_session("to_timestamp")
+
+    return Column(
+        expression.ConvertTimezone(
+            timestamp=to_timestamp(Column.ensure_col(sourceTs)).expression,
+            source_tz=sourceTz.expression if sourceTz else None,
+            target_tz=Column.ensure_col(targetTz).expression,
+        )
+    )
 
 
 @meta(unsupported_engines="postgres")

{sqlframe-2.0.0 → sqlframe-2.2.0}/sqlframe/base/session.py

@@ -27,6 +27,7 @@ from sqlframe.base.catalog import _BaseCatalog
 from sqlframe.base.dataframe import _BaseDataFrame
 from sqlframe.base.normalize import normalize_dict
 from sqlframe.base.readerwriter import _BaseDataFrameReader, _BaseDataFrameWriter
+from sqlframe.base.udf import _BaseUDFRegistration
 from sqlframe.base.util import (
     get_column_mapping_from_schema_input,
     normalize_string,
@@ -64,16 +65,18 @@ CATALOG = t.TypeVar("CATALOG", bound=_BaseCatalog)
 READER = t.TypeVar("READER", bound=_BaseDataFrameReader)
 WRITER = t.TypeVar("WRITER", bound=_BaseDataFrameWriter)
 DF = t.TypeVar("DF", bound=_BaseDataFrame)
+UDF_REGISTRATION = t.TypeVar("UDF_REGISTRATION", bound=_BaseUDFRegistration)
 
 _MISSING = "MISSING"
 
 
-class _BaseSession(t.Generic[CATALOG, READER, WRITER, DF, CONN]):
+class _BaseSession(t.Generic[CATALOG, READER, WRITER, DF, CONN, UDF_REGISTRATION]):
     _instance = None
     _reader: t.Type[READER]
     _writer: t.Type[WRITER]
     _catalog: t.Type[CATALOG]
     _df: t.Type[DF]
+    _udf_registration: t.Type[UDF_REGISTRATION]
 
     SANITIZE_COLUMN_NAMES = False
 
@@ -81,7 +84,6 @@ class _BaseSession(t.Generic[CATALOG, READER, WRITER, DF, CONN]):
         self,
         conn: t.Optional[CONN] = None,
         schema: t.Optional[MappingSchema] = None,
-        case_sensitive: bool = False,
         *args,
         **kwargs,
     ):
@@ -91,11 +93,6 @@ class _BaseSession(t.Generic[CATALOG, READER, WRITER, DF, CONN]):
             self.execution_dialect: Dialect = Dialect.get_or_raise(
                 self.builder.DEFAULT_EXECUTION_DIALECT
             )
-            self.case_sensitive: bool = case_sensitive
-            if self.case_sensitive:
-                self.input_dialect.NORMALIZATION_STRATEGY = NormalizationStrategy.CASE_SENSITIVE
-                self.output_dialect.NORMALIZATION_STRATEGY = NormalizationStrategy.CASE_SENSITIVE
-                self.execution_dialect.NORMALIZATION_STRATEGY = NormalizationStrategy.CASE_SENSITIVE
             self.known_ids: t.Set[str] = set()
             self.known_branch_ids: t.Set[str] = set()
             self.known_sequence_ids: t.Set[str] = set()
@@ -176,6 +173,13 @@ class _BaseSession(t.Generic[CATALOG, READER, WRITER, DF, CONN]):
     def _has_connection(self) -> bool:
         return hasattr(self, "_connection") and bool(self._connection)
 
+    @property
+    def udf(self) -> UDF_REGISTRATION:
+        return self._udf_registration(self)
+
+    def getActiveSession(self) -> Self:
+        return self
+
     def range(self, *args):
         start = 0
         step = 1
@@ -573,53 +577,10 @@ class _BaseSession(t.Generic[CATALOG, READER, WRITER, DF, CONN]):
             converted_values.append(cls._to_value(value))
         return _create_row(columns, converted_values)
 
-    @property
-    def _is_standalone(self) -> bool:
-        from sqlframe.standalone.session import StandaloneSession
-
-        return isinstance(self, StandaloneSession)
-
-    @property
-    def _is_duckdb(self) -> bool:
-        from sqlframe.duckdb.session import DuckDBSession
-
-        return isinstance(self, DuckDBSession)
-
-    @property
-    def _is_postgres(self) -> bool:
-        from sqlframe.postgres.session import PostgresSession
-
-        return isinstance(self, PostgresSession)
-
-    @property
-    def _is_spark(self) -> bool:
-        from sqlframe.spark.session import SparkSession
-
-        return isinstance(self, SparkSession)
-
-    @property
-    def _is_bigquery(self) -> bool:
-        from sqlframe.bigquery.session import BigQuerySession
-
-        return isinstance(self, BigQuerySession)
-
-    @property
-    def _is_redshift(self) -> bool:
-        from sqlframe.redshift.session import RedshiftSession
-
-        return isinstance(self, RedshiftSession)
-
-    @property
-    def _is_snowflake(self) -> bool:
-        from sqlframe.snowflake.session import SnowflakeSession
-
-        return isinstance(self, SnowflakeSession)
-
     class Builder:
         SQLFRAME_INPUT_DIALECT_KEY = "sqlframe.input.dialect"
         SQLFRAME_OUTPUT_DIALECT_KEY = "sqlframe.output.dialect"
         SQLFRAME_EXECUTION_DIALECT_KEY = "sqlframe.execution.dialect"
-        SQLFRAME_CASE_SENSITIVE_KEY = "spark.sql.caseSensitive"
         SQLFRAME_CONN_KEY = "sqlframe.conn"
         SQLFRAME_SCHEMA_KEY = "sqlframe.schema"
         DEFAULT_INPUT_DIALECT = "spark"
@@ -665,8 +626,6 @@ class _BaseSession(t.Generic[CATALOG, READER, WRITER, DF, CONN]):
                     self._session_kwargs["conn"] = value
                 elif key == self.SQLFRAME_SCHEMA_KEY:
                     self._session_kwargs["schema"] = value
-                elif key == self.SQLFRAME_CASE_SENSITIVE_KEY:
-                    self._session_kwargs["case_sensitive"] = value
                 else:
                     self._session_kwargs[key] = value
             if map:
@@ -676,8 +635,6 @@ class _BaseSession(t.Generic[CATALOG, READER, WRITER, DF, CONN]):
                     self.output_dialect = map[self.SQLFRAME_OUTPUT_DIALECT_KEY]
                 if self.SQLFRAME_EXECUTION_DIALECT_KEY in map:
                     self.execution_dialect = map[self.SQLFRAME_EXECUTION_DIALECT_KEY]
-                if self.SQLFRAME_CASE_SENSITIVE_KEY in map:
-                    self._session_kwargs["case_sensitive"] = map[self.SQLFRAME_CASE_SENSITIVE_KEY]
                 if self.SQLFRAME_CONN_KEY in map:
                     self._session_kwargs["conn"] = map[self.SQLFRAME_CONN_KEY]
                 if self.SQLFRAME_SCHEMA_KEY in map:
@@ -700,15 +657,5 @@ class _BaseSession(t.Generic[CATALOG, READER, WRITER, DF, CONN]):
             self.session.execution_dialect = Dialect.get_or_raise(self.execution_dialect)
             if hasattr(self.session, "_connection") and not self.session._connection:
                 self.session._connection = self._conn
-            if self.session.case_sensitive:
-                self.session.input_dialect.NORMALIZATION_STRATEGY = (
-                    NormalizationStrategy.CASE_SENSITIVE
-                )
-                self.session.output_dialect.NORMALIZATION_STRATEGY = (
-                    NormalizationStrategy.CASE_SENSITIVE
-                )
-                self.session.execution_dialect.NORMALIZATION_STRATEGY = (
-                    NormalizationStrategy.CASE_SENSITIVE
-                )
 
     builder = Builder()

sqlframe-2.2.0/sqlframe/base/udf.py

@@ -0,0 +1,36 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+from __future__ import annotations
+
+import typing as t
+
+if t.TYPE_CHECKING:
+    from sqlframe.base._typing import DataTypeOrString, UserDefinedFunctionLike
+    from sqlframe.base.session import _BaseSession
+
+    SESSION = t.TypeVar("SESSION", bound=_BaseSession)
+else:
+    SESSION = t.TypeVar("SESSION")
+
+
+class _BaseUDFRegistration(t.Generic[SESSION]):
+    def __init__(self, sparkSession: SESSION):
+        self.sparkSession = sparkSession
+
+    def register(
+        self,
+        name: str,
+        f: t.Union[t.Callable[..., t.Any], UserDefinedFunctionLike],
+        returnType: t.Optional[DataTypeOrString] = None,
+    ) -> UserDefinedFunctionLike:
+        raise NotImplementedError
+
+    def registerJavaFunction(
+        self,
+        name: str,
+        javaClassName: str,
+        returnType: t.Optional[DataTypeOrString] = None,
+    ) -> None:
+        raise NotImplementedError
+
+    def registerJavaUDAF(self, name: str, javaClassName: str) -> None:
+        raise NotImplementedError

{sqlframe-2.0.0 → sqlframe-2.2.0}/sqlframe/bigquery/session.py

@@ -9,6 +9,7 @@ from sqlframe.bigquery.readwriter import (
     BigQueryDataFrameReader,
     BigQueryDataFrameWriter,
 )
+from sqlframe.bigquery.udf import BigQueryUDFRegistration
 
 if t.TYPE_CHECKING:
     from google.cloud.bigquery.client import Client as BigQueryClient
@@ -25,12 +26,14 @@ class BigQuerySession(
         BigQueryDataFrameWriter,
         BigQueryDataFrame,
         BigQueryConnection,
+        BigQueryUDFRegistration,
     ],
 ):
     _catalog = BigQueryCatalog
     _reader = BigQueryDataFrameReader
     _writer = BigQueryDataFrameWriter
     _df = BigQueryDataFrame
+    _udf_registration = BigQueryUDFRegistration
 
     QUALIFY_INFO_SCHEMA_WITH_DATABASE = True
     SANITIZE_COLUMN_NAMES = True

sqlframe-2.2.0/sqlframe/bigquery/udf.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlframe.base.udf import _BaseUDFRegistration
+
+if t.TYPE_CHECKING:
+    from sqlframe.bigquery.session import BigQuerySession
+
+
+class BigQueryUDFRegistration(_BaseUDFRegistration["BigQuerySession"]): ...

{sqlframe-2.0.0 → sqlframe-2.2.0}/sqlframe/duckdb/session.py

@@ -11,6 +11,7 @@ from sqlframe.duckdb.readwriter import (
     DuckDBDataFrameReader,
     DuckDBDataFrameWriter,
 )
+from sqlframe.duckdb.udf import DuckDBUDFRegistration
 
 if t.TYPE_CHECKING:
     from duckdb import DuckDBPyConnection
@@ -26,12 +27,14 @@ class DuckDBSession(
         DuckDBDataFrameWriter,
         DuckDBDataFrame,
         DuckDBPyConnection,
+        DuckDBUDFRegistration,
     ]
 ):
     _catalog = DuckDBCatalog
     _reader = DuckDBDataFrameReader
     _writer = DuckDBDataFrameWriter
     _df = DuckDBDataFrame
+    _udf_registration = DuckDBUDFRegistration
 
     def __init__(self, conn: t.Optional[DuckDBPyConnection] = None, *args, **kwargs):
         import duckdb

sqlframe-2.2.0/sqlframe/duckdb/udf.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlframe.base.udf import _BaseUDFRegistration
+
+if t.TYPE_CHECKING:
+    from sqlframe.base._typing import DataTypeOrString, UserDefinedFunctionLike
+    from sqlframe.duckdb.session import DuckDBSession
+
+
+class DuckDBUDFRegistration(_BaseUDFRegistration["DuckDBSession"]):
+    def register(  # type: ignore
+        self,
+        name: str,
+        f: t.Union[t.Callable[..., t.Any], UserDefinedFunctionLike],
+        returnType: t.Optional[DataTypeOrString] = None,
+    ) -> UserDefinedFunctionLike:
+        self.sparkSession._conn.create_function(name, f, return_type=returnType)  # type: ignore

{sqlframe-2.0.0 → sqlframe-2.2.0}/sqlframe/postgres/session.py

@@ -11,6 +11,7 @@ from sqlframe.postgres.readwriter import (
     PostgresDataFrameReader,
     PostgresDataFrameWriter,
 )
+from sqlframe.postgres.udf import PostgresUDFRegistration
 
 if t.TYPE_CHECKING:
     from psycopg2.extensions import connection as psycopg2_connection
@@ -27,12 +28,14 @@ class PostgresSession(
         PostgresDataFrameWriter,
         PostgresDataFrame,
         psycopg2_connection,
+        PostgresUDFRegistration,
     ],
 ):
     _catalog = PostgresCatalog
     _reader = PostgresDataFrameReader
     _writer = PostgresDataFrameWriter
     _df = PostgresDataFrame
+    _udf_registration = PostgresUDFRegistration
 
     def __init__(self, conn: t.Optional[psycopg2_connection] = None):
         if not hasattr(self, "_conn"):

sqlframe-2.2.0/sqlframe/postgres/udf.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlframe.base.udf import _BaseUDFRegistration
+
+if t.TYPE_CHECKING:
+    from sqlframe.postgres.session import PostgresSession
+
+
+class PostgresUDFRegistration(_BaseUDFRegistration["PostgresSession"]): ...

{sqlframe-2.0.0 → sqlframe-2.2.0}/sqlframe/redshift/session.py

@@ -10,6 +10,7 @@ from sqlframe.redshift.readwriter import (
     RedshiftDataFrameReader,
     RedshiftDataFrameWriter,
 )
+from sqlframe.redshift.udf import RedshiftUDFRegistration
 
 if t.TYPE_CHECKING:
     from redshift_connector.core import Connection as RedshiftConnection
@@ -24,12 +25,14 @@ class RedshiftSession(
         RedshiftDataFrameWriter,
         RedshiftDataFrame,
         RedshiftConnection,
+        RedshiftUDFRegistration,
     ],
 ):
     _catalog = RedshiftCatalog
     _reader = RedshiftDataFrameReader
     _writer = RedshiftDataFrameWriter
     _df = RedshiftDataFrame
+    _udf_registration = RedshiftUDFRegistration
 
     def __init__(self, conn: t.Optional[RedshiftConnection] = None):
         warnings.warn(

sqlframe-2.2.0/sqlframe/redshift/udf.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlframe.base.udf import _BaseUDFRegistration
+
+if t.TYPE_CHECKING:
+    from sqlframe.redshift.session import RedshiftSession
+
+
+class RedshiftUDFRegistration(_BaseUDFRegistration["RedshiftSession"]): ...

{sqlframe-2.0.0 → sqlframe-2.2.0}/sqlframe/snowflake/session.py

@@ -4,6 +4,8 @@ import json
 import typing as t
 import warnings
 
+from sqlframe.snowflake.udf import SnowflakeUDFRegistration
+
 try:
     from snowflake.connector.converter import SnowflakeConverter
 except ImportError:
@@ -50,12 +52,14 @@ class SnowflakeSession(
         SnowflakeDataFrameWriter,
         SnowflakeDataFrame,
         SnowflakeConnection,
+        SnowflakeUDFRegistration,
     ],
 ):
     _catalog = SnowflakeCatalog
     _reader = SnowflakeDataFrameReader
     _writer = SnowflakeDataFrameWriter
     _df = SnowflakeDataFrame
+    _udf_registration = SnowflakeUDFRegistration
 
     def __init__(self, conn: t.Optional[SnowflakeConnection] = None):
         import snowflake

sqlframe-2.2.0/sqlframe/snowflake/udf.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlframe.base.udf import _BaseUDFRegistration
+
+if t.TYPE_CHECKING:
+    from sqlframe.snowflake.session import SnowflakeSession
+
+
+class SnowflakeUDFRegistration(_BaseUDFRegistration["SnowflakeSession"]): ...

{sqlframe-2.0.0 → sqlframe-2.2.0}/sqlframe/spark/session.py

@@ -14,6 +14,7 @@ from sqlframe.spark.readwriter import (
     SparkDataFrameWriter,
 )
 from sqlframe.spark.types import Row
+from sqlframe.spark.udf import SparkUDFRegistration
 
 if t.TYPE_CHECKING:
     import pandas as pd
@@ -32,12 +33,14 @@ class SparkSession(
         SparkDataFrameWriter,
         SparkDataFrame,
         PySparkSession,
+        SparkUDFRegistration,
     ],
 ):
     _catalog = SparkCatalog
     _reader = SparkDataFrameReader
     _writer = SparkDataFrameWriter
     _df = SparkDataFrame
+    _udf_registration = SparkUDFRegistration
 
     def __init__(self, conn: t.Optional[PySparkSession] = None, *args, **kwargs):
         from pyspark.sql.session import DataFrame, SparkSession