Repository: ray-project/ray
File: python/ray/data/datasource/datasource.py (Python)

Code before:
        clone = copy.copy(self)

        # Combine with existing predicate using AND
        clone._predicate_expr = (
            predicate_expr
            if clone._predicate_expr is None
            else clone._predicate_expr & predicate_expr
        )

        return clone


@PublicAPI
class Datasource(_DatasourceProjectionPushdownMixin, _DatasourcePredicatePushdownMixin):
    """Interface for defining a custom :class:`~ray.data.Dataset` datasource.

    User may subclass this class to implement a custom datasource. The subclass should
    implement :meth:`.get_read_tasks` and
    :meth:`.estimate_inmemory_data_size` to read the data and estimate the in-memory data size, respectively.

    To read a datasource into a dataset, use :meth:`~ray.data.read_datasource`.

    Example:
        >>> class MyDatasource(Datasource):
        ...     def get_read_tasks(self, parallelism: int, **read_args) -> List["ReadTask"]:
        ...         return [ReadTask(lambda: [pa.Table.from_pydict({"data": range(10)})], BlockMetadata(10, 80, None, None))]
        ...     def estimate_inmemory_data_size(self) -> Optional[int]:
        ...         return 80
        >>> ds = MyDatasource()
        >>> tasks = ds.get_read_tasks(parallelism=1)
        >>> len(tasks) == 1
        True
        >>> tasks[0].metadata.num_rows == 10
        True
    """  # noqa: E501

    def __init__(self):
        """Initialize the datasource and its mixins."""
        _DatasourcePredicatePushdownMixin.__init__(self)

    @Deprecated
    def create_reader(self, **read_args) -> "Reader":
        """
        Deprecated: Implement :meth:`~ray.data.Datasource.get_read_tasks` and
        :meth:`~ray.data.Datasource.estimate_inmemory_data_size` instead.
        """
        return _LegacyDatasourceReader(self, **read_args)

    @Deprecated
    def prepare_read(self, parallelism: int, **read_args) -> List["ReadTask"]:
        """

Review comment (suggestion):
I think this example might confuse readers because it ignores the specified `parallelism`. Could you revise the example to make it simple but still valid?

Also, I think this example would be more illustrative if we explicitly listed the keyword-parameters rather than using **read_args.

Code after:
        clone = copy.copy(self)

        # Combine with existing predicate using AND
        clone._predicate_expr = (
            predicate_expr
            if clone._predicate_expr is None
            else clone._predicate_expr & predicate_expr
        )

        return clone


@PublicAPI
class Datasource(_DatasourceProjectionPushdownMixin, _DatasourcePredicatePushdownMixin):
    """Interface for defining a custom :class:`~ray.data.Dataset` datasource.

    User may subclass this class to implement a custom datasource. The subclass should
    implement :meth:`.get_read_tasks` and
    :meth:`.estimate_inmemory_data_size` to read the data and estimate the in-memory data size, respectively.

    To read a datasource into a dataset, use :meth:`~ray.data.read_datasource`.

    Example:
        >>> from ray.data.context import DataContext
        >>> class MyDatasource(Datasource):
        ...     def __init__(self, num_rows: int = 100):
        ...         super().__init__()
        ...         self.num_rows = num_rows
        ...     def get_read_tasks(
        ...         self,
        ...         parallelism: int,
        ...         per_task_row_limit: int | None = None,
        ...         data_context: DataContext | None = None,
        ...     ) -> List["ReadTask"]:
        ...         # Split num_rows across parallelism tasks
        ...         rows_per_task = self.num_rows // parallelism
        ...         return [
        ...             ReadTask(
        ...                 lambda: [pa.Table.from_pydict({"data": range(rows_per_task)})],
        ...                 BlockMetadata(rows_per_task, rows_per_task * 8, None, None),
        ...             ) for _ in range(parallelism)
        ...         ]
        ...     def estimate_inmemory_data_size(self) -> Optional[int]:
        ...         # Return total size for all data (independent of parallelism)
        ...         return self.num_rows * 8
        >>> ds = MyDatasource(num_rows=100)
        >>> tasks = ds.get_read_tasks(parallelism=5)
        >>> len(tasks) == 5
        True
        >>> tasks[0].metadata.num_rows == 20
        True