sqlframe 3.4.0__tar.gz → 3.5.0__tar.gz

{sqlframe-3.4.0 → sqlframe-3.5.0}/PKG-INFO

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

{sqlframe-3.4.0 → sqlframe-3.5.0}/docs/duckdb.md

@@ -443,6 +443,7 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [sort_array](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.sort_array.html)
 * [soundex](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.soundex.html)
 * [split](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.split.html)
+* [split_part](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.split_part.html)
 * [sqrt](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.sqrt.html)
 * [stddev](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.stddev.html)
 * [stddev_pop](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.stddev_pop.html)

{sqlframe-3.4.0 → sqlframe-3.5.0}/docs/snowflake.md

@@ -484,6 +484,7 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [soundex](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.soundex.html)
 * [split](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.split.html)
   * Regular expressions not supported
+* [split_part](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.split_part.html)
 * [sqrt](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.sqrt.html)
 * [stddev](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.stddev.html)
 * [stddev_pop](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.stddev_pop.html)

{sqlframe-3.4.0 → sqlframe-3.5.0}/setup.py

@@ -20,7 +20,7 @@ setup(
     python_requires=">=3.8",
     install_requires=[
         "prettytable<3.11.1",
-        "sqlglot>=24.0.0,<25.25",
+        "sqlglot>=24.0.0,<25.26",
         "typing_extensions>=4.8,<5",
     ],
     extras_require={
@@ -31,8 +31,8 @@ setup(
         "dev": [
             "duckdb>=0.9,<1.2",
             "findspark>=2,<3",
-            "mypy>=1.10.0,<1.12",
-            "openai>=1.30,<1.52",
+            "mypy>=1.10.0,<1.13",
+            "openai>=1.30,<1.53",
             "pandas>=2,<3",
             "pandas-stubs>=2,<3",
             "psycopg>=3.1,<4",
@@ -43,7 +43,7 @@ setup(
             "pytest-xdist>=3.6,<3.7",
             "pre-commit>=3.5;python_version=='3.8'",
             "pre-commit>=3.7,<4.1;python_version>='3.9'",
-            "ruff>=0.4.4,<0.7",
+            "ruff>=0.4.4,<0.8",
             "types-psycopg2>=2.9,<3",
         ],
         "docs": [
@@ -58,7 +58,7 @@ setup(
             "pandas>=2,<3",
         ],
         "openai": [
-            "openai>=1.30,<1.52",
+            "openai>=1.30,<1.53",
         ],
         "pandas": [
             "pandas>=2,<3",

{sqlframe-3.4.0 → sqlframe-3.5.0}/sqlframe/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.4.0'
-__version_tuple__ = version_tuple = (3, 4, 0)
+__version__ = version = '3.5.0'
+__version_tuple__ = version_tuple = (3, 5, 0)

{sqlframe-3.4.0 → sqlframe-3.5.0}/sqlframe/base/column.py

@@ -338,6 +338,9 @@ class Column:
         new_expression = exp.Not(this=exp.Is(this=self.column_expression, expression=exp.Null()))
         return Column(new_expression)
 
+    def eqNullSafe(self, other: ColumnOrLiteral) -> Column:
+        return self.binary_op(exp.NullSafeEQ, other)
+
     def cast(
         self,
         dataType: t.Union[str, DataType, exp.DataType, exp.DataType.Type],
@@ -385,7 +388,7 @@ class Column:
 
     def isin(self, *cols: t.Union[ColumnOrLiteral, t.Iterable[ColumnOrLiteral]]):
         columns = flatten(cols) if isinstance(cols[0], (list, set, tuple)) else cols  # type: ignore
-        expressions = [self._lit(x).expression for x in columns]
+        expressions = [self._lit(x).expression for x in columns]  # type: ignore
         return Column.invoke_expression_over_column(self, exp.In, expressions=expressions)  # type: ignore
 
     def between(

{sqlframe-3.4.0 → sqlframe-3.5.0}/sqlframe/base/dataframe.py

@@ -15,7 +15,7 @@ from dataclasses import dataclass
 
 import sqlglot
 from prettytable import PrettyTable
-from sqlglot import Dialect
+from sqlglot import Dialect, maybe_parse
 from sqlglot import expressions as exp
 from sqlglot import lineage as sqlglot_lineage
 from sqlglot.helper import ensure_list, flatten, object_to_dict, seq_get
@@ -460,16 +460,40 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         df.expression.ctes[-1].set("cache_storage_level", storage_level)
         return df
 
-    @classmethod
-    def _add_ctes_to_expression(cls, expression: exp.Select, ctes: t.List[exp.CTE]) -> exp.Select:
+    def _add_ctes_to_expression(self, expression: exp.Select, ctes: t.List[exp.CTE]) -> exp.Select:
         expression = expression.copy()
         with_expression = expression.args.get("with")
         if with_expression:
             existing_ctes = with_expression.expressions
-            existsing_cte_names = {x.alias_or_name for x in existing_ctes}
+            existing_cte_counts = {x.alias_or_name: 0 for x in existing_ctes}
+            replaced_cte_names = {}  # type: ignore
             for cte in ctes:
-                if cte.alias_or_name not in existsing_cte_names:
-                    existing_ctes.append(cte)
+                if replaced_cte_names:
+                    cte = cte.transform(replace_id_value, replaced_cte_names)  # type: ignore
+                if cte.alias_or_name in existing_cte_counts:
+                    existing_cte_counts[cte.alias_or_name] += 10
+                    cte.set(
+                        "this",
+                        cte.this.where(
+                            exp.EQ(
+                                this=exp.Literal.number(existing_cte_counts[cte.alias_or_name]),
+                                expression=exp.Literal.number(
+                                    existing_cte_counts[cte.alias_or_name]
+                                ),
+                            )
+                        ),
+                    )
+                    new_cte_alias = self._create_hash_from_expression(cte.this)
+                    replaced_cte_names[cte.args["alias"].this] = maybe_parse(
+                        new_cte_alias, dialect=self.session.input_dialect, into=exp.Identifier
+                    )
+                    cte.set(
+                        "alias",
+                        maybe_parse(
+                            new_cte_alias, dialect=self.session.input_dialect, into=exp.TableAlias
+                        ),
+                    )
+                existing_ctes.append(cte)
         else:
             existing_ctes = ctes
         expression.set("with", exp.With(expressions=existing_ctes))
@@ -843,11 +867,11 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
             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
-        join_expression = self.expression.join(
-            other_df.latest_cte_name, join_type=how.replace("_", " ")
+        join_expression = join_expression.join(
+            join_expression.ctes[-1].alias, join_type=how.replace("_", " ")
         )
-        join_expression = self._add_ctes_to_expression(join_expression, other_df.expression.ctes)
         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)

{sqlframe-3.4.0 → sqlframe-3.5.0}/sqlframe/base/function_alternatives.py

@@ -1320,8 +1320,8 @@ def flatten_using_array_flatten(col: ColumnOrName) -> Column:
 def map_concat_using_map_cat(*cols: t.Union[ColumnOrName, t.Iterable[ColumnOrName]]) -> Column:
     columns = list(flatten(cols)) if not isinstance(cols[0], (str, Column)) else cols  # type: ignore
     if len(columns) == 1:
-        return Column.invoke_anonymous_function(columns[0], "MAP_CAT")
-    return Column.invoke_anonymous_function(columns[0], "MAP_CAT", *columns[1:])
+        return Column.invoke_anonymous_function(columns[0], "MAP_CAT")  # type: ignore
+    return Column.invoke_anonymous_function(columns[0], "MAP_CAT", *columns[1:])  # type: ignore
 
 
 def sequence_from_generate_series(

{sqlframe-3.4.0 → sqlframe-3.5.0}/sqlframe/base/functions.py

@@ -1652,8 +1652,8 @@ def arrays_zip(*cols: ColumnOrName) -> Column:
 def map_concat(*cols: t.Union[ColumnOrName, t.Iterable[ColumnOrName]]) -> Column:
     columns = list(flatten(cols)) if not isinstance(cols[0], (str, Column)) else cols  # type: ignore
     if len(columns) == 1:
-        return Column.invoke_anonymous_function(columns[0], "MAP_CONCAT")
-    return Column.invoke_anonymous_function(columns[0], "MAP_CONCAT", *columns[1:])
+        return Column.invoke_anonymous_function(columns[0], "MAP_CONCAT")  # type: ignore
+    return Column.invoke_anonymous_function(columns[0], "MAP_CONCAT", *columns[1:])  # type: ignore
 
 
 @meta(unsupported_engines="postgres")
@@ -4697,7 +4697,7 @@ def spark_partition_id() -> Column:
     return Column.invoke_anonymous_function(None, "spark_partition_id")
 
 
-@meta(unsupported_engines="*")
+@meta(unsupported_engines=["bigquery", "postgres"])
 def split_part(src: ColumnOrName, delimiter: ColumnOrName, partNum: ColumnOrName) -> Column:
     """
     Splits `str` by delimiter and return requested part of the split (1-based).
@@ -4723,7 +4723,9 @@ def split_part(src: ColumnOrName, delimiter: ColumnOrName, partNum: ColumnOrName
     >>> df.select(split_part(df.a, df.b, df.c).alias('r')).collect()
     [Row(r='13')]
     """
-    return Column.invoke_anonymous_function(src, "split_part", delimiter, partNum)
+    return Column.invoke_expression_over_column(
+        src, expression.SplitPart, delimiter=delimiter, part_index=partNum
+    )
 
 
 @meta()

{sqlframe-3.4.0 → sqlframe-3.5.0}/sqlframe/duckdb/functions.pyi

@@ -177,6 +177,7 @@ from sqlframe.base.functions import (
     skewness as skewness,
     sort_array as sort_array,
     soundex as soundex,
+    split_part as split_part,
     sqrt as sqrt,
     startswith as startswith,
     stddev as stddev,

{sqlframe-3.4.0 → sqlframe-3.5.0}/sqlframe/snowflake/functions.pyi

@@ -187,6 +187,7 @@ from sqlframe.base.functions import (
     sinh as sinh,
     size as size,
     soundex as soundex,
+    split_part as split_part,
     sqrt as sqrt,
     startswith as startswith,
     stddev as stddev,

{sqlframe-3.4.0 → sqlframe-3.5.0}/sqlframe.egg-info/PKG-INFO

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

{sqlframe-3.4.0 → sqlframe-3.5.0}/sqlframe.egg-info/requires.txt

@@ -1,5 +1,5 @@
 prettytable<3.11.1
-sqlglot<25.25,>=24.0.0
+sqlglot<25.26,>=24.0.0
 typing_extensions<5,>=4.8
 
 [bigquery]
@@ -9,8 +9,8 @@ google-cloud-bigquery[pandas]<4,>=3
 [dev]
 duckdb<1.2,>=0.9
 findspark<3,>=2
-mypy<1.12,>=1.10.0
-openai<1.52,>=1.30
+mypy<1.13,>=1.10.0
+openai<1.53,>=1.30
 pandas-stubs<3,>=2
 pandas<3,>=2
 psycopg<4,>=3.1
@@ -19,7 +19,7 @@ pyspark<3.6,>=2
 pytest-postgresql<7,>=6
 pytest-xdist<3.7,>=3.6
 pytest<8.4,>=8.2.0
-ruff<0.7,>=0.4.4
+ruff<0.8,>=0.4.4
 types-psycopg2<3,>=2.9
 
 [dev:python_version == "3.8"]
@@ -40,7 +40,7 @@ duckdb<1.2,>=0.9
 pandas<3,>=2
 
 [openai]
-openai<1.52,>=1.30
+openai<1.53,>=1.30
 
 [pandas]
 pandas<3,>=2

{sqlframe-3.4.0 → sqlframe-3.5.0}/tests/integration/engines/test_engine_column.py

@@ -5,6 +5,7 @@ import typing as t
 import pytest
 
 from sqlframe.base import types
+from sqlframe.base.types import Row
 from sqlframe.bigquery import BigQuerySession
 from sqlframe.postgres import PostgresSession
 
@@ -14,6 +15,23 @@ if t.TYPE_CHECKING:
 pytest_plugins = ["tests.integration.fixtures"]
 
 
+def test_columnt_eq_null_safe(get_session: t.Callable[[], _BaseSession], get_func):
+    session = get_session()
+    df1 = session.createDataFrame([Row(id=1, value="foo"), Row(id=2, value=None)])
+    lit = get_func("lit", session)
+    values = df1.select(
+        df1["value"] == "foo",
+        df1["value"].eqNullSafe("foo"),
+        df1["value"].eqNullSafe(None),
+        df1["id"].eqNullSafe(1),
+        df1["id"] == lit(None),
+        df1["id"].eqNullSafe(None),
+    ).collect()
+
+    assert values[0] == (True, True, False, True, None, False)
+    assert values[1] == (None, False, True, False, None, False)
+
+
 def test_column_get_item_array(get_session: t.Callable[[], _BaseSession], get_func):
     session = get_session()
     lit = get_func("lit", session)

{sqlframe-3.4.0 → sqlframe-3.5.0}/tests/integration/test_int_dataframe.py

@@ -231,6 +231,17 @@ def test_where_clause_single(
     compare_frames(df_employee, dfs_employee)
 
 
+def test_where_clause_eq_nullsafe(
+    pyspark_employee: PySparkDataFrame,
+    get_df: t.Callable[[str], _BaseDataFrame],
+    compare_frames: t.Callable,
+):
+    employee = get_df("employee")
+    df_employee = pyspark_employee.where(F.col("age").eqNullSafe(F.lit(37)))
+    dfs_employee = employee.where(SF.col("age") == SF.lit(37))
+    compare_frames(df_employee, dfs_employee)
+
+
 def test_where_clause_multiple_and(
     pyspark_employee: PySparkDataFrame,
     get_df: t.Callable[[str], _BaseDataFrame],
@@ -2180,3 +2191,59 @@ def test_chained_join_common_key(
     dfs = dfs.join(dfs_height, how="left", on="name").join(dfs_location, how="left", on="name")
 
     compare_frames(df, dfs, compare_schema=False)
+
+
+# https://github.com/eakmanrq/sqlframe/issues/185
+def test_chaining_joins_with_selects(
+    pyspark_employee: PySparkDataFrame,
+    pyspark_store: PySparkDataFrame,
+    pyspark_district: PySparkDataFrame,
+    get_df: t.Callable[[str], _BaseDataFrame],
+    compare_frames: t.Callable,
+    is_spark: t.Callable,
+):
+    if is_spark():
+        pytest.skip(
+            "This test is not supported in Spark. This is related to how duplicate columns are handled in Spark"
+        )
+    df = (
+        pyspark_employee.alias("employee")
+        .join(
+            pyspark_store.filter(F.col("store_name") != "test").alias("store"),
+            on=F.col("employee.employee_id") == F.col("store.store_id"),
+        )
+        .join(
+            pyspark_district.alias("district"),
+            on=F.col("store.store_id") == F.col("district.district_id"),
+        )
+        .join(
+            pyspark_district.alias("district2"),
+            on=(F.col("store.store_id") + 1) == F.col("district2.district_id"),
+            how="left",
+        )
+        .select("*")
+    )
+
+    employee = get_df("employee")
+    store = get_df("store")
+    district = get_df("district")
+
+    dfs = (
+        employee.alias("employee")
+        .join(
+            store.filter(SF.col("store_name") != "test").alias("store"),
+            on=SF.col("employee.employee_id") == SF.col("store.store_id"),
+        )
+        .join(
+            district.alias("district"),
+            on=SF.col("store.store_id") == SF.col("district.district_id"),
+        )
+        .join(
+            district.alias("district2"),
+            on=(SF.col("store.store_id") + 1) == SF.col("district2.district_id"),
+            how="left",
+        )
+        .select("*")
+    )
+
+    compare_frames(df, dfs, compare_schema=False)