sqlframe 3.39.2__tar.gz → 3.39.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 3.39.2
+Version: 3.39.4
 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.4}/setup.py

@@ -20,8 +20,9 @@ setup(
     python_requires=">=3.9",
     install_requires=[
         "prettytable<4",
-        "sqlglot>=24.0.0,<27.9",
+        "sqlglot>=24.0.0,<27.10",
         "typing_extensions",
+        "more-itertools",
     ],
     extras_require={
         "bigquery": [

{sqlframe-3.39.2 → sqlframe-3.39.4}/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.4'
+__version_tuple__ = version_tuple = (3, 39, 4)
 
-__commit_id__ = commit_id = 'g772b3a6bf'
+__commit_id__ = commit_id = 'g7103a1e73'

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

@@ -31,6 +31,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,
@@ -1633,14 +1634,30 @@ class BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
 
     @operation(Operation.SELECT)
     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.4}/sqlframe/base/functions.py

@@ -37,9 +37,7 @@ def _get_session() -> _BaseSession:
 
 @meta()
 def col(column_name: t.Union[ColumnOrName, t.Any]) -> Column:
-    from sqlframe.base.session import _BaseSession
-
-    dialect = _BaseSession().input_dialect
+    dialect = _get_session().input_dialect
     if isinstance(column_name, str):
         col_expression = expression.to_column(column_name, dialect=dialect).transform(
             dialect.normalize_identifier
@@ -662,9 +660,7 @@ def grouping_id(*cols: ColumnOrName) -> Column:
 
 @meta()
 def input_file_name() -> Column:
-    from sqlframe.base.session import _BaseSession
-
-    return Column(expression.Literal.string(_BaseSession()._last_loaded_file or ""))
+    return Column(expression.Literal.string(_get_session()._last_loaded_file or ""))
 
 
 @meta()
@@ -959,12 +955,10 @@ def current_timestamp() -> Column:
 
 @meta()
 def date_format(col: ColumnOrName, format: str) -> Column:
-    from sqlframe.base.session import _BaseSession
-
     return Column.invoke_expression_over_column(
         Column(expression.TimeStrToTime(this=Column.ensure_col(col).column_expression)),
         expression.TimeToStr,
-        format=_BaseSession().format_time(format),
+        format=_get_session().format_time(format),
     )
 
 
@@ -1450,6 +1444,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)
 
@@ -3375,10 +3372,9 @@ def get(col: ColumnOrName, index: t.Union[ColumnOrName, int]) -> Column:
 def get_active_spark_context() -> SparkContext:
     """Raise RuntimeError if SparkContext is not initialized,
     otherwise, returns the active SparkContext."""
-    from sqlframe.base.session import _BaseSession
     from sqlframe.spark.session import SparkSession
 
-    session: _BaseSession = _BaseSession()
+    session = _get_session()
     if not isinstance(session, SparkSession):
         raise RuntimeError("This function is only available in SparkSession.")
     return session.spark_session.sparkContext
@@ -6341,7 +6337,8 @@ def to_unix_timestamp(
     session = _get_session()
 
     if session._is_duckdb:
-        format = format or _BaseSession().default_time_format
+        format = format or session.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.4}/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.4}/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.4}/sqlframe.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 3.39.2
+Version: 3.39.4
 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.4}/sqlframe.egg-info/requires.txt

@@ -1,5 +1,6 @@
+more-itertools
 prettytable<4
-sqlglot<27.9,>=24.0.0
+sqlglot<27.10,>=24.0.0
 typing_extensions
 
 [bigquery]

sqlframe-3.39.4/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.4}/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.4}/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"