sqlframe 3.8.2__py3-none-any.whl → 3.9.1__py3-none-any.whl

sqlframe/__init__.py

@@ -15,6 +15,7 @@ ENGINE_TO_PREFIX = {
     "redshift": "Redshift",
     "snowflake": "Snowflake",
     "spark": "Spark",
+    "databricks": "Databricks",
     "standalone": "Standalone",
 }
 

sqlframe/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.8.2'
-__version_tuple__ = version_tuple = (3, 8, 2)
+__version__ = version = '3.9.1'
+__version_tuple__ = version_tuple = (3, 9, 1)

sqlframe/base/dataframe.py

@@ -629,10 +629,10 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
                 # We will drop the "view" if it exists before running the cache table
                 output_expressions.append(exp.Drop(this=cache_table, exists=True, kind="VIEW"))
             elif expression_type == exp.Create:
-                expression = df.output_expression_container.copy()
+                expression = df.output_expression_container.copy()  # type: ignore
                 expression.set("expression", select_expression)
             elif expression_type == exp.Insert:
-                expression = df.output_expression_container.copy()
+                expression = df.output_expression_container.copy()  # type: ignore
                 select_without_ctes = select_expression.copy()
                 select_without_ctes.set("with", None)
                 expression.set("expression", select_without_ctes)
@@ -866,6 +866,7 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         if on is None:
             logger.warning("Got no value for on. This appears to change the join to a cross join.")
             how = "cross"
+
         other_df = other_df._convert_leaf_to_cte()
         join_expression = self._add_ctes_to_expression(self.expression, other_df.expression.ctes)
         # We will determine actual "join on" expression later so we don't provide it at first
@@ -875,6 +876,13 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         self_columns = self._get_outer_select_columns(join_expression)
         other_columns = self._get_outer_select_columns(other_df.expression)
         join_columns = self._ensure_and_normalize_cols(on)
+        # If the two dataframes being joined come from the same branch and reference the same table,
+        # then really one of the tables was supposed to reference the other dataframe so we update it to do that.
+        if self.branch_id == other_df.branch_id:
+            for col in join_columns:
+                for eq in col.expression.find_all(exp.EQ):
+                    if eq.this.table == eq.expression.table:
+                        eq.expression.set("table", exp.to_identifier(other_df.latest_cte_name))
         # Determines the join clause and select columns to be used passed on what type of columns were provided for
         # the join. The columns returned changes based on how the on expression is provided.
         if how != "cross":

sqlframe/databricks/__init__.py

@@ -0,0 +1,32 @@
+from sqlframe.databricks.catalog import DatabricksCatalog
+from sqlframe.databricks.column import Column
+from sqlframe.databricks.dataframe import (
+    DatabricksDataFrame,
+    DatabricksDataFrameNaFunctions,
+    DatabricksDataFrameStatFunctions,
+)
+from sqlframe.databricks.group import DatabricksGroupedData
+from sqlframe.databricks.readwriter import (
+    DatabricksDataFrameReader,
+    DatabricksDataFrameWriter,
+)
+from sqlframe.databricks.session import DatabricksSession
+from sqlframe.databricks.types import Row
+from sqlframe.databricks.udf import DatabricksUDFRegistration
+from sqlframe.databricks.window import Window, WindowSpec
+
+__all__ = [
+    "Column",
+    "Row",
+    "DatabricksCatalog",
+    "DatabricksDataFrame",
+    "DatabricksDataFrameNaFunctions",
+    "DatabricksGroupedData",
+    "DatabricksDataFrameReader",
+    "DatabricksDataFrameWriter",
+    "DatabricksSession",
+    "DatabricksDataFrameStatFunctions",
+    "DatabricksUDFRegistration",
+    "Window",
+    "WindowSpec",
+]

sqlframe/databricks/catalog.py

@@ -0,0 +1,302 @@
+# 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 fnmatch
+import json
+import typing as t
+
+from sqlglot import exp, parse_one
+
+from sqlframe.base.catalog import Column, Function, _BaseCatalog
+from sqlframe.base.mixins.catalog_mixins import (
+    GetCurrentCatalogFromFunctionMixin,
+    GetCurrentDatabaseFromFunctionMixin,
+    ListCatalogsFromInfoSchemaMixin,
+    ListDatabasesFromInfoSchemaMixin,
+    ListTablesFromInfoSchemaMixin,
+    SetCurrentCatalogFromUseMixin,
+    SetCurrentDatabaseFromUseMixin,
+)
+from sqlframe.base.util import normalize_string, schema_, to_schema
+
+if t.TYPE_CHECKING:
+    from sqlframe.databricks.session import DatabricksSession  # noqa
+    from sqlframe.databricks.dataframe import DatabricksDataFrame  # noqa
+
+
+class DatabricksCatalog(
+    SetCurrentCatalogFromUseMixin["DatabricksSession", "DatabricksDataFrame"],
+    GetCurrentCatalogFromFunctionMixin["DatabricksSession", "DatabricksDataFrame"],
+    GetCurrentDatabaseFromFunctionMixin["DatabricksSession", "DatabricksDataFrame"],
+    ListDatabasesFromInfoSchemaMixin["DatabricksSession", "DatabricksDataFrame"],
+    ListCatalogsFromInfoSchemaMixin["DatabricksSession", "DatabricksDataFrame"],
+    SetCurrentDatabaseFromUseMixin["DatabricksSession", "DatabricksDataFrame"],
+    ListTablesFromInfoSchemaMixin["DatabricksSession", "DatabricksDataFrame"],
+    _BaseCatalog["DatabricksSession", "DatabricksDataFrame"],
+):
+    CURRENT_CATALOG_EXPRESSION: exp.Expression = exp.func("current_catalog")
+    UPPERCASE_INFO_SCHEMA = True
+
+    def listFunctions(
+        self, dbName: t.Optional[str] = None, pattern: t.Optional[str] = None
+    ) -> t.List[Function]:
+        """
+        Returns a t.List of functions registered in the specified database.
+
+        .. versionadded:: 3.4.0
+
+        Parameters
+        ----------
+        dbName : str
+            name of the database to t.List the functions.
+            ``dbName`` can be qualified with catalog name.
+        pattern : str
+            The pattern that the function name needs to match.
+
+            .. versionchanged: 3.5.0
+                Adds ``pattern`` argument.
+
+        Returns
+        -------
+        t.List
+            A t.List of :class:`Function`.
+
+        Notes
+        -----
+        If no database is specified, the current database and catalog
+        are used. This API includes all temporary functions.
+
+        Examples
+        --------
+        >>> spark.catalog.t.listFunctions()
+        [Function(name=...
+
+        >>> spark.catalog.t.listFunctions(pattern="to_*")
+        [Function(name=...
+
+        >>> spark.catalog.t.listFunctions(pattern="*not_existing_func*")
+        []
+        """
+        if dbName is None:
+            schema = schema_(
+                db=exp.parse_identifier(
+                    self.currentDatabase(), dialect=self.session.output_dialect
+                ),
+                catalog=exp.parse_identifier(
+                    self.currentCatalog(), dialect=self.session.output_dialect
+                ),
+            )
+        else:
+            dbName = normalize_string(dbName, from_dialect="input", is_schema=True)
+            schema = to_schema(dbName, dialect=self.session.input_dialect)
+            if not schema.catalog:
+                schema.set(
+                    "catalog",
+                    exp.parse_identifier(
+                        self.currentCatalog(), dialect=self.session.output_dialect
+                    ),
+                )
+        current_catalog = self.currentCatalog()
+        self.session._collect(f"USE CATALOG {schema.catalog}")
+
+        query = parse_one(
+            f"""SHOW USER FUNCTIONS IN {schema.sql(dialect=self.session.input_dialect)}""",
+            dialect=self.session.input_dialect,
+        )
+        functions = [
+            Function(
+                name=normalize_string(x["function"], from_dialect="execution", to_dialect="output"),
+                catalog=normalize_string(
+                    schema.catalog, from_dialect="execution", to_dialect="output"
+                ),
+                namespace=[
+                    normalize_string(schema.db, from_dialect="execution", to_dialect="output")
+                ],
+                description=None,
+                className="",
+                isTemporary=False,
+            )
+            for x in self.session._collect(query)
+        ]
+        if pattern:
+            normalized_pattern = normalize_string(
+                pattern, from_dialect="input", to_dialect="output", is_pattern=True
+            )
+            functions = [x for x in functions if fnmatch.fnmatch(x.name, normalized_pattern)]
+        self.session._collect(f"USE CATALOG {current_catalog}")
+        return functions
+
+    def get_columns(self, table: exp.Table | str) -> t.Dict[str, exp.DataType]:
+        table = (
+            normalize_string(table, from_dialect="input", is_table=True)
+            if isinstance(table, str)
+            else table
+        )
+        table = exp.to_table(table, dialect=self.session.input_dialect)
+        if not table.catalog:
+            table.set(
+                "catalog",
+                exp.parse_identifier(
+                    normalize_string(
+                        self.currentCatalog(), from_dialect="output", to_dialect="input"
+                    ),
+                    dialect=self.session.input_dialect,
+                ),
+            )
+        if not table.db:
+            table.set(
+                "db",
+                exp.parse_identifier(
+                    normalize_string(
+                        self.currentDatabase(), from_dialect="output", to_dialect="input"
+                    ),
+                    dialect=self.session.input_dialect,
+                ),
+            )
+        sql = f"DESCRIBE TABLE {table.sql(dialect=self.session.input_dialect)}"
+        results = self.session._collect(sql)
+        return {
+            normalize_string(
+                row["col_name"],
+                from_dialect="execution",
+                to_dialect="output",
+                is_column=True,
+            ): exp.DataType.build(
+                normalize_string(
+                    row["data_type"],
+                    from_dialect="execution",
+                    to_dialect="output",
+                    is_datatype=True,
+                ),
+                dialect=self.session.output_dialect,
+                udt=True,
+            )
+            for row in results
+            if row["data_type"] != "" and row["data_type"] != "data_type"
+        }
+
+    def listColumns(
+        self, tableName: str, dbName: t.Optional[str] = None, include_temp: bool = False
+    ) -> t.List[Column]:
+        """Returns a t.List of columns for the given table/view in the specified database.
+
+        .. versionadded:: 2.0.0
+
+        Parameters
+        ----------
+        tableName : str
+            name of the table to t.List columns.
+
+            .. versionchanged:: 3.4.0
+               Allow ``tableName`` to be qualified with catalog name when ``dbName`` is None.
+
+        dbName : str, t.Optional
+            name of the database to find the table to t.List columns.
+
+        Returns
+        -------
+        t.List
+            A t.List of :class:`Column`.
+
+        Notes
+        -----
+        The order of arguments here is different from that of its JVM counterpart
+        because Python does not support method overloading.
+
+        If no database is specified, the current database and catalog
+        are used. This API includes all temporary views.
+
+        Examples
+        --------
+        >>> _ = spark.sql("DROP TABLE IF EXISTS tbl1")
+        >>> _ = spark.sql("CREATE TABLE tblA (name STRING, age INT) USING parquet")
+        >>> spark.catalog.t.listColumns("tblA")
+        [Column(name='name', description=None, dataType='string', nullable=True, ...
+        >>> _ = spark.sql("DROP TABLE tblA")
+        """
+        tableName = normalize_string(
+            tableName, from_dialect=self.session.input_dialect, is_table=True
+        )
+        dbName = (
+            normalize_string(dbName, from_dialect=self.session.input_dialect, is_schema=True)
+            if dbName
+            else None
+        )
+        if df := self.session.temp_views.get(tableName):
+            return [
+                Column(
+                    name=x,
+                    description=None,
+                    dataType="",
+                    nullable=True,
+                    isPartition=False,
+                    isBucket=False,
+                )
+                for x in df.columns
+            ]
+
+        table = exp.to_table(tableName, dialect=self.session.input_dialect)
+        schema = to_schema(dbName, dialect=self.session.input_dialect) if dbName else None
+        if not table.db:
+            if schema and schema.db:
+                table.set("db", schema.args["db"])
+            else:
+                current_database = normalize_string(
+                    self.currentDatabase(),
+                    from_dialect=self.session.output_dialect,
+                    to_dialect=self.session.input_dialect,
+                )
+                table.set(
+                    "db",
+                    exp.parse_identifier(current_database, dialect=self.session.input_dialect),
+                )
+        if not table.catalog:
+            if schema and schema.catalog:
+                table.set("catalog", schema.args["catalog"])
+            else:
+                current_catalog = normalize_string(
+                    self.currentCatalog(),
+                    from_dialect=self.session.output_dialect,
+                    to_dialect=self.session.input_dialect,
+                )
+                table.set(
+                    "catalog",
+                    exp.parse_identifier(current_catalog, dialect=self.session.input_dialect),
+                )
+        sql = f"DESCRIBE TABLE {'.'.join(part.name for part in table.parts)}"
+        results = self.session._collect(sql)
+
+        is_partition = False
+        partitions = set([])
+        for row in results:
+            if row["col_name"] == "# Partition Information":
+                is_partition = True
+            if is_partition and row["data_type"] != "" and row["data_type"] != "data_type":
+                partitions.add(row["col_name"])
+
+        columns = []
+        for row in results:
+            if row["data_type"] == "" or row["data_type"] == "data_type":
+                break
+            columns.append(
+                Column(
+                    name=normalize_string(
+                        row["col_name"],
+                        from_dialect=self.session.execution_dialect,
+                        to_dialect=self.session.output_dialect,
+                    ),
+                    description=row["comment"],
+                    dataType=normalize_string(
+                        row["data_type"],
+                        from_dialect=self.session.execution_dialect,
+                        to_dialect=self.session.output_dialect,
+                        is_datatype=True,
+                    ),
+                    nullable=True,
+                    isPartition=True if row["col_name"] in partitions else False,
+                    isBucket=False,
+                )
+            )
+
+        return columns

sqlframe/databricks/column.py

@@ -0,0 +1 @@
+from sqlframe.base.column import Column

sqlframe/databricks/dataframe.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import logging
+import sys
+import typing as t
+
+from sqlframe.base.catalog import Column as CatalogColumn
+from sqlframe.base.dataframe import (
+    _BaseDataFrame,
+    _BaseDataFrameNaFunctions,
+    _BaseDataFrameStatFunctions,
+)
+from sqlframe.base.mixins.dataframe_mixins import NoCachePersistSupportMixin
+from sqlframe.base.util import normalize_string
+from sqlframe.databricks.group import DatabricksGroupedData
+
+if t.TYPE_CHECKING:
+    from sqlframe.databricks.readwriter import DatabricksDataFrameWriter
+    from sqlframe.databricks.session import DatabricksSession
+
+
+logger = logging.getLogger(__name__)
+
+
+class DatabricksDataFrameNaFunctions(_BaseDataFrameNaFunctions["DatabricksDataFrame"]):
+    pass
+
+
+class DatabricksDataFrameStatFunctions(_BaseDataFrameStatFunctions["DatabricksDataFrame"]):
+    pass
+
+
+class DatabricksDataFrame(
+    NoCachePersistSupportMixin,
+    _BaseDataFrame[
+        "DatabricksSession",
+        "DatabricksDataFrameWriter",
+        "DatabricksDataFrameNaFunctions",
+        "DatabricksDataFrameStatFunctions",
+        "DatabricksGroupedData",
+    ],
+):
+    _na = DatabricksDataFrameNaFunctions
+    _stat = DatabricksDataFrameStatFunctions
+    _group_data = DatabricksGroupedData
+
+    @property
+    def _typed_columns(self) -> t.List[CatalogColumn]:
+        sql = self.session._to_sql(self.expression)
+        columns = []
+        for row in self.session._collect(r"DESCRIBE QUERY ({sql})".format(sql=sql)):
+            columns.append(
+                CatalogColumn(
+                    name=normalize_string(
+                        row.col_name, from_dialect="execution", to_dialect="output"
+                    ),
+                    dataType=normalize_string(
+                        row.data_type,
+                        from_dialect="execution",
+                        to_dialect="output",
+                        is_datatype=True,
+                    ),
+                    nullable=True,
+                    description=row.comment,
+                    isPartition=False,
+                    isBucket=False,
+                )
+            )
+        return columns

sqlframe/databricks/functions.py

@@ -0,0 +1,22 @@
+import inspect
+import sys
+
+import sqlframe.base.functions  # noqa
+
+module = sys.modules["sqlframe.base.functions"]
+globals().update(
+    {
+        name: func
+        for name, func in inspect.getmembers(module, inspect.isfunction)
+        if hasattr(func, "unsupported_engines") and "databricks" not in func.unsupported_engines
+    }
+)
+
+
+from sqlframe.base.function_alternatives import (  # noqa
+    percentile_without_disc as percentile,
+    add_months_by_multiplication as add_months,
+    arrays_overlap_renamed as arrays_overlap,
+    _is_string_using_typeof_string_lcase as _is_string,
+    try_element_at_zero_based as try_element_at,
+)

sqlframe/databricks/functions.pyi

@@ -0,0 +1,416 @@
+from sqlframe.base.function_alternatives import (  # noqa
+    percentile_without_disc as percentile,
+    add_months_by_multiplication as add_months,
+    arrays_overlap_renamed as arrays_overlap,
+    try_element_at_zero_based as try_element_at,
+)
+from sqlframe.base.functions import (
+    abs as abs,
+    acos as acos,
+    acosh as acosh,
+    aes_decrypt as aes_decrypt,
+    aes_encrypt as aes_encrypt,
+    aggregate as aggregate,
+    any_value as any_value,
+    approxCountDistinct as approxCountDistinct,
+    approx_count_distinct as approx_count_distinct,
+    approx_percentile as approx_percentile,
+    array as array,
+    array_agg as array_agg,
+    array_append as array_append,
+    array_compact as array_compact,
+    array_contains as array_contains,
+    array_distinct as array_distinct,
+    array_except as array_except,
+    array_insert as array_insert,
+    array_intersect as array_intersect,
+    array_join as array_join,
+    array_max as array_max,
+    array_min as array_min,
+    array_position as array_position,
+    array_prepend as array_prepend,
+    array_remove as array_remove,
+    array_repeat as array_repeat,
+    array_size as array_size,
+    array_sort as array_sort,
+    array_union as array_union,
+    arrays_zip as arrays_zip,
+    asc as asc,
+    asc_nulls_first as asc_nulls_first,
+    asc_nulls_last as asc_nulls_last,
+    ascii as ascii,
+    asin as asin,
+    asinh as asinh,
+    assert_true as assert_true,
+    atan as atan,
+    atan2 as atan2,
+    atanh as atanh,
+    avg as avg,
+    base64 as base64,
+    bin as bin,
+    bit_and as bit_and,
+    bit_count as bit_count,
+    bit_get as bit_get,
+    bit_length as bit_length,
+    bit_or as bit_or,
+    bit_xor as bit_xor,
+    bitmap_bit_position as bitmap_bit_position,
+    bitmap_bucket_number as bitmap_bucket_number,
+    bitmap_construct_agg as bitmap_construct_agg,
+    bitmap_count as bitmap_count,
+    bitmap_or_agg as bitmap_or_agg,
+    bitwiseNOT as bitwiseNOT,
+    bitwise_not as bitwise_not,
+    bool_and as bool_and,
+    bool_or as bool_or,
+    broadcast as broadcast,
+    bround as bround,
+    btrim as btrim,
+    bucket as bucket,
+    call_function as call_function,
+    cardinality as cardinality,
+    cbrt as cbrt,
+    ceil as ceil,
+    ceiling as ceiling,
+    char as char,
+    char_length as char_length,
+    character_length as character_length,
+    coalesce as coalesce,
+    col as col,
+    collect_list as collect_list,
+    collect_set as collect_set,
+    concat as concat,
+    concat_ws as concat_ws,
+    contains as contains,
+    conv as conv,
+    convert_timezone as convert_timezone,
+    corr as corr,
+    cos as cos,
+    cosh as cosh,
+    cot as cot,
+    count as count,
+    countDistinct as countDistinct,
+    count_distinct as count_distinct,
+    count_if as count_if,
+    count_min_sketch as count_min_sketch,
+    covar_pop as covar_pop,
+    covar_samp as covar_samp,
+    crc32 as crc32,
+    create_map as create_map,
+    csc as csc,
+    cume_dist as cume_dist,
+    curdate as curdate,
+    current_catalog as current_catalog,
+    current_database as current_database,
+    current_date as current_date,
+    current_schema as current_schema,
+    current_timestamp as current_timestamp,
+    current_timezone as current_timezone,
+    current_user as current_user,
+    date_add as date_add,
+    date_diff as date_diff,
+    date_format as date_format,
+    date_from_unix_date as date_from_unix_date,
+    date_part as date_part,
+    date_sub as date_sub,
+    date_trunc as date_trunc,
+    dateadd as dateadd,
+    datediff as datediff,
+    datepart as datepart,
+    day as day,
+    dayofmonth as dayofmonth,
+    dayofweek as dayofweek,
+    dayofyear as dayofyear,
+    days as days,
+    decode as decode,
+    degrees as degrees,
+    dense_rank as dense_rank,
+    desc as desc,
+    desc_nulls_first as desc_nulls_first,
+    desc_nulls_last as desc_nulls_last,
+    e as e,
+    element_at as element_at,
+    elt as elt,
+    encode as encode,
+    endswith as endswith,
+    equal_null as equal_null,
+    every as every,
+    exists as exists,
+    exp as exp,
+    explode as explode,
+    explode_outer as explode_outer,
+    expm1 as expm1,
+    expr as expr,
+    extract as extract,
+    factorial as factorial,
+    filter as filter,
+    find_in_set as find_in_set,
+    first as first,
+    first_value as first_value,
+    flatten as flatten,
+    floor as floor,
+    forall as forall,
+    format_number as format_number,
+    format_string as format_string,
+    from_csv as from_csv,
+    from_json as from_json,
+    from_unixtime as from_unixtime,
+    from_utc_timestamp as from_utc_timestamp,
+    get as get,
+    get_active_spark_context as get_active_spark_context,
+    get_json_object as get_json_object,
+    getbit as getbit,
+    greatest as greatest,
+    grouping as grouping,
+    grouping_id as grouping_id,
+    hash as hash,
+    hex as hex,
+    histogram_numeric as histogram_numeric,
+    hll_sketch_agg as hll_sketch_agg,
+    hll_sketch_estimate as hll_sketch_estimate,
+    hll_union as hll_union,
+    hll_union_agg as hll_union_agg,
+    hour as hour,
+    hours as hours,
+    hypot as hypot,
+    ifnull as ifnull,
+    ilike as ilike,
+    initcap as initcap,
+    inline as inline,
+    inline_outer as inline_outer,
+    input_file_name as input_file_name,
+    instr as instr,
+    isnan as isnan,
+    isnotnull as isnotnull,
+    isnull as isnull,
+    java_method as java_method,
+    json_array_length as json_array_length,
+    json_object_keys as json_object_keys,
+    json_tuple as json_tuple,
+    kurtosis as kurtosis,
+    lag as lag,
+    last as last,
+    last_day as last_day,
+    last_value as last_value,
+    lcase as lcase,
+    lead as lead,
+    least as least,
+    left as left,
+    length as length,
+    levenshtein as levenshtein,
+    like as like,
+    lit as lit,
+    ln as ln,
+    localtimestamp as localtimestamp,
+    locate as locate,
+    log as log,
+    log10 as log10,
+    log1p as log1p,
+    log2 as log2,
+    lower as lower,
+    lpad as lpad,
+    ltrim as ltrim,
+    make_date as make_date,
+    make_dt_interval as make_dt_interval,
+    make_interval as make_interval,
+    make_timestamp as make_timestamp,
+    make_timestamp_ltz as make_timestamp_ltz,
+    make_timestamp_ntz as make_timestamp_ntz,
+    make_ym_interval as make_ym_interval,
+    map_concat as map_concat,
+    map_contains_key as map_contains_key,
+    map_entries as map_entries,
+    map_filter as map_filter,
+    map_from_arrays as map_from_arrays,
+    map_from_entries as map_from_entries,
+    map_keys as map_keys,
+    map_values as map_values,
+    map_zip_with as map_zip_with,
+    mask as mask,
+    max as max,
+    max_by as max_by,
+    md5 as md5,
+    mean as mean,
+    median as median,
+    min as min,
+    min_by as min_by,
+    minute as minute,
+    mode as mode,
+    monotonically_increasing_id as monotonically_increasing_id,
+    month as month,
+    months as months,
+    months_between as months_between,
+    named_struct as named_struct,
+    nanvl as nanvl,
+    negate as negate,
+    negative as negative,
+    next_day as next_day,
+    now as now,
+    nth_value as nth_value,
+    ntile as ntile,
+    nullif as nullif,
+    nvl as nvl,
+    nvl2 as nvl2,
+    octet_length as octet_length,
+    overlay as overlay,
+    parse_url as parse_url,
+    percent_rank as percent_rank,
+    percentile_approx as percentile_approx,
+    pi as pi,
+    pmod as pmod,
+    posexplode as posexplode,
+    posexplode_outer as posexplode_outer,
+    position as position,
+    positive as positive,
+    pow as pow,
+    power as power,
+    printf as printf,
+    quarter as quarter,
+    radians as radians,
+    raise_error as raise_error,
+    rand as rand,
+    randn as randn,
+    rank as rank,
+    reduce as reduce,
+    reflect as reflect,
+    regexp as regexp,
+    regexp_count as regexp_count,
+    regexp_extract as regexp_extract,
+    regexp_extract_all as regexp_extract_all,
+    regexp_instr as regexp_instr,
+    regexp_like as regexp_like,
+    regexp_replace as regexp_replace,
+    regexp_substr as regexp_substr,
+    regr_avgx as regr_avgx,
+    regr_avgy as regr_avgy,
+    regr_count as regr_count,
+    regr_intercept as regr_intercept,
+    regr_r2 as regr_r2,
+    regr_slope as regr_slope,
+    regr_sxx as regr_sxx,
+    regr_sxy as regr_sxy,
+    regr_syy as regr_syy,
+    repeat as repeat,
+    replace as replace,
+    reverse as reverse,
+    right as right,
+    rint as rint,
+    rlike as rlike,
+    round as round,
+    row_number as row_number,
+    rpad as rpad,
+    rtrim as rtrim,
+    schema_of_csv as schema_of_csv,
+    schema_of_json as schema_of_json,
+    sec as sec,
+    second as second,
+    sentences as sentences,
+    sequence as sequence,
+    sha as sha,
+    sha1 as sha1,
+    sha2 as sha2,
+    shiftLeft as shiftLeft,
+    shiftRight as shiftRight,
+    shiftRightUnsigned as shiftRightUnsigned,
+    shiftleft as shiftleft,
+    shiftright as shiftright,
+    shiftrightunsigned as shiftrightunsigned,
+    shuffle as shuffle,
+    sign as sign,
+    signum as signum,
+    sin as sin,
+    sinh as sinh,
+    size as size,
+    skewness as skewness,
+    slice as slice,
+    some as some,
+    sort_array as sort_array,
+    soundex as soundex,
+    spark_partition_id as spark_partition_id,
+    split as split,
+    split_part as split_part,
+    sqrt as sqrt,
+    stack as stack,
+    startswith as startswith,
+    std as std,
+    stddev as stddev,
+    stddev_pop as stddev_pop,
+    stddev_samp as stddev_samp,
+    str_to_map as str_to_map,
+    struct as struct,
+    substr as substr,
+    substring as substring,
+    substring_index as substring_index,
+    sum as sum,
+    sumDistinct as sumDistinct,
+    sum_distinct as sum_distinct,
+    tan as tan,
+    tanh as tanh,
+    timestamp_micros as timestamp_micros,
+    timestamp_millis as timestamp_millis,
+    timestamp_seconds as timestamp_seconds,
+    toDegrees as toDegrees,
+    toRadians as toRadians,
+    to_binary as to_binary,
+    to_char as to_char,
+    to_csv as to_csv,
+    to_date as to_date,
+    to_json as to_json,
+    to_number as to_number,
+    to_timestamp as to_timestamp,
+    to_timestamp_ltz as to_timestamp_ltz,
+    to_timestamp_ntz as to_timestamp_ntz,
+    to_unix_timestamp as to_unix_timestamp,
+    to_utc_timestamp as to_utc_timestamp,
+    to_varchar as to_varchar,
+    transform as transform,
+    transform_keys as transform_keys,
+    transform_values as transform_values,
+    translate as translate,
+    trim as trim,
+    trunc as trunc,
+    try_add as try_add,
+    try_aes_decrypt as try_aes_decrypt,
+    try_avg as try_avg,
+    try_divide as try_divide,
+    try_multiply as try_multiply,
+    try_subtract as try_subtract,
+    try_sum as try_sum,
+    try_to_binary as try_to_binary,
+    try_to_number as try_to_number,
+    try_to_timestamp as try_to_timestamp,
+    typeof as typeof,
+    ucase as ucase,
+    unbase64 as unbase64,
+    unhex as unhex,
+    unix_date as unix_date,
+    unix_micros as unix_micros,
+    unix_millis as unix_millis,
+    unix_seconds as unix_seconds,
+    unix_timestamp as unix_timestamp,
+    upper as upper,
+    url_decode as url_decode,
+    url_encode as url_encode,
+    user as user,
+    var_pop as var_pop,
+    var_samp as var_samp,
+    variance as variance,
+    version as version,
+    weekday as weekday,
+    weekofyear as weekofyear,
+    when as when,
+    width_bucket as width_bucket,
+    xpath as xpath,
+    xpath_boolean as xpath_boolean,
+    xpath_double as xpath_double,
+    xpath_float as xpath_float,
+    xpath_int as xpath_int,
+    xpath_long as xpath_long,
+    xpath_number as xpath_number,
+    xpath_short as xpath_short,
+    xpath_string as xpath_string,
+    xxhash64 as xxhash64,
+    year as year,
+    years as years,
+    zip_with as zip_with,
+)

sqlframe/databricks/group.py

@@ -0,0 +1,14 @@
+# 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
+
+from sqlframe.base.group import _BaseGroupedData
+
+if t.TYPE_CHECKING:
+    from sqlframe.databricks.dataframe import DatabricksDataFrame
+
+
+class DatabricksGroupedData(_BaseGroupedData["DatabricksDataFrame"]):
+    pass

sqlframe/databricks/readwriter.py

@@ -0,0 +1,96 @@
+# 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 sys
+import typing as t
+
+import sqlglot as sg
+from sqlglot import exp
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+from sqlframe.base.mixins.readwriter_mixins import PandasLoaderMixin, PandasWriterMixin
+from sqlframe.base.readerwriter import (
+    _BaseDataFrameReader,
+    _BaseDataFrameWriter,
+)
+from sqlframe.base.util import normalize_string
+
+if t.TYPE_CHECKING:
+    from sqlframe.databricks.session import DatabricksSession  # noqa
+    from sqlframe.databricks.dataframe import DatabricksDataFrame  # noqa
+
+
+class DatabricksDataFrameReader(
+    PandasLoaderMixin["DatabricksSession", "DatabricksDataFrame"],
+    _BaseDataFrameReader["DatabricksSession", "DatabricksDataFrame"],
+):
+    pass
+
+
+class DatabricksDataFrameWriter(
+    PandasWriterMixin["DatabricksSession", "DatabricksDataFrame"],
+    _BaseDataFrameWriter["DatabricksSession", "DatabricksDataFrame"],
+):
+    def saveAsTable(
+        self,
+        name: str,
+        format: t.Optional[str] = None,
+        mode: t.Optional[str] = None,
+        partitionBy: t.Optional[t.Union[str, t.List[str]]] = None,
+        clusterBy: t.Optional[t.Union[str, t.List[str]]] = None,
+        **options,
+    ) -> Self:
+        if format is not None:
+            raise NotImplementedError("Providing Format in the save as table is not supported")
+        exists, replace, mode = None, None, mode or str(self._mode)
+        if mode == "append":
+            return self.insertInto(name)
+        if mode == "ignore":
+            exists = True
+        if mode == "overwrite":
+            replace = True
+        name = normalize_string(name, from_dialect="input", is_table=True)
+
+        properties: t.List[exp.Expression] = []
+        if partitionBy is not None:
+            if isinstance(partitionBy, str):
+                partition_by = [partitionBy]
+            else:
+                partition_by = partitionBy
+            properties.append(
+                exp.PartitionedByProperty(
+                    this=exp.Tuple(expressions=list(map(sg.to_identifier, partition_by)))
+                )
+            )
+        if clusterBy is not None:
+            if isinstance(clusterBy, str):
+                cluster_by = [clusterBy]
+            else:
+                cluster_by = clusterBy
+            properties.append(
+                exp.Cluster(
+                    expressions=[exp.Tuple(expressions=list(map(sg.to_identifier, cluster_by)))]
+                )
+            )
+
+        properties.extend(
+            exp.Property(this=sg.to_identifier(name), value=exp.convert(value))
+            for name, value in (options or {}).items()
+        )
+
+        output_expression_container = exp.Create(
+            this=exp.to_table(name, dialect=self._session.input_dialect),
+            kind="TABLE",
+            exists=exists,
+            replace=replace,
+            properties=exp.Properties(expressions=properties),
+        )
+        df = self._df.copy(output_expression_container=output_expression_container)
+        if self._session._has_connection:
+            df.collect()
+        return self.copy(_df=df)

sqlframe/databricks/session.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import typing as t
+import warnings
+
+from sqlframe.base.session import _BaseSession
+from sqlframe.databricks.catalog import DatabricksCatalog
+from sqlframe.databricks.dataframe import DatabricksDataFrame
+from sqlframe.databricks.readwriter import (
+    DatabricksDataFrameReader,
+    DatabricksDataFrameWriter,
+)
+from sqlframe.databricks.udf import DatabricksUDFRegistration
+
+if t.TYPE_CHECKING:
+    from databricks.sql.client import Connection as DatabricksConnection
+else:
+    DatabricksConnection = t.Any
+
+
+class DatabricksSession(
+    _BaseSession[  # type: ignore
+        DatabricksCatalog,
+        DatabricksDataFrameReader,
+        DatabricksDataFrameWriter,
+        DatabricksDataFrame,
+        DatabricksConnection,
+        DatabricksUDFRegistration,
+    ],
+):
+    _catalog = DatabricksCatalog
+    _reader = DatabricksDataFrameReader
+    _writer = DatabricksDataFrameWriter
+    _df = DatabricksDataFrame
+    _udf_registration = DatabricksUDFRegistration
+
+    def __init__(
+        self,
+        conn: t.Optional[DatabricksConnection] = None,
+        server_hostname: t.Optional[str] = None,
+        http_path: t.Optional[str] = None,
+        access_token: t.Optional[str] = None,
+    ):
+        from databricks import sql
+
+        if not hasattr(self, "_conn"):
+            super().__init__(conn or sql.connect(server_hostname, http_path, access_token))
+
+    class Builder(_BaseSession.Builder):
+        DEFAULT_EXECUTION_DIALECT = "databricks"
+
+        @property
+        def session(self) -> DatabricksSession:
+            return DatabricksSession(**self._session_kwargs)
+
+        def getOrCreate(self) -> DatabricksSession:
+            return super().getOrCreate()  # type: ignore
+
+    builder = Builder()

sqlframe/databricks/types.py

@@ -0,0 +1 @@
+from sqlframe.base.types import *

sqlframe/databricks/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.databricks.session import DatabricksSession
+
+
+class DatabricksUDFRegistration(_BaseUDFRegistration["DatabricksSession"]): ...

sqlframe/databricks/window.py

@@ -0,0 +1 @@
+from sqlframe.base.window import *

{sqlframe-3.8.2.dist-info → sqlframe-3.9.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 3.8.2
+Version: 3.9.1
 Summary: Turning PySpark Into a Universal DataFrame API
 Home-page: https://github.com/eakmanrq/sqlframe
 Author: Ryan Eakman
@@ -18,16 +18,18 @@ Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: prettytable (<3.12.1)
-Requires-Dist: sqlglot (<25.32,>=24.0.0)
+Requires-Dist: sqlglot (<25.33,>=24.0.0)
 Requires-Dist: typing-extensions (<5,>=4.8)
 Provides-Extra: bigquery
 Requires-Dist: google-cloud-bigquery-storage (<3,>=2) ; extra == 'bigquery'
 Requires-Dist: google-cloud-bigquery[pandas] (<4,>=3) ; extra == 'bigquery'
+Provides-Extra: databricks
+Requires-Dist: databricks-sql-connector (<4,>=3.6) ; extra == 'databricks'
 Provides-Extra: dev
 Requires-Dist: duckdb (<1.2,>=0.9) ; extra == 'dev'
 Requires-Dist: findspark (<3,>=2) ; extra == 'dev'
 Requires-Dist: mypy (<1.14,>=1.10.0) ; extra == 'dev'
-Requires-Dist: openai (<1.55,>=1.30) ; extra == 'dev'
+Requires-Dist: openai (<1.56,>=1.30) ; extra == 'dev'
 Requires-Dist: pandas-stubs (<3,>=2) ; extra == 'dev'
 Requires-Dist: pandas (<3,>=2) ; extra == 'dev'
 Requires-Dist: psycopg (<4,>=3.1) ; extra == 'dev'
@@ -36,7 +38,7 @@ Requires-Dist: pyspark (<3.6,>=2) ; extra == 'dev'
 Requires-Dist: pytest-postgresql (<7,>=6) ; extra == 'dev'
 Requires-Dist: pytest-xdist (<3.7,>=3.6) ; extra == 'dev'
 Requires-Dist: pytest (<8.4,>=8.2.0) ; extra == 'dev'
-Requires-Dist: ruff (<0.8,>=0.4.4) ; extra == 'dev'
+Requires-Dist: ruff (<0.9,>=0.4.4) ; extra == 'dev'
 Requires-Dist: types-psycopg2 (<3,>=2.9) ; extra == 'dev'
 Requires-Dist: pre-commit (>=3.5) ; (python_version == "3.8") and extra == 'dev'
 Requires-Dist: pre-commit (<4.1,>=3.7) ; (python_version >= "3.9") and extra == 'dev'
@@ -50,7 +52,7 @@ Provides-Extra: duckdb
 Requires-Dist: duckdb (<1.2,>=0.9) ; extra == 'duckdb'
 Requires-Dist: pandas (<3,>=2) ; extra == 'duckdb'
 Provides-Extra: openai
-Requires-Dist: openai (<1.55,>=1.30) ; extra == 'openai'
+Requires-Dist: openai (<1.56,>=1.30) ; extra == 'openai'
 Provides-Extra: pandas
 Requires-Dist: pandas (<3,>=2) ; extra == 'pandas'
 Provides-Extra: postgres
@@ -76,6 +78,11 @@ SQLFrame currently supports the following engines (many more in development):
 * [Snowflake](https://sqlframe.readthedocs.io/en/stable/snowflake)
 * [Spark](https://sqlframe.readthedocs.io/en/stable/spark)
 
+There are also two engines in development. These engines lack test coverage and robust documentation, but are available for early testing:
+
+* [Redshift](https://sqlframe.readthedocs.io/en/stable/redshift)
+* [Databricks](https://sqlframe.readthedocs.io/en/stable/databricks)
+
 SQLFrame also has a "Standalone" session that be used to generate SQL without any connection to a database engine.
 
 * [Standalone](https://sqlframe.readthedocs.io/en/stable/standalone)
@@ -100,6 +107,10 @@ pip install "sqlframe[postgres]"
 pip install "sqlframe[snowflake]"
 # Spark
 pip install "sqlframe[spark]"
+# Redshift (in development)
+pip install "sqlframe[redshift]"
+# Databricks (in development)
+pip install "sqlframe[databricks]"
 # Standalone
 pip install sqlframe
 ```

{sqlframe-3.8.2.dist-info → sqlframe-3.9.1.dist-info}/RECORD

@@ -1,10 +1,10 @@
-sqlframe/__init__.py,sha256=E3qCJ4PSEgKz6Hg3ves6LWt3JrQOV8c9HVLSIUOzKNc,3106
-sqlframe/_version.py,sha256=bV13UGjgKKBESjgpeQicXd0eKXwIF_qqxftysbl4rqs,411
+sqlframe/__init__.py,sha256=wfqm98eLoLid9oV_FzzpG5loKC6LxOhj2lXpfN7SARo,3138
+sqlframe/_version.py,sha256=juEjthQzCmK7Szp2cv6Mn7hKjzVyOZ5uyw49uPvQDw8,411
 sqlframe/base/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sqlframe/base/_typing.py,sha256=b2clI5HI1zEZKB_3Msx3FeAJQyft44ubUifJwQRVXyQ,1298
 sqlframe/base/catalog.py,sha256=SzFQalTWdhWzxUY-4ut1f9TfOECp_JmJEgNPfrRKCe0,38457
 sqlframe/base/column.py,sha256=06fhVZ2nCn2QLxnfjdK-oYKeTFJC_smgSxu7u2UYlVg,17878
-sqlframe/base/dataframe.py,sha256=coeUwntwYbT1g6YKVwk3ZfWMfJqAzd1ECYabBSsNsV0,72892
+sqlframe/base/dataframe.py,sha256=sPSffZHmFUA_uwQXnwhUhwx4xtXorRXRipATmklSXOw,73442
 sqlframe/base/decorators.py,sha256=Jy4bf8MhZ-AJ6CWTj59bBJRqamtLbPC0USUMFrY6g0w,449
 sqlframe/base/exceptions.py,sha256=9Uwvqn2eAkDpqm4BrRgbL61qM-GMCbJEMAW8otxO46s,370
 sqlframe/base/function_alternatives.py,sha256=jofb2-nweefqcjUsd4xVqfRmJSZ-T_0Iq5roW2pL0OA,50768
@@ -35,6 +35,18 @@ sqlframe/bigquery/session.py,sha256=uSiEWWiDEryq3gIJJUmsu1DIalRGomNiymVulxt439c,
 sqlframe/bigquery/types.py,sha256=KwNyuXIo-2xVVd4bZED3YrQOobKCtemlxGrJL7DrTC8,34
 sqlframe/bigquery/udf.py,sha256=ZZ1-P1zWZhQqmhBqwAxfNeKl31nDkkZgkuz7Dn28P_0,264
 sqlframe/bigquery/window.py,sha256=6GKPzuxeSapJakBaKBeT9VpED1ACdjggDv9JRILDyV0,35
+sqlframe/databricks/__init__.py,sha256=BkB_eO1UYwcf8j6x7bi4BWmDCMkfn0CUMwossWgwaG4,993
+sqlframe/databricks/catalog.py,sha256=LpAn2UahzyIC3VKvQ8_6OVevaxFS_mf2HM2K4e4azZ8,11309
+sqlframe/databricks/column.py,sha256=E1tUa62Y5HajkhgFuebU9zohrGyieudcHzTT8gfalio,40
+sqlframe/databricks/dataframe.py,sha256=2VbzMs6ehrXUzgFCKNCARNMVO-FNnmEeku8hPIv948c,2141
+sqlframe/databricks/functions.py,sha256=rDTnOyj_SKm6mzFtL4GraDNXRnbYyfpijtn8dOQDmog,640
+sqlframe/databricks/functions.pyi,sha256=jiZr-EGGuXEodEeKq56MwibcXBk4Lpy4H3brFd0DDVA,11628
+sqlframe/databricks/group.py,sha256=dU3g0DVLRlfOSCamKchQFXRd1WTFbdxoXkpEX8tPD6Y,399
+sqlframe/databricks/readwriter.py,sha256=Lter5V3y6YqDSyv6FOrLF9JRfM6f5kIgB8AC-4nfJJo,3285
+sqlframe/databricks/session.py,sha256=BOpYMy2bgtYZ7XvEhQD2gqzl5XHEvFHt3R2V40y2isI,1752
+sqlframe/databricks/types.py,sha256=KwNyuXIo-2xVVd4bZED3YrQOobKCtemlxGrJL7DrTC8,34
+sqlframe/databricks/udf.py,sha256=3rmxv_6zSLfIxH8P8P050ZO-ki0aqBb9wWuUQBtl4m8,272
+sqlframe/databricks/window.py,sha256=6GKPzuxeSapJakBaKBeT9VpED1ACdjggDv9JRILDyV0,35
 sqlframe/duckdb/__init__.py,sha256=KAw_uZEhFMwi3D9Wj6AgHAKqLNk-EAx2uDIYu56oL44,872
 sqlframe/duckdb/catalog.py,sha256=YYYVmetLUaJOdObKw4AJ7L0P-msshkta4xHlcZQ9zEA,4795
 sqlframe/duckdb/column.py,sha256=E1tUa62Y5HajkhgFuebU9zohrGyieudcHzTT8gfalio,40
@@ -107,8 +119,8 @@ sqlframe/standalone/udf.py,sha256=azmgtUjHNIPs0WMVNId05SHwiYn41MKVBhKXsQJ5dmY,27
 sqlframe/standalone/window.py,sha256=6GKPzuxeSapJakBaKBeT9VpED1ACdjggDv9JRILDyV0,35
 sqlframe/testing/__init__.py,sha256=VVCosQhitU74A3NnE52O4mNtGZONapuEXcc20QmSlnQ,132
 sqlframe/testing/utils.py,sha256=9DDYVuocO7tygee3RaajuJNZ24sJwf_LY556kKg7kTw,13011
-sqlframe-3.8.2.dist-info/LICENSE,sha256=VZu79YgW780qxaFJMr0t5ZgbOYEh04xWoxaWOaqIGWk,1068
-sqlframe-3.8.2.dist-info/METADATA,sha256=h-5QphoAVwjsAkYc7nS9tHBxFDLHxWDmHKcCNGqN89A,8639
-sqlframe-3.8.2.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-sqlframe-3.8.2.dist-info/top_level.txt,sha256=T0_RpoygaZSF6heeWwIDQgaP0varUdSK1pzjeJZRjM8,9
-sqlframe-3.8.2.dist-info/RECORD,,
+sqlframe-3.9.1.dist-info/LICENSE,sha256=VZu79YgW780qxaFJMr0t5ZgbOYEh04xWoxaWOaqIGWk,1068
+sqlframe-3.9.1.dist-info/METADATA,sha256=IKeqsU6mGTjVsFUwUVMY4xDpzIEmIhNC3ox0uWaDkQs,9142
+sqlframe-3.9.1.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+sqlframe-3.9.1.dist-info/top_level.txt,sha256=T0_RpoygaZSF6heeWwIDQgaP0varUdSK1pzjeJZRjM8,9
+sqlframe-3.9.1.dist-info/RECORD,,