sqlframe 3.35.0__tar.gz → 3.36.0__tar.gz

{sqlframe-3.35.0 → sqlframe-3.36.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 3.35.0
+Version: 3.36.0
 Summary: Turning PySpark Into a Universal DataFrame API
 Home-page: https://github.com/eakmanrq/sqlframe
 Author: Ryan Eakman

{sqlframe-3.35.0 → sqlframe-3.36.0}/docs/databricks.md

@@ -103,6 +103,29 @@ A DatabricksSession, which implements the PySpark Session API, is created by pas
     session = SparkSession.builder.getOrCreate()
     ```
 
+### Creating Session with Idle Connections
+
+The Databricks SQL Connector for Python will automatically close connections that have been idle for a while. 
+This will cause errors when using SQLFrame since it will retry to use a closed connection.
+To avoid this, you can have SQLFrame create the connection for you and it will automatically reconnect when needed.
+Note that this will not work with the `activate` function since it requires a `databricks.sql.client.Connection` object.
+
+```python
+import os
+
+from sqlframe.databricks import DatabricksSession
+
+session = DatabricksSession(
+    server_hostname="dbc-xxxxxxxx-xxxx.cloud.databricks.com",
+    http_path="/sql/1.0/warehouses/xxxxxxxxxxxxxxxx",
+    access_token=os.environ["ACCESS_TOKEN"],  # Replace this with how you get your databricks access token
+    auth_type="access_token",
+    catalog="catalog",
+    schema="schema",
+)
+```
+
+
 ## Example Usage
 
 ```python

{sqlframe-3.35.0 → sqlframe-3.36.0}/docs/snowflake.md

@@ -544,6 +544,7 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [mean](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.GroupedData.mean.html)
 * [min](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.GroupedData.min.html)
 * [pivot](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.GroupedData.pivot.html)
+  * Doesn't support multiple aggregate functions on a single pivot
 * [sum](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.GroupedData.sum.html)
 
 ### DataFrameReader Class

{sqlframe-3.35.0 → sqlframe-3.36.0}/mkdocs.yml

@@ -3,14 +3,14 @@ repo_url: https://github.com/eakmanrq/sqlframe
 repo_name: eakmanrq/sqlframe
 nav:
   - "Overview": index.md
+  - "Configuration": configuration.md
   - "BigQuery": bigquery.md
+  - "Databricks": databricks.md
   - "DuckDB": duckdb.md
   - "Postgres": postgres.md
   - "Spark": spark.md
   - "Standalone": standalone.md
-  - "Configuration": configuration.md
   - "Redshift (In-Development)": redshift.md
-  - "Databricks (In-Development)": databricks.md
 theme:
   name: material
   logo: images/SF.png

{sqlframe-3.35.0 → sqlframe-3.36.0}/setup.py

@@ -20,7 +20,7 @@ setup(
     python_requires=">=3.9",
     install_requires=[
         "prettytable<4",
-        "sqlglot>=24.0.0,<26.25",
+        "sqlglot>=24.0.0,<26.32",
         "typing_extensions",
     ],
     extras_require={
@@ -38,12 +38,12 @@ setup(
             "psycopg>=3.1,<4",
             "pyarrow>=10,<21",
             "pyspark>=2,<3.6",
-            "pytest>=8.2.0,<8.4",
+            "pytest>=8.2.0,<8.5",
             "pytest-forked",
             "pytest-postgresql>=6,<8",
             "pytest-xdist>=3.6,<3.8",
             "pre-commit>=3.7,<5",
-            "ruff>=0.4.4,<0.12",
+            "ruff>=0.4.4,<0.13",
             "types-psycopg2>=2.9,<3",
         ],
         "docs": [

{sqlframe-3.35.0 → sqlframe-3.36.0}/sqlframe/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.35.0'
-__version_tuple__ = version_tuple = (3, 35, 0)
+__version__ = version = '3.36.0'
+__version_tuple__ = version_tuple = (3, 36, 0)

{sqlframe-3.35.0 → sqlframe-3.36.0}/sqlframe/base/function_alternatives.py

@@ -999,27 +999,6 @@ def element_at_using_brackets(col: ColumnOrName, value: ColumnOrLiteral) -> Colu
     )
 
 
-def array_remove_using_filter(col: ColumnOrName, value: ColumnOrLiteral) -> Column:
-    lit = get_func_from_session("lit")
-    col_func = get_func_from_session("col")
-
-    value = value if isinstance(value, Column) else lit(value)
-    return Column(
-        expression.Anonymous(
-            this="LIST_FILTER",
-            expressions=[
-                col_func(col).column_expression,
-                expression.Lambda(
-                    this=expression.NEQ(
-                        this=expression.Identifier(this="x"), expression=value.column_expression
-                    ),
-                    expressions=[expression.Identifier(this="x")],
-                ),
-            ],
-        )
-    )
-
-
 def array_union_using_list_concat(col1: ColumnOrName, col2: ColumnOrName) -> Column:
     col_func = get_func_from_session("col")
 
@@ -1321,10 +1300,6 @@ def day_with_try_to_timestamp(col: ColumnOrName) -> Column:
     )
 
 
-def endswith_with_underscore(str: ColumnOrName, suffix: ColumnOrName) -> Column:
-    return Column.invoke_anonymous_function(str, "ENDS_WITH", suffix)
-
-
 def endswith_using_like(str: ColumnOrName, suffix: ColumnOrName) -> Column:
     concat = get_func_from_session("concat")
     lit = get_func_from_session("lit")
@@ -1664,41 +1639,6 @@ def array_position_bgutil(col: ColumnOrName, value: ColumnOrLiteral) -> Column:
     )
 
 
-def array_remove_bgutil(col: ColumnOrName, value: ColumnOrLiteral) -> Column:
-    lit = get_func_from_session("lit")
-
-    value_col = value if isinstance(value, Column) else lit(value)
-
-    filter_subquery = expression.select(
-        "*",
-    ).from_(
-        expression.Unnest(
-            expressions=[Column.ensure_col(col).column_expression],
-            alias=expression.TableAlias(
-                columns=[expression.to_identifier("x")],
-            ),
-        )
-    )
-
-    agg_subquery = (
-        expression.select(
-            expression.Anonymous(
-                this="ARRAY_AGG",
-                expressions=[expression.column("x")],
-            ),
-        )
-        .from_(filter_subquery.subquery("t"))
-        .where(
-            expression.NEQ(
-                this=expression.column("x", "t"),
-                expression=value_col.column_expression,
-            )
-        )
-    )
-
-    return Column(agg_subquery.subquery())
-
-
 def array_distinct_bgutil(col: ColumnOrName) -> Column:
     return Column(
         expression.Anonymous(

{sqlframe-3.35.0 → sqlframe-3.36.0}/sqlframe/base/functions.py

@@ -2268,21 +2268,10 @@ def element_at(col: ColumnOrName, value: ColumnOrLiteral) -> Column:
 
 @meta()
 def array_remove(col: ColumnOrName, value: ColumnOrLiteral) -> Column:
-    from sqlframe.base.function_alternatives import (
-        array_remove_bgutil,
-        array_remove_using_filter,
-    )
-
-    session = _get_session()
-
-    if session._is_bigquery:
-        return array_remove_bgutil(col, value)
-
-    if session._is_duckdb:
-        return array_remove_using_filter(col, value)
-
     value_col = value if isinstance(value, Column) else lit(value)
-    return Column.invoke_anonymous_function(col, "ARRAY_REMOVE", value_col)
+    return Column.invoke_expression_over_column(
+        col, expression.ArrayRemove, expression=value_col.column_expression
+    )
 
 
 @meta(unsupported_engines="postgres")
@@ -2299,14 +2288,14 @@ def array_distinct(col: ColumnOrName) -> Column:
 
 @meta(unsupported_engines=["bigquery", "postgres"])
 def array_intersect(col1: ColumnOrName, col2: ColumnOrName) -> Column:
-    from sqlframe.base.function_alternatives import array_intersect_using_intersection
-
-    session = _get_session()
-
-    if session._is_snowflake:
-        return array_intersect_using_intersection(col1, col2)
-
-    return Column.invoke_anonymous_function(col1, "ARRAY_INTERSECT", Column.ensure_col(col2))
+    return Column(
+        expression.ArrayIntersect(
+            expressions=[
+                Column.ensure_col(col1).column_expression,
+                Column.ensure_col(col2).column_expression,
+            ]
+        )
+    )
 
 
 @meta(unsupported_engines=["postgres"])
@@ -3237,18 +3226,16 @@ def elt(*inputs: ColumnOrName) -> Column:
 def endswith(str: ColumnOrName, suffix: ColumnOrName) -> Column:
     from sqlframe.base.function_alternatives import (
         endswith_using_like,
-        endswith_with_underscore,
     )
 
     session = _get_session()
 
-    if session._is_bigquery or session._is_duckdb:
-        return endswith_with_underscore(str, suffix)
-
     if session._is_postgres:
         return endswith_using_like(str, suffix)
 
-    return Column.invoke_anonymous_function(str, "endswith", suffix)
+    return Column.invoke_expression_over_column(
+        str, expression.EndsWith, expression=Column.ensure_col(suffix).column_expression
+    )
 
 
 @meta(unsupported_engines="*")
@@ -5666,10 +5653,9 @@ def replace(
     ):
         replace = expression.Literal.string("")  # type: ignore
 
-    if replace is not None:
-        return Column.invoke_anonymous_function(src, "replace", search, replace)
-    else:
-        return Column.invoke_anonymous_function(src, "replace", search)
+    return Column.invoke_expression_over_column(
+        src, expression.Replace, expression=search, replacement=replace
+    )
 
 
 @meta()

sqlframe-3.36.0/sqlframe/base/group.py

@@ -0,0 +1,227 @@
+# 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
+
+from sqlframe.base.operations import Operation, group_operation, operation
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+if t.TYPE_CHECKING:
+    from sqlframe.base.column import Column
+    from sqlframe.base.session import DF
+else:
+    DF = t.TypeVar("DF")
+
+
+# https://spark.apache.org/docs/latest/sql-ref-syntax-qry-select-groupby.html
+# https://stackoverflow.com/questions/37975227/what-is-the-difference-between-cube-rollup-and-groupby-operators
+class _BaseGroupedData(t.Generic[DF]):
+    last_op: Operation
+
+    def __init__(
+        self,
+        df: DF,
+        group_by_cols: t.Union[t.List[Column], t.List[t.List[Column]]],
+        last_op: Operation,
+    ):
+        self._df = df.copy()
+        self.session = df.session
+        self.last_op = last_op
+        self.group_by_cols = group_by_cols
+        self.pivot_col: t.Optional[str] = None
+        self.pivot_values: t.Optional[t.List[t.Any]] = None
+
+    def _get_function_applied_columns(
+        self, func_name: str, cols: t.Tuple[str, ...]
+    ) -> t.List[Column]:
+        from sqlframe.base import functions as F
+
+        func_name = func_name.lower()
+        return [
+            getattr(F, func_name)(name).alias(
+                self.session._sanitize_column_name(f"{func_name}({name})")
+            )
+            for name in cols
+        ]
+
+    @group_operation(Operation.SELECT)
+    def agg(self, *exprs: t.Union[Column, t.Dict[str, str]]) -> DF:
+        from sqlframe.base.column import Column
+
+        columns = (
+            [
+                self._get_function_applied_columns(agg_func, (column_name,))[0]
+                for column_name, agg_func in exprs[0].items()
+            ]
+            if isinstance(exprs[0], dict)
+            else exprs
+        )
+        cols = self._df._ensure_and_normalize_cols(columns)
+
+        # Handle pivot transformation
+        if self.pivot_col is not None and self.pivot_values is not None:
+            from sqlglot import exp
+
+            from sqlframe.base import functions as F
+
+            # Build the pivot expression
+            # First, we need to convert the DataFrame to include the pivot logic
+            df = self._df.copy()
+
+            # Create the base query with group by columns, pivot column, and aggregation columns
+            select_cols = []
+            # Add group by columns
+            for col in self.group_by_cols:
+                select_cols.append(col.expression)  # type: ignore
+            # Add pivot column
+            select_cols.append(Column.ensure_col(self.pivot_col).expression)
+            # Add the value columns that will be aggregated
+            for agg_col in cols:
+                # Extract the column being aggregated from the aggregation function
+                # For example, from SUM(earnings), we want to extract 'earnings'
+                if (
+                    isinstance(agg_col.column_expression, exp.AggFunc)
+                    and agg_col.column_expression.this
+                ):
+                    if agg_col.column_expression.this not in select_cols:
+                        select_cols.append(agg_col.column_expression.this)
+
+            # Create the base query
+            base_query = df.expression.select(*select_cols, append=False)
+
+            # Build pivot expression
+            pivot_expressions = []
+            for agg_col in cols:
+                if isinstance(agg_col.column_expression, exp.AggFunc):
+                    # Clone the aggregation function
+                    # Snowflake doesn't support alias in the pivot, so we need to use the column_expression
+                    agg_func = (
+                        agg_col.column_expression.copy()
+                        if self.session._is_snowflake
+                        else agg_col.expression.copy()
+                    )
+                    pivot_expressions.append(agg_func)
+
+            # Create the IN clause with pivot values
+            in_values = []
+            for v in self.pivot_values:
+                if isinstance(v, str):
+                    in_values.append(exp.Literal.string(v))
+                else:
+                    in_values.append(exp.Literal.number(v))
+
+            # Build the pivot node with the fields parameter
+            pivot = exp.Pivot(
+                expressions=pivot_expressions,
+                fields=[
+                    exp.In(
+                        this=Column.ensure_col(self.pivot_col).column_expression,
+                        expressions=in_values,
+                    )
+                ],
+            )
+
+            # Create a subquery with the pivot attached
+            subquery = base_query.subquery()
+            subquery.set("pivots", [pivot])
+
+            # Create the final select from the pivoted subquery
+            expression = exp.select("*").from_(subquery)
+
+            return self._df.copy(expression=expression)
+
+        # Original non-pivot logic
+        if not self.group_by_cols or not isinstance(self.group_by_cols[0], (list, tuple, set)):
+            expression = self._df.expression.group_by(
+                # User column_expression for group by to avoid alias in group by
+                *[x.column_expression for x in self.group_by_cols]  # type: ignore
+            ).select(*[x.expression for x in self.group_by_cols + cols], append=False)  # type: ignore
+            group_by_cols = self.group_by_cols
+        else:
+            from sqlglot import exp
+
+            expression = self._df.expression
+            all_grouping_sets = []
+            group_by_cols = []
+            for grouping_set in self.group_by_cols:
+                all_grouping_sets.append(
+                    exp.Tuple(expressions=[x.column_expression for x in grouping_set])  # type: ignore
+                )
+                group_by_cols.extend(grouping_set)  # type: ignore
+            group_by_cols = list(dict.fromkeys(group_by_cols))
+            group_by = exp.Group(grouping_sets=[exp.GroupingSets(expressions=all_grouping_sets)])
+            expression.set("group", group_by)
+        for col in cols:
+            # Spark supports having an empty grouping_id which means all of the columns but other dialects
+            # like duckdb don't support this so we expand the grouping_id to include all of the columns
+            if col.column_expression.this == "GROUPING_ID":
+                col.column_expression.set("expressions", [x.expression for x in group_by_cols])  # type: ignore
+        expression = expression.select(*[x.expression for x in group_by_cols + cols], append=False)  # type: ignore
+        return self._df.copy(expression=expression)
+
+    def count(self) -> DF:
+        from sqlframe.base import functions as F
+
+        return self.agg(F.count("*").alias("count"))
+
+    def mean(self, *cols: str) -> DF:
+        return self.avg(*cols)
+
+    def avg(self, *cols: str) -> DF:
+        return self.agg(*self._get_function_applied_columns("avg", cols))
+
+    def max(self, *cols: str) -> DF:
+        return self.agg(*self._get_function_applied_columns("max", cols))
+
+    def min(self, *cols: str) -> DF:
+        return self.agg(*self._get_function_applied_columns("min", cols))
+
+    def sum(self, *cols: str) -> DF:
+        return self.agg(*self._get_function_applied_columns("sum", cols))
+
+    def pivot(self, pivot_col: str, values: t.Optional[t.List[t.Any]] = None) -> Self:
+        """
+        Pivots a column of the current DataFrame and perform the specified aggregation.
+
+        There are two versions of the pivot function: one that requires the caller
+        to specify the list of distinct values to pivot on, and one that does not.
+        The latter is more concise but less efficient, because Spark needs to first
+        compute the list of distinct values internally.
+
+        Parameters
+        ----------
+        pivot_col : str
+            Name of the column to pivot.
+        values : list, optional
+            List of values that will be translated to columns in the output DataFrame.
+
+        Returns
+        -------
+        GroupedData
+            Returns self to allow chaining with aggregation methods.
+        """
+        if self.session._is_postgres:
+            raise NotImplementedError(
+                "Pivot operation is not supported in Postgres. Please create an issue if you would like a workaround implemented."
+            )
+
+        self.pivot_col = pivot_col
+
+        if values is None:
+            # Eagerly compute distinct values
+            from sqlframe.base.column import Column
+
+            distinct_df = self._df.select(pivot_col).distinct()
+            distinct_rows = distinct_df.collect()
+            # Sort to make the results deterministic
+            self.pivot_values = sorted([row[0] for row in distinct_rows])
+        else:
+            self.pivot_values = values
+
+        return self

{sqlframe-3.35.0 → sqlframe-3.36.0}/sqlframe/base/types.py

@@ -104,22 +104,22 @@ class FloatType(DataType):
 
 
 class ByteType(DataType):
-    def __str__(self) -> str:
+    def simpleString(self) -> str:
         return "tinyint"
 
 
 class IntegerType(DataType):
-    def __str__(self) -> str:
+    def simpleString(self) -> str:
         return "int"
 
 
 class LongType(DataType):
-    def __str__(self) -> str:
+    def simpleString(self) -> str:
         return "bigint"
 
 
 class ShortType(DataType):
-    def __str__(self) -> str:
+    def simpleString(self) -> str:
         return "smallint"
 
 

{sqlframe-3.35.0 → sqlframe-3.36.0}/sqlframe/databricks/session.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
+import logging
 import typing as t
-import warnings
 
 from sqlframe.base.session import _BaseSession
 from sqlframe.databricks.catalog import DatabricksCatalog
@@ -19,6 +19,9 @@ else:
     DatabricksConnection = t.Any
 
 
+logger = logging.getLogger(__name__)
+
+
 class DatabricksSession(
     _BaseSession[  # type: ignore
         DatabricksCatalog,
@@ -43,14 +46,60 @@ class DatabricksSession(
         server_hostname: t.Optional[str] = None,
         http_path: t.Optional[str] = None,
         access_token: t.Optional[str] = None,
+        **kwargs: t.Any,
     ):
         from databricks import sql
 
+        self._conn_kwargs = (
+            {}
+            if conn
+            else {
+                "server_hostname": server_hostname,
+                "http_path": http_path,
+                "access_token": access_token,
+                "disable_pandas": True,
+                **kwargs,
+            }
+        )
+
         if not hasattr(self, "_conn"):
             super().__init__(
-                conn or sql.connect(server_hostname, http_path, access_token, disable_pandas=True)
+                conn or sql.connect(**self._conn_kwargs),
             )
 
+    def _execute(self, sql: str) -> None:
+        from databricks.sql import connect
+        from databricks.sql.exc import DatabaseError, RequestError
+
+        try:
+            super()._execute(sql)
+        except (DatabaseError, RequestError) as e:
+            logger.warning("Failed to execute query")
+            if not self._is_session_expired_error(e):
+                logger.error("Error is not related to session expiration, re-raising")
+                raise e
+            if self._conn_kwargs:
+                logger.info("Attempting to reconnect with provided connection parameters")
+                self._connection = connect(**self._conn_kwargs)
+                # Clear the cached cursor
+                if hasattr(self, "_cur"):
+                    delattr(self, "_cur")
+                super()._execute(sql)
+            else:
+                logger.error("No connection parameters provided so could not reconnect")
+                raise
+
+    def _is_session_expired_error(self, error: Exception) -> bool:
+        error_str = str(error).lower()
+        session_keywords = [
+            "invalid sessionhandle",
+            "session is closed",
+            "session expired",
+            "session not found",
+            "sessionhandle",
+        ]
+        return any(keyword in error_str for keyword in session_keywords)
+
     @classmethod
     def _try_get_map(cls, value: t.Any) -> t.Optional[t.Dict[str, t.Any]]:
         if (

{sqlframe-3.35.0 → sqlframe-3.36.0}/sqlframe.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 3.35.0
+Version: 3.36.0
 Summary: Turning PySpark Into a Universal DataFrame API
 Home-page: https://github.com/eakmanrq/sqlframe
 Author: Ryan Eakman

{sqlframe-3.35.0 → sqlframe-3.36.0}/sqlframe.egg-info/requires.txt

@@ -1,5 +1,5 @@
 prettytable<4
-sqlglot<26.25,>=24.0.0
+sqlglot<26.32,>=24.0.0
 typing_extensions
 
 [bigquery]
@@ -23,8 +23,8 @@ pyspark<3.6,>=2
 pytest-forked
 pytest-postgresql<8,>=6
 pytest-xdist<3.8,>=3.6
-pytest<8.4,>=8.2.0
-ruff<0.12,>=0.4.4
+pytest<8.5,>=8.2.0
+ruff<0.13,>=0.4.4
 types-psycopg2<3,>=2.9
 
 [docs]

{sqlframe-3.35.0 → sqlframe-3.36.0}/tests/integration/engines/test_int_functions.py

@@ -3409,6 +3409,10 @@ def test_bitmap_or_agg(get_session_and_func, get_func):
 
 def test_any_value(get_session_and_func):
     session, any_value = get_session_and_func("any_value")
+    if isinstance(session, PostgresSession):
+        pytest.skip(
+            "any_value is supported in SQLGlot for Postgres but by default assumes Postgres 16+. Tests run against 15. Therefore skipping but should remove this if SQLFrame addss the ability to define Postgres version."
+        )
     df = session.createDataFrame(
         [("c", None), ("a", 2), ("a", 3), ("b", 8), ("b", 2)], ["c1", "c2"]
     )
@@ -3419,9 +3423,9 @@ def test_any_value(get_session_and_func):
         assert non_ignore_nulls == [Row(value="c", value2=2)]
         assert ignore_nulls == [Row(value="c", value2=2)]
     # SQLGlot converts any_value to max
-    elif isinstance(session, PostgresSession):
-        assert non_ignore_nulls == [Row(value="c", value2=8)]
-        assert ignore_nulls == [Row(value="c", value2=8)]
+    # elif isinstance(session, PostgresSession):
+    #     assert non_ignore_nulls == [Row(value="c", value2=8)]
+    #     assert ignore_nulls == [Row(value="c", value2=8)]
     # Always includes nulls
     elif isinstance(session, SnowflakeSession):
         assert non_ignore_nulls == [Row(value="c", value2=None)]