sqlframe 3.36.1__tar.gz → 3.36.3__tar.gz

{sqlframe-3.36.1 → sqlframe-3.36.3}/Makefile

@@ -11,7 +11,7 @@ fast-test:
 	pytest -n auto tests/unit
 
 local-test:
-	pytest -n auto -m "fast or local"
+	pytest -n auto -m "fast or local" --reruns 5
 
 %-test:
 	pytest -n auto -m "${*}"

{sqlframe-3.36.1 → sqlframe-3.36.3}/PKG-INFO

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

{sqlframe-3.36.1 → sqlframe-3.36.3}/docs/bigquery.md

@@ -507,6 +507,7 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [sum_distinct](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.sum_distinct.html)
 * [tan](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.tan.html)
 * [tanh](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.tanh.html)
+* [timestamp_add](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.timestamp_add.html)
 * [timestamp_seconds](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.timestamp_seconds.html)
 * [toDegrees](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.toDegrees.html)
 * [toRadians](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.toRadians.html)

{sqlframe-3.36.1 → sqlframe-3.36.3}/docs/duckdb.md

@@ -468,6 +468,7 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [sumDistinct](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.sumDistinct.html)
 * [sum_distinct](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.sum_distinct.html)
 * [tan](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.tan.html)
+* [timestamp_add](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.timestamp_add.html)
 * [timestamp_seconds](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.timestamp_seconds.html)
 * [toDegrees](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.toDegrees.html)
 * [toRadians](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.toRadians.html)

{sqlframe-3.36.1 → sqlframe-3.36.3}/docs/postgres.md

@@ -455,6 +455,8 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [sum_distinct](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.sum_distinct.html)
 * [tan](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.tan.html)
 * [tanh](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.tanh.html)
+* [timestamp_add](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.timestamp_add.html)
+  * The quantity argument must be literal, not a column
 * [timestamp_seconds](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.timestamp_seconds.html)
 * [toDegrees](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.toDegrees.html)
 * [toRadians](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.toRadians.html)

{sqlframe-3.36.1 → sqlframe-3.36.3}/docs/snowflake.md

@@ -507,6 +507,7 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [sum_distinct](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.sum_distinct.html)
 * [tan](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.tan.html)
 * [tanh](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.tanh.html)
+* [timestamp_add](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.timestamp_add.html)
 * [timestamp_seconds](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.timestamp_seconds.html)
 * [toDegrees](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.toDegrees.html)
 * [toRadians](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.toRadians.html)

{sqlframe-3.36.1 → sqlframe-3.36.3}/setup.py

@@ -20,7 +20,7 @@ setup(
     python_requires=">=3.9",
     install_requires=[
         "prettytable<4",
-        "sqlglot>=24.0.0,<26.32",
+        "sqlglot>=24.0.0,<26.34",
         "typing_extensions",
     ],
     extras_require={
@@ -41,7 +41,8 @@ setup(
             "pytest>=8.2.0,<8.5",
             "pytest-forked",
             "pytest-postgresql>=6,<8",
-            "pytest-xdist>=3.6,<3.8",
+            "pytest-rerunfailures",
+            "pytest-xdist>=3.6,<3.9",
             "pre-commit>=3.7,<5",
             "ruff>=0.4.4,<0.13",
             "types-psycopg2>=2.9,<3",
@@ -70,7 +71,7 @@ setup(
             "redshift_connector>=2.1.1,<2.2.0",
         ],
         "snowflake": [
-            "snowflake-connector-python[secure-local-storage]>=3.10.0,<3.16",
+            "snowflake-connector-python[secure-local-storage]>=3.10.0,<3.17",
         ],
         "spark": [
             "pyspark>=2,<3.6",

{sqlframe-3.36.1 → sqlframe-3.36.3}/sqlframe/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.36.1'
-__version_tuple__ = version_tuple = (3, 36, 1)
+__version__ = version = '3.36.3'
+__version_tuple__ = version_tuple = (3, 36, 3)

{sqlframe-3.36.1 → sqlframe-3.36.3}/sqlframe/base/functions.py

@@ -1475,6 +1475,32 @@ def timestamp_seconds(col: ColumnOrName) -> Column:
     return Column.invoke_expression_over_column(col, expression.UnixToTime)
 
 
+@meta()
+def timestamp_add(unit: str, quantity: ColumnOrName, ts: ColumnOrName) -> Column:
+    session = _get_session()
+
+    if session._is_duckdb or session._is_postgres:
+        quantity = lit(quantity) if isinstance(quantity, int) else quantity
+        if (
+            isinstance(quantity, Column)
+            and isinstance(quantity.expression, expression.Literal)
+            and quantity.expression.is_number
+            and int(quantity.expression.this) < 0
+        ):
+            # If quantity is a negative literal, we use DateSub
+            expr = expression.DateSub
+            quantity.expression.set("this", str(-int(quantity.expression.this)))
+        else:
+            expr = expression.DateAdd  # type: ignore
+        return Column.invoke_expression_over_column(
+            ts, expr, expression=quantity, unit=expression.Var(this=unit.upper())
+        )
+
+    return Column.invoke_expression_over_column(
+        ts, expression.TimestampAdd, expression=quantity, unit=expression.Var(this=unit.upper())
+    )
+
+
 @meta(unsupported_engines=["*", "spark"])
 def window(
     timeColumn: ColumnOrName,
@@ -1796,7 +1822,9 @@ def substring_index(str: ColumnOrName, delim: str, count: int) -> Column:
     if session._is_bigquery:
         return substring_index_bgutil(str, delim, count)
 
-    return Column.invoke_anonymous_function(str, "SUBSTRING_INDEX", lit(delim), lit(count))
+    return Column.invoke_expression_over_column(
+        str, expression.SubstringIndex, delimiter=lit(delim), count=lit(count)
+    )
 
 
 @meta(unsupported_engines="bigquery")
@@ -2205,7 +2233,9 @@ def slice(
 
     start_col = lit(start) if isinstance(start, int) else start
     length_col = lit(length) if isinstance(length, int) else length
-    return Column.invoke_anonymous_function(x, "SLICE", start_col, length_col)
+    return Column.invoke_expression_over_column(
+        x, expression.ArraySlice, start=start_col, end=length_col
+    )
 
 
 @meta()
@@ -2748,7 +2778,7 @@ def typeof(col: ColumnOrName) -> Column:
     if session._is_snowflake:
         return typeof_from_variant(col)
 
-    return Column.invoke_anonymous_function(col, "TYPEOF")
+    return Column.invoke_expression_over_column(col, expression.Typeof)
 
 
 @meta()

{sqlframe-3.36.1 → sqlframe-3.36.3}/sqlframe/base/group.py

@@ -70,6 +70,11 @@ class _BaseGroupedData(t.Generic[DF]):
 
             from sqlframe.base import functions as F
 
+            if self.session._is_snowflake and len(cols) > 1:
+                raise ValueError(
+                    "Snowflake does not support multiple aggregation functions in a single group by operation."
+                )
+
             # Build the pivot expression
             # First, we need to convert the DataFrame to include the pivot logic
             df = self._df.copy()
@@ -132,11 +137,55 @@ class _BaseGroupedData(t.Generic[DF]):
             subquery.set("pivots", [pivot])
 
             # Create the final select from the pivoted subquery
-            expression = exp.select("*").from_(subquery)
+            final_select_in_values = []
+            for col in in_values:  # type: ignore
+                for agg_col in cols:
+                    original_name = col.alias_or_name  # type: ignore
+                    if self.session._is_snowflake:
+                        # Snowflake takes the provided values, like 'Java', and creates the column as "'Java'"
+                        # Therefore the user to select the column would need to use "'Java'"
+                        # This does not conform to the PySpark API, nor is it very user-friendly.
+                        # Therefore, we select the column as expected, and tell SQLFrame it is case-sensitive, but then
+                        # alias is to case-insensitive "Java" so that the user can select it without quotes.
+                        # This has a downside that if a user really needed case-sensitive column names then it wouldn't work.
+                        new_col = exp.to_column(
+                            col.alias_or_name,  # type: ignore
+                            quoted=True,
+                            dialect=self.session.execution_dialect,
+                        )
+                        new_col.this.set("this", f"'{new_col.this.this}'")
+                        new_col = exp.alias_(new_col, original_name)
+                        new_col.unalias()._meta = {"case_sensitive": True}
+                    elif self.session._is_bigquery:
+                        # BigQuery flips the alias order to <alias>_<value> instead of <value>_<alias>
+                        new_col = exp.to_column(
+                            f"{agg_col.alias_or_name}_{original_name}",
+                            dialect=self.session.execution_dialect,
+                        )
+                        new_col = (
+                            exp.alias_(new_col, original_name)
+                            if len(cols) == 1
+                            else exp.alias_(new_col, f"{original_name}_{agg_col.alias_or_name}")
+                        )
+                    elif self.session._is_duckdb:
+                        # DuckDB always respects the alias if if num_cols == 1
+                        new_col = exp.column(f"{original_name}_{agg_col.expression.alias_or_name}")
+                        if len(cols) == 1:
+                            new_col = exp.alias_(new_col, original_name)
+                    else:
+                        new_col = (
+                            exp.column(original_name)
+                            if len(cols) == 1
+                            else exp.column(f"{original_name}_{agg_col.expression.alias_or_name}")
+                        )
+                    final_select_in_values.append(new_col)
+
+            expression = exp.select(
+                *[x.column_expression for x in self.group_by_cols] + final_select_in_values  # type: ignore
+            ).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

{sqlframe-3.36.1 → sqlframe-3.36.3}/sqlframe.egg-info/PKG-INFO

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

{sqlframe-3.36.1 → sqlframe-3.36.3}/sqlframe.egg-info/requires.txt

@@ -1,5 +1,5 @@
 prettytable<4
-sqlglot<26.32,>=24.0.0
+sqlglot<26.34,>=24.0.0
 typing_extensions
 
 [bigquery]
@@ -22,7 +22,8 @@ pyarrow<21,>=10
 pyspark<3.6,>=2
 pytest-forked
 pytest-postgresql<8,>=6
-pytest-xdist<3.8,>=3.6
+pytest-rerunfailures
+pytest-xdist<3.9,>=3.6
 pytest<8.5,>=8.2.0
 ruff<0.13,>=0.4.4
 types-psycopg2<3,>=2.9
@@ -51,7 +52,7 @@ psycopg2<3,>=2.8
 redshift_connector<2.2.0,>=2.1.1
 
 [snowflake]
-snowflake-connector-python[secure-local-storage]<3.16,>=3.10.0
+snowflake-connector-python[secure-local-storage]<3.17,>=3.10.0
 
 [spark]
 pyspark<3.6,>=2

{sqlframe-3.36.1 → sqlframe-3.36.3}/tests/integration/engines/test_int_functions.py

@@ -1531,6 +1531,48 @@ def test_timestamp_seconds(get_session_and_func):
     )
 
 
+def test_timestamp_add(get_session_and_func, get_func):
+    session, timestamp_add = get_session_and_func("timestamp_add")
+    lit = get_func("lit", session)
+
+    # Test data from PySpark examples
+    df = session.createDataFrame(
+        [(datetime.datetime(2016, 3, 11, 9, 0, 7), 2), (datetime.datetime(2024, 4, 2, 9, 0, 7), 3)],
+        ["ts", "quantity"],
+    )
+
+    # Test adding years
+    if not session._is_postgres:
+        result = df.select(timestamp_add("year", "quantity", "ts")).collect()
+        expected_years = [
+            datetime.datetime(2018, 3, 11, 9, 0, 7),
+            datetime.datetime(2027, 4, 2, 9, 0, 7),
+        ]
+        for i, row in enumerate(result):
+            actual = row[0].replace(tzinfo=None) if row[0] else None
+            assert actual == expected_years[i], f"Year addition failed for row {i}"
+
+    # Test adding weeks
+    result = df.select(timestamp_add("WEEK", lit(5), "ts")).collect()
+    expected_weeks = [
+        datetime.datetime(2016, 4, 15, 9, 0, 7),
+        datetime.datetime(2024, 5, 7, 9, 0, 7),
+    ]
+    for i, row in enumerate(result):
+        actual = row[0].replace(tzinfo=None) if row[0] else None
+        assert actual == expected_weeks[i], f"Week addition failed for row {i}"
+
+    # Test subtracting days
+    result = df.select(timestamp_add("day", lit(-5), "ts")).collect()
+    expected_days = [
+        datetime.datetime(2016, 3, 6, 9, 0, 7),
+        datetime.datetime(2024, 3, 28, 9, 0, 7),
+    ]
+    for i, row in enumerate(result):
+        actual = row[0].replace(tzinfo=None) if row[0] else None
+        assert actual == expected_days[i], f"Day subtraction failed for row {i}"
+
+
 def test_window(get_session_and_func, get_func):
     session, window = get_session_and_func("window")
     sum = get_func("sum", session)

{sqlframe-3.36.1 → sqlframe-3.36.3}/tests/integration/test_int_grouped_data.py

@@ -370,3 +370,64 @@ def test_pivot_multiple_aggregations(
     )
 
     compare_frames(df_pivot, dfs_pivot)
+
+
+def test_pivot_without_values_and_selects(
+    pyspark_employee: DataFrame,
+    compare_frames: t.Callable,
+    get_session: t.Callable,
+    is_postgres: t.Callable,
+):
+    """Test pivot without values (auto-detect)"""
+    sqlf_spark = get_session()
+    if is_postgres():
+        pytest.skip("Pivot operation is not supported in Postgres")
+    spark = pyspark_employee.sparkSession
+
+    # Create test data based on PySpark documentation example
+    df1 = spark.createDataFrame(
+        [
+            Row(course="dotNET", year=2012, earnings=10000),
+            Row(course="Java", year=2012, earnings=20000),
+            Row(course="dotNET", year=2012, earnings=5000),
+            Row(course="dotNET", year=2013, earnings=48000),
+            Row(course="Java", year=2013, earnings=30000),
+        ]
+    )
+
+    # Create the same DataFrame in SQLFrame
+    dfs1 = sqlf_spark.createDataFrame(
+        [
+            {
+                "course": "dotNET",
+                "year": 2012,
+                "earnings": 10000,
+            },
+            {
+                "course": "Java",
+                "year": 2012,
+                "earnings": 20000,
+            },
+            {
+                "course": "dotNET",
+                "year": 2012,
+                "earnings": 5000,
+            },
+            {
+                "course": "dotNET",
+                "year": 2013,
+                "earnings": 48000,
+            },
+            {
+                "course": "Java",
+                "year": 2013,
+                "earnings": 30000,
+            },
+        ]
+    )
+
+    # Test pivot without values (auto-detect)
+    df_pivot = df1.groupBy("year").pivot("course").sum("earnings").select("Java")
+    dfs_pivot = dfs1.groupBy("year").pivot("course").sum("earnings").select("Java")
+
+    compare_frames(df_pivot, dfs_pivot)

{sqlframe-3.36.1 → sqlframe-3.36.3}/tests/unit/standalone/test_functions.py

@@ -1544,6 +1544,24 @@ def test_timestamp_seconds(expression, expected):
     assert expression.column_expression.sql(dialect="spark") == expected
 
 
+@pytest.mark.parametrize(
+    "expression, expected",
+    [
+        (SF.timestamp_add("year", "quantity", "ts"), "DATE_ADD(YEAR, quantity, ts)"),
+        (SF.timestamp_add("WEEK", SF.lit(5), "ts"), "DATE_ADD(WEEK, 5, ts)"),
+        (SF.timestamp_add("day", SF.lit(-5), "ts"), "DATE_ADD(DAY, -5, ts)"),
+        (
+            SF.timestamp_add("hour", SF.col("quantity"), SF.col("ts")),
+            "DATE_ADD(HOUR, quantity, ts)",
+        ),
+        (SF.timestamp_add("second", SF.lit(120), "ts"), "DATE_ADD(SECOND, 120, ts)"),
+        (SF.timestamp_add("month", "quantity", "ts"), "DATE_ADD(MONTH, quantity, ts)"),
+    ],
+)
+def test_timestamp_add(expression, expected):
+    assert expression.column_expression.sql(dialect="spark") == expected
+
+
 @pytest.mark.parametrize(
     "expression, expected",
     [