[data] Add a parameter to allow overriding LanceDB scanner options (#…

…46975)

The default `batch_size` for the LanceDatasource is too big for some use
cases. This PR adds a new `scanner_options` argument to allow overriding
the LanceDB scanner options.

---------

Signed-off-by: Hao Chen <[email protected]>

python/ray/data/_internal/datasource/lance_datasource.py

@@ -1,5 +1,5 @@
 import logging
-from typing import TYPE_CHECKING, Dict, Iterator, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional
 
 import numpy as np
 
@@ -23,14 +23,18 @@ def __init__(
         columns: Optional[List[str]] = None,
         filter: Optional[str] = None,
         storage_options: Optional[Dict[str, str]] = None,
+        scanner_options: Optional[Dict[str, Any]] = None,
     ):
         _check_import(self, module="lance", package="pylance")
 
         import lance
 
         self.uri = uri
-        self.columns = columns
-        self.filter = filter
+        self.scanner_options = scanner_options or {}
+        if columns is not None:
+            self.scanner_options["columns"] = columns
+        if filter is not None:
+            self.scanner_options["filter"] = filter
         self.storage_options = storage_options
         self.lance_ds = lance.dataset(uri=uri, storage_options=storage_options)
 
@@ -54,14 +58,11 @@ def get_read_tasks(self, parallelism: int) -> List[ReadTask]:
                 size_bytes=None,
                 exec_stats=None,
             )
-            columns = self.columns
-            row_filter = self.filter
+            scanner_options = self.scanner_options
             lance_ds = self.lance_ds
 
             read_task = ReadTask(
-                lambda f=fragment_ids: _read_fragments(
-                    f, lance_ds, columns, row_filter
-                ),
+                lambda f=fragment_ids: _read_fragments(f, lance_ds, scanner_options),
                 metadata,
             )
             read_tasks.append(read_task)
@@ -74,7 +75,9 @@ def estimate_inmemory_data_size(self) -> Optional[int]:
 
 
 def _read_fragments(
-    fragment_ids, lance_ds, columns, row_filter
+    fragment_ids,
+    lance_ds,
+    scanner_options,
 ) -> Iterator["pyarrow.Table"]:
     """Read Lance fragments in batches.
 
@@ -84,6 +87,7 @@ def _read_fragments(
     import pyarrow
 
     fragments = [lance_ds.get_fragment(id) for id in fragment_ids]
-    scanner = lance_ds.scanner(columns, filter=row_filter, fragments=fragments)
+    scanner_options["fragments"] = fragments
+    scanner = lance_ds.scanner(**scanner_options)
     for batch in scanner.to_reader():
         yield pyarrow.Table.from_batches([batch])

python/ray/data/read_api.py

@@ -3007,6 +3007,7 @@ def read_lance(
     columns: Optional[List[str]] = None,
     filter: Optional[str] = None,
     storage_options: Optional[Dict[str, str]] = None,
+    scanner_options: Optional[Dict[str, Any]] = None,
     ray_remote_args: Optional[Dict[str, Any]] = None,
     concurrency: Optional[int] = None,
     override_num_blocks: Optional[int] = None,
@@ -3033,6 +3034,10 @@ def read_lance(
             connection. This is used to store connection parameters like credentials,
             endpoint, etc. For more information, see `Object Store Configuration <https\
                 ://lancedb.github.io/lance/read_and_write.html#object-store-configuration>`_.
+        scanner_options: Additional options to configure the `LanceDataset.scanner()`
+            method, such as `batch_size`. For more information,
+            see `LanceDB API doc <https://lancedb.github.io\
+            /lance/api/python/lance.html#lance.dataset.LanceDataset.scanner>`_
         ray_remote_args: kwargs passed to :meth:`~ray.remote` in the read tasks.
         concurrency: The maximum number of Ray tasks to run concurrently. Set this
             to control number of tasks to run concurrently. This doesn't change the
@@ -3051,6 +3056,7 @@ def read_lance(
         columns=columns,
         filter=filter,
         storage_options=storage_options,
+        scanner_options=scanner_options,
     )
 
     return read_datasource(

python/ray/data/tests/test_lance.py

@@ -28,7 +28,11 @@
         ),
     ],
 )
-def test_lance_read_basic(fs, data_path):
+@pytest.mark.parametrize(
+    "batch_size",
+    [None, 100],
+)
+def test_lance_read_basic(fs, data_path, batch_size):
     # NOTE: Lance only works with PyArrow 12 or above.
     pyarrow_version = _get_pyarrow_version()
     if pyarrow_version is not None:
@@ -51,7 +55,10 @@ def test_lance_read_basic(fs, data_path):
     )
     ds_lance.merge(df2, "one")
 
-    ds = ray.data.read_lance(path)
+    if batch_size is None:
+        ds = ray.data.read_lance(path)
+    else:
+        ds = ray.data.read_lance(path, scanner_options={"batch_size": batch_size})
 
     # Test metadata-only ops.
     assert ds.count() == 6