chore: import upstream snapshot with attribution

python/ray/data/_internal/datasource_v2/datasource_v2.py

@@ -0,0 +1,191 @@
+"""
+DataSourceV2 API - Unified Abstraction for Reading Data Sources
+
+This module defines a unified, extensible API for reading data from diverse sources
+in Ray Data. The API provides a common abstraction layer that enables datasources to
+declaratively expose their capabilities—such as filter pushdown, projection pruning,
+and parallel reads—while allowing the execution engine to leverage these capabilities
+transparently.
+
+Core Principles:
+- Modularity: Separate concerns (indexing, scanning, reading)
+- Expressivity: Declarative capability exposure via mixins
+- Extensibility: Easy to add new datasources with custom optimizations
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Generic,
+    Optional,
+)
+
+import pyarrow as pa
+
+from ray.data._internal.datasource_v2 import InputSplit
+from ray.data._internal.datasource_v2.listing.file_indexer import FileIndexer
+from ray.util.annotations import DeveloperAPI
+
+if TYPE_CHECKING:
+    from pyarrow.fs import FileSystem
+
+    from ray.data._internal.datasource_v2.readers.in_memory_size_estimator import (
+        InMemorySizeEstimator,
+    )
+    from ray.data._internal.datasource_v2.scanners.scanner import Scanner
+
+
+@DeveloperAPI
+class DatasourceCategory(Enum):
+    """Categories of datasources with different capability profiles.
+
+    Each category has a distinct set of applicable optimizations:
+    - FILE_BASED: Local/cloud files (parquet, csv, json, images)
+    - DATABASE: SQL databases (postgres, mysql, snowflake)
+    - DATA_LAKE: Table formats (iceberg, delta, hudi)
+    - IN_MEMORY: In-process data (pandas, numpy, arrow)
+    - SYNTHETIC: Generated data (range, range_tensor)
+    - STREAMING: Unbounded sources (kafka, kinesis)
+    """
+
+    FILE_BASED = "file_based"
+    DATABASE = "database"
+    DATA_LAKE = "data_lake"
+    IN_MEMORY = "in_memory"
+    SYNTHETIC = "synthetic"
+    STREAMING = "streaming"
+
+
+@DeveloperAPI
+class DataSourceV2(ABC, Generic[InputSplit]):
+    """Abstract base class for V2 datasources.
+
+    DataSourceV2 is the entry point for reading data from a source. It provides:
+    1. File listing (for file-based sources) - via _get_file_indexer()
+    2. Schema inference
+    3. Size estimation
+    4. Scanner creation
+
+    Subclasses should implement the abstract methods and can optionally
+    override _get_file_indexer() and get_size_estimator() for file-based sources.
+
+    Example::
+
+        datasource = ParquetDatasourceV2()
+        indexer = datasource._get_file_indexer()
+        # List files with optional sampling
+        for manifest in indexer.list_files(paths, filesystem=fs):
+            schema = datasource.infer_schema(manifest)
+            break  # Just need first manifest for schema
+        scanner = datasource.create_scanner(schema)
+        scanner = scanner.prune_columns(["col1", "col2"])
+        reader = scanner.create_reader()
+        for table in reader.read(manifest):
+            process(table)
+    """
+
+    def __init__(self, name: str, category: DatasourceCategory):
+        """Initialize the datasource.
+
+        Args:
+            name: Human-readable name for this datasource.
+            category: Category of this datasource.
+        """
+        self._name = name
+        self._category = category
+        # File-based subclasses set this to ``False`` in their ``__init__``
+        # when the user-supplied paths are in the ``local://`` scheme —
+        # the driver node is the only one that can read those files.
+        # ``_read_datasource_v2`` consults the flag to decide whether to
+        # pin read tasks via a ``label_selector``.
+        self._supports_distributed_reads: bool = True
+
+    @property
+    def name(self) -> str:
+        """Human-readable name for this datasource."""
+        return self._name
+
+    @property
+    def category(self) -> DatasourceCategory:
+        """Category of this datasource."""
+        return self._category
+
+    @property
+    def supports_distributed_reads(self) -> bool:
+        """Whether read tasks may run on any cluster node.
+
+        Defaults to ``True``. File-based subclasses (e.g.
+        :class:`ParquetDatasourceV2`) flip this to ``False`` when the
+        user supplies ``local://``-scheme paths so ``_read_datasource_v2``
+        can pin reads to the driver node via a ``ray.io/node-id``
+        label selector. Mirrors V1 ``Datasource.supports_distributed_reads``.
+        """
+        return self._supports_distributed_reads
+
+    def _get_file_indexer(self) -> Optional[FileIndexer]:
+        """Return FileIndexer component if applicable.
+
+        Override this for file-based datasources to provide file discovery.
+
+        Returns:
+            FileIndexer instance, or None for non-file-based sources.
+        """
+        return None
+
+    def get_size_estimator(self) -> Optional[InMemorySizeEstimator]:
+        """Return size estimator for this datasource.
+
+        Override this to provide format-specific size estimation.
+
+        Returns:
+            InMemorySizeEstimator instance, or None if not supported.
+        """
+        return None
+
+    @abstractmethod
+    def infer_schema(self, sample: InputSplit) -> pa.Schema:
+        """Infer schema from a sample of data.
+
+        Args:
+            sample: Sample data to infer schema from.
+
+        Returns:
+            PyArrow Schema inferred from the sample.
+
+        Raises:
+            ValueError: If schema cannot be inferred from the sample.
+        """
+        ...
+
+    @abstractmethod
+    def create_scanner(
+        self,
+        schema: pa.Schema,
+        filesystem: Optional["FileSystem"] = None,
+        **options: Any,
+    ) -> Scanner[InputSplit]:
+        """Create a Scanner for reading data.
+
+        Args:
+            schema: Schema for the data to read.
+            filesystem: Optional filesystem for file-based sources.
+            **options: Additional datasource-specific options.
+
+        Returns:
+            Configured Scanner instance.
+        """
+        ...
+
+    def resolve_partitioning(self, sample: InputSplit) -> Optional[Any]:
+        """Return a partitioning descriptor derived from ``sample``, or ``None``.
+
+        Override this for file-based sources whose partition keys must be
+        discovered from a sample path (e.g. hive layouts where field names
+        are not known up front). The resolved descriptor is passed into
+        :meth:`create_scanner`.
+        """
+        return None