sqlframe 3.9.1__tar.gz → 3.9.3__tar.gz

{sqlframe-3.9.1 → sqlframe-3.9.3}/.github/workflows/publish.workflow.yaml

@@ -5,17 +5,20 @@ on:
       - 'v[0-9]+.[0-9]+.[0-9]+'
 permissions:
   contents: write
+  id-token: write
 jobs:
   deploy:
     runs-on: ubuntu-latest
-    env:
-      TWINE_USERNAME: ${{ secrets.TWINE_USERNAME }}
-      TWINE_PASSWORD: ${{ secrets.TWINE_PASSWORD }}
+    environment:
+      name: publish
+      url: https://pypi.org/p/sqlframe
     steps:
     - name: Checkout
       uses: actions/checkout@v4
+    - name: Package
+      run: make package
     - name: Publish
-      run: make publish
+      uses: pypa/gh-action-pypi-publish@release/v1
     - name: Create release
       env:
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 3.9.1
+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.1 → sqlframe-3.9.3}/docs/databricks.md

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

{sqlframe-3.9.1 → 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.1'
-__version_tuple__ = version_tuple = (3, 9, 1)
+__version__ = version = '3.9.3'
+__version_tuple__ = version_tuple = (3, 9, 3)

{sqlframe-3.9.1 → sqlframe-3.9.3}/sqlframe/base/dataframe.py

@@ -12,6 +12,7 @@ import typing as t
 import zlib
 from copy import copy
 from dataclasses import dataclass
+from uuid import uuid4
 
 import sqlglot
 from prettytable import PrettyTable
@@ -208,6 +209,8 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         expression: exp.Select,
         branch_id: t.Optional[str] = None,
         sequence_id: t.Optional[str] = None,
+        join_on_uuid: t.Optional[str] = None,
+        known_uuids: t.Optional[t.Set[str]] = None,
         last_op: Operation = Operation.INIT,
         pending_hints: t.Optional[t.List[exp.Expression]] = None,
         output_expression_container: t.Optional[OutputExpressionContainer] = None,
@@ -217,6 +220,9 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         self.expression: exp.Select = expression
         self.branch_id = branch_id or self.session._random_branch_id
         self.sequence_id = sequence_id or self.session._random_sequence_id
+        self.join_on_uuid = join_on_uuid or str(uuid4())
+        self.known_uuids = known_uuids or set()
+        self.known_uuids.add(self.join_on_uuid)
         self.last_op = last_op
         self.pending_hints = pending_hints or []
         self.output_expression_container = output_expression_container or exp.Select()
@@ -228,10 +234,12 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
     def __getitem__(self, column_name: str) -> Column:
         from sqlframe.base.util import get_func_from_session
 
-        col = get_func_from_session("col", self.session)
+        col_func = get_func_from_session("col", self.session)
 
         column_name = f"{self.branch_id}.{column_name}"
-        return col(column_name)
+        col = col_func(column_name)
+        col.expression.meta["join_on_uuid"] = self.join_on_uuid
+        return col
 
     def __copy__(self):
         return self.copy()
@@ -715,6 +723,7 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         return ";\n".join(results)
 
     def copy(self, **kwargs) -> Self:
+        kwargs["join_on_uuid"] = str(uuid4())
         return self.__class__(**object_to_dict(self, **kwargs))
 
     @operation(Operation.SELECT)
@@ -876,13 +885,21 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         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)
-        # If the two dataframes being joined come from the same branch and reference the same table,
-        # then really one of the tables was supposed to reference the other dataframe so we update it to do that.
+        # If the two dataframes being joined come from the same branch, we then check if they have any columns that
+        # were created using the "branch_id" (df["column_name"]). If so, we know that we need to differentiate
+        # the two columns since they would end up with the same table name. We do this by checking for the unique
+        # uuids in the other df and finding columns that have metadata on them that match the uuids. If so, we know
+        # it comes from the other df and we change the table name to the other df's table name.
+        # See `test_self_join` for an example of this.
         if self.branch_id == other_df.branch_id:
+            other_df_unique_uuids = other_df.known_uuids - self.known_uuids
             for col in join_columns:
-                for eq in col.expression.find_all(exp.EQ):
-                    if eq.this.table == eq.expression.table:
-                        eq.expression.set("table", exp.to_identifier(other_df.latest_cte_name))
+                for col_expr in col.expression.find_all(exp.Column):
+                    if (
+                        "join_on_uuid" in col_expr.meta
+                        and col_expr.meta["join_on_uuid"] in other_df_unique_uuids
+                    ):
+                        col_expr.set("table", exp.to_identifier(other_df.latest_cte_name))
         # Determines the join clause and select columns to be used passed on what type of columns were provided for
         # the join. The columns returned changes based on how the on expression is provided.
         if how != "cross":

{sqlframe-3.9.1 → 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.1 → 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.1 → 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.1 → 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.1 → 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.1 → 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.1 → sqlframe-3.9.3}/sqlframe.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 3.9.1
+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.1 → 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.1 → 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 [