Repository: ibis-project/ibis

Issue:
feat: support PyArrow Dataset and Table in datafusion register function
### Is your feature request related to a problem?

DataFusion Python supports registering PyArrow datasets through the `register_dataset()` method. It would be nice if that was passed through the `con.register()` method. We support this in the DuckDB backend, although that might be by accident 😉.

### Describe the solution you'd like

Handle `source` types other than string, including `pyarrow.dataset.Dataset` and `pyarrow.Table`.

https://github.com/ibis-project/ibis/blob/63db64333e0a7a178008663657fec7fa2425c439/ibis/backends/datafusion/__init__.py#L103-L105

### What version of ibis are you running?

5.1.0

### What backend(s) are you using, if any?

Datafusion

### Code of Conduct

- [X] I agree to follow this project's Code of Conduct


Discussion:
I agree, it should be fairly easy to implement. Thanks @wjones127 for the suggestion! 
Supporting pyarrow Datasets and Tables in the DuckDB was never an accident :)

This would be a nice improvement for sure.

Ideally this issue doesn't stay open forever.
> Supporting pyarrow Datasets and Tables in the DuckDB was never an accident :)

Oh good to know. I assumed datasets weren't since they aren't included in the type hint here

https://github.com/ibis-project/ibis/blob/3a590b75bfc71822e8bf441eacf0cbbd84372f05/ibis/backends/duckdb/__init__.py#L398-L402

If Krisztián isn't already looking at it, I'll look at tackling this in the next week or two.

Resolving patch:
diff --git a/ibis/backends/datafusion/__init__.py b/ibis/backends/datafusion/__init__.py
--- a/ibis/backends/datafusion/__init__.py
+++ b/ibis/backends/datafusion/__init__.py
@@ -3,7 +3,7 @@
 import re
 from functools import lru_cache
 from pathlib import Path
-from typing import Any, Mapping
+from typing import TYPE_CHECKING, Any, Mapping
 
 import pyarrow as pa
 
@@ -23,6 +23,9 @@
 
 import datafusion
 
+if TYPE_CHECKING:
+    import pandas as pd
+
 
 class Backend(BaseBackend):
     name = 'datafusion'
@@ -101,21 +104,65 @@ def table(self, name: str, schema: sch.Schema | None = None) -> ir.Table:
         return ops.DatabaseTable(name, schema, self).to_expr()
 
     def register(
-        self, source: str | Path, table_name: str | None = None, **kwargs: Any
+        self,
+        source: str | Path | pa.Table | pa.RecordBatch | pa.Dataset | pd.DataFrame,
+        table_name: str | None = None,
+        **kwargs: Any,
     ) -> ir.Table:
-        """Register a CSV or Parquet file with `table_name` located at `source`.
+        """Register a data set with `table_name` located at `source`.
 
         Parameters
         ----------
         source
-            The path to the file
+            The data source(s). May be a path to a file or directory of
+            parquet/csv files, a pandas dataframe, or a pyarrow table, dataset
+            or record batch.
         table_name
             The name of the table
         kwargs
             Datafusion-specific keyword arguments
+
+        Examples
+        --------
+        Register a csv:
+
+        >>> import ibis
+        >>> conn = ibis.datafusion.connect(config)
+        >>> conn.register("path/to/data.csv", "my_table")
+        >>> conn.table("my_table")
+
+        Register a PyArrow table:
+
+        >>> import pyarrow as pa
+        >>> tab = pa.table({"x": [1, 2, 3]})
+        >>> conn.register(tab, "my_table")
+        >>> conn.table("my_table")
+
+        Register a PyArrow dataset:
+
+        >>> import pyarrow.dataset as ds
+        >>> dataset = ds.dataset("path/to/table")
+        >>> conn.register(dataset, "my_table")
+        >>> conn.table("my_table")
         """
+        import pandas as pd
+
         if isinstance(source, (str, Path)):
             first = str(source)
+        elif isinstance(source, pa.Table):
+            self._context.deregister_table(table_name)
+            self._context.register_record_batches(table_name, [source.to_batches()])
+            return self.table(table_name)
+        elif isinstance(source, pa.RecordBatch):
+            self._context.deregister_table(table_name)
+            self._context.register_record_batches(table_name, [[source]])
+            return self.table(table_name)
+        elif isinstance(source, pa.dataset.Dataset):
+            self._context.deregister_table(table_name)
+            self._context.register_dataset(table_name, source)
+            return self.table(table_name)
+        elif isinstance(source, pd.DataFrame):
+            return self.register(pa.Table.from_pandas(source), table_name, **kwargs)
         else:
             raise ValueError("`source` must be either a string or a pathlib.Path")
 
@@ -207,14 +254,14 @@ def _get_frame(
             return self.compile(expr, params, **kwargs)
         elif isinstance(expr, ir.Column):
             # expression must be named for the projection
-            expr = expr.name('tmp').as_table()
+            expr = expr.as_table()
             return self.compile(expr, params, **kwargs)
         elif isinstance(expr, ir.Scalar):
             if an.find_immediate_parent_tables(expr.op()):
                 # there are associated datafusion tables so convert the expr
                 # to a selection which we can directly convert to a datafusion
                 # plan
-                expr = expr.name('tmp').as_table()
+                expr = expr.as_table()
                 frame = self.compile(expr, params, **kwargs)
             else:
                 # doesn't have any tables associated so create a plan from a


Test patch:
diff --git a/ibis/backends/datafusion/tests/test_register.py b/ibis/backends/datafusion/tests/test_register.py
new file mode 100644
--- /dev/null
+++ b/ibis/backends/datafusion/tests/test_register.py
@@ -0,0 +1,46 @@
+import pandas as pd
+import pyarrow as pa
+import pyarrow.dataset as ds
+import pytest
+
+import ibis
+
+
+@pytest.fixture
+def conn():
+    return ibis.datafusion.connect()
+
+
+def test_read_csv(conn, data_directory):
+    t = conn.read_csv(data_directory / "csv" / "functional_alltypes.csv")
+    assert t.count().execute()
+
+
+def test_read_parquet(conn, data_directory):
+    t = conn.read_parquet(data_directory / "parquet" / "functional_alltypes.parquet")
+    assert t.count().execute()
+
+
+def test_register_table(conn):
+    tab = pa.table({'x': [1, 2, 3]})
+    conn.register(tab, "my_table")
+    assert conn.table("my_table").x.sum().execute() == 6
+
+
+def test_register_pandas(conn):
+    df = pd.DataFrame({'x': [1, 2, 3]})
+    conn.register(df, "my_table")
+    assert conn.table("my_table").x.sum().execute() == 6
+
+
+def test_register_batches(conn):
+    batch = pa.record_batch([pa.array([1, 2, 3])], names=["x"])
+    conn.register(batch, "my_table")
+    assert conn.table("my_table").x.sum().execute() == 6
+
+
+def test_register_dataset(conn):
+    tab = pa.table({'x': [1, 2, 3]})
+    dataset = ds.InMemoryDataset(tab)
+    conn.register(dataset, "my_table")
+    assert conn.table("my_table").x.sum().execute() == 6
diff --git a/ibis/backends/tests/test_export.py b/ibis/backends/tests/test_export.py
--- a/ibis/backends/tests/test_export.py
+++ b/ibis/backends/tests/test_export.py
@@ -95,7 +95,6 @@ def test_empty_column_to_pyarrow(limit, awards_players):
     assert len(array) == 0
 
 
-@pytest.mark.notyet(["datafusion"], reason="DataFusion backend doesn't support sum")
 @pytest.mark.parametrize("limit", no_limit)
 def test_empty_scalar_to_pyarrow(limit, awards_players):
     expr = awards_players.filter(awards_players.awardID == "DEADBEEF").yearID.sum()
@@ -103,7 +102,6 @@ def test_empty_scalar_to_pyarrow(limit, awards_players):
     assert isinstance(array, pa.Scalar)
 
 
-@pytest.mark.notyet(["datafusion"], reason="DataFusion backend doesn't support sum")
 @pytest.mark.parametrize("limit", no_limit)
 def test_scalar_to_pyarrow_scalar(limit, awards_players):
     scalar = awards_players.yearID.sum().to_pyarrow(limit=limit)