[data] Add LanceDB Datasource

python/ray/data/__init__.py

@@ -48,6 +48,7 @@
     read_datasource,
     read_images,
     read_json,
+    read_lance,
     read_mongo,
     read_numpy,
     read_parquet,
@@ -128,6 +129,7 @@
     "read_datasource",
     "read_images",
     "read_json",
+    "read_lance",
     "read_numpy",
     "read_mongo",
     "read_parquet",

python/ray/data/datasource/__init__.py

@@ -1,3 +1,5 @@
+from python.ray.data.datasource.lance_datasource import LanceDatasource
+
 from ray.data.datasource.avro_datasource import AvroDatasource
 from ray.data.datasource.bigquery_datasink import _BigQueryDatasink
 from ray.data.datasource.bigquery_datasource import BigQueryDatasource
@@ -91,6 +93,7 @@
     "ImageDatasource",
     "_JSONDatasink",
     "JSONDatasource",
+    "LanceDatasource",
     "_NumpyDatasink",
     "NumpyDatasource",
     "ParquetBaseDatasource",

python/ray/data/datasource/lance_datasource.py

@@ -0,0 +1,78 @@
+import logging
+from typing import List, Optional
+
+import lance
+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
+
+import pyarrow as pa
+
+logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+class LanceDatasource(Datasource):
+    """Lance Datasource
+    Read a Lance table as a Ray Dataset
+
+    Parameters
+    ----------
+    uri : str
+        The base URI of the Lance 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: Optional[List[str]] = None,
+        filter: Optional[str] = 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 table from the fragment
+            batches = fragment.to_batches(columns=self.columns, filter=self.filter)
+            for batch in batches:
+                yield pa.Table.from_batches([batch])
+
+        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,59 @@ def from_torch(
     )
 
 
+@PublicAPI
+def read_lance(
+    *,
+    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 Lance dataset. The dataset to read from
+    is specified using a fully qualified ```uri```. Using Lance, 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_lance(
+        ...     uri="./db_name.lance",
+        ...     columns=["column_name", "label"],
+        ...     filter="label = 2 AND column_name IS NOT NULL",
+        ...     parallelism=64
+        ... )
+
+    Tip:
+
+        For more details about these Lance concepts, see the following:
+        - URI: https://lancedb.github.io/lance/read_and_write.html#object-store-configuration
+
+    Args:
+        uri: The URI of the source Lance dataset to read from.
+            Currently supports local file paths, S3, GCS, and AZ 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 Lance 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,