sqlframe 3.3.1__tar.gz → 3.4.1__tar.gz

{sqlframe-3.3.1 → sqlframe-3.4.1}/PKG-INFO

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

{sqlframe-3.3.1 → sqlframe-3.4.1}/docs/bigquery.md

@@ -234,7 +234,6 @@ See something that you would like to see supported? [Open an issue](https://gith
 * sql 
     * SQLFrame Specific: Get the SQL representation of a given column
 * [startswith](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.Column.startswith.html)
-* [substr](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.Column.substr.html)
 * [when](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.Column.when.html)
 
 ### DataFrame Class

{sqlframe-3.3.1 → sqlframe-3.4.1}/docs/duckdb.md

@@ -198,7 +198,6 @@ See something that you would like to see supported? [Open an issue](https://gith
 * sql 
     * SQLFrame Specific: Get the SQL representation of a given column
 * [startswith](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.Column.startswith.html)
-* [substr](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.Column.substr.html)
 * [when](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.Column.when.html)
 
 ### DataFrame Class

{sqlframe-3.3.1 → sqlframe-3.4.1}/docs/snowflake.md

@@ -230,7 +230,6 @@ See something that you would like to see supported? [Open an issue](https://gith
 * sql 
     * SQLFrame Specific: Get the SQL representation of a given column
 * [startswith](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.Column.startswith.html)
-* [substr](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.Column.substr.html)
 * [when](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.Column.when.html)
 
 ### DataFrame Class

{sqlframe-3.3.1 → sqlframe-3.4.1}/setup.py

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

{sqlframe-3.3.1 → sqlframe-3.4.1}/sqlframe/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.3.1'
-__version_tuple__ = version_tuple = (3, 3, 1)
+__version__ = version = '3.4.1'
+__version_tuple__ = version_tuple = (3, 4, 1)

{sqlframe-3.3.1 → sqlframe-3.4.1}/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.3.1 → sqlframe-3.4.1}/sqlframe/base/functions.py

@@ -1129,11 +1129,15 @@ def overlay(
     pos: t.Union[ColumnOrName, int],
     len: t.Optional[t.Union[ColumnOrName, int]] = None,
 ) -> Column:
-    pos_value = lit(pos) if isinstance(pos, int) else pos
-    if len is not None:
-        len_value = lit(len) if isinstance(len, int) else len
-        return Column.invoke_anonymous_function(src, "OVERLAY", replace, pos_value, len_value)
-    return Column.invoke_anonymous_function(src, "OVERLAY", replace, pos_value)
+    return Column.invoke_expression_over_column(
+        src,
+        expression.Overlay,
+        **{
+            "expression": Column(replace).expression,
+            "from": lit(pos).expression,
+            "for": lit(len).expression if len is not None else None,
+        },
+    )
 
 
 @meta(unsupported_engines=["bigquery", "duckdb", "postgres", "snowflake"])
@@ -4834,7 +4838,7 @@ def str_to_map(
     )
 
 
-@meta(unsupported_engines="*")
+@meta(unsupported_engines="postgres")
 def substr(str: ColumnOrName, pos: ColumnOrName, len: t.Optional[ColumnOrName] = None) -> Column:
     """
     Returns the substring of `str` that starts at `pos` and is of length `len`,
@@ -4873,10 +4877,7 @@ def substr(str: ColumnOrName, pos: ColumnOrName, len: t.Optional[ColumnOrName] =
     |                   k SQL|
     +------------------------+
     """
-    if len is not None:
-        return Column.invoke_anonymous_function(str, "substr", pos, len)
-    else:
-        return Column.invoke_anonymous_function(str, "substr", pos)
+    return Column.invoke_expression_over_column(str, expression.Substring, start=pos, length=len)
 
 
 @meta(unsupported_engines="*")

{sqlframe-3.3.1 → sqlframe-3.4.1}/sqlframe/bigquery/functions.pyi

@@ -251,6 +251,7 @@ from sqlframe.base.functions import stddev as stddev
 from sqlframe.base.functions import stddev_pop as stddev_pop
 from sqlframe.base.functions import stddev_samp as stddev_samp
 from sqlframe.base.functions import struct as struct
+from sqlframe.base.functions import substr as substr
 from sqlframe.base.functions import substring as substring
 from sqlframe.base.functions import sum as sum
 from sqlframe.base.functions import sum_distinct as sum_distinct

{sqlframe-3.3.1 → sqlframe-3.4.1}/sqlframe/duckdb/functions.pyi

@@ -184,6 +184,7 @@ from sqlframe.base.functions import (
     stddev_samp as stddev_samp,
     struct as struct,
     substring as substring,
+    substr as substr,
     sum as sum,
     sumDistinct as sumDistinct,
     sum_distinct as sum_distinct,

{sqlframe-3.3.1 → sqlframe-3.4.1}/sqlframe/duckdb/session.py

@@ -52,6 +52,10 @@ class DuckDBSession(
             super().__init__(conn, *args, **kwargs)
             self._last_result = None
 
+    @cached_property
+    def _cur(self) -> DuckDBPyConnection:  # type: ignore
+        return self._conn
+
     @classmethod
     def _try_get_map(cls, value: t.Any) -> t.Optional[t.Dict[str, t.Any]]:
         if value and isinstance(value, dict):

{sqlframe-3.3.1 → sqlframe-3.4.1}/sqlframe/snowflake/functions.pyi

@@ -193,6 +193,7 @@ from sqlframe.base.functions import (
     stddev_pop as stddev_pop,
     stddev_samp as stddev_samp,
     substring as substring,
+    substr as substr,
     sum as sum,
     sumDistinct as sumDistinct,
     sum_distinct as sum_distinct,

{sqlframe-3.3.1 → sqlframe-3.4.1}/sqlframe/spark/functions.py

@@ -18,4 +18,5 @@ from sqlframe.base.function_alternatives import (  # noqa
     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-3.3.1 → sqlframe-3.4.1}/sqlframe/spark/functions.pyi

@@ -2,6 +2,7 @@ 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,
@@ -372,7 +373,6 @@ from sqlframe.base.functions import (
     try_aes_decrypt as try_aes_decrypt,
     try_avg as try_avg,
     try_divide as try_divide,
-    try_element_at as try_element_at,
     try_multiply as try_multiply,
     try_subtract as try_subtract,
     try_sum as try_sum,

{sqlframe-3.3.1 → sqlframe-3.4.1}/sqlframe.egg-info/PKG-INFO

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

{sqlframe-3.3.1 → sqlframe-3.4.1}/sqlframe.egg-info/requires.txt

@@ -1,5 +1,5 @@
 prettytable<3.11.1
-sqlglot<25.23,>=24.0.0
+sqlglot<25.25,>=24.0.0
 typing_extensions<5,>=4.8
 
 [bigquery]
@@ -10,7 +10,7 @@ google-cloud-bigquery[pandas]<4,>=3
 duckdb<1.2,>=0.9
 findspark<3,>=2
 mypy<1.12,>=1.10.0
-openai<1.48,>=1.30
+openai<1.52,>=1.30
 pandas-stubs<3,>=2
 pandas<3,>=2
 psycopg<4,>=3.1
@@ -26,7 +26,7 @@ types-psycopg2<3,>=2.9
 pre-commit>=3.5
 
 [dev:python_version >= "3.9"]
-pre-commit<3.9,>=3.7
+pre-commit<4.1,>=3.7
 
 [docs]
 mkdocs-include-markdown-plugin==6.0.6
@@ -40,7 +40,7 @@ duckdb<1.2,>=0.9
 pandas<3,>=2
 
 [openai]
-openai<1.48,>=1.30
+openai<1.52,>=1.30
 
 [pandas]
 pandas<3,>=2

{sqlframe-3.3.1 → sqlframe-3.4.1}/tests/common_fixtures.py

@@ -117,8 +117,11 @@ def spark_session(pyspark_session: PySparkSession) -> SparkSession:
 
 @pytest.fixture(scope="function")
 def duckdb_session() -> DuckDBSession:
-    connection = duckdb.connect(config={"TimeZone": "UTC"})
-    return DuckDBSession(conn=connection)
+    connector = duckdb.connect()
+    connector.execute("set TimeZone = 'UTC'")
+    connector.execute("SELECT * FROM duckdb_settings() WHERE name = 'TimeZone'")
+    assert connector.fetchone()[1] == "UTC"  # type: ignore
+    return DuckDBSession(conn=connector)
 
 
 @pytest.fixture

{sqlframe-3.3.1 → sqlframe-3.4.1}/tests/integration/engines/duck/test_duckdb_catalog.py

@@ -8,7 +8,10 @@ from sqlframe.duckdb.session import DuckDBSession
 def duckdb_session() -> DuckDBSession:
     import duckdb
 
-    connector = duckdb.connect(config={"TimeZone": "UTC"})
+    connector = duckdb.connect()
+    connector.execute("set TimeZone = 'UTC'")
+    connector.execute("SELECT * FROM duckdb_settings() WHERE name = 'TimeZone'")
+    assert connector.fetchone()[1] == "UTC"  # type: ignore
     session = DuckDBSession(connector)
     session._execute('''ATTACH ':memory:' AS "default"''')
     session._execute('''ATTACH ':memory:' AS "catalog1"''')

{sqlframe-3.3.1 → sqlframe-3.4.1}/tests/integration/engines/test_int_functions.py

@@ -18,7 +18,7 @@ from sqlframe.base.util import (
     get_func_from_session as get_func_from_session_without_fallback,
 )
 from sqlframe.bigquery import BigQuerySession
-from sqlframe.duckdb import DuckDBSession
+from sqlframe.duckdb import DuckDBCatalog, DuckDBSession
 from sqlframe.postgres import PostgresDataFrame, PostgresSession
 from sqlframe.snowflake import SnowflakeSession
 from sqlframe.spark.session import SparkSession

{sqlframe-3.3.1 → sqlframe-3.4.1}/tests/integration/test_int_dataframe.py

@@ -2180,3 +2180,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)

{sqlframe-3.3.1 → sqlframe-3.4.1}/tests/unit/standalone/test_functions.py

@@ -1823,12 +1823,12 @@ def test_instr(expression, expected):
 @pytest.mark.parametrize(
     "expression, expected",
     [
-        (SF.overlay("cola", "colb", 3, 7), "OVERLAY(cola, colb, 3, 7)"),
+        (SF.overlay("cola", "colb", 3, 7), "OVERLAY(cola PLACING colb FROM 3 FOR 7)"),
         (
             SF.overlay(SF.col("cola"), SF.col("colb"), SF.lit(3), SF.lit(7)),
-            "OVERLAY(cola, colb, 3, 7)",
+            "OVERLAY(cola PLACING colb FROM 3 FOR 7)",
         ),
-        (SF.overlay("cola", "colb", 3), "OVERLAY(cola, colb, 3)"),
+        (SF.overlay("cola", "colb", 3), "OVERLAY(cola PLACING colb FROM 3)"),
     ],
 )
 def test_overlay(expression, expected):
@@ -2263,8 +2263,8 @@ def test_array_contains(expression, expected):
 @pytest.mark.parametrize(
     "expression, expected",
     [
-        (SF.arrays_overlap("cola", "colb"), "ARRAY_OVERLAPS(cola, colb)"),
-        (SF.arrays_overlap(SF.col("cola"), SF.col("colb")), "ARRAY_OVERLAPS(cola, colb)"),
+        (SF.arrays_overlap("cola", "colb"), "cola && colb"),
+        (SF.arrays_overlap(SF.col("cola"), SF.col("colb")), "cola && colb"),
     ],
 )
 def test_arrays_overlap(expression, expected):
@@ -4495,9 +4495,9 @@ def test_str_to_map(expression, expected):
 @pytest.mark.parametrize(
     "expression, expected",
     [
-        (SF.substr("cola", "colb"), "SUBSTR(cola, colb)"),
-        (SF.substr(SF.col("cola"), SF.col("colb")), "SUBSTR(cola, colb)"),
-        (SF.substr("cola", "colb", "colc"), "SUBSTR(cola, colb, colc)"),
+        (SF.substr("cola", "colb"), "SUBSTRING(cola, colb)"),
+        (SF.substr(SF.col("cola"), SF.col("colb")), "SUBSTRING(cola, colb)"),
+        (SF.substr("cola", "colb", "colc"), "SUBSTRING(cola, colb, colc)"),
     ],
 )
 def test_substr(expression, expected):