chore: import upstream snapshot with attribution

doc/source/data/doc_code/custom_datasource_example.py

@@ -0,0 +1,79 @@
+# flake8: noqa
+# fmt: off
+
+from typing import Iterator, Union, List
+
+import pyarrow
+
+from ray.data.block import Block
+
+# __datasource_constructor_start__
+from ray.data.datasource import FileBasedDatasource
+
+class ImageDatasource(FileBasedDatasource):
+    def __init__(self, paths: Union[str, List[str]], *, mode: str):
+        super().__init__(
+            paths,
+            file_extensions=["png", "jpg", "jpeg", "bmp", "gif", "tiff"],
+        )
+
+        self.mode = mode  # Specify read options in the constructor
+# __datasource_constructor_end__
+
+# __read_stream_start__
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        import io
+        import numpy as np
+        from PIL import Image
+        from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+
+        data = f.readall()
+        image = Image.open(io.BytesIO(data))
+        image = image.convert(self.mode)
+
+        # Each block contains one row
+        builder = DelegatingBlockBuilder()
+        array = np.asarray(image)
+        item = {"image": array}
+        builder.add(item)
+        yield builder.build()
+# __read_stream_end__
+
+# __read_datasource_start__
+import ray
+
+ds = ray.data.read_datasource(
+    ImageDatasource("s3://anonymous@ray-example-data/batoidea", mode="RGB")
+)
+# __read_datasource_end__
+
+
+from typing import Any, Dict
+import pyarrow
+
+# __datasink_constructor_start__
+from ray.data.datasource import RowBasedFileDatasink
+
+class ImageDatasink(RowBasedFileDatasink):
+    def __init__(self, path: str, column: str, file_format: str):
+        super().__init__(path, file_format=file_format)
+
+        self.column = column
+        self.file_format = file_format  # Specify write options in the constructor
+# __datasink_constructor_end__
+
+# __write_row_to_file_start__
+    def write_row_to_file(self, row: Dict[str, Any], file: pyarrow.NativeFile):
+        import io
+        from PIL import Image
+
+        # PIL can't write to a NativeFile, so we have to write to a buffer first.
+        image = Image.fromarray(row[self.column])
+        buffer = io.BytesIO()
+        image.save(buffer, format=self.file_format)
+        file.write(buffer.getvalue())
+# __write_row_to_file_end__
+
+# __write_datasink_start__
+ds.write_datasink(ImageDatasink("/tmp/results", column="image", file_format="png"))
+# __write_datasink_end__