chore: import upstream snapshot with attribution

python/ray/data/_internal/datasource/audio_datasource.py

@@ -0,0 +1,57 @@
+import io
+from typing import TYPE_CHECKING, Iterator, List, Union
+
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.util import _check_import
+from ray.data.block import Block
+from ray.data.datasource.file_based_datasource import FileBasedDatasource
+
+if TYPE_CHECKING:
+    import pyarrow
+
+
+class AudioDatasource(FileBasedDatasource):
+    _FILE_EXTENSIONS = [
+        "mp3",
+        "wav",
+        "aac",
+        "flac",
+        "ogg",
+        "m4a",
+        "wma",
+        "alac",
+        "aiff",
+        "pcm",
+        "amr",
+        "opus",
+        "ra",
+        "rm",
+        "au",
+        "mid",
+        "midi",
+        "caf",
+    ]
+
+    def __init__(
+        self,
+        paths: Union[str, List[str]],
+        **file_based_datasource_kwargs,
+    ):
+        super().__init__(paths, **file_based_datasource_kwargs)
+
+        _check_import(self, module="soundfile", package="soundfile")
+
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        import soundfile
+
+        # `soundfile` doesn't support reading from a `pyarrow.NativeFile` directly, so
+        # we need to read the file into memory first.
+        stream = io.BytesIO(f.read())
+        amplitude, sample_rate = soundfile.read(stream, always_2d=True, dtype="float32")
+
+        # (amplitude, channels) -> (channels, amplitude)
+        amplitude = amplitude.transpose((1, 0))
+
+        builder = DelegatingBlockBuilder()
+        builder.add({"amplitude": amplitude, "sample_rate": sample_rate})
+        yield builder.build()