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

Code before:


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

    To read a datasource into a dataset, use :meth:`~ray.data.read_datasource`.
    """  # 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"]:
        """
        Deprecated: Implement :meth:`~ray.data.Datasource.get_read_tasks` and
        :meth:`~ray.data.Datasource.estimate_inmemory_data_size` instead.
        """
        raise NotImplementedError

    def get_name(self) -> str:
        """Return a human-readable name for this datasource.
        This will be used as the names of the read tasks.
        """
        name = type(self).__name__
        datasource_suffix = "Datasource"
        if name.endswith(datasource_suffix):
            name = name[: -len(datasource_suffix)]
        return name

    def estimate_inmemory_data_size(self) -> Optional[int]:
        """Return an estimate of the in-memory data size, or None if unknown.

        Note that the in-memory data size may be larger than the on-disk data size.
        """
        raise NotImplementedError

    def get_read_tasks(
        self,
        parallelism: int,
        per_task_row_limit: Optional[int] = None,
        data_context: Optional["DataContext"] = None,

Review comment (bug):
Nit: I think this can produce the wrong number of rows when `num_rows` isn't divisible by `parallelism`.

Code after:


@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
        >>> ds.estimate_inmemory_data_size() == sum(t.metadata.size_bytes for t in tasks)
        True
    """  # noqa: E501

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

    @Deprecated
    def create_reader(self, **read_args) -> "Reader":