feat(python): Support running Polars SQL queries against any objects implementing the PyCapsule interface (pola-rs#22235)

Author: alexander-beedie

py-polars/polars/sql/context.py

@@ -12,6 +12,7 @@
 
 from polars._typing import FrameType
 from polars._utils.deprecation import deprecate_renamed_parameter
+from polars._utils.pycapsule import is_pycapsule
 from polars._utils.unstable import issue_unstable_warning
 from polars._utils.various import _get_stack_locals
 from polars._utils.wrap import wrap_ldf
@@ -52,14 +53,14 @@
         pa.RecordBatch,
     ]
 
-
 __all__ = ["SQLContext"]
 
 
 def _compatible_frame(obj: Any) -> bool:
     """Check if the object can be converted to DataFrame."""
     return (
-        isinstance(obj, (DataFrame, LazyFrame, Series))
+        is_pycapsule(obj)
+        or isinstance(obj, LazyFrame)
         or (_check_for_pandas(obj) and isinstance(obj, (pd.DataFrame, pd.Series)))
         or (_check_for_pyarrow(obj) and isinstance(obj, (pa.Table, pa.RecordBatch)))
     )
@@ -68,14 +69,16 @@ def _compatible_frame(obj: Any) -> bool:
 def _ensure_lazyframe(obj: Any) -> LazyFrame:
     """Return LazyFrame from compatible input."""
     if isinstance(obj, (DataFrame, LazyFrame)):
-        return obj if isinstance(obj, LazyFrame) else obj.lazy()
+        return obj.lazy()
     elif isinstance(obj, Series):
         return obj.to_frame().lazy()
     elif _check_for_pandas(obj) and isinstance(obj, (pd.DataFrame, pd.Series)):
         if isinstance(frame := from_pandas(obj), Series):
             frame = frame.to_frame()
         return frame.lazy()
-    elif _check_for_pyarrow(obj) and isinstance(obj, (pa.Table, pa.RecordBatch)):
+    elif is_pycapsule(obj) or (
+        _check_for_pyarrow(obj) and isinstance(obj, (pa.Table, pa.RecordBatch))
+    ):
         return from_arrow(obj).lazy()  # type: ignore[union-attr]
     else:
         msg = f"Unrecognised frame type: {type(obj)}"

py-polars/tests/unit/sql/test_miscellaneous.py

@@ -9,6 +9,7 @@
 import polars as pl
 from polars.exceptions import ColumnNotFoundError, SQLInterfaceError, SQLSyntaxError
 from polars.testing import assert_frame_equal
+from tests.unit.utils.pycapsule_utils import PyCapsuleStreamHolder
 
 if TYPE_CHECKING:
     from polars.datatypes import DataType
@@ -243,6 +244,7 @@ def test_sql_on_compatible_frame_types() -> None:
     dfp = df.to_pandas()
     dfa = df.to_arrow()
     dfb = dfa.to_batches()[0]  # noqa: F841
+    dfo = PyCapsuleStreamHolder(df)  # noqa: F841
 
     # run polars sql query against all frame types
     for dfs in (  # noqa: B007
@@ -256,14 +258,15 @@ def test_sql_on_compatible_frame_types() -> None:
                 UNION ALL SELECT * FROM dfp  -- pandas frame
                 UNION ALL SELECT * FROM dfa  -- pyarrow table
                 UNION ALL SELECT * FROM dfb  -- pyarrow record batch
+                UNION ALL SELECT * FROM dfo  -- arbitrary pycapsule object
             ) tbl
             INNER JOIN dfs ON dfs.c == tbl.b -- join on pandas/polars series
             GROUP BY "a", "b"
             ORDER BY "a", "b"
             """
         ).collect()
 
-        expected = pl.DataFrame({"a": [1, 3], "b": [4, 6], "cc": [16, 24]})
+        expected = pl.DataFrame({"a": [1, 3], "b": [4, 6], "cc": [20, 30]})
         assert_frame_equal(left=expected, right=res)
 
     # register and operate on non-polars frames
@@ -400,60 +403,60 @@ def test_select_output_heights_20058_21084(filter_expr: str, order_expr: str) ->
     )
 
     assert_frame_equal(
-        df.sql(f"""\
-    SELECT 1 + 1 as a, 1 as b FROM self {filter_expr} {order_expr}
-    """).cast(pl.Int64),
+        df.sql(f"SELECT 1 + 1 as a, 1 as b FROM self {filter_expr} {order_expr}").cast(
+            pl.Int64
+        ),
         pl.DataFrame({"a": [2, 2, 2], "b": [1, 1, 1]}),
     )
 
     # Queries that aggregate to unit height
 
     assert_frame_equal(
-        df.sql(f"""\
-    SELECT COUNT(*) as a FROM self {filter_expr} {order_expr}
-    """).cast(pl.Int64),
+        df.sql(f"SELECT COUNT(*) as a FROM self {filter_expr} {order_expr}").cast(
+            pl.Int64
+        ),
         pl.DataFrame({"a": 3}),
     )
 
     assert_frame_equal(
-        df.sql(f"""\
-    SELECT COUNT(*) as a, 1 as b FROM self {filter_expr} {order_expr}
-    """).cast(pl.Int64),
+        df.sql(
+            f"SELECT COUNT(*) as a, 1 as b FROM self {filter_expr} {order_expr}"
+        ).cast(pl.Int64),
         pl.DataFrame({"a": 3, "b": 1}),
     )
 
     assert_frame_equal(
-        df.sql(f"""\
-    SELECT FIRST(a) as a, 1 as b FROM self {filter_expr} {order_expr}
-    """).cast(pl.Int64),
+        df.sql(
+            f"SELECT FIRST(a) as a, 1 as b FROM self {filter_expr} {order_expr}"
+        ).cast(pl.Int64),
         pl.DataFrame({"a": 1, "b": 1}),
     )
 
     assert_frame_equal(
-        df.sql(f"""\
-    SELECT SUM(a) as a, 1 as b FROM self {filter_expr} {order_expr}
-    """).cast(pl.Int64),
+        df.sql(f"SELECT SUM(a) as a, 1 as b FROM self {filter_expr} {order_expr}").cast(
+            pl.Int64
+        ),
         pl.DataFrame({"a": 6, "b": 1}),
     )
 
     assert_frame_equal(
-        df.sql(f"""\
-    SELECT FIRST(1) as a, 1 as b FROM self {filter_expr} {order_expr}
-    """).cast(pl.Int64),
+        df.sql(
+            f"SELECT FIRST(1) as a, 1 as b FROM self {filter_expr} {order_expr}"
+        ).cast(pl.Int64),
         pl.DataFrame({"a": 1, "b": 1}),
     )
 
     assert_frame_equal(
-        df.sql(f"""\
-    SELECT FIRST(1) + 1 as a, 1 as b FROM self {filter_expr} {order_expr}
-    """).cast(pl.Int64),
+        df.sql(
+            f"SELECT FIRST(1) + 1 as a, 1 as b FROM self {filter_expr} {order_expr}"
+        ).cast(pl.Int64),
         pl.DataFrame({"a": 2, "b": 1}),
     )
 
     assert_frame_equal(
-        df.sql(f"""\
-    SELECT FIRST(1 + 1) as a, 1 as b FROM self {filter_expr} {order_expr}
-    """).cast(pl.Int64),
+        df.sql(
+            f"SELECT FIRST(1 + 1) as a, 1 as b FROM self {filter_expr} {order_expr}"
+        ).cast(pl.Int64),
         pl.DataFrame({"a": 2, "b": 1}),
     )
 
@@ -473,48 +476,34 @@ def test_select_explode_height_filter_order_by() -> None:
     # extended with NULLs.
 
     assert_frame_equal(
-        df.sql(
-            """\
-    SELECT UNNEST(list_long) as list FROM self ORDER BY sort_key
-        """
-        ),
+        df.sql("SELECT UNNEST(list_long) as list FROM self ORDER BY sort_key"),
         pl.Series("list", [2, 1, 3, 4, 5, 6]).to_frame(),
     )
 
     assert_frame_equal(
         df.sql(
-            """\
-    SELECT UNNEST(list_long) as list FROM self ORDER BY sort_key NULLS FIRST
-        """
+            "SELECT UNNEST(list_long) as list FROM self ORDER BY sort_key NULLS FIRST"
         ),
         pl.Series("list", [3, 4, 5, 6, 2, 1]).to_frame(),
     )
 
     # Literals are broadcasted to output height of UNNEST:
     assert_frame_equal(
-        df.sql(
-            """\
-    SELECT UNNEST(list_long) as list, 1 as x FROM self ORDER BY sort_key
-        """
-        ),
+        df.sql("SELECT UNNEST(list_long) as list, 1 as x FROM self ORDER BY sort_key"),
         pl.select(pl.Series("list", [2, 1, 3, 4, 5, 6]), x=1),
     )
 
     # Note: Filter applies before projections in SQL
     assert_frame_equal(
         df.sql(
-            """\
-    SELECT UNNEST(list_long) as list FROM self WHERE filter_mask ORDER BY sort_key
-        """
+            "SELECT UNNEST(list_long) as list FROM self WHERE filter_mask ORDER BY sort_key"
         ),
         pl.Series("list", [4, 5, 6]).to_frame(),
     )
 
     assert_frame_equal(
         df.sql(
-            """\
-    SELECT UNNEST(list_long) as list FROM self WHERE filter_mask_all_true ORDER BY sort_key
-        """
+            "SELECT UNNEST(list_long) as list FROM self WHERE filter_mask_all_true ORDER BY sort_key"
         ),
         pl.Series("list", [2, 1, 3, 4, 5, 6]).to_frame(),
     )