sqlframe 3.9.2__tar.gz → 3.9.3__tar.gz

{sqlframe-3.9.2 → sqlframe-3.9.3}/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.9.3}/PKG-INFO

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

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

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

{sqlframe-3.9.2 → sqlframe-3.9.3}/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.9.3'
+__version_tuple__ = version_tuple = (3, 9, 3)

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

@@ -1220,6 +1220,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.9.3}/sqlframe/base/functions.py

@@ -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.9.3}/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.9.3}/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.9.3}/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.9.3}/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.9.3}/sqlframe.egg-info/PKG-INFO

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

{sqlframe-3.9.2 → sqlframe-3.9.3}/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.9.3}/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,
     )
@@ -231,6 +233,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.3/tests/integration/engines/databricks/test_databricks_catalog.py

@@ -0,0 +1,338 @@
+import typing as t
+
+import pytest
+
+from sqlframe.base.catalog import CatalogMetadata, Column, Database, Function, Table
+from sqlframe.databricks.session import DatabricksSession
+
+pytest_plugins = ["tests.integration.fixtures"]
+pytestmark = [
+    pytest.mark.databricks,
+    pytest.mark.xdist_group("databricks_tests"),
+]
+
+
+@pytest.fixture
+def reset_catalog(databricks_session: DatabricksSession) -> t.Iterator[None]:
+    yield
+    databricks_session.catalog.setCurrentCatalog("sqlframe")
+    databricks_session.catalog.setCurrentDatabase("db1")
+
+
+@pytest.fixture
+def reset_database(databricks_session: DatabricksSession) -> t.Iterator[None]:
+    yield
+    databricks_session.catalog.setCurrentDatabase("db1")
+
+
+def test_current_catalog(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.currentCatalog() == "sqlframe"
+
+
+def test_set_current_catalog(databricks_session: DatabricksSession, reset_catalog):
+    assert databricks_session.catalog.currentCatalog() == "sqlframe"
+    databricks_session.catalog.setCurrentCatalog("catalog1")
+    assert databricks_session.catalog.currentCatalog() == "catalog1"
+
+
+def test_list_catalogs(databricks_session: DatabricksSession):
+    assert sorted(databricks_session.catalog.listCatalogs(), key=lambda x: x.name) == [
+        CatalogMetadata(name="sqlframe", description=None)
+    ]
+
+
+def test_current_database(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.currentDatabase() == "db1"
+
+
+def test_set_current_database(databricks_session: DatabricksSession, reset_database):
+    assert databricks_session.catalog.currentDatabase() == "db1"
+    databricks_session.catalog.setCurrentDatabase("default")
+    assert databricks_session.catalog.currentDatabase() == "default"
+
+
+def test_list_databases(databricks_session: DatabricksSession):
+    assert sorted(
+        databricks_session.catalog.listDatabases(), key=lambda x: (x.catalog, x.name)
+    ) == [
+        Database(name="db1", catalog="sqlframe", description=None, locationUri=""),
+        Database(name="default", catalog="sqlframe", description=None, locationUri=""),
+        Database(name="information_schema", catalog="sqlframe", description=None, locationUri=""),
+    ]
+
+
+def test_list_databases_pattern(databricks_session: DatabricksSession):
+    assert sorted(
+        databricks_session.catalog.listDatabases("db*"), key=lambda x: (x.catalog, x.name)
+    ) == [
+        Database(name="db1", catalog="sqlframe", description=None, locationUri=""),
+    ]
+
+
+def test_get_database_no_match(databricks_session: DatabricksSession):
+    with pytest.raises(ValueError):
+        assert databricks_session.catalog.getDatabase("nonexistent")
+
+
+def test_get_database_name_only(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.getDatabase("db1") == Database(
+        name="db1", catalog="sqlframe", description=None, locationUri=""
+    )
+
+
+def test_get_database_name_and_catalog(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.getDatabase("sqlframe.db1") == Database(
+        name="db1", catalog="sqlframe", description=None, locationUri=""
+    )
+
+
+def test_database_exists_does_exist(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.databaseExists("db1") is True
+
+
+def test_database_exists_does_not_exist(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.databaseExists("nonexistent") is False
+
+
+def test_list_tables_no_args(databricks_session: DatabricksSession):
+    assert sorted(
+        databricks_session.catalog.listTables(), key=lambda x: (x.catalog, x.database, x.name)
+    ) == [
+        Table(
+            name="table1",
+            catalog="sqlframe",
+            namespace=["db1"],
+            description=None,
+            tableType="MANAGED",
+            isTemporary=False,
+        )
+    ]
+
+
+def test_list_tables_db_no_catalog(databricks_session: DatabricksSession):
+    assert sorted(
+        databricks_session.catalog.listTables("db1"), key=lambda x: (x.catalog, x.database, x.name)
+    ) == [
+        Table(
+            name="table1",
+            catalog="sqlframe",
+            namespace=["db1"],
+            description=None,
+            tableType="MANAGED",
+            isTemporary=False,
+        )
+    ]
+
+
+def test_list_tables_db_and_catalog(databricks_session: DatabricksSession):
+    assert sorted(
+        databricks_session.catalog.listTables("sqlframe.db1"),
+        key=lambda x: (x.catalog, x.database, x.name),
+    ) == [
+        Table(
+            name="table1",
+            catalog="sqlframe",
+            namespace=["db1"],
+            description=None,
+            tableType="MANAGED",
+            isTemporary=False,
+        )
+    ]
+
+
+def test_list_tables_pattern(databricks_session: DatabricksSession):
+    assert Table(
+        name="table1",
+        catalog="sqlframe",
+        namespace=["db1"],
+        description=None,
+        tableType="MANAGED",
+        isTemporary=False,
+    ) in databricks_session.catalog.listTables(pattern="tab*")
+
+
+def test_get_table(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.getTable("sqlframe.db1.table1") == Table(
+        name="table1",
+        catalog="sqlframe",
+        namespace=["db1"],
+        description=None,
+        tableType="MANAGED",
+        isTemporary=False,
+    )
+
+
+def test_get_table_not_exists(databricks_session: DatabricksSession):
+    with pytest.raises(ValueError):
+        assert databricks_session.catalog.getTable("dev.db1.nonexistent")
+
+
+def test_list_functions(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.listFunctions() == [
+        Function(
+            name="add",
+            catalog="sqlframe",
+            namespace=["db1"],
+            description=None,
+            className="",
+            isTemporary=False,
+        )
+    ]
+
+
+def test_list_functions_pattern(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.listFunctions(dbName="db1", pattern="ad*") == [
+        Function(
+            name="add",
+            catalog="sqlframe",
+            namespace=["db1"],
+            description=None,
+            className="",
+            isTemporary=False,
+        )
+    ]
+
+
+def test_function_exists_does_exist(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.functionExists("add", dbName="sqlframe.db1") is True
+
+
+def test_function_exists_does_not_exist(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.functionExists("nonexistent") is False
+
+
+def test_get_function_exists(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.getFunction("sqlframe.db1.add") == Function(
+        name="add",
+        catalog="sqlframe",
+        namespace=["db1"],
+        description=None,
+        className="",
+        isTemporary=False,
+    )
+
+
+def test_get_function_not_exists(databricks_session: DatabricksSession):
+    with pytest.raises(ValueError):
+        assert databricks_session.catalog.getFunction("sqlframe.db1.nonexistent")
+
+
+def test_list_columns(databricks_session: DatabricksSession):
+    assert sorted(
+        databricks_session.catalog.listColumns("sqlframe.db1.table1"), key=lambda x: x.name
+    ) == [
+        Column(
+            name="id",
+            description=None,
+            dataType="INT",
+            nullable=True,
+            isPartition=False,
+            isBucket=False,
+        ),
+        Column(
+            name="name",
+            description=None,
+            dataType="VARCHAR(100)",
+            nullable=True,
+            isPartition=False,
+            isBucket=False,
+        ),
+    ]
+
+
+def test_list_columns_use_db_name(databricks_session: DatabricksSession):
+    assert sorted(
+        databricks_session.catalog.listColumns("table1", dbName="sqlframe.db1"),
+        key=lambda x: x.name,
+    ) == [
+        Column(
+            name="id",
+            description=None,
+            dataType="INT",
+            nullable=True,
+            isPartition=False,
+            isBucket=False,
+        ),
+        Column(
+            name="name",
+            description=None,
+            dataType="VARCHAR(100)",
+            nullable=True,
+            isPartition=False,
+            isBucket=False,
+        ),
+    ]
+
+
+def test_table_exists_table_name_only(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.tableExists("sqlframe.db1.table1") is True
+
+
+def test_table_exists_table_name_and_db_name(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.tableExists("table1", dbName="sqlframe.db1") is True
+
+
+def test_table_not_exists(databricks_session: DatabricksSession):
+    assert databricks_session.catalog.tableExists("nonexistent") is False
+
+
+def test_create_external_table(databricks_session: DatabricksSession):
+    with pytest.raises(NotImplementedError):
+        databricks_session.catalog.createExternalTable(
+            "table1", "sqlframe.default", "path/to/table"
+        )
+
+
+def test_create_table(databricks_session: DatabricksSession):
+    with pytest.raises(NotImplementedError):
+        databricks_session.catalog.createTable("table1", "sqlframe.default")
+
+
+def test_drop_temp_view(databricks_session: DatabricksSession):
+    with pytest.raises(NotImplementedError):
+        databricks_session.catalog.dropTempView("view1")
+
+
+def test_drop_global_temp_view(databricks_session: DatabricksSession):
+    with pytest.raises(NotImplementedError):
+        databricks_session.catalog.dropGlobalTempView("view1")
+
+
+def test_register_function(databricks_session: DatabricksSession):
+    with pytest.raises(NotImplementedError):
+        databricks_session.catalog.registerFunction("function1", lambda x: x)
+
+
+def test_is_cached(databricks_session: DatabricksSession):
+    with pytest.raises(NotImplementedError):
+        databricks_session.catalog.isCached("table1")
+
+
+def test_cache_table(databricks_session: DatabricksSession):
+    with pytest.raises(NotImplementedError):
+        databricks_session.catalog.cacheTable("table1")
+
+
+def test_uncache_table(databricks_session: DatabricksSession):
+    with pytest.raises(NotImplementedError):
+        databricks_session.catalog.uncacheTable("table1")
+
+
+def test_clear_cache(databricks_session: DatabricksSession):
+    with pytest.raises(NotImplementedError):
+        databricks_session.catalog.clearCache()
+
+
+def test_refresh_table(databricks_session: DatabricksSession):
+    with pytest.raises(NotImplementedError):
+        databricks_session.catalog.refreshTable("table1")
+
+
+def test_recover_partitions(databricks_session: DatabricksSession):
+    with pytest.raises(NotImplementedError):
+        databricks_session.catalog.recoverPartitions("table1")
+
+
+def test_refresh_by_path(databricks_session: DatabricksSession):
+    with pytest.raises(NotImplementedError):
+        databricks_session.catalog.refreshByPath("path/to/table")