sqlframe 3.29.1__tar.gz → 3.31.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 3.29.1
+Version: 3.31.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.31.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.31.0}/docs/duckdb.md

@@ -379,6 +379,7 @@ See something that you would like to see supported? [Open an issue](https://gith
 * [kurtosis](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.kurtosis.html)
 * [lag](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.lag.html)
 * [last](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.last.html)
+* [last_value](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.last_value.html)
 * [last_day](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.last_day.html)
 * [lcase](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.lcase.html)
 * [lead](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.functions.lead.html)

{sqlframe-3.29.1 → sqlframe-3.31.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.17",
         "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.31.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.31.0'
+__version_tuple__ = version_tuple = (3, 31, 0)

{sqlframe-3.29.1 → sqlframe-3.31.0}/sqlframe/base/functions.py

@@ -678,7 +678,22 @@ def isnull(col: ColumnOrName) -> Column:
 def last(col: ColumnOrName, ignorenulls: t.Optional[bool] = None) -> Column:
     this = Column.invoke_expression_over_column(col, expression.Last)
     if ignorenulls:
-        return Column.invoke_expression_over_column(this, expression.IgnoreNulls)
+        if _get_session()._is_duckdb:
+            return Column(
+                expression.Filter(
+                    this=this.expression,
+                    expression=expression.Where(
+                        this=expression.Not(
+                            this=expression.Is(
+                                this=Column.ensure_col(col).expression,
+                                expression=expression.Null(),
+                            )
+                        )
+                    ),
+                )
+            )
+        else:
+            return Column.invoke_expression_over_column(this, expression.IgnoreNulls)
     return this
 
 
@@ -3872,7 +3887,7 @@ def json_object_keys(col: ColumnOrName) -> Column:
     return Column.invoke_anonymous_function(col, "json_object_keys")
 
 
-@meta(unsupported_engines="*")
+@meta(unsupported_engines=["bigquery", "postgres", "snowflake"])
 def last_value(col: ColumnOrName, ignoreNulls: t.Optional[t.Union[bool, Column]] = None) -> Column:
     """Returns the last value of `col` for a group of rows. It will return the last non-null
     value it sees when `ignoreNulls` is set to true. If all values are null, then null is returned.
@@ -3913,6 +3928,11 @@ def last_value(col: ColumnOrName, ignoreNulls: t.Optional[t.Union[bool, Column]]
     |            b|            2|
     +-------------+-------------+
     """
+    session = _get_session()
+
+    if session._is_duckdb:
+        return last(col, ignoreNulls)  # type: ignore
+
     column = Column.invoke_expression_over_column(col, expression.LastValue)
 
     if ignoreNulls:

{sqlframe-3.29.1 → sqlframe-3.31.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.31.0}/sqlframe.egg-info/PKG-INFO

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

@@ -1,5 +1,5 @@
 prettytable<4
-sqlglot<26.15,>=24.0.0
+sqlglot<26.17,>=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.31.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.31.0}/tests/unit/standalone/test_functions.py

@@ -1223,8 +1223,8 @@ def test_date_format(expression, expected):
 @pytest.mark.parametrize(
     "expression, expected",
     [
-        (SF.year("cola"), "YEAR(TO_DATE(cola))"),
-        (SF.year(SF.col("cola")), "YEAR(TO_DATE(cola))"),
+        (SF.year("cola"), "YEAR(cola)"),
+        (SF.year(SF.col("cola")), "YEAR(cola)"),
     ],
 )
 def test_year(expression, expected):
@@ -1245,8 +1245,8 @@ def test_quarter(expression, expected):
 @pytest.mark.parametrize(
     "expression, expected",
     [
-        (SF.month("cola"), "MONTH(TO_DATE(cola))"),
-        (SF.month(SF.col("cola")), "MONTH(TO_DATE(cola))"),
+        (SF.month("cola"), "MONTH(cola)"),
+        (SF.month(SF.col("cola")), "MONTH(cola)"),
     ],
 )
 def test_month(expression, expected):

{sqlframe-3.29.1 → sqlframe-3.31.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)