sqlframe 3.9.2__tar.gz → 3.10.0__tar.gz

{sqlframe-3.9.2 → sqlframe-3.10.0}/Makefile

@@ -22,6 +22,9 @@ duckdb-test:
 snowflake-test:
 	pytest -n auto -m "snowflake"
 
+databricks-test:
+	pytest -n auto -m "databricks"
+
 style:
 	pre-commit run --all-files
 

{sqlframe-3.9.2 → sqlframe-3.10.0}/PKG-INFO

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

{sqlframe-3.9.2 → sqlframe-3.10.0}/docs/databricks.md

@@ -1,5 +1,3 @@
-from test import auth_type
-
 # Databricks (In Development)
 
 ## Installation

{sqlframe-3.9.2 → sqlframe-3.10.0}/setup.py

@@ -20,7 +20,7 @@ setup(
     python_requires=">=3.8",
     install_requires=[
         "prettytable<3.12.1",
-        "sqlglot>=24.0.0,<25.33",
+        "sqlglot>=24.0.0,<26.1",
         "typing_extensions>=4.8,<5",
     ],
     extras_require={
@@ -32,13 +32,14 @@ setup(
             "duckdb>=0.9,<1.2",
             "findspark>=2,<3",
             "mypy>=1.10.0,<1.14",
-            "openai>=1.30,<1.56",
+            "openai>=1.30,<1.58",
             "pandas>=2,<3",
             "pandas-stubs>=2,<3",
             "psycopg>=3.1,<4",
             "pyarrow>=10,<19",
             "pyspark>=2,<3.6",
             "pytest>=8.2.0,<8.4",
+            "pytest-forked",
             "pytest-postgresql>=6,<7",
             "pytest-xdist>=3.6,<3.7",
             "pre-commit>=3.5;python_version=='3.8'",
@@ -58,7 +59,7 @@ setup(
             "pandas>=2,<3",
         ],
         "openai": [
-            "openai>=1.30,<1.56",
+            "openai>=1.30,<1.58",
         ],
         "pandas": [
             "pandas>=2,<3",

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.9.2'
-__version_tuple__ = version_tuple = (3, 9, 2)
+__version__ = version = '3.10.0'
+__version_tuple__ = version_tuple = (3, 10, 0)

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe/base/function_alternatives.py

@@ -64,6 +64,39 @@ def first_always_ignore_nulls(col: ColumnOrName, ignorenulls: t.Optional[bool] =
     return first(col)
 
 
+def to_timestamp_with_time_zone(col: ColumnOrName, format: t.Optional[str] = None) -> Column:
+    from sqlframe.base.session import _BaseSession
+
+    if format is not None:
+        return Column.invoke_expression_over_column(
+            col, expression.StrToTime, format=_BaseSession().format_time(format)
+        )
+
+    return Column.ensure_col(col).cast("timestamp with time zone", dialect="postgres")
+
+
+def to_timestamp_tz(col: ColumnOrName, format: t.Optional[str] = None) -> Column:
+    from sqlframe.base.session import _BaseSession
+
+    if format is not None:
+        return Column.invoke_expression_over_column(
+            col, expression.StrToTime, format=_BaseSession().format_time(format)
+        )
+
+    return Column.ensure_col(col).cast("timestamptz", dialect="duckdb")
+
+
+def to_timestamp_just_timestamp(col: ColumnOrName, format: t.Optional[str] = None) -> Column:
+    from sqlframe.base.session import _BaseSession
+
+    if format is not None:
+        return Column.invoke_expression_over_column(
+            col, expression.StrToTime, format=_BaseSession().format_time(format)
+        )
+
+    return Column.ensure_col(col).cast("datetime", dialect="bigquery")
+
+
 def bitwise_not_from_bitnot(col: ColumnOrName) -> Column:
     return Column.invoke_anonymous_function(col, "BITNOT")
 
@@ -1220,6 +1253,11 @@ def get_json_object_cast_object(col: ColumnOrName, path: str) -> Column:
     return get_json_object(col_func(col).cast("variant"), path)
 
 
+def get_json_object_using_function(col: ColumnOrName, path: str) -> Column:
+    lit = get_func_from_session("lit")
+    return Column.invoke_anonymous_function(col, "GET_JSON_OBJECT", lit(path))
+
+
 def create_map_with_cast(*cols: t.Union[ColumnOrName, t.Iterable[ColumnOrName]]) -> Column:
     from sqlframe.base.functions import create_map
 

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe/base/functions.py

@@ -900,7 +900,7 @@ def to_timestamp(col: ColumnOrName, format: t.Optional[str] = None) -> Column:
             col, expression.StrToTime, format=_BaseSession().format_time(format)
         )
 
-    return Column.ensure_col(col).cast("timestamp")
+    return Column.ensure_col(col).cast("timestampltz")
 
 
 @meta()
@@ -2173,7 +2173,7 @@ def current_database() -> Column:
 current_schema = current_database
 
 
-@meta(unsupported_engines="*")
+@meta(unsupported_engines=["*", "databricks"])
 def current_timezone() -> Column:
     return Column.invoke_anonymous_function(None, "current_timezone")
 
@@ -2261,7 +2261,7 @@ def get(col: ColumnOrName, index: t.Union[ColumnOrName, int]) -> Column:
     return Column.invoke_anonymous_function(col, "get", index)
 
 
-@meta(unsupported_engines="*")
+@meta(unsupported_engines=["*", "databricks"])
 def get_active_spark_context() -> SparkContext:
     """Raise RuntimeError if SparkContext is not initialized,
     otherwise, returns the active SparkContext."""
@@ -2778,7 +2778,7 @@ def isnotnull(col: ColumnOrName) -> Column:
     return Column.invoke_anonymous_function(col, "isnotnull")
 
 
-@meta(unsupported_engines="*")
+@meta(unsupported_engines=["*", "databricks"])
 def java_method(*cols: ColumnOrName) -> Column:
     """
     Calls a method with reflection.
@@ -3050,7 +3050,7 @@ def ln(col: ColumnOrName) -> Column:
     return Column.invoke_expression_over_column(col, expression.Ln)
 
 
-@meta(unsupported_engines="*")
+@meta(unsupported_engines=["*", "databricks"])
 def localtimestamp() -> Column:
     """
     Returns the current timestamp without time zone at the start of query evaluation
@@ -3080,7 +3080,7 @@ def localtimestamp() -> Column:
     return Column.invoke_anonymous_function(None, "localtimestamp")
 
 
-@meta(unsupported_engines="*")
+@meta(unsupported_engines=["*", "databricks"])
 def make_dt_interval(
     days: t.Optional[ColumnOrName] = None,
     hours: t.Optional[ColumnOrName] = None,
@@ -3227,7 +3227,7 @@ def make_timestamp(
         )
 
 
-@meta(unsupported_engines="*")
+@meta(unsupported_engines=["*", "databricks"])
 def make_timestamp_ltz(
     years: ColumnOrName,
     months: ColumnOrName,
@@ -3354,7 +3354,7 @@ def make_timestamp_ntz(
     )
 
 
-@meta(unsupported_engines="*")
+@meta(unsupported_engines=["*", "databricks"])
 def make_ym_interval(
     years: t.Optional[ColumnOrName] = None,
     months: t.Optional[ColumnOrName] = None,
@@ -3922,7 +3922,7 @@ def printf(format: ColumnOrName, *cols: ColumnOrName) -> Column:
     return Column.invoke_anonymous_function(format, "printf", *cols)
 
 
-@meta(unsupported_engines=["*", "spark"])
+@meta(unsupported_engines=["*", "spark", "databricks"])
 def product(col: ColumnOrName) -> Column:
     """
     Aggregate function: returns the product of the values in a group.
@@ -3961,7 +3961,7 @@ def product(col: ColumnOrName) -> Column:
 reduce = aggregate
 
 
-@meta(unsupported_engines="*")
+@meta(unsupported_engines=["*", "databricks"])
 def reflect(*cols: ColumnOrName) -> Column:
     """
     Calls a method with reflection.
@@ -5046,7 +5046,7 @@ def to_str(value: t.Any) -> t.Optional[str]:
         return str(value)
 
 
-@meta(unsupported_engines="*")
+@meta(unsupported_engines=["*", "databricks"])
 def to_timestamp_ltz(
     timestamp: ColumnOrName,
     format: t.Optional[ColumnOrName] = None,

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe/base/session.py

@@ -570,6 +570,8 @@ class _BaseSession(t.Generic[CATALOG, READER, WRITER, DF, CONN, UDF_REGISTRATION
             return cls._to_row(list(value.keys()), list(value.values()))
         elif isinstance(value, (list, set, tuple)) and value:
             return [cls._to_value(x) for x in value]
+        elif isinstance(value, datetime.datetime):
+            return value.replace(tzinfo=None)
         return value
 
     @classmethod

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe/bigquery/functions.py

@@ -73,6 +73,7 @@ from sqlframe.base.function_alternatives import (  # noqa
     _is_string_using_typeof_string as _is_string,
     array_append_using_array_cat as array_append,
     endswith_with_underscore as endswith,
+    to_timestamp_just_timestamp as to_timestamp,
 )
 
 

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe/databricks/catalog.py

@@ -26,7 +26,6 @@ if t.TYPE_CHECKING:
 
 
 class DatabricksCatalog(
-    SetCurrentCatalogFromUseMixin["DatabricksSession", "DatabricksDataFrame"],
     GetCurrentCatalogFromFunctionMixin["DatabricksSession", "DatabricksDataFrame"],
     GetCurrentDatabaseFromFunctionMixin["DatabricksSession", "DatabricksDataFrame"],
     ListDatabasesFromInfoSchemaMixin["DatabricksSession", "DatabricksDataFrame"],
@@ -38,6 +37,15 @@ class DatabricksCatalog(
     CURRENT_CATALOG_EXPRESSION: exp.Expression = exp.func("current_catalog")
     UPPERCASE_INFO_SCHEMA = True
 
+    def setCurrentCatalog(self, catalogName: str) -> None:
+        self.session._collect(
+            exp.Use(
+                kind=exp.Var(this=exp.to_identifier("CATALOG")),
+                this=exp.parse_identifier(catalogName, dialect=self.session.input_dialect),
+            ),
+            quote_identifiers=False,
+        )
+
     def listFunctions(
         self, dbName: t.Optional[str] = None, pattern: t.Optional[str] = None
     ) -> t.List[Function]:
@@ -106,7 +114,9 @@ class DatabricksCatalog(
         )
         functions = [
             Function(
-                name=normalize_string(x["function"], from_dialect="execution", to_dialect="output"),
+                name=normalize_string(
+                    x["function"].split(".")[-1], from_dialect="execution", to_dialect="output"
+                ),
                 catalog=normalize_string(
                     schema.catalog, from_dialect="execution", to_dialect="output"
                 ),

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe/databricks/dataframe.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import logging
-import sys
 import typing as t
 
 from sqlframe.base.catalog import Column as CatalogColumn
@@ -52,7 +51,9 @@ class DatabricksDataFrame(
             columns.append(
                 CatalogColumn(
                     name=normalize_string(
-                        row.col_name, from_dialect="execution", to_dialect="output"
+                        row.col_name,
+                        from_dialect="execution",
+                        to_dialect="output",
                     ),
                     dataType=normalize_string(
                         row.data_type,

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe/databricks/functions.py

@@ -19,4 +19,5 @@ from sqlframe.base.function_alternatives import (  # noqa
     arrays_overlap_renamed as arrays_overlap,
     _is_string_using_typeof_string_lcase as _is_string,
     try_element_at_zero_based as try_element_at,
+    get_json_object_using_function as get_json_object,
 )

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe/databricks/session.py

@@ -44,7 +44,20 @@ class DatabricksSession(
         from databricks import sql
 
         if not hasattr(self, "_conn"):
-            super().__init__(conn or sql.connect(server_hostname, http_path, access_token))
+            super().__init__(
+                conn or sql.connect(server_hostname, http_path, access_token, disable_pandas=True)
+            )
+
+    @classmethod
+    def _try_get_map(cls, value: t.Any) -> t.Optional[t.Dict[str, t.Any]]:
+        if (
+            value
+            and isinstance(value, list)
+            and all(isinstance(item, tuple) for item in value)
+            and all(len(item) == 2 for item in value)
+        ):
+            return dict(value)
+        return None
 
     class Builder(_BaseSession.Builder):
         DEFAULT_EXECUTION_DIALECT = "databricks"

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe/duckdb/functions.py

@@ -53,4 +53,5 @@ from sqlframe.base.function_alternatives import (  # noqa
     endswith_with_underscore as endswith,
     last_day_with_cast as last_day,
     regexp_replace_global_option as regexp_replace,
+    to_timestamp_tz as to_timestamp,
 )

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe/duckdb/functions.pyi

@@ -29,6 +29,7 @@ from sqlframe.base.function_alternatives import (  # noqa
     try_element_at_zero_based as try_element_at,
     to_unix_timestamp_include_default_format as to_unix_timestamp,
     regexp_replace_global_option as regexp_replace,
+    to_timestamp_tz as to_timestamp,
 )
 from sqlframe.base.functions import (
     abs as abs,

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe/postgres/functions.py

@@ -70,4 +70,5 @@ from sqlframe.base.function_alternatives import (  # noqa
     endswith_using_like as endswith,
     last_day_with_cast as last_day,
     regexp_replace_global_option as regexp_replace,
+    to_timestamp_with_time_zone as to_timestamp,
 )

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe.egg-info/PKG-INFO

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

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe.egg-info/SOURCES.txt

@@ -316,6 +316,10 @@ tests/integration/engines/bigquery/__init__.py
 tests/integration/engines/bigquery/test_bigquery_catalog.py
 tests/integration/engines/bigquery/test_bigquery_dataframe.py
 tests/integration/engines/bigquery/test_bigquery_session.py
+tests/integration/engines/databricks/__init__.py
+tests/integration/engines/databricks/test_databricks_catalog.py
+tests/integration/engines/databricks/test_databricks_dataframe.py
+tests/integration/engines/databricks/test_databricks_session.py
 tests/integration/engines/duck/__init__.py
 tests/integration/engines/duck/test_duckdb_activate.py
 tests/integration/engines/duck/test_duckdb_catalog.py
@@ -345,6 +349,8 @@ tests/unit/test_activate.py
 tests/unit/test_util.py
 tests/unit/bigquery/__init__.py
 tests/unit/bigquery/test_activate.py
+tests/unit/databricks/__init__.py
+tests/unit/databricks/test_activate.py
 tests/unit/duck/__init__.py
 tests/unit/duck/test_activate.py
 tests/unit/postgres/__init__.py

{sqlframe-3.9.2 → sqlframe-3.10.0}/sqlframe.egg-info/requires.txt

@@ -1,5 +1,5 @@
 prettytable<3.12.1
-sqlglot<25.33,>=24.0.0
+sqlglot<26.1,>=24.0.0
 typing_extensions<5,>=4.8
 
 [bigquery]
@@ -13,12 +13,13 @@ databricks-sql-connector<4,>=3.6
 duckdb<1.2,>=0.9
 findspark<3,>=2
 mypy<1.14,>=1.10.0
-openai<1.56,>=1.30
+openai<1.58,>=1.30
 pandas-stubs<3,>=2
 pandas<3,>=2
 psycopg<4,>=3.1
 pyarrow<19,>=10
 pyspark<3.6,>=2
+pytest-forked
 pytest-postgresql<7,>=6
 pytest-xdist<3.7,>=3.6
 pytest<8.4,>=8.2.0
@@ -43,7 +44,7 @@ duckdb<1.2,>=0.9
 pandas<3,>=2
 
 [openai]
-openai<1.56,>=1.30
+openai<1.58,>=1.30
 
 [pandas]
 pandas<3,>=2

{sqlframe-3.9.2 → sqlframe-3.10.0}/tests/common_fixtures.py

@@ -12,6 +12,7 @@ from pytest_postgresql.janitor import DatabaseJanitor
 
 from sqlframe.base.session import _BaseSession
 from sqlframe.bigquery.session import BigQuerySession
+from sqlframe.databricks.session import DatabricksSession
 from sqlframe.duckdb.session import DuckDBSession
 from sqlframe.postgres.session import PostgresSession
 from sqlframe.redshift.session import RedshiftSession
@@ -22,6 +23,7 @@ from sqlframe.standalone.dataframe import StandaloneDataFrame
 from sqlframe.standalone.session import StandaloneSession
 
 if t.TYPE_CHECKING:
+    from databricks.sql import Connection as DatabricksConnection
     from google.cloud.bigquery.dbapi.connection import (
         Connection as BigQueryConnection,
     )
@@ -90,7 +92,7 @@ def pyspark_session(tmp_path_factory, gen_tpcds: t.List[Path]) -> PySparkSession
         .config("spark.sql.warehouse.dir", data_dir)
         .config("spark.driver.extraJavaOptions", f"-Dderby.system.home={derby_dir}")
         .config("spark.sql.shuffle.partitions", 1)
-        .config("spark.sql.session.timeZone", "America/Los_Angeles")
+        .config("spark.sql.session.timeZone", "UTC")
         .master("local[1]")
         .appName("Unit-tests")
         .getOrCreate()
@@ -223,6 +225,7 @@ def snowflake_connection() -> SnowflakeConnection:
 @pytest.fixture
 def snowflake_session(snowflake_connection: SnowflakeConnection) -> SnowflakeSession:
     session = SnowflakeSession(snowflake_connection)
+    session._execute("ALTER SESSION SET TIMEZONE = 'UTC'")
     session._execute("CREATE SCHEMA IF NOT EXISTS db1")
     session._execute("CREATE TABLE IF NOT EXISTS db1.table1 (id INTEGER, name VARCHAR(100))")
     session._execute(
@@ -231,6 +234,31 @@ def snowflake_session(snowflake_connection: SnowflakeConnection) -> SnowflakeSes
     return session
 
 
+@pytest.fixture(scope="session")
+def databricks_connection() -> DatabricksConnection:
+    from databricks.sql import connect
+
+    conn = connect(
+        server_hostname=os.environ["SQLFRAME_DATABRICKS_SERVER_HOSTNAME"],
+        http_path=os.environ["SQLFRAME_DATABRICKS_HTTP_PATH"],
+        access_token=os.environ["SQLFRAME_DATABRICKS_ACCESS_TOKEN"],
+        auth_type="access_token",
+        catalog=os.environ["SQLFRAME_DATABRICKS_CATALOG"],
+        schema=os.environ["SQLFRAME_DATABRICKS_SCHEMA"],
+        _disable_pandas=True,
+    )
+    return conn
+
+
+@pytest.fixture
+def databricks_session(databricks_connection: DatabricksConnection) -> DatabricksSession:
+    session = DatabricksSession(databricks_connection)
+    session._execute("CREATE SCHEMA IF NOT EXISTS db1")
+    session._execute("CREATE TABLE IF NOT EXISTS db1.table1 (id INTEGER, name VARCHAR(100))")
+    session._execute("CREATE OR REPLACE FUNCTION db1.add(x INT, y INT) RETURNS INT RETURN x + y")
+    return session
+
+
 @pytest.fixture(scope="module")
 def _employee_data() -> EmployeeData:
     return [

{sqlframe-3.9.2 → sqlframe-3.10.0}/tests/conftest.py

@@ -16,7 +16,7 @@ def pytest_collection_modifyitems(items, *args, **kwargs):
 def set_tz():
     import os
 
-    os.environ["TZ"] = "US/Pacific"
+    os.environ["TZ"] = "UTC"
     time.tzset()
     yield
     del os.environ["TZ"]