[data] Add LanceDB Datasource

python/ray/data/__init__.py

@@ -48,6 +48,7 @@
     read_datasource,
     read_images,
     read_json,
+    read_lancedb,
     read_mongo,
     read_numpy,
     read_parquet,

python/ray/data/datasource/__init__.py

@@ -37,6 +37,7 @@
 from ray.data.datasource.image_datasource import ImageDatasource
 from ray.data.datasource.json_datasink import _JSONDatasink
 from ray.data.datasource.json_datasource import JSONDatasource
+from ray.data.datasource.lancedb_datasource import LanceDBDatasource
 from ray.data.datasource.mongo_datasink import _MongoDatasink
 from ray.data.datasource.mongo_datasource import MongoDatasource
 from ray.data.datasource.numpy_datasink import _NumpyDatasink
@@ -91,6 +92,7 @@
     "ImageDatasource",
     "_JSONDatasink",
     "JSONDatasource",
+    "LanceDBDatasource",
     "_NumpyDatasink",
     "NumpyDatasource",
     "ParquetBaseDatasource",

python/ray/data/datasource/lancedb_datasource.py

@@ -0,0 +1,87 @@
+"""
+Module to read an LanceDB table into a Ray Dataset, by using the Ray Datasource API.
+"""
+
+import logging
+from typing import List, Optional, Union
+
+import lance
+import pyarrow as pa
+from lance import LanceFragment
+
+from ray.data import ReadTask
+from ray.data.block import Block, BlockMetadata
+from ray.data.datasource import Datasource
+from ray.util.annotations import DeveloperAPI
+
+logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+class LanceDBDatasource(Datasource):
+    """LanceDB Datasource
+    Read a LanceDB table as a Ray Dataset
+
+    Parameters
+    ----------
+    uri : str
+        The base URI of the LanceDB dataset.
+    columns: list
+        A list of columns to return from the dataset.
+    filter: str
+        A standard SQL expressions as predicates for dataset filtering.
+    """
+
+    def __init__(
+        self,
+        uri: str,
+        columns: Union[list, None] = None,
+        filter: Union[str, None] = None,
+    ):
+        self.uri = uri
+        self.columns = columns
+        self.filter = filter
+
+        self.lance_ds = lance.dataset(uri)
+        self.fragments = self.lance_ds.get_fragments()
+
+    def get_read_tasks(self, parallelism: int) -> List[ReadTask]:
+        # To begin with, read one Fragment at a time
+        # Each Ray Data Block contains a Pandas RecordBatch
+        def _read_single_fragment(fragment: LanceFragment) -> Block:
+            # Fetch batches from the fragment
+            batches = fragment.to_batches(columns=self.columns, filter=self.filter)
+
+            # Convert the generator of RecordBatch objects to a list
+            batches_list = list(batches)
+
+            # Convert the list of RecordBatch objects to a Table
+            table = pa.Table.from_batches(batches_list)
+
+            return table
+
+        read_tasks = []
+        for fragment in self.fragments:
+            data_files = ", ".join(
+                [data_file.path() for data_file in fragment.data_files()]
+            )
+
+            metadata = BlockMetadata(
+                num_rows=fragment.count_rows(),
+                size_bytes=None,
+                schema=fragment.schema,
+                input_files=[data_files],
+                exec_stats=None,
+            )
+
+            read_task = ReadTask(
+                lambda fragment=fragment: [_read_single_fragment(fragment)],
+                metadata,
+            )
+            read_tasks.append(read_task)
+
+        return read_tasks
+
+    def estimate_inmemory_data_size(self) -> Optional[int]:
+        # TODO: Add memory size estimation to improve auto-tune of parallelism.
+        return None

python/ray/data/read_api.py

@@ -54,6 +54,7 @@
     Datasource,
     ImageDatasource,
     JSONDatasource,
+    LanceDBDatasource,
     MongoDatasource,
     NumpyDatasource,
     ParquetBaseDatasource,
@@ -2900,6 +2901,54 @@ def from_torch(
     )
 
 
+@PublicAPI
+def from_lancedb(
+    *,
+    uri: str,
+    columns: Optional[List[str]] = None,
+    filter: Optional[str] = None,
+    parallelism: int = -1,
+    ray_remote_args: Optional[Dict[str, Any]] = None,
+) -> Dataset:
+    """
+    Create a :class:`~ray.data.Dataset` from a LanceDB dataset. The dataset to read from
+    is specified using a fully qualified ```uri```. Using LanceDB, any
+    intended columns or filters are applied, and the files that satisfy
+    the query are distributed across Ray read tasks. The number of tasks is
+    determined by ``parallelism`` which can be requested from this interface or
+    automatically chosen if unspecified (see the``parallelism`` arg below).
+
+    Examples:
+        >>> import ray
+        >>> ds = ray.data.read_lancedb(
+        ...     uri="./db_name.lance",
+        ...     columns=["column_name", "label"],
+        ...     filter="label = 2 AND column_name IS NOT NULL",
+        ...     parallelism=64
+        ... )
+
+    Args:
+        uri: The URI of the source LanceDB dataset to read from.
+            Currently supports local file paths, S3, or GCS URIs are supported.
+        columns: The columns to read from the dataset.
+            If not specified, all columns are read.
+        filter: The filter to apply to the dataset.
+            If not specified, no filter is applied.
+        parallelism: Degree of parallelism to use for the Dataset
+        ray_remote_args: Optional arguments to pass to `ray.remote` in the read tasks
+
+    Returns:
+        A :class:`~ray.data.Dataset` the LanceDB dataset from the results of executing the read.
+    """  # noqa: E501
+    datasource = LanceDBDatasource(uri=uri, columns=columns, filter=filter)
+
+    dataset = read_datasource(
+        datasource=datasource, parallelism=parallelism, ray_remote_args=ray_remote_args
+    )
+
+    return dataset
+
+
 def _get_datasource_or_legacy_reader(
     ds: Datasource,
     ctx: DataContext,