sqlframe 3.27.1__tar.gz → 3.28.1__tar.gz

{sqlframe-3.27.1 → sqlframe-3.28.1}/PKG-INFO

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

{sqlframe-3.27.1 → sqlframe-3.28.1}/setup.py

@@ -20,7 +20,7 @@ setup(
     python_requires=">=3.9",
     install_requires=[
         "prettytable<4",
-        "sqlglot>=24.0.0,<26.12",
+        "sqlglot>=24.0.0,<26.13",
         "typing_extensions",
     ],
     extras_require={

{sqlframe-3.27.1 → sqlframe-3.28.1}/sqlframe/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.27.1'
-__version_tuple__ = version_tuple = (3, 27, 1)
+__version__ = version = '3.28.1'
+__version_tuple__ = version_tuple = (3, 28, 1)

{sqlframe-3.27.1 → sqlframe-3.28.1}/sqlframe/base/dataframe.py

@@ -910,6 +910,8 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
 
     @operation(Operation.GROUP_BY)
     def groupBy(self, *cols, **kwargs) -> GROUP_DATA:
+        if cols and isinstance(cols[0], list):
+            cols = cols[0]  # type: ignore
         columns = self._ensure_and_normalize_cols(cols)
         return self._group_data(self, columns, self.last_op)
 
@@ -1084,7 +1086,7 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
                     for left_col, right_col in join_column_pairs
                 ]
                 # To match spark behavior only the join clause gets deduplicated and it gets put in the front of the column list
-                select_column_names = [
+                select_column_names: list[str | Column] = [
                     (
                         column.alias_or_name
                         if not isinstance(column.expression.this, exp.Star)

{sqlframe-3.27.1 → sqlframe-3.28.1}/sqlframe/base/decorators.py

@@ -1,19 +1,23 @@
-from __future__ import annotations
-
+import functools
 import re
 import typing as t
 
 from sqlglot import exp
 from sqlglot.helper import ensure_list
+from typing_extensions import ParamSpec
 
 from sqlframe.base.column import Column
 
-CALLING_CLASS = t.TypeVar("CALLING_CLASS")
+P = ParamSpec("P")
+T = t.TypeVar("T")
 
 
-def func_metadata(unsupported_engines: t.Optional[t.Union[str, t.List[str]]] = None) -> t.Callable:
-    def _metadata(func: t.Callable) -> t.Callable:
-        def wrapper(*args, **kwargs):
+def func_metadata(
+    unsupported_engines: t.Optional[t.Union[str, t.List[str]]] = None,
+) -> t.Callable[[t.Callable[P, T]], t.Callable[P, T]]:
+    def _metadata(func: t.Callable[P, T]) -> t.Callable[P, T]:
+        @functools.wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
             funcs_to_not_auto_alias = [
                 "posexplode",
                 "explode_outer",
@@ -34,14 +38,17 @@ def func_metadata(unsupported_engines: t.Optional[t.Union[str, t.List[str]]] = N
                 and not isinstance(result.expression, exp.Alias)
                 and func.__name__ not in funcs_to_not_auto_alias
             ):
-                col_name = result.column_expression.find(exp.Identifier)
-                if col_name:
-                    col_name = col_name.name
+                col_name = ""
+                col_name_exp: t.Optional[exp.Expression] = result.column_expression.find(
+                    exp.Identifier
+                )
+                if col_name_exp:
+                    col_name = col_name_exp.name
                 else:
-                    col_name = result.column_expression.find(exp.Literal)
-                    if col_name:
-                        col_name = col_name.this
-                alias_name = f"{func.__name__}__{col_name or ''}__"
+                    col_name_exp = result.column_expression.find(exp.Literal)
+                    if col_name_exp:
+                        col_name = col_name_exp.this
+                alias_name = f"{func.__name__}__{col_name}__"
                 # BigQuery has restrictions on alias names so we constrain it to alphanumeric characters and underscores
                 return result.alias(re.sub(r"\W", "_", alias_name))  # type: ignore
             return result

{sqlframe-3.27.1 → sqlframe-3.28.1}/sqlframe/base/functions.py

@@ -3481,12 +3481,9 @@ def hll_union(
     +------------+
     """
     if allowDifferentLgConfigK is not None:
-        allowDifferentLgConfigK = (
-            lit(allowDifferentLgConfigK)
-            if isinstance(allowDifferentLgConfigK, bool)
-            else allowDifferentLgConfigK
+        return Column.invoke_anonymous_function(
+            col1, "hll_union", col2, lit(allowDifferentLgConfigK)
         )
-        return Column.invoke_anonymous_function(col1, "hll_union", col2, allowDifferentLgConfigK)  # type: ignore
     else:
         return Column.invoke_anonymous_function(col1, "hll_union", col2)
 

{sqlframe-3.27.1 → sqlframe-3.28.1}/sqlframe/base/group.py

@@ -47,7 +47,10 @@ class _BaseGroupedData(t.Generic[DF]):
         from sqlframe.base.column import Column
 
         columns = (
-            [Column(f"{agg_func}({column_name})") for column_name, agg_func in exprs[0].items()]
+            [
+                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
         )
@@ -55,7 +58,8 @@ class _BaseGroupedData(t.Generic[DF]):
 
         if not self.group_by_cols or not isinstance(self.group_by_cols[0], (list, tuple, set)):
             expression = self._df.expression.group_by(
-                *[x.expression for x in self.group_by_cols]  # type: ignore
+                # 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:
@@ -66,7 +70,7 @@ class _BaseGroupedData(t.Generic[DF]):
             group_by_cols = []
             for grouping_set in self.group_by_cols:
                 all_grouping_sets.append(
-                    exp.Tuple(expressions=[x.expression for x in grouping_set])  # type: ignore
+                    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))

{sqlframe-3.27.1 → sqlframe-3.28.1}/sqlframe.egg-info/PKG-INFO

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

{sqlframe-3.27.1 → sqlframe-3.28.1}/sqlframe.egg-info/requires.txt

@@ -1,5 +1,5 @@
 prettytable<4
-sqlglot<26.12,>=24.0.0
+sqlglot<26.13,>=24.0.0
 typing_extensions
 
 [bigquery]

{sqlframe-3.27.1 → sqlframe-3.28.1}/tests/integration/test_int_dataframe.py

@@ -2634,3 +2634,41 @@ def test_rows_between_positive_both_start_end(
     dfs = dfs.withColumn("rolling_price", SF.mean("price").over(swindow))
 
     compare_frames(df, dfs, compare_schema=False)
+
+
+# https://github.com/eakmanrq/sqlframe/issues/356
+def test_array_of_grouping_columns(
+    pyspark_employee: PySparkDataFrame,
+    get_df: t.Callable[[str], BaseDataFrame],
+    compare_frames: t.Callable,
+):
+    df = pyspark_employee.sparkSession.createDataFrame(
+        [(2, "Alice"), (2, "Bob"), (2, "Bob"), (5, "Bob")], schema=["age", "name"]
+    )
+    df = df.groupBy(["name", df.age]).count().sort("name", "age")  # type: ignore
+
+    dfs = get_df("employee").sparkSession.createDataFrame(
+        [(2, "Alice"), (2, "Bob"), (2, "Bob"), (5, "Bob")], schema=["age", "name"]
+    )
+    dfs = dfs.groupBy(["name", dfs.age]).count().sort("name", "age")
+
+    compare_frames(df, dfs, compare_schema=False, sort=True)
+
+
+# https://github.com/eakmanrq/sqlframe/issues/356
+def test_alias_group_by_column(
+    pyspark_employee: PySparkDataFrame,
+    get_df: t.Callable[[str], BaseDataFrame],
+    compare_frames: t.Callable,
+):
+    df = pyspark_employee.sparkSession.createDataFrame(
+        [(2, "Alice"), (2, "Bob"), (2, "Bob"), (5, "Bob")], schema=["age", "name"]
+    )
+    df = df.groupBy(F.col("name").alias("Firstname")).count()
+
+    dfs = get_df("employee").sparkSession.createDataFrame(
+        [(2, "Alice"), (2, "Bob"), (2, "Bob"), (5, "Bob")], schema=["age", "name"]
+    )
+    dfs = dfs.groupBy(SF.col("name").alias("Firstname")).count()
+
+    compare_frames(df, dfs, compare_schema=False, sort=True)

{sqlframe-3.27.1 → sqlframe-3.28.1}/tests/unit/standalone/test_dataframe.py

@@ -156,3 +156,11 @@ def test_unquoted_identifiers(standalone_employee: StandaloneDataFrame):
         standalone_employee.sql(dialect="snowflake", pretty=False, quote_identifiers=False)
         == """SELECT A1.EMPLOYEE_ID AS "employee_id", CAST(A1.FNAME AS TEXT) AS "fname", CAST(A1.LNAME AS TEXT) AS "lname", A1.AGE AS "age", A1.STORE_ID AS "store_id" FROM (VALUES (1, 'Jack', 'Shephard', 37, 1), (2, 'John', 'Locke', 65, 1), (3, 'Kate', 'Austen', 37, 2), (4, 'Claire', 'Littleton', 27, 2), (5, 'Hugo', 'Reyes', 29, 100)) AS A1(EMPLOYEE_ID, FNAME, LNAME, AGE, STORE_ID)"""
     )
+
+
+# https://github.com/eakmanrq/sqlframe/issues/356
+def test_aliased_dictionary_agg(standalone_employee: StandaloneDataFrame):
+    assert (
+        standalone_employee.groupBy("fname").agg({"age": "avg", "lname": "count"}).sql(pretty=False)
+        == "SELECT CAST(`a1`.`fname` AS STRING) AS `fname`, AVG(`a1`.`age`) AS `avg(age)`, COUNT(CAST(`a1`.`lname` AS STRING)) AS `count(lname)` FROM VALUES (1, 'Jack', 'Shephard', 37, 1), (2, 'John', 'Locke', 65, 1), (3, 'Kate', 'Austen', 37, 2), (4, 'Claire', 'Littleton', 27, 2), (5, 'Hugo', 'Reyes', 29, 100) AS `a1`(`employee_id`, `fname`, `lname`, `age`, `store_id`) GROUP BY CAST(`a1`.`fname` AS STRING)"
+    )

{sqlframe-3.27.1 → sqlframe-3.28.1}/tests/unit/standalone/test_functions.py

@@ -17,7 +17,16 @@ def test_invoke_anonymous(name, func):
     # make_interval - SQLGlot doesn't support week
     # to_char - convert to a cast that ignores the format provided
     # ltrim/rtrim - don't seem to convert correctly on some engines
-    ignore_funcs = {"array_size", "exists", "make_interval", "to_char", "ltrim", "rtrim"}
+    ignore_funcs = {
+        "array_size",
+        "exists",
+        "make_interval",
+        "to_char",
+        "ltrim",
+        "rtrim",
+        "ascii",
+        "current_schema",
+    }
     if "invoke_anonymous_function" in inspect.getsource(func) and name not in ignore_funcs:
         func = parse_one(f"{name}()", read="spark", error_level=ErrorLevel.IGNORE)
         assert isinstance(func, exp.Anonymous)
@@ -3332,7 +3341,7 @@ def test_contains(expression, expected):
             "CONVERT_TIMEZONE(cola, CAST(colb AS TIMESTAMP_LTZ))",
         ),
         (
-            SF.convert_timezone(SF.col("colc"), "cola", "colb"),
+            SF.convert_timezone(SF.col("colc"), SF.col("cola"), "colb"),
             "CONVERT_TIMEZONE(colc, cola, CAST(colb AS TIMESTAMP_LTZ))",
         ),
     ],
@@ -3632,7 +3641,7 @@ def test_histogram_numeric(expression, expected):
         (SF.hll_sketch_agg("cola"), "HLL_SKETCH_AGG(cola)"),
         (SF.hll_sketch_agg(SF.col("cola")), "HLL_SKETCH_AGG(cola)"),
         (SF.hll_sketch_agg("cola", 12), "HLL_SKETCH_AGG(cola, 12)"),
-        (SF.hll_sketch_agg("cola", "colb"), "HLL_SKETCH_AGG(cola, colb)"),
+        (SF.hll_sketch_agg("cola", SF.lit(12)), "HLL_SKETCH_AGG(cola, 12)"),
     ],
 )
 def test_hll_sketch_agg(expression, expected):
@@ -3655,7 +3664,6 @@ def test_hll_sketch_estimate(expression, expected):
     [
         (SF.hll_union("cola", "colb"), "HLL_UNION(cola, colb)"),
         (SF.hll_union(SF.col("cola"), SF.col("colb")), "HLL_UNION(cola, colb)"),
-        (SF.hll_union("cola", "colb", "colc"), "HLL_UNION(cola, colb, colc)"),
         (SF.hll_union("cola", "colb", False), "HLL_UNION(cola, colb, FALSE)"),
     ],
 )
@@ -3668,7 +3676,7 @@ def test_hll_union(expression, expected):
     [
         (SF.hll_union_agg("cola"), "HLL_UNION_AGG(cola)"),
         (SF.hll_union_agg(SF.col("cola")), "HLL_UNION_AGG(cola)"),
-        (SF.hll_union_agg("cola", "colb"), "HLL_UNION_AGG(cola, colb)"),
+        (SF.hll_union_agg("cola", SF.lit(True)), "HLL_UNION_AGG(cola, TRUE)"),
         (SF.hll_union_agg("cola", False), "HLL_UNION_AGG(cola, FALSE)"),
     ],
 )