sqlframe 3.39.2__tar.gz → 3.39.3__tar.gz

{sqlframe-3.39.2 → sqlframe-3.39.3}/PKG-INFO

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

{sqlframe-3.39.2 → sqlframe-3.39.3}/setup.py

@@ -22,6 +22,7 @@ setup(
         "prettytable<4",
         "sqlglot>=24.0.0,<27.9",
         "typing_extensions",
+        "more-itertools",
     ],
     extras_require={
         "bigquery": [

{sqlframe-3.39.2 → sqlframe-3.39.3}/sqlframe/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '3.39.2'
-__version_tuple__ = version_tuple = (3, 39, 2)
+__version__ = version = '3.39.3'
+__version_tuple__ = version_tuple = (3, 39, 3)
 
-__commit_id__ = commit_id = 'g772b3a6bf'
+__commit_id__ = commit_id = 'g9d915cb1e'

{sqlframe-3.39.2 → sqlframe-3.39.3}/sqlframe/base/dataframe.py

@@ -16,6 +16,7 @@ from dataclasses import dataclass
 from uuid import uuid4
 
 import sqlglot
+from more_itertools import partition
 from prettytable import PrettyTable
 from sqlglot import Dialect, maybe_parse
 from sqlglot import expressions as exp
@@ -31,6 +32,7 @@ from sqlframe.base.util import (
     get_func_from_session,
     get_tables_from_expression_with_join,
     normalize_string,
+    partition_to,
     quote_preserving_alias_or_name,
     sqlglot_to_spark,
     verify_openai_installed,
@@ -540,16 +542,23 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         expression.set("with", exp.With(expressions=existing_ctes))
         return expression
 
+    @classmethod
+    def _get_outer_select_expressions(
+        cls, item: exp.Expression
+    ) -> t.List[t.Union[exp.Column, exp.Alias]]:
+        outer_select = item.find(exp.Select)
+        if outer_select:
+            return outer_select.expressions
+        return []
+
     @classmethod
     def _get_outer_select_columns(cls, item: exp.Expression) -> t.List[Column]:
         from sqlframe.base.session import _BaseSession
 
         col = get_func_from_session("col", _BaseSession())
 
-        outer_select = item.find(exp.Select)
-        if outer_select:
-            return [col(quote_preserving_alias_or_name(x)) for x in outer_select.expressions]
-        return []
+        outer_expressions = cls._get_outer_select_expressions(item)
+        return [col(quote_preserving_alias_or_name(x)) for x in outer_expressions]
 
     def _create_hash_from_expression(self, expression: exp.Expression) -> str:
         from sqlframe.base.session import _BaseSession
@@ -1503,20 +1512,23 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         """
         return func(self, *args, **kwargs)  # type: ignore
 
-    @operation(Operation.SELECT)
+    @operation(Operation.SELECT_CONSTRAINED)
     def withColumn(self, colName: str, col: Column) -> Self:
         return self.withColumns.__wrapped__(self, {colName: col})  # type: ignore
 
-    @operation(Operation.SELECT)
+    @operation(Operation.SELECT_CONSTRAINED)
     def withColumnRenamed(self, existing: str, new: str) -> Self:
+        col_func = get_func_from_session("col", self.session)
         expression = self.expression.copy()
         existing = self.session._normalize_string(existing)
-        columns = self._get_outer_select_columns(expression)
+        outer_expressions = self._get_outer_select_expressions(expression)
         results = []
         found_match = False
-        for column in columns:
-            if column.alias_or_name == existing:
-                column = column.alias(new)
+        for expr in outer_expressions:
+            column = col_func(expr.copy())
+            if existing == quote_preserving_alias_or_name(expr):
+                if isinstance(column.expression, exp.Alias):
+                    column.expression.set("alias", exp.to_identifier(new))
                 self._update_display_name_mapping([column], [new])
                 found_match = True
             results.append(column)
@@ -1524,7 +1536,7 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
             raise ValueError("Tried to rename a column that doesn't exist")
         return self.select.__wrapped__(self, *results, skip_update_display_name_mapping=True)  # type: ignore
 
-    @operation(Operation.SELECT)
+    @operation(Operation.SELECT_CONSTRAINED)
     def withColumnsRenamed(self, colsMap: t.Dict[str, str]) -> Self:
         """
         Returns a new :class:`DataFrame` by renaming multiple columns. If a non-existing column is
@@ -1570,7 +1582,7 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
 
         return self.select.__wrapped__(self, *results, skip_update_display_name_mapping=True)  # type: ignore
 
-    @operation(Operation.SELECT)
+    @operation(Operation.SELECT_CONSTRAINED)
     def withColumns(self, *colsMap: t.Dict[str, Column]) -> Self:
         """
         Returns a new :class:`DataFrame` by adding multiple columns or replacing the
@@ -1608,13 +1620,14 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         """
         if len(colsMap) != 1:
             raise ValueError("Only a single map is supported")
+        col_func = get_func_from_session("col")
         col_map = {
             self._ensure_and_normalize_col(k): (self._ensure_and_normalize_col(v), k)
             for k, v in colsMap[0].items()
         }
-        existing_cols = self._get_outer_select_columns(self.expression)
-        existing_col_names = [x.alias_or_name for x in existing_cols]
-        select_columns = existing_cols
+        existing_expr = self._get_outer_select_expressions(self.expression)
+        existing_col_names = [x.alias_or_name for x in existing_expr]
+        select_columns = [col_func(x) for x in existing_expr]
         for col, (col_value, display_name) in col_map.items():
             column_name = col.alias_or_name
             existing_col_index = (
@@ -1631,16 +1644,32 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         )
         return self.select.__wrapped__(self, *select_columns, skip_update_display_name_mapping=True)  # type: ignore
 
-    @operation(Operation.SELECT)
+    @operation(Operation.SELECT_CONSTRAINED)
     def drop(self, *cols: t.Union[str, Column]) -> Self:
-        all_columns = self._get_outer_select_columns(self.expression)
-        drop_cols = self._ensure_and_normalize_cols(cols)
-        new_columns = [
-            col
-            for col in all_columns
-            if col.alias_or_name not in [drop_column.alias_or_name for drop_column in drop_cols]
-        ]
-        return self.copy().select(*new_columns, append=False)
+        # Separate string column names from Column objects for different handling
+        column_objs, column_names = partition_to(lambda x: isinstance(x, str), cols, list, set)
+
+        # Normalize only the Column objects (strings will be handled as unqualified)
+        drop_cols = self._ensure_and_normalize_cols(column_objs) if column_objs else []
+
+        # Work directly with the expression's select columns to preserve table qualifiers
+        current_expressions = self.expression.expressions
+        drop_sql = {drop_col.expression.sql() for drop_col in drop_cols}
+
+        # Create a more sophisticated matching function that considers table qualifiers
+        def should_drop_expression(expr: exp.Expression) -> bool:
+            # Check against fully qualified Column objects and
+            # Check against unqualified string column names (drop ALL columns with this name)
+            if expr.sql() in drop_sql or (
+                isinstance(expr, exp.Column) and expr.alias_or_name in column_names
+            ):
+                return True
+            return False
+
+        new_expressions = [expr for expr in current_expressions if not should_drop_expression(expr)]
+        return self.select.__wrapped__(  # type: ignore
+            self, *new_expressions, skip_update_display_name_mapping=True
+        )
 
     @operation(Operation.LIMIT)
     def limit(self, num: int) -> Self:

{sqlframe-3.39.2 → sqlframe-3.39.3}/sqlframe/base/functions.py

@@ -1450,6 +1450,9 @@ def unix_timestamp(
 
     session = _get_session()
 
+    if session._is_duckdb or session._is_postgres or session._is_snowflake or session._is_bigquery:
+        timestamp = Column.ensure_col(timestamp).cast("string")
+
     if session._is_bigquery:
         return unix_timestamp_bgutil(timestamp, format)
 
@@ -6342,6 +6345,7 @@ def to_unix_timestamp(
 
     if session._is_duckdb:
         format = format or _BaseSession().default_time_format
+        timestamp = Column.ensure_col(timestamp).cast("string")
 
     if format is not None:
         return Column.invoke_expression_over_column(

{sqlframe-3.39.2 → sqlframe-3.39.3}/sqlframe/base/operations.py

@@ -27,9 +27,10 @@ class Operation(IntEnum):
     WHERE = 2
     GROUP_BY = 3
     HAVING = 4
-    SELECT = 5
-    ORDER_BY = 6
-    LIMIT = 7
+    SELECT_CONSTRAINED = 5
+    SELECT = 6
+    ORDER_BY = 7
+    LIMIT = 8
 
 
 # We want to decorate a function (self: DF, *args, **kwargs) -> T

{sqlframe-3.39.2 → sqlframe-3.39.3}/sqlframe/base/session.py

@@ -179,7 +179,7 @@ class _BaseSession(t.Generic[CATALOG, READER, WRITER, DF, TABLE, CONN, UDF_REGIS
         return self._table(self, *args, **kwargs)
 
     def __new__(cls, *args, **kwargs):
-        if _BaseSession._instance is None:
+        if _BaseSession._instance is None or not isinstance(_BaseSession._instance, cls):
             _BaseSession._instance = super().__new__(cls)
         return _BaseSession._instance
 
@@ -194,6 +194,11 @@ class _BaseSession(t.Generic[CATALOG, READER, WRITER, DF, TABLE, CONN, UDF_REGIS
     def getActiveSession(self) -> Self:
         return self
 
+    def stop(self) -> None:
+        if connection := getattr(self, "_connection", None):
+            connection.close()
+        _BaseSession._instance = None
+
     def range(
         self,
         start: int,

{sqlframe-3.39.2 → sqlframe-3.39.3}/sqlframe/base/util.py

@@ -6,6 +6,7 @@ import string
 import typing as t
 import unicodedata
 
+from more_itertools import partition
 from sqlglot import expressions as exp
 from sqlglot import parse_one, to_table
 from sqlglot.dialects import DuckDB
@@ -537,3 +538,17 @@ def is_relativedelta_like(value: t.Any) -> bool:
         and hasattr(value, "weeks")
         and hasattr(value, "leapdays")
     )
+
+
+T = t.TypeVar("T")
+R1 = t.TypeVar("R1")
+R2 = t.TypeVar("R2")
+
+
+def partition_to(
+    pred: t.Callable[[T], bool],
+    iterable: t.Iterable[T],
+    result1: t.Type[R1],
+    result2: t.Type[R2],
+) -> tuple[R1, R2]:
+    return (lambda x, y: (result1(x), result2(y)))(*partition(pred, iterable))  # type: ignore

{sqlframe-3.39.2 → sqlframe-3.39.3}/sqlframe.egg-info/PKG-INFO

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

{sqlframe-3.39.2 → sqlframe-3.39.3}/sqlframe.egg-info/requires.txt

@@ -1,3 +1,4 @@
+more-itertools
 prettytable<4
 sqlglot<27.9,>=24.0.0
 typing_extensions

sqlframe-3.39.3/tests/integration/engines/duck/test_duckdb_session.py

@@ -0,0 +1,38 @@
+import pytest
+from duckdb.duckdb import ConnectionException
+from sqlglot import exp
+
+from sqlframe.base.types import Row
+from sqlframe.duckdb.session import DuckDBSession
+
+pytest_plugins = ["tests.common_fixtures"]
+
+
+def test_session_from_config():
+    import duckdb
+
+    conn = duckdb.connect()
+    conn.execute("CREATE TABLE test_table (cola INT, colb STRING)")
+    session = DuckDBSession.builder.config("sqlframe.conn", conn).getOrCreate()
+    columns = session.catalog.get_columns("test_table")
+    assert columns == {"cola": exp.DataType.build("INT"), "colb": exp.DataType.build("TEXT")}
+    assert session.execution_dialect_name == "duckdb"
+
+
+@pytest.mark.forked
+def test_session_stop(duckdb_session: DuckDBSession):
+    assert duckdb_session.range(1, 2).collect() == [Row(id=1)]
+    duckdb_session.stop()
+    with pytest.raises(ConnectionException):
+        duckdb_session.range(1, 10).collect()
+
+
+@pytest.mark.forked
+def test_session_new_session(duckdb_session: DuckDBSession):
+    # Remove old session
+    assert duckdb_session.range(1, 2).collect() == [Row(id=1)]
+    duckdb_session.stop()
+    new_session = DuckDBSession.builder.getOrCreate()
+    assert new_session is not duckdb_session
+    assert isinstance(new_session, DuckDBSession)
+    assert new_session.range(1, 2).collect() == [Row(id=1)]

{sqlframe-3.39.2 → sqlframe-3.39.3}/tests/integration/engines/test_int_functions.py

@@ -1502,11 +1502,15 @@ def test_from_unixtime(get_session_and_func):
     assert df.select(from_unixtime("unix_time").alias("ts")).first()[0] == expected
 
 
-def test_unix_timestamp(get_session_and_func):
+def test_unix_timestamp(get_session_and_func, get_func):
     session, unix_timestamp = get_session_and_func("unix_timestamp")
     df = session.createDataFrame([("2015-04-08",)], ["dt"])
     result = df.select(unix_timestamp("dt", "yyyy-MM-dd").alias("unix_time")).first()[0]
     assert result == 1428451200
+    ts_type = "TIMESTAMP" if isinstance(session, PySparkSession) else "TIMESTAMPNTZ"
+    df = session.createDataFrame([(datetime.datetime(2015, 4, 8),)], schema=f"ts {ts_type}")
+    result = df.select(unix_timestamp("ts").alias("unix_time")).first()[0]
+    assert result == 1428451200
 
 
 def test_from_utc_timestamp(get_session_and_func):
@@ -5003,6 +5007,12 @@ def test_to_unix_timestamp(get_session_and_func, get_func):
     else:
         df = session.createDataFrame([("2016-04-08",)], ["e"])
         assert df.select(to_unix_timestamp(df.e).alias("r")).collect() == [Row(r=None)]
+    ts_type = "TIMESTAMP" if isinstance(session, PySparkSession) else "TIMESTAMPNTZ"
+    df = session.createDataFrame([(datetime.datetime(2015, 4, 8),)], schema=f"ts {ts_type}")
+    result = df.select(
+        to_unix_timestamp("ts", lit("yyyy-MM-dd HH:mm:ss")).alias("unix_time")
+    ).first()[0]
+    assert result == 1428451200
 
 
 def test_to_varchar(get_session_and_func, get_func):

{sqlframe-3.39.2 → sqlframe-3.39.3}/tests/integration/test_int_dataframe.py

@@ -1892,6 +1892,97 @@ def test_drop_column_reference_join(
     compare_frames(df, dfs, sort=True)
 
 
+# TODO: This test exposes a bug in Spark where the column order is not preserved after a join
+# for some reason spark flips foo/bar to bar/foo
+# def test_left_join_column_order(
+#     pyspark_employee: PySparkDataFrame,
+#     get_df: t.Callable[[str], BaseDataFrame],
+#     compare_frames: t.Callable,
+# ):
+#     df1 = pyspark_employee.sparkSession.createDataFrame([{"foo": 0, "bar": "a"}, {"foo": 1, "bar": "b"}]).alias("df1")
+#     df2 = pyspark_employee.sparkSession.createDataFrame([{"foo": 0, "baz": 1.5}]).alias("df2")
+#     df_joined = df1.join(df2, on=F.col("df1.foo") == F.col("df2.foo"), how="left")
+#
+#     employee = get_df("employee")
+#     dfs1 = employee.sparkSession.createDataFrame([{"foo": 0, "bar": "a"}, {"foo": 1, "bar": "b"}]).alias("dfs1")
+#     dfs2 = employee.sparkSession.createDataFrame([{"foo": 0, "baz": 1.5}]).alias("dfs2")
+#     dfs_joined = dfs1.join(dfs2, on=SF.col("dfs1.foo") == SF.col("dfs2.foo"), how="left")
+#
+#     compare_frames(df_joined, dfs_joined, sort=True)
+
+
+# def test_drop_column_join_column(
+#     pyspark_employee: PySparkDataFrame,
+#     get_df: t.Callable[[str], BaseDataFrame],
+#     compare_frames: t.Callable,
+# ):
+#     df1 = pyspark_employee.sparkSession.createDataFrame([{"foo": 0, "bar": "a"}, {"foo": 1, "bar": "b"}]).alias("df1")
+#     df2 = pyspark_employee.sparkSession.createDataFrame([{"foo": 0, "baz": 1.5}]).alias("df2")
+#     df_joined = df1.join(df2, on=F.col("df1.foo") == F.col("df2.foo"), how="left").drop(F.col("df1.foo"))
+#
+#     employee = get_df("employee")
+#     dfs1 = employee.sparkSession.createDataFrame([{"foo": 0, "bar": "a"}, {"foo": 1, "bar": "b"}]).alias("dfs1")
+#     dfs2 = employee.sparkSession.createDataFrame([{"foo": 0, "baz": 1.5}]).alias("dfs2")
+#     dfs_joined = dfs1.join(dfs2, on=SF.col("dfs1.foo") == SF.col("dfs2.foo"), how="left").drop(SF.col("dfs1.foo"))
+#
+#     compare_frames(df_joined, dfs_joined, sort=True)
+
+
+def test_drop_column_join_column_df_reference(
+    pyspark_employee: PySparkDataFrame,
+    get_df: t.Callable[[str], BaseDataFrame],
+    compare_frames: t.Callable,
+):
+    df1 = pyspark_employee.sparkSession.createDataFrame(  # type: ignore
+        [{"foo": 0, "bar": "a"}, {"foo": 1, "bar": "b"}]
+    ).alias("df1")
+    df2 = pyspark_employee.sparkSession.createDataFrame([{"foo": 0, "baz": 1.5}]).alias("df2")  # type: ignore
+    df_joined = (
+        df1.join(df2, on=F.col("df1.foo") == F.col("df2.foo"), how="left")
+        .drop(df1.foo)
+        .select(df1.bar, df2.baz, df2.foo)
+    )
+
+    employee = get_df("employee")
+    dfs1 = employee.sparkSession.createDataFrame(
+        [{"foo": 0, "bar": "a"}, {"foo": 1, "bar": "b"}]
+    ).alias("dfs1")
+    dfs2 = employee.sparkSession.createDataFrame([{"foo": 0, "baz": 1.5}]).alias("dfs2")
+    dfs_joined = (
+        dfs1.join(dfs2, on=SF.col("dfs1.foo") == SF.col("dfs2.foo"), how="left")
+        .drop(dfs1.foo)
+        .select(dfs1.bar, dfs2.baz, dfs2.foo)
+    )
+
+    compare_frames(df_joined, dfs_joined, sort=True, compare_schema=False)
+
+
+def test_drop_join_column_unqualified(
+    pyspark_employee: PySparkDataFrame,
+    get_df: t.Callable[[str], BaseDataFrame],
+    compare_frames: t.Callable,
+):
+    df1 = pyspark_employee.sparkSession.createDataFrame(  # type: ignore
+        [{"foo": 0, "bar": "a"}, {"foo": 1, "bar": "b"}]
+    ).alias("df1")
+    df2 = pyspark_employee.sparkSession.createDataFrame([{"foo": 0, "baz": 1.5}]).alias("df2")  # type: ignore
+    df_joined = (
+        df1.join(df2, on=F.col("df1.foo") == F.col("df2.foo"), how="left").drop("foo")
+        # select the columns to work around column order bug
+    )
+
+    employee = get_df("employee")
+    dfs1 = employee.sparkSession.createDataFrame(
+        [{"foo": 0, "bar": "a"}, {"foo": 1, "bar": "b"}]
+    ).alias("dfs1")
+    dfs2 = employee.sparkSession.createDataFrame([{"foo": 0, "baz": 1.5}]).alias("dfs2")
+    dfs_joined = dfs1.join(dfs2, on=SF.col("dfs1.foo") == SF.col("dfs2.foo"), how="left").drop(
+        "foo"
+    )
+
+    compare_frames(df_joined, dfs_joined, sort=True, compare_schema=False)
+
+
 def test_limit(
     pyspark_employee: PySparkDataFrame,
     get_df: t.Callable[[str], BaseDataFrame],
@@ -2167,10 +2258,10 @@ def test_transform(
     get_df: t.Callable[[str], BaseDataFrame],
     compare_frames: t.Callable,
 ):
-    def cast_all_to_int_pyspark(input_df):
+    def cast_all_to_string_pyspark(input_df):
         return input_df.select([F.col(col_name).cast("string") for col_name in input_df.columns])
 
-    def cast_all_to_int_sqlframe(input_df):
+    def cast_all_to_string_sqlframe(input_df):
         return input_df.select([SF.col(col_name).cast("string") for col_name in input_df.columns])
 
     def sort_columns_asc(input_df):
@@ -2178,8 +2269,8 @@ def test_transform(
 
     employee = get_df("employee")
 
-    df = pyspark_employee.transform(cast_all_to_int_pyspark).transform(sort_columns_asc)
-    dfs = employee.transform(cast_all_to_int_sqlframe).transform(sort_columns_asc)
+    df = pyspark_employee.transform(cast_all_to_string_pyspark).transform(sort_columns_asc)
+    dfs = employee.transform(cast_all_to_string_sqlframe).transform(sort_columns_asc)
     compare_frames(df, dfs)
 
 

{sqlframe-3.39.2 → sqlframe-3.39.3}/tests/unit/standalone/test_session.py

@@ -136,24 +136,3 @@ def test_sql_insert(standalone_session: StandaloneSession, compare_sql: t.Callab
 
 def test_session_create_builder_patterns():
     assert StandaloneSession.builder.appName("abc").getOrCreate() == StandaloneSession()
-
-
-# @pytest.mark.parametrize(
-#     "input, expected",
-#     [
-#         (
-#             StandaloneSession._to_row(["a"], [1]),
-#             types.Row(a=1),
-#         ),
-#         (
-#             StandaloneSession._to_row(["a", "b"], [1, 2]),
-#             types.Row(a=1, b=2),
-#         ),
-#         (
-#             StandaloneSession._to_row(["a", "a"], [1, 2]),
-#             types.Row(a=1, a=2),
-#         ),
-#     ],
-# )
-# def test_to_row(input, expected):
-#     assert input == expected

sqlframe-3.39.2/tests/integration/engines/duck/test_duckdb_session.py

@@ -1,14 +0,0 @@
-from sqlglot import exp
-
-from sqlframe.duckdb.session import DuckDBSession
-
-
-def test_session_from_config():
-    import duckdb
-
-    conn = duckdb.connect()
-    conn.execute("CREATE TABLE test_table (cola INT, colb STRING)")
-    session = DuckDBSession.builder.config("sqlframe.conn", conn).getOrCreate()
-    columns = session.catalog.get_columns("test_table")
-    assert columns == {"cola": exp.DataType.build("INT"), "colb": exp.DataType.build("TEXT")}
-    assert session.execution_dialect_name == "duckdb"