feat(datafusion): enable reading multiple paths in read_csv (#10317)

Author: mesejo

ibis/backends/datafusion/__init__.py

@@ -28,7 +28,7 @@
 from ibis.common.dispatch import lazy_singledispatch
 from ibis.expr.operations.udf import InputType
 from ibis.formats.pyarrow import PyArrowSchema, PyArrowType
-from ibis.util import deprecated, gen_name, normalize_filename
+from ibis.util import deprecated, gen_name, normalize_filename, normalize_filenames
 
 try:
     from datafusion import ExecutionContext as SessionContext
@@ -414,13 +414,16 @@ def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
         self.con.from_arrow(op.data.to_pyarrow(op.schema), op.name)
 
     def read_csv(
-        self, path: str | Path, table_name: str | None = None, **kwargs: Any
+        self,
+        source_list: str | Path | list[str | Path] | tuple[str | Path],
+        table_name: str | None = None,
+        **kwargs: Any,
     ) -> ir.Table:
         """Register a CSV file as a table in the current database.
 
         Parameters
         ----------
-        path
+        source_list
             The data source. A string or Path to the CSV file.
         table_name
             An optional name to use for the created table. This defaults to
@@ -434,9 +437,9 @@ def read_csv(
             The just-registered table
 
         """
-        path = normalize_filename(path)
+        path = normalize_filenames(source_list)
         table_name = table_name or gen_name("read_csv")
-        # Our other backends support overwriting views / tables when reregistering
+        # Our other backends support overwriting views / tables when re-registering
         self.con.deregister_table(table_name)
         self.con.register_csv(table_name, path, **kwargs)
         return self.table(table_name)

ibis/backends/datafusion/tests/test_register.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import pathlib
+
 import pandas as pd
 import pyarrow as pa
 import pytest
@@ -17,6 +19,20 @@ def test_read_csv(conn, data_dir):
     assert t.count().execute()
 
 
+@pytest.mark.parametrize(
+    "function",
+    [pathlib.Path, str],
+)
+def test_read_csv_path_list(conn, data_dir, function):
+    path = data_dir / "csv" / "functional_alltypes.csv"
+
+    t = conn.read_csv(path, table_name="alltypes1")
+    t2 = conn.read_csv([function(path), function(path)], table_name="alltypes2")
+
+    assert t2.schema() == t.schema()
+    assert t2.count().execute() == 2 * t.count().execute()
+
+
 def test_read_parquet(conn, data_dir):
     t = conn.read_parquet(data_dir / "parquet" / "functional_alltypes.parquet")
     assert t.count().execute()

ibis/backends/tests/test_signatures.py

@@ -107,7 +107,7 @@ def _scrape_methods(modules, params):
     "read_csv": pytest.param(
         BaseBackend,
         "read_csv",
-        marks=pytest.mark.notyet(["duckdb", "flink", "pyspark"]),
+        marks=pytest.mark.notyet(["duckdb", "flink", "pyspark", "datafusion"]),
     ),
     "read_delta": pytest.param(
         BaseBackend,