sqlframe 3.29.1__tar.gz → 3.30.0__tar.gz

{sqlframe-3.29.1 → sqlframe-3.30.0}/PKG-INFO

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

{sqlframe-3.29.1 → sqlframe-3.30.0}/docs/databricks.md

@@ -361,3 +361,17 @@ Output:
 | 6  | Sun-Hwa |   Kwon    |  27 |    5     |
 +----+---------+-----------+-----+----------+
 ```
+
+### toArrow Batching
+
+SQLFrame supports the `toArrow` method in order to return a `pyarrow.Table` object. 
+```python
+arrow_table = df.toArrow()
+```
+
+Unique to SQLFrame though is a `batch_size` argument. If provided, a class similar to `pyarrow.RecordBatchReader` is returned which can be used to iterate over the batches.
+`StopIteration` will be raised when there are no more batches to read.
+```python
+record_batch_reader = df.toArrow(batch_size=1)
+first_batch_arrow_table = record_batch_reader.read_next_batch()
+```

{sqlframe-3.29.1 → sqlframe-3.30.0}/setup.py

@@ -20,7 +20,7 @@ setup(
     python_requires=">=3.9",
     install_requires=[
         "prettytable<4",
-        "sqlglot>=24.0.0,<26.15",
+        "sqlglot>=24.0.0,<26.16",
         "typing_extensions",
     ],
     extras_require={
@@ -76,7 +76,7 @@ setup(
             "pyspark>=2,<3.6",
         ],
         "databricks": [
-            "databricks-sql-connector>=3.6,<5",
+            "databricks-sql-connector[pyarrow]>=3.6,<5",
         ],
     },
     classifiers=[

{sqlframe-3.29.1 → sqlframe-3.30.0}/sqlframe/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.29.1'
-__version_tuple__ = version_tuple = (3, 29, 1)
+__version__ = version = '3.30.0'
+__version_tuple__ = version_tuple = (3, 30, 0)

{sqlframe-3.29.1 → sqlframe-3.30.0}/sqlframe/databricks/dataframe.py

@@ -14,6 +14,9 @@ from sqlframe.base.util import normalize_string
 from sqlframe.databricks.group import DatabricksGroupedData
 
 if t.TYPE_CHECKING:
+    from databricks.sql.client import Cursor
+    from pyarrow import Table as ArrowTable
+
     from sqlframe.databricks.readwriter import DatabricksDataFrameWriter
     from sqlframe.databricks.session import DatabricksSession
 
@@ -21,6 +24,18 @@ if t.TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+class RecordBatchReaderFacade:
+    def __init__(self, cur: Cursor, batch_size: int):
+        self.cur = cur
+        self.batch_size = batch_size
+
+    def read_next_batch(self) -> ArrowTable:
+        result = self.cur.fetchmany_arrow(self.batch_size)
+        if result.num_rows == 0:
+            raise StopIteration
+        return result
+
+
 class DatabricksDataFrameNaFunctions(_BaseDataFrameNaFunctions["DatabricksDataFrame"]):
     pass
 
@@ -68,3 +83,17 @@ class DatabricksDataFrame(
                 )
             )
         return columns
+
+    @t.overload  # type: ignore
+    def toArrow(self) -> ArrowTable: ...
+
+    @t.overload
+    def toArrow(self, batch_size: int) -> RecordBatchReaderFacade: ...
+
+    def toArrow(
+        self, batch_size: t.Optional[int] = None
+    ) -> t.Union[ArrowTable, RecordBatchReaderFacade]:
+        self._collect(skip_rows=True)
+        if not batch_size:
+            return self.session._cur.fetchall_arrow()  # type: ignore
+        return RecordBatchReaderFacade(self.session._cur, batch_size)  # type: ignore

{sqlframe-3.29.1 → sqlframe-3.30.0}/sqlframe.egg-info/PKG-INFO

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

{sqlframe-3.29.1 → sqlframe-3.30.0}/sqlframe.egg-info/requires.txt

@@ -1,5 +1,5 @@
 prettytable<4
-sqlglot<26.15,>=24.0.0
+sqlglot<26.16,>=24.0.0
 typing_extensions
 
 [bigquery]
@@ -7,7 +7,7 @@ google-cloud-bigquery-storage<3,>=2
 google-cloud-bigquery[pandas]<4,>=3
 
 [databricks]
-databricks-sql-connector<5,>=3.6
+databricks-sql-connector[pyarrow]<5,>=3.6
 
 [dev]
 duckdb<1.3,>=1.2

{sqlframe-3.29.1 → sqlframe-3.30.0}/tests/integration/engines/databricks/test_databricks_dataframe.py

@@ -175,3 +175,80 @@ def test_explain(databricks_employee: DatabricksDataFrame, capsys):
     assert "== Physical Plan ==" in output
     assert "LocalTableScan" in output
     assert "== Photon Explanation ==" in output
+
+
+def test_to_arrow(databricks_employee: DatabricksDataFrame):
+    arrow_table = databricks_employee.toArrow()
+    assert arrow_table.num_rows == 5
+    assert arrow_table.num_columns == 5
+    assert arrow_table.column_names == [
+        "employee_id",
+        "fname",
+        "lname",
+        "age",
+        "store_id",
+    ]
+    assert arrow_table.column(0).to_pylist() == [1, 2, 3, 4, 5]
+    assert arrow_table.column(1).to_pylist() == ["Jack", "John", "Kate", "Claire", "Hugo"]
+    assert arrow_table.column(2).to_pylist() == [
+        "Shephard",
+        "Locke",
+        "Austen",
+        "Littleton",
+        "Reyes",
+    ]
+    assert arrow_table.column(3).to_pylist() == [37, 65, 37, 27, 29]
+    assert arrow_table.column(4).to_pylist() == [1, 1, 2, 2, 100]
+
+
+def test_to_arrow_batch(databricks_employee: DatabricksDataFrame):
+    record_batch_reader = databricks_employee.toArrow(batch_size=1)
+    first_batch = record_batch_reader.read_next_batch()
+    assert first_batch.num_rows == 1
+    assert first_batch.num_columns == 5
+    assert first_batch.column_names == [
+        "employee_id",
+        "fname",
+        "lname",
+        "age",
+        "store_id",
+    ]
+    assert first_batch.column(0).to_pylist() == [1]
+    assert first_batch.column(1).to_pylist() == ["Jack"]
+    assert first_batch.column(2).to_pylist() == ["Shephard"]
+    assert first_batch.column(3).to_pylist() == [37]
+    assert first_batch.column(4).to_pylist() == [1]
+    second_batch = record_batch_reader.read_next_batch()
+    assert second_batch.num_rows == 1
+    assert second_batch.num_columns == 5
+    assert second_batch.column(0).to_pylist() == [2]
+    assert second_batch.column(1).to_pylist() == ["John"]
+    assert second_batch.column(2).to_pylist() == ["Locke"]
+    assert second_batch.column(3).to_pylist() == [65]
+    assert second_batch.column(4).to_pylist() == [1]
+    third_batch = record_batch_reader.read_next_batch()
+    assert third_batch.num_rows == 1
+    assert third_batch.num_columns == 5
+    assert third_batch.column(0).to_pylist() == [3]
+    assert third_batch.column(1).to_pylist() == ["Kate"]
+    assert third_batch.column(2).to_pylist() == ["Austen"]
+    assert third_batch.column(3).to_pylist() == [37]
+    assert third_batch.column(4).to_pylist() == [2]
+    fourth_batch = record_batch_reader.read_next_batch()
+    assert fourth_batch.num_rows == 1
+    assert fourth_batch.num_columns == 5
+    assert fourth_batch.column(0).to_pylist() == [4]
+    assert fourth_batch.column(1).to_pylist() == ["Claire"]
+    assert fourth_batch.column(2).to_pylist() == ["Littleton"]
+    assert fourth_batch.column(3).to_pylist() == [27]
+    assert fourth_batch.column(4).to_pylist() == [2]
+    fifth_batch = record_batch_reader.read_next_batch()
+    assert fifth_batch.num_rows == 1
+    assert fifth_batch.num_columns == 5
+    assert fifth_batch.column(0).to_pylist() == [5]
+    assert fifth_batch.column(1).to_pylist() == ["Hugo"]
+    assert fifth_batch.column(2).to_pylist() == ["Reyes"]
+    assert fifth_batch.column(3).to_pylist() == [29]
+    assert fifth_batch.column(4).to_pylist() == [100]
+    with pytest.raises(StopIteration):
+        record_batch_reader.read_next_batch()

{sqlframe-3.29.1 → sqlframe-3.30.0}/tests/unit/standalone/test_session.py

@@ -84,16 +84,14 @@ def test_typed_schema_nested_map(standalone_session: StandaloneSession, compare_
         ]
     )
     df = standalone_session.createDataFrame([[{"sub_cola": 1, "sub_colb": "test"}]], schema)  # type: ignore
-    expected = "SELECT `a1`.`cola` AS `cola` FROM VALUES (MAP('sub_cola', 1, 'sub_colb', 'test')) AS `a1`(`cola`)"
+    expected = "SELECT CAST(`a1`.`cola` AS MAP<INT, STRING>) AS `cola` FROM VALUES (MAP('sub_cola', 1, 'sub_colb', 'test')) AS `a1`(`cola`)"
 
     compare_sql(df, expected)
 
 
 def test_nested_struct(standalone_session: StandaloneSession, compare_sql: t.Callable):
     df = standalone_session.createDataFrame([[types.Row(cola=1, colb="test")]])  # type: ignore
-    expected = (
-        "SELECT `a1`.`_1` AS `_1` FROM VALUES (STRUCT(1 AS `cola`, 'test' AS `colb`)) AS `a1`(`_1`)"
-    )
+    expected = "SELECT CAST(`a1`.`_1` AS STRUCT<`cola`: BIGINT, `colb`: STRING>) AS `_1` FROM VALUES (STRUCT(1 AS `cola`, 'test' AS `colb`)) AS `a1`(`_1`)"
     compare_sql(df, expected)