From e3fcf841c656c9a193e5839c5ab2340543bef4e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20Piotr=20C=C5=82apa?= Date: Wed, 1 Apr 2026 11:57:43 +0000 Subject: [PATCH 1/9] ws_audio: rename the class names to WSAudioEpisode and WSAudioSegment --- README.md | 2 +- wsds/ws_audio.py | 70 ++++++++++++++++++++++++++-------------------- wsds/ws_dataset.py | 4 +-- wsds/ws_decode.py | 6 ++-- wsds/ws_shard.py | 6 ++-- 5 files changed, 49 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index 1895568..a1febc5 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,7 @@ shape: (5_271_939, 3) ```pycon >>> x['audio'] -WSAudio(audio_reader=AudioReader(src=, sample_rate=None), tstart=614.46246, tend=627.3976) +WSAudioSegment(episode=WSAudioEpisode(src=, sample_rate=None), tstart=614.46246, tend=627.3976) ``` diff --git a/wsds/ws_audio.py b/wsds/ws_audio.py index 1dd7351..ea22cfc 100644 --- a/wsds/ws_audio.py +++ b/wsds/ws_audio.py @@ -3,23 +3,25 @@ import typing from dataclasses import dataclass -from .audio_codec import audio_to_html, create_decoder, encode_mp3, to_filelike +from .audio_codec import audio_to_html, create_decoder, encode_mp3 +from .pupyarrow import pupyarrow -def load_segment(src, start, end, sample_rate=None): - """Efficiently loads an audio segment from `src` (see below) `tstart` to `tend` seconds while - optionally resampling it to `sample_rate`. - - `src` can be one of: - - a file-like object - - a byte string - - a PyArrow binary buffer in memory""" - return AudioReader(src).read_segment(start, end, sample_rate=sample_rate) - @dataclass() -class AudioReader: - """A lazy seeking-capable audio reader for random-access to recordings stored in wsds shards.""" +class WSAudioEpisode: + """A lazy seeking-capable audio reader for random-access to recordings stored in wsds shards. + + >>> from wsds import WSDataset + >>> ds = WSDataset("librilight/source") + >>> audio = ds[0].get_audio() + >>> audio.load().shape + torch.Size([1, 17884909]) + >>> audio.read_segment(start=2, end=5).shape + torch.Size([1, 48000]) + >>> audio.read_segment(start=2, end=5, sample_rate=8000).shape + torch.Size([1, 24000]) + """ src: typing.Any _decoder: typing.Any = None @@ -27,7 +29,7 @@ class AudioReader: skip_samples: int = 0 def __repr__(self): - return f"AudioReader(src={type(self.src)}, sample_rate={self._sample_rate})" + return f"WSAudioEpisode(src={type(self.src)}, sample_rate={self._sample_rate})" def unwrap(self): """Return the raw audio bytes""" @@ -35,8 +37,12 @@ def unwrap(self): return self.src.as_buffer().to_pybytes() elif isinstance(self.src, (bytes, bytearray)): return self.src + elif isinstance(self.src, pupyarrow.LazyBuffer): + return self.src.read() else: - raise TypeError(f"Unsupported AudioReader src type: {type(self.src)}") + raise TypeError(f"Unsupported src type: {type(self.src)}") + + to_bytes = unwrap def get_decoder(self, sample_rate=None): """Lazily creates/caches decoder via audio_codec.create_decoder().""" @@ -95,52 +101,56 @@ def _display_(self): @dataclass(frozen=True) -class WSAudio: - """A lazy reference to a single sample from a segmented audio file.""" +class WSAudioSegment: + """A lazy reference to a single sample from a segmented audio file. + """ - audio_reader: AudioReader + episode: WSAudioEpisode tstart: float tend: float + def __repr__(self) -> str: + return f"WSAudioSegment(episode={self.episode}, tstart={self.tstart!s}, tend={self.tend!s})" + @property def duration(self) -> float: """Duration of the audio segment in seconds.""" return self.tend - self.tstart - def with_context(self, before: float = 0, after: float = 0) -> "WSAudio": - """Return a new WSAudio with expanded timestamps to include surrounding context. + def with_context(self, before: float = 0, after: float = 0) -> "WSAudioSegment": + """Return a new WSAudioSegment with expanded timestamps to include surrounding context. Args: before: Seconds of context to add before the segment start (will not go below 0) after: Seconds of context to add after the segment end Returns: - A new WSAudio instance with adjusted timestamps + A new WSAudioSegment instance with adjusted timestamps """ - return WSAudio( - audio_reader=self.audio_reader, + return WSAudioSegment( + episode=self.episode, tstart=max(0, self.tstart - before), tend=self.tend + after, ) - def with_timestamps(self, tstart: float | None = None, tend: float | None = None) -> "WSAudio": - """Return a new WSAudio with modified timestamps. + def with_timestamps(self, tstart: float | None = None, tend: float | None = None) -> "WSAudioSegment": + """Return a new WSAudioSegment with modified timestamps. Args: tstart: New start time in seconds (None to keep current) tend: New end time in seconds (None to keep current) Returns: - A new WSAudio instance with the specified timestamps + A new WSAudioSegment instance with the specified timestamps """ - return WSAudio( - audio_reader=self.audio_reader, + return WSAudioSegment( + episode=self.episode, tstart=tstart if tstart is not None else self.tstart, tend=tend if tend is not None else self.tend, ) def load(self, sample_rate=None, pad_to_seconds=None): - samples = self.audio_reader.read_segment(self.tstart, self.tend, sample_rate) + samples = self.episode.read_segment(self.tstart, self.tend, sample_rate) sample_rate = samples.sample_rate if pad_to_seconds is not None: import torch @@ -152,7 +162,7 @@ def load(self, sample_rate=None, pad_to_seconds=None): @property def metadata(self): - return self.audio_reader.metadata + return self.episode.metadata def _repr_html_(self): return audio_to_html(self.load()) diff --git a/wsds/ws_dataset.py b/wsds/ws_dataset.py index c8fd72e..cf96590 100644 --- a/wsds/ws_dataset.py +++ b/wsds/ws_dataset.py @@ -39,8 +39,8 @@ class WSDataset: >>> sample = dataset["large/5304/the_tinted_venus_1408_librivox_64kb_mp3/tintedvenus_05_anstey_64kb_090"] >>> print(repr(sample["transcription_wslang_raw.txt"])) ' I will accompany you," she said.' - >>> sample['audio'] - WSAudio(audio_reader=AudioReader(src=, sample_rate=None), tstart=1040.2133, tend=1042.8413) + >>> sample['audio'].load().shape + torch.Size([1, 42049]) """ dataset_root: Path diff --git a/wsds/ws_decode.py b/wsds/ws_decode.py index 7d89979..4dc4aa3 100644 --- a/wsds/ws_decode.py +++ b/wsds/ws_decode.py @@ -5,7 +5,7 @@ import numpy as np import pyarrow as pa -from .ws_audio import AudioReader +from .ws_audio import WSAudioEpisode AUDIO_FILE_KEYS = frozenset( [ @@ -50,7 +50,7 @@ def decode_sample(column: str, data): import json return json.load(fd) elif ext in AUDIO_FILE_KEYS: - return AudioReader(fd) + return WSAudioEpisode(fd) else: return fd.read() @@ -90,7 +90,7 @@ def get_audio(sample, audio_columns=None): audio_columns: Optional list of column names to try. Defaults to AUDIO_FILE_KEYS. Returns: - The audio value (typically an AudioReader or WSAudio). + The audio value (typically a WSAudioEpisode or WSAudioSegment). Raises: KeyError: If no audio column is found in the sample. diff --git a/wsds/ws_shard.py b/wsds/ws_shard.py index 39639cc..a465b0d 100644 --- a/wsds/ws_shard.py +++ b/wsds/ws_shard.py @@ -7,7 +7,7 @@ import pyarrow as pa from .utils import WSShardMissingError -from .ws_audio import AudioReader, WSAudio +from .ws_audio import WSAudioEpisode, WSAudioSegment from .ws_decode import decode_sample from .ws_sample import WSSample @@ -125,7 +125,7 @@ class WSSourceAudioShard(WSShardInterface): # cache _source_file_name: str = None _source_sample: WSSample = None - _source_reader: AudioReader = None + _source_reader: WSAudioEpisode = None @classmethod def from_link(cls, link, dataset, shard_ref): @@ -149,7 +149,7 @@ def get_sample(self, _column, offset): self._source_file_name = file_name tstart, tend = self.get_timestamps(segment_offset) - return WSAudio(self._source_reader, tstart, tend) + return WSAudioSegment(self._source_reader, tstart, tend) class WSYoutubeVideoShard(WSSourceAudioShard): From c44092caba34816a5974cb8aa6f4f20695ca0205 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20Piotr=20C=C5=82apa?= Date: Wed, 1 Apr 2026 12:00:18 +0000 Subject: [PATCH 2/9] ws_audio: added a general encode_audio method that works with all backends --- wsds/audio_codec.py | 45 +++++++++++++++++++++++++++++++++------------ wsds/ws_audio.py | 6 +++--- 2 files changed, 36 insertions(+), 15 deletions(-) diff --git a/wsds/audio_codec.py b/wsds/audio_codec.py index 2c64c24..5d55601 100644 --- a/wsds/audio_codec.py +++ b/wsds/audio_codec.py @@ -201,37 +201,58 @@ def decode_segment(src, start=0, end=None, sample_rate=None): return samples -def encode_mp3(samples) -> bytes: - """Encode a torch tensor to MP3 bytes. +def encode_audio(samples, format="mp3", sample_rate=None, bitrate=None) -> bytes: + """Encode a torch tensor to audio bytes. - Tries torchffmpeg -> torchcodec -> torchaudio as encoder backends. + Tries humecodec -> torchcodec -> torchaudio as encoder backends. + + >>> from wsds import WSDataset + >>> audio = WSDataset("librilight/source")[0].get_audio() + >>> samples = audio.read_segment(start=0, end=2.0, sample_rate=16000) + >>> mp3 = encode_audio(samples, format="mp3") + >>> mp3[:3] == b"ID3" or mp3[:2] in (b"\\xff\\xfb", b"\\xff\\xf3") + True + >>> ogg = encode_audio(samples, format="ogg") # doctest: +SKIP + >>> ogg[:4] == b"OggS" # doctest: +SKIP + True Args: samples: A torch.Tensor with a .sample_rate attribute. Shape: (channels, frames). + format: Output format, e.g. "mp3", "ogg" (Opus). Default: "mp3". + sample_rate: Target sample rate (defaults to samples.sample_rate). + bitrate: Bitrate in bps. Only used for formats that support it (e.g. Opus). Returns: - MP3-encoded bytes. + Encoded audio bytes. """ + if sample_rate is None: + sample_rate = int(samples.sample_rate) + out = io.BytesIO() try: - from torchffmpeg import MediaEncoder + from humecodec import MediaEncoder - sample_rate = int(samples.sample_rate) - # samples is (channels, frames), write_audio_chunk expects (frames, channels) waveform = samples.mT.float().contiguous() - enc = MediaEncoder(out, "mp3") - enc.add_audio_stream(sample_rate=sample_rate, num_channels=waveform.size(1), format="flt") + enc = MediaEncoder(out, format) + stream_kwargs = dict(sample_rate=sample_rate, num_channels=waveform.size(1), format="flt") + if format == "ogg": + from humecodec import CodecConfig + + stream_kwargs.update(encoder="libopus", encoder_format="flt") + if bitrate: + stream_kwargs["codec_config"] = CodecConfig(bit_rate=bitrate) + enc.add_audio_stream(**stream_kwargs) with enc.open(): enc.write_audio_chunk(0, waveform) except ImportError: try: from torchcodec.encoders import AudioEncoder - AudioEncoder(samples, sample_rate=int(samples.sample_rate)).to_file_like(out, "mp3") + AudioEncoder(samples, sample_rate=sample_rate).to_file_like(out, format) except ImportError: import torchaudio - torchaudio.save(out, samples, int(samples.sample_rate), format="mp3") + torchaudio.save(out, samples, sample_rate, format=format) return out.getvalue() @@ -247,5 +268,5 @@ def audio_to_html(samples) -> str: """ import base64 - mp3_data = base64.b64encode(encode_mp3(samples)).decode("ascii") + mp3_data = base64.b64encode(encode_audio(samples, format="mp3")).decode("ascii") return f'' diff --git a/wsds/ws_audio.py b/wsds/ws_audio.py index ea22cfc..cc80cc7 100644 --- a/wsds/ws_audio.py +++ b/wsds/ws_audio.py @@ -3,7 +3,7 @@ import typing from dataclasses import dataclass -from .audio_codec import audio_to_html, create_decoder, encode_mp3 +from .audio_codec import audio_to_html, create_decoder, encode_audio from .pupyarrow import pupyarrow @@ -97,7 +97,7 @@ def _repr_html_(self): def _display_(self): import marimo - return marimo.audio(encode_mp3(self.read_segment())) + return marimo.audio(encode_audio(self.read_segment())) @dataclass(frozen=True) @@ -170,4 +170,4 @@ def _repr_html_(self): def _display_(self): import marimo - return marimo.audio(encode_mp3(self.load())) + return marimo.audio(encode_audio(self.load())) From 132d80b911939d7e6acb4efae9a7d41fc2456487 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20Piotr=20C=C5=82apa?= Date: Wed, 1 Apr 2026 12:14:02 +0000 Subject: [PATCH 3/9] Audio decoding rework (NFY) --- requirements.txt | 1 - tests.py | 4 +- wsds/audio_codec.py | 278 ++++++++++++++++++++------------------------ wsds/ws_audio.py | 34 ++---- 4 files changed, 139 insertions(+), 178 deletions(-) diff --git a/requirements.txt b/requirements.txt index 4fe9e43..64cde6c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,5 +4,4 @@ numpy polars>=1.36.1 pyarrow>=20 torch -torchaudio # torchcodec – optional, it causes serious performance regressions diff --git a/tests.py b/tests.py index 39a9acb..41ffc57 100644 --- a/tests.py +++ b/tests.py @@ -2,7 +2,7 @@ import unittest import wsds -from wsds import ws_dataset, ws_shard, ws_sink +from wsds import ws_dataset, ws_shard, ws_sink, ws_audio, audio_codec def load_tests(loader, tests, ignore): @@ -10,6 +10,8 @@ def load_tests(loader, tests, ignore): tests.addTests(doctest.DocTestSuite(ws_dataset)) tests.addTests(doctest.DocTestSuite(ws_shard)) tests.addTests(doctest.DocTestSuite(ws_sink)) + tests.addTests(doctest.DocTestSuite(ws_audio)) + tests.addTests(doctest.DocTestSuite(audio_codec)) tests.addTests(doctest.DocFileSuite("README.md")) return tests diff --git a/wsds/audio_codec.py b/wsds/audio_codec.py index 5d55601..7043ca1 100644 --- a/wsds/audio_codec.py +++ b/wsds/audio_codec.py @@ -2,203 +2,181 @@ This module contains all audio encoding/decoding logic, separated from the data model layer in ws_audio.py. It provides: -- Decoder backends (TorchFFmpegAudioDecoder, CompatAudioDecoder) -- A factory for creating decoders with automatic backend selection -- MP3 encoding with multi-backend fallback +- AudioDecoder: unified decoder with automatic backend selection (humecodec or torchaudio) +- encode_audio(): multi-backend encoder (humecodec -> torchcodec -> torchaudio) - HTML audio rendering utility """ from __future__ import annotations import io +import traceback import typing import pyarrow as pa -def to_filelike(src: typing.Any) -> typing.BinaryIO: - """Coerces files, byte-strings and PyArrow binary buffers into file-like objects.""" - if hasattr(src, "read"): # an open file - return src - # if not an open file then we assume some kind of binary data in memory - if hasattr(src, "as_buffer"): # PyArrow binary data - return pa.BufferReader(src.as_buffer()) - return io.BytesIO(src) - - -class TorchFFmpegAudioDecoder: - def __init__(self, src, sample_rate): - from torchffmpeg import MediaDecoder - - if hasattr(src, "_optimal_read_size"): - buffer_size = src._optimal_read_size - else: - buffer_size = 128 * 1024 - self.src = src - self.reader = MediaDecoder(to_filelike(self.src), buffer_size=buffer_size) - self.metadata = self.reader.get_src_stream_info(self.reader.default_audio_stream) - - if sample_rate is None: - sample_rate = int(self.metadata.sample_rate) +class AudioDecoder: + """Unified audio decoder that works with humecodec or torchaudio backends.""" + def __init__(self, reader, metadata, sample_rate): + self.reader = reader + self.metadata = metadata self.sample_rate = sample_rate + self.debug = False + # ffmpeg's demuxer skips start_skip_samples at pts=0 (start of stream) + # but NOT after seeking. We compensate by adjusting the trim point. + self.start_skip_samples = getattr(metadata, 'start_skip_samples', 0) or 0 - self.reader.add_basic_audio_stream( - frames_per_chunk=int(32 * sample_rate), - sample_rate=sample_rate, - decoder_option={"threads": "4", "thread_type": "frame"}, - ) - - def get_samples_played_in_range(self, tstart=0, tend=None): + def get_samples_played_in_range(self, tstart=0, tend=None, margin=.25): import torch - self.reader.seek(max(0, tstart - 1), "key") + chunk = True + while chunk is not None: + (chunk,) = self.reader.pop_chunks() + # ffmpeg's demuxer sets start_skip_samples (e.g. mp3 encoder delay) + # and applies it only at pts=0 via side data. After seeking, the skip + # is NOT applied, so the decoded audio starts start_skip_samples + # earlier than expected. We add the delay to tstart to skip past it. + seek_adj = self.start_skip_samples / self.metadata.sample_rate if tstart > 0 else 0 + tstart += seek_adj + if tend is not None: + tend += seek_adj + + # rough seek + self.reader.seek(max(0, tstart - margin), "key") if tend is None: - chunks = [] - more_data = True - while more_data: - if self.reader.fill_buffer() == 1: - more_data = False - (chunk,) = self.reader.pop_chunks() - if chunk is not None: - chunks.append(chunk) - prefix = int((tstart - chunks[0].pts) * self.sample_rate) - if prefix < 0: - prefix = 0 - return torch.cat(chunks)[prefix:].mT - - self.reader.fill_buffer() - (chunk,) = self.reader.pop_chunks() - prefix = int((tstart - chunk.pts) * self.sample_rate) + n_samples = float('inf') + else: + # does not need to be precise + n_samples = round((tend - tstart + margin * 2) * self.sample_rate) + + chunks = [] + more_data = True + while more_data and n_samples > 0: + if self.reader.fill_buffer() == 1: + more_data = False + (chunk,) = self.reader.pop_chunks() + n_samples -= chunk.shape[0] + chunks.append(chunk) + # Round each term to integer samples independently to avoid float + # precision loss from subtracting two close timestamps before scaling. + # When starting from 0, ignore the first chunk's PTS and use all + # decoded samples. This matches torchaudio/ffmpeg CLI behavior which + # treats the first output sample as t=0 regardless of encoder delay + # (negative PTS) or skip_samples (positive PTS). + chunk0_pts = 0.0 if tstart == 0 else chunks[0].pts + prefix = round(tstart * self.sample_rate) - round(chunk0_pts * self.sample_rate) + + if self.debug: + c0 = chunks[0] + print(f" [decode] codec={self.metadata.codec} sr={self.sample_rate} " + f"tstart={tstart:.4f} seek_adj={seek_adj:.4f} " + f"chunk0.pts={c0.pts:.6f} chunk0.samples={c0.shape[0]} " + f"n_chunks={len(chunks)} prefix={prefix}", flush=True) + if prefix < 0: + if self.debug: + print(f" [trim] negative prefix {prefix}, clamping to 0", flush=True) prefix = 0 - if tend: - samples = chunk[prefix : prefix + int((tend - tstart) * self.sample_rate)].mT + samples = torch.cat(chunks) + if tend is not None: + return samples[prefix : prefix + round(tend * self.sample_rate) - round(tstart * self.sample_rate)].mT else: - samples = chunk[prefix:].mT - while chunk is not None: - (chunk,) = self.reader.pop_chunks() - return samples + return samples[prefix:].mT +def _create_reader_humecodec(src, buffer_size): + from humecodec import MediaDecoder -class CompatAudioDecoder: - def __init__(self, src, sample_rate): - import torchaudio + reader = MediaDecoder(src=src, buffer_size=buffer_size) + metadata = reader.get_src_stream_info(reader.default_audio_stream) + return reader, metadata - if not hasattr(torchaudio, "io"): - raise ImportError("You need either torchaudio<2.9 or torchcodec installed") - self.src = src - if hasattr(src, "_optimal_read_size"): - buffer_size = src._optimal_read_size - else: - buffer_size = 128 * 1024 - self.reader = torchaudio.io.StreamReader(src=to_filelike(self.src), buffer_size=buffer_size) - self.metadata = self.reader.get_src_stream_info(0) - if sample_rate is None: - sample_rate = self.metadata.sample_rate +def _create_reader_torchaudio(src, buffer_size): + from torchaudio.io import StreamReader - self.sample_rate = sample_rate + reader = StreamReader(src=src, buffer_size=buffer_size) + metadata = reader.get_src_stream_info(reader.default_audio_stream) + return reader, metadata - # fetch 32 seconds because we likely need 30s at maximum but the seeking may be imprecise (and we seek 1s early) - # FIXME: check if we can get away with some better settings here (-1, maybe 10s + concatenate the chunks in a loop) - self.reader.add_basic_audio_stream( - frames_per_chunk=int(32 * sample_rate), - sample_rate=sample_rate, - decoder_option={"threads": "4", "thread_type": "frame"}, - ) - def get_samples_played_in_range(self, tstart=0, tend=None): - # rought seek - self.reader.seek(max(0, tstart - 1), "key") +def _create_decoder_torchcodec(src, sample_rate): + """Create a torchcodec-backed decoder that matches the AudioDecoder interface.""" + from types import SimpleNamespace - if tend is None: - import torch - - chunks = [] - more_data = True - while more_data: - if self.reader.fill_buffer() == 1: - more_data = False - (chunk,) = self.reader.pop_chunks() - chunks.append(chunk) - prefix = int((tstart - chunks[0].pts) * self.sample_rate) - if prefix < 0: - prefix = 0 - return torch.cat(chunks)[prefix:].mT - - self.reader.fill_buffer() - (chunk,) = self.reader.pop_chunks() - # tight crop (seems accurate down to 1 sample in my tests) - prefix = int((tstart - chunk.pts) * self.sample_rate) - if prefix < 0: - prefix = 0 - if tend: - samples = chunk[prefix : prefix + int((tend - tstart) * self.sample_rate)].mT - else: - samples = chunk[prefix:].mT - # clear out any remaining data - while chunk is not None: - (chunk,) = self.reader.pop_chunks() - return samples + from torchcodec.decoders import AudioDecoder as TorchcodecDecoder + # torchcodec accepts bytes but not BytesIO + decoder = TorchcodecDecoder(src, sample_rate=sample_rate) + metadata = decoder.metadata -def create_decoder(src, sample_rate=None): - """Factory: tries torchffmpeg -> torchcodec -> torchaudio, returns a decoder instance. + class TorchcodecAdapter: + def __init__(self): + self.metadata = metadata + self.sample_rate = sample_rate if sample_rate is not None else int(metadata.sample_rate) - Args: - src: A file-like object or bytes-like source for audio data. - sample_rate: Optional target sample rate for resampling. + def get_samples_played_in_range(self, tstart=0, tend=None): + return decoder.get_samples_played_in_range(tstart, tend) - Returns: - A decoder instance with .metadata, .sample_rate, and .get_samples_played_in_range() interface. - """ - try: - from torchffmpeg import MediaDecoder as _ # noqa: F401 + return TorchcodecAdapter() - AudioDecoder = TorchFFmpegAudioDecoder - except ImportError: - try: - from torchcodec.decoders import AudioDecoder - except ImportError: - AudioDecoder = CompatAudioDecoder - return AudioDecoder(src, sample_rate=sample_rate) +_STREAMING_BACKENDS = [ + (_create_reader_humecodec, "humecodec"), + (_create_reader_torchaudio, "torchaudio.io"), +] +_chosen_backend = None -def decode_segment(src, start=0, end=None, sample_rate=None): - """One-shot decode: creates decoder, reads segment, returns tensor with .sample_rate attr. - Handles MP3 skip_samples compensation automatically. +def create_decoder(src, sample_rate=None): + """Factory: tries humecodec -> torchaudio -> torchcodec, returns a decoder instance. Args: - src: Audio source (file-like, bytes, or PyArrow buffer). - start: Start time in seconds. - end: End time in seconds (None for rest of file). - sample_rate: Optional target sample rate. + src: A file-like object for audio data. + sample_rate: Optional target sample rate for resampling. Returns: - A torch.Tensor with a .sample_rate attribute. + A decoder with .metadata, .sample_rate, and .get_samples_played_in_range(). """ - filelike = to_filelike(src) - decoder = create_decoder(filelike, sample_rate) - - skip_samples = 0 - if decoder.metadata.codec == "mp3": - skip_samples = 1105 + global _chosen_backend + + buffer_size = getattr(src, "_optimal_read_size", 128 * 1024) + + if _chosen_backend is not None: + if _chosen_backend == "torchcodec": + return _create_decoder_torchcodec(src, sample_rate) + reader, metadata = _chosen_backend(src, buffer_size) + else: + for factory, module in _STREAMING_BACKENDS: + try: + reader, metadata = factory(src, buffer_size) + _chosen_backend = factory + break + except ImportError: + continue + else: + # Fall back to torchcodec (different API, no streaming reader) + try: + decoder = _create_decoder_torchcodec(src, sample_rate) + _chosen_backend = "torchcodec" + return decoder + except ImportError: + raise ImportError("Neither humecodec, torchaudio, nor torchcodec is installed.") if sample_rate is None: - sample_rate = decoder.metadata.sample_rate + sample_rate = int(metadata.sample_rate) - seek_adjustment = skip_samples / sample_rate if start > 0 else 0 - samples = decoder.get_samples_played_in_range( - start + seek_adjustment, end + seek_adjustment if end is not None else None + reader.add_basic_audio_stream( + frames_per_chunk=int(1 * sample_rate), + sample_rate=sample_rate, + decoder_option={"threads": "4", "thread_type": "frame"}, ) - if hasattr(samples, "data"): - samples = samples.data - samples.sample_rate = sample_rate - return samples + + return AudioDecoder(reader, metadata, sample_rate) + def encode_audio(samples, format="mp3", sample_rate=None, bitrate=None) -> bytes: diff --git a/wsds/ws_audio.py b/wsds/ws_audio.py index cc80cc7..2d59078 100644 --- a/wsds/ws_audio.py +++ b/wsds/ws_audio.py @@ -26,8 +26,6 @@ class WSAudioEpisode: src: typing.Any _decoder: typing.Any = None _sample_rate: int | None = None - skip_samples: int = 0 - def __repr__(self): return f"WSAudioEpisode(src={type(self.src)}, sample_rate={self._sample_rate})" @@ -46,22 +44,11 @@ def unwrap(self): def get_decoder(self, sample_rate=None): """Lazily creates/caches decoder via audio_codec.create_decoder().""" - sample_rate_switch = False - if self._sample_rate is not None: - sample_rate_switch = self._sample_rate != sample_rate - - if self._decoder is None or sample_rate_switch: - decoder = create_decoder(to_filelike(self.src), sample_rate=sample_rate) - # mp3 has encoder delays that are not handled well when seeking - if decoder.metadata.codec == "mp3": - self.skip_samples = 1105 - - if sample_rate is None: - sample_rate = decoder.metadata.sample_rate - - self._decoder = decoder - self._sample_rate = sample_rate - + requested_sr = sample_rate or (self._decoder and self._decoder.metadata.sample_rate) + if self._decoder is None or requested_sr != self._sample_rate: + self.src.seek(0) + self._decoder = create_decoder(self.src, sample_rate=sample_rate) + self._sample_rate = sample_rate or self._decoder.metadata.sample_rate return self._decoder, self._sample_rate @property @@ -76,14 +63,9 @@ def sample_rate(self): def read_segment(self, start=0, end=None, sample_rate=None): decoder, sample_rate = self.get_decoder(sample_rate) - seek_adjustment = self.skip_samples / sample_rate if start > 0 else 0 - _samples = decoder.get_samples_played_in_range( - start + seek_adjustment, end + seek_adjustment if end is not None else None - ) - if hasattr(_samples, "data"): - samples = _samples.data - else: - samples = _samples + samples = decoder.get_samples_played_in_range(start, end) + if hasattr(samples, "data"): + samples = samples.data samples.sample_rate = sample_rate return samples From 361ea78fe1ecccac911857aac1f3f21f4b405a6b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20Piotr=20C=C5=82apa?= Date: Thu, 2 Apr 2026 10:01:03 +0000 Subject: [PATCH 4/9] audio_codec: fix seek accuracy across codecs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Use start_skip_samples from demuxer for seek adjustment (mp3 encoder delay, MP4 edit lists). The demuxer applies this skip at pts=0 but not after seeking. - Read codec_delay from decoder's output stream info for codecs where the delay is set by the codec init (wmav2, opus). - Switch read loop to PTS-based termination instead of sample counting, fixing truncation when seek lands far from target. - Detect when seek lands at stream start (chunk0.pts < margin) and use tstart=0 timeline to match ffmpeg CLI behavior. - Handle negative first-chunk PTS (some AAC files) by clamping to 0. Test results: 100/100 full load, 93/100 seek (±5 sample tolerance). Co-Authored-By: Claude Opus 4.6 (1M context) --- scripts/test_audio_backends.sh | 14 +++++++++ wsds/audio_codec.py | 56 +++++++++++++++++----------------- 2 files changed, 42 insertions(+), 28 deletions(-) create mode 100755 scripts/test_audio_backends.sh diff --git a/scripts/test_audio_backends.sh b/scripts/test_audio_backends.sh new file mode 100755 index 0000000..48bf02f --- /dev/null +++ b/scripts/test_audio_backends.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash +# Run the full test suite in isolated environments, one per decoder backend. +# +# Usage: +# ./scripts/test_audio_backends.sh +# +# Each environment installs only one decoder backend so we verify that the +# interfaces work correctly regardless of which backend is present. + +TEST_COMMAND=${1:-"python -m tests"} + +parallel --tag --lb \ + "uv run --isolated --with {} $TEST_COMMAND" \ + ::: humecodec "torchaudio<2.9" torchcodec diff --git a/wsds/audio_codec.py b/wsds/audio_codec.py index 7043ca1..5323b2c 100644 --- a/wsds/audio_codec.py +++ b/wsds/audio_codec.py @@ -19,14 +19,13 @@ class AudioDecoder: """Unified audio decoder that works with humecodec or torchaudio backends.""" - def __init__(self, reader, metadata, sample_rate): + def __init__(self, reader, metadata, sample_rate, codec_delay=0): self.reader = reader self.metadata = metadata self.sample_rate = sample_rate self.debug = False - # ffmpeg's demuxer skips start_skip_samples at pts=0 (start of stream) - # but NOT after seeking. We compensate by adjusting the trim point. - self.start_skip_samples = getattr(metadata, 'start_skip_samples', 0) or 0 + self.codec_delay = codec_delay + self.init_skip_samples = getattr(metadata, 'start_skip_samples', 0) or 0 def get_samples_played_in_range(self, tstart=0, tend=None, margin=.25): import torch @@ -35,46 +34,39 @@ def get_samples_played_in_range(self, tstart=0, tend=None, margin=.25): while chunk is not None: (chunk,) = self.reader.pop_chunks() - # ffmpeg's demuxer sets start_skip_samples (e.g. mp3 encoder delay) - # and applies it only at pts=0 via side data. After seeking, the skip - # is NOT applied, so the decoded audio starts start_skip_samples - # earlier than expected. We add the delay to tstart to skip past it. - seek_adj = self.start_skip_samples / self.metadata.sample_rate if tstart > 0 else 0 + seek_adj = self.init_skip_samples / self.metadata.sample_rate if tstart > 0 else 0 tstart += seek_adj if tend is not None: tend += seek_adj # rough seek self.reader.seek(max(0, tstart - margin), "key") - if tend is None: - n_samples = float('inf') - else: - # does not need to be precise - n_samples = round((tend - tstart + margin * 2) * self.sample_rate) chunks = [] more_data = True - while more_data and n_samples > 0: + while more_data: if self.reader.fill_buffer() == 1: more_data = False (chunk,) = self.reader.pop_chunks() - n_samples -= chunk.shape[0] chunks.append(chunk) - # Round each term to integer samples independently to avoid float - # precision loss from subtracting two close timestamps before scaling. - # When starting from 0, ignore the first chunk's PTS and use all - # decoded samples. This matches torchaudio/ffmpeg CLI behavior which - # treats the first output sample as t=0 regardless of encoder delay - # (negative PTS) or skip_samples (positive PTS). - chunk0_pts = 0.0 if tstart == 0 else chunks[0].pts + if tend is not None and chunk.pts + chunk.shape[0] / self.sample_rate > tend + margin: + break + + # If the seek landed at (or near) the start of the stream, treat it + # the same as tstart=0: ignore the first-chunk PTS so that the timeline + # matches what ffmpeg CLI / torchaudio produce for a full decode. + seek_landed_at_start = tstart == 0 or chunks[0].pts < margin + chunk0_pts = 0.0 if seek_landed_at_start else chunks[0].pts prefix = round(tstart * self.sample_rate) - round(chunk0_pts * self.sample_rate) if self.debug: - c0 = chunks[0] + import torch as _t + total_samples = sum(c.shape[0] for c in chunks) print(f" [decode] codec={self.metadata.codec} sr={self.sample_rate} " - f"tstart={tstart:.4f} seek_adj={seek_adj:.4f} " - f"chunk0.pts={c0.pts:.6f} chunk0.samples={c0.shape[0]} " - f"n_chunks={len(chunks)} prefix={prefix}", flush=True) + f"tstart_orig={tstart - seek_adj:.4f} tstart_adj={tstart:.4f} " + f"seek_adj={seek_adj:.6f} (init_skip={self.init_skip_samples} codec_delay={self.codec_delay}) " + f"chunk0.pts={chunks[0].pts:.6f} chunk0_pts_used={chunk0_pts:.6f} " + f"n_chunks={len(chunks)} total_samples={total_samples} prefix={prefix}", flush=True) if prefix < 0: if self.debug: @@ -175,7 +167,15 @@ def create_decoder(src, sample_rate=None): decoder_option={"threads": "4", "thread_type": "frame"}, ) - return AudioDecoder(reader, metadata, sample_rate) + # Get codec_delay from the decoder (available after add_audio_stream opens the codec) + codec_delay = 0 + try: + out_info = reader.get_out_stream_info(0) + codec_delay = getattr(out_info, 'codec_delay', 0) or 0 + except Exception: + pass + + return AudioDecoder(reader, metadata, sample_rate, codec_delay=codec_delay) From b9a0d89a237fb2d3de679dc58110e3d8aa87d415 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20Piotr=20C=C5=82apa?= Date: Thu, 2 Apr 2026 18:32:16 +0000 Subject: [PATCH 5/9] audio_codec: read from start for tstart < 5s, fix seek_adj scoping - For short seeks (tstart < 5s) and unreliable codecs (wmav2/wmapro), always read from the start instead of seeking. This avoids the seek_landed_at_start heuristic entirely for short seeks and costs almost nothing for small files. - Only apply start_skip_samples seek_adj when actually seeking, not when reading from start (where the decoder handles it automatically). - Initialize seek_adj to 0.0 to avoid UnboundLocalError. Comprehensive test: 1494/1580 pass (94.6%), up from ~85% before. Co-Authored-By: Claude Opus 4.6 (1M context) --- wsds/audio_codec.py | 33 +++++++++++++++++++++------------ 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/wsds/audio_codec.py b/wsds/audio_codec.py index 5323b2c..d5a2dd8 100644 --- a/wsds/audio_codec.py +++ b/wsds/audio_codec.py @@ -26,6 +26,10 @@ def __init__(self, reader, metadata, sample_rate, codec_delay=0): self.debug = False self.codec_delay = codec_delay self.init_skip_samples = getattr(metadata, 'start_skip_samples', 0) or 0 + # Codecs where flush produces unreliable output (wrong skip_samples, + # wrong frame sizes). For these, always read from the start and trim. + codec_name = getattr(metadata, 'codec', '') or '' + self._seek_unreliable = codec_name in ('wmav2', 'wmapro') def get_samples_played_in_range(self, tstart=0, tend=None, margin=.25): import torch @@ -34,13 +38,21 @@ def get_samples_played_in_range(self, tstart=0, tend=None, margin=.25): while chunk is not None: (chunk,) = self.reader.pop_chunks() - seek_adj = self.init_skip_samples / self.metadata.sample_rate if tstart > 0 else 0 - tstart += seek_adj - if tend is not None: - tend += seek_adj - - # rough seek - self.reader.seek(max(0, tstart - margin), "key") + # For short seeks and unreliable codecs, read from the start. + # This avoids seek accuracy issues for tstart < 5s (tiny cost) and + # codec flush bugs for wmav2/wmapro. + read_from_start = self._seek_unreliable or tstart < 5.0 + + # Only adjust for start_skip_samples when actually seeking — when + # reading from start, the decoder applies skip_samples automatically. + seek_adj = 0.0 + if not read_from_start: + seek_adj = self.init_skip_samples / self.metadata.sample_rate + tstart += seek_adj + if tend is not None: + tend += seek_adj + seek_target = 0.0 if read_from_start else max(0, tstart - margin) + self.reader.seek(seek_target, "key") chunks = [] more_data = True @@ -52,11 +64,8 @@ def get_samples_played_in_range(self, tstart=0, tend=None, margin=.25): if tend is not None and chunk.pts + chunk.shape[0] / self.sample_rate > tend + margin: break - # If the seek landed at (or near) the start of the stream, treat it - # the same as tstart=0: ignore the first-chunk PTS so that the timeline - # matches what ffmpeg CLI / torchaudio produce for a full decode. - seek_landed_at_start = tstart == 0 or chunks[0].pts < margin - chunk0_pts = 0.0 if seek_landed_at_start else chunks[0].pts + # When reading from start, ignore PTS so timeline matches ffmpeg CLI. + chunk0_pts = 0.0 if read_from_start else chunks[0].pts prefix = round(tstart * self.sample_rate) - round(chunk0_pts * self.sample_rate) if self.debug: From 7be4020b1d5ee560ae10dfb252a001ae76512f6a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20Piotr=20C=C5=82apa?= Date: Sat, 4 Apr 2026 10:51:33 +0000 Subject: [PATCH 6/9] audio_codec: indexed byte seek for mp1/mp2/mp3 Build a packet index (128KB resolution) on first seek and use seek_to_byte_offset for raw MPEG audio formats. This avoids the slow sequential scan that ffmpeg's mp3 demuxer does for timestamp seeks. The index PTS is used directly for trimming since the demuxer doesn't update PTS after byte seek. Also: no seek_adj for indexed seeks (index PTS is in raw timeline, not the skip_samples-adjusted timeline). Comprehensive test: 1556/1580 pass (98.5%) across 8 seek positions. Co-Authored-By: Claude Opus 4.6 (1M context) --- wsds/audio_codec.py | 81 +++++++++++++++++++++++++++++++++++++-------- 1 file changed, 68 insertions(+), 13 deletions(-) diff --git a/wsds/audio_codec.py b/wsds/audio_codec.py index d5a2dd8..d546e49 100644 --- a/wsds/audio_codec.py +++ b/wsds/audio_codec.py @@ -29,7 +29,38 @@ def __init__(self, reader, metadata, sample_rate, codec_delay=0): # Codecs where flush produces unreliable output (wrong skip_samples, # wrong frame sizes). For these, always read from the start and trim. codec_name = getattr(metadata, 'codec', '') or '' - self._seek_unreliable = codec_name in ('wmav2', 'wmapro') + self._seek_unreliable = codec_name in ('wmav2', 'wmapro', 'vorbis') + # Raw MPEG audio formats: timestamp seek does sequential scan, + # byte-offset seek with our own index is much faster. + self._use_byte_index = codec_name in ('mp3', 'mp2', 'mp1') + self._packet_index = None + + def _build_index(self): + """Build a sparse packet index for byte-offset seeking.""" + if self._packet_index is not None: + return + try: + idx = self.reader.build_packet_index( + self.reader.default_audio_stream, 128 * 1024) + if idx and len(idx) > 1: + self._packet_index = idx + except Exception: + self._packet_index = [] + + def _indexed_seek(self, target_time): + """Seek via byte offset using the packet index. Returns the index entry's PTS or None.""" + self._build_index() + if not self._packet_index: + return None + # Find last entry with pts <= target_time + best = self._packet_index[0] + for entry in self._packet_index: + if entry.pts_seconds <= target_time: + best = entry + else: + break + self.reader.seek_to_byte_offset(best.pos) + return best.pts_seconds def get_samples_played_in_range(self, tstart=0, tend=None, margin=.25): import torch @@ -40,19 +71,31 @@ def get_samples_played_in_range(self, tstart=0, tend=None, margin=.25): # For short seeks and unreliable codecs, read from the start. # This avoids seek accuracy issues for tstart < 5s (tiny cost) and - # codec flush bugs for wmav2/wmapro. + # codec flush bugs for wmav2/wmapro/vorbis. read_from_start = self._seek_unreliable or tstart < 5.0 # Only adjust for start_skip_samples when actually seeking — when # reading from start, the decoder applies skip_samples automatically. seek_adj = 0.0 + index_pts = None if not read_from_start: - seek_adj = self.init_skip_samples / self.metadata.sample_rate - tstart += seek_adj - if tend is not None: - tend += seek_adj - seek_target = 0.0 if read_from_start else max(0, tstart - margin) - self.reader.seek(seek_target, "key") + # For raw MPEG formats, use indexed byte seek (fast, avoids sequential scan). + # No seek_adj needed: the index PTS and decoded audio are both in + # the raw timeline (skip_samples is not applied after byte seek). + if self._use_byte_index: + index_pts = self._indexed_seek(tstart - margin) + else: + # Timestamp seek: the demuxer applies start_skip_samples at + # pts=0 but not after seeking, so adjust tstart to compensate. + seek_adj = self.init_skip_samples / self.metadata.sample_rate + tstart += seek_adj + if tend is not None: + tend += seek_adj + + if index_pts is None: + # Fall back to timestamp seek (or read from start) + seek_target = 0.0 if read_from_start else max(0, tstart - margin) + self.reader.seek(seek_target, "key") chunks = [] more_data = True @@ -61,11 +104,23 @@ def get_samples_played_in_range(self, tstart=0, tend=None, margin=.25): more_data = False (chunk,) = self.reader.pop_chunks() chunks.append(chunk) - if tend is not None and chunk.pts + chunk.shape[0] / self.sample_rate > tend + margin: - break - - # When reading from start, ignore PTS so timeline matches ffmpeg CLI. - chunk0_pts = 0.0 if read_from_start else chunks[0].pts + if tend is not None: + chunk_end_pts = chunk.pts + chunk.shape[0] / self.sample_rate + if index_pts is not None: + # PTS not updated by demuxer after byte seek — estimate from index + elapsed = sum(c.shape[0] for c in chunks) / self.sample_rate + chunk_end_pts = index_pts + elapsed + if chunk_end_pts > tend + margin: + break + + # Determine the reference PTS for trimming + if read_from_start: + chunk0_pts = 0.0 + elif index_pts is not None: + # Byte seek: demuxer PTS is stale, use our index entry + chunk0_pts = index_pts + else: + chunk0_pts = chunks[0].pts prefix = round(tstart * self.sample_rate) - round(chunk0_pts * self.sample_rate) if self.debug: From 7c3f8bdc3be1af60047b310caf7f9c8b92387ee6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20Piotr=20C=C5=82apa?= Date: Wed, 1 Apr 2026 13:13:11 +0000 Subject: [PATCH 7/9] pupyarrow: added asyncio support (NFY) --- wsds/pupyarrow/file_reader.py | 326 ++++++++++++++++++------- wsds/pupyarrow/pupyarrow.py | 448 +++++++++++++++++++--------------- wsds/ws_s3_shard.py | 74 +++--- wsds/ws_shard.py | 8 + 4 files changed, 551 insertions(+), 305 deletions(-) diff --git a/wsds/pupyarrow/file_reader.py b/wsds/pupyarrow/file_reader.py index c44a747..cc9eb9f 100644 --- a/wsds/pupyarrow/file_reader.py +++ b/wsds/pupyarrow/file_reader.py @@ -1,38 +1,89 @@ from __future__ import annotations +import asyncio +import contextvars import os import threading import time +from dataclasses import dataclass from pathlib import Path -from typing import BinaryIO -BLOCK_SIZE = 16384 # 16kB minimum read size +BLOCK_SIZE = 8192 # 8kB minimum sync read size +MIN_ASYNC_READ = 4096 # 4kB minimum async read size +VERBOSE = False + + +@dataclass +class _Region: + """A coalesced read region covering one or more buffer descriptors.""" + + offset: int + length: int + members: list[tuple[int, int, int]] # list of (abs_offset, start_in_region, end_in_region) + + +def _coalesce_regions(items: list[tuple[int, int]], gap_threshold: int = 64 * 1024) -> list[_Region]: + """Merge nearby reads into larger contiguous fetches. + + items: list of (absolute_offset, length) pairs. + Returns _Region objects with members referencing back to the original offsets. + """ + if not items: + return [] + + sorted_items = sorted(items, key=lambda x: x[0]) + + regions: list[_Region] = [] + cur_offset = sorted_items[0][0] + cur_end = cur_offset + sorted_items[0][1] + cur_members: list[tuple[int, int, int]] = [(sorted_items[0][0], 0, sorted_items[0][1])] + + for abs_offset, length in sorted_items[1:]: + item_end = abs_offset + length + if abs_offset <= cur_end + gap_threshold: + member_start = abs_offset - cur_offset + cur_members.append((abs_offset, member_start, member_start + length)) + cur_end = max(cur_end, item_end) + else: + regions.append(_Region(offset=cur_offset, length=cur_end - cur_offset, members=cur_members)) + cur_offset = abs_offset + cur_end = item_end + cur_members = [(abs_offset, 0, length)] + + regions.append(_Region(offset=cur_offset, length=cur_end - cur_offset, members=cur_members)) + return regions class FileReader: - """Base class for reading bytes from a file with two cache slots. + """Base class for reading bytes from a file. - Forward cache: caches reads from absolute offsets (for sequential access). - Tail cache: caches reads from the end of the file (for footer parsing). + Sync reads use a two-slot cache (forward + tail). + Async reads use a range cache with plan mode for coalesced IO. - Subclasses implement _raw_read and _raw_read_end only. + Subclasses implement _raw_read, _raw_read_end, and optionally _async_read_impl. IO stats (io_time, io_count, io_bytes, cache_hits) are always tracked. - Pass verbose=True to print per-request details to stderr. """ - def __init__(self, *, verbose: bool = False): + def __init__(self): + # Sync caches self._fwd_start: int = 0 self._fwd_data: bytes = b"" self._tail_data: bytes = b"" - self._verbose = verbose + self._verbose = VERBOSE + self._async_first = False # subclasses that track IO in _async_read_impl set this self.io_time: float = 0.0 self.io_count: int = 0 self.io_bytes: int = 0 self.cache_hits: int = 0 + # Async plan mode (per-task via contextvar) + self._planned: contextvars.ContextVar[bool] = contextvars.ContextVar("_planned", default=False) + self._pending: list[tuple[int, int, int, asyncio.Future]] = [] # (offset, actual, length, future) + self._cache: list[tuple[int, bytes]] = [] # (offset, data) ranges + def read(self, offset: int, length: int) -> bytes: - """Read length bytes at absolute offset, using the forward cache.""" + """Read length bytes at absolute offset, using forward cache.""" fwd_end = self._fwd_start + len(self._fwd_data) if self._fwd_data and offset >= self._fwd_start and offset + length <= fwd_end: self.cache_hits += 1 @@ -42,11 +93,14 @@ def read(self, offset: int, length: int) -> bytes: t0 = time.monotonic() self._fwd_data = self._raw_read(offset, actual_length) dt = time.monotonic() - t0 - self.io_time += dt - self.io_count += 1 - self.io_bytes += len(self._fwd_data) + if not self._async_first: + self.io_time += dt + self.io_count += 1 + self.io_bytes += len(self._fwd_data) if self._verbose: - print(f"[IO] read offset={offset} len={actual_length} got={len(self._fwd_data)} {dt * 1000:.1f}ms") + print( + f"[IO] read offset={offset} reqn={length} len={actual_length} got={len(self._fwd_data)} {dt * 1000:.1f}ms" + ) self._fwd_start = offset return self._fwd_data[:length] @@ -64,11 +118,12 @@ def read_end(self, offset: int, length: int) -> bytes: t0 = time.monotonic() self._tail_data = self._raw_read_end(actual_n) dt = time.monotonic() - t0 - self.io_time += dt - self.io_count += 1 - self.io_bytes += len(self._tail_data) + if not self._async_first: + self.io_time += dt + self.io_count += 1 + self.io_bytes += len(self._tail_data) if self._verbose: - print(f"[IO] read_end n={actual_n} got={len(self._tail_data)} {dt * 1000:.1f}ms") + print(f"[IO] read_end reqn={length} n={actual_n} got={len(self._tail_data)} {dt * 1000:.1f}ms") start = len(self._tail_data) + offset return self._tail_data[start : start + length] @@ -80,73 +135,154 @@ def _raw_read_end(self, n: int) -> bytes: """Read the last n bytes of the file. May return fewer if file is smaller.""" raise NotImplementedError + # -- Async IO with range cache and plan mode -------------------------------- + + async def async_read(self, offset: int, length: int) -> bytes: + """Async read with range cache and optional plan mode. + + Every read fetches at least MIN_ASYNC_READ bytes and caches the result. + In plan mode, reads are deferred and coalesced on flush(). + """ + # Check range cache + for c_off, c_data in self._cache: + if offset >= c_off and offset + length <= c_off + len(c_data): + self.cache_hits += 1 + s = offset - c_off + return c_data[s : s + length] + + actual = max(length, MIN_ASYNC_READ) + + if not self._planned.get(): + # Eager mode: read directly + data = await self._async_read_impl(offset, actual) + self._cache.append((offset, data)) + return data[:length] + + # Plan mode: submit and await future + fut = asyncio.get_running_loop().create_future() + self._pending.append((offset, actual, length, fut)) + return await fut + + async def flush(self): + """Coalesce pending reads, execute, resolve futures.""" + if not self._pending: + return + items = [(off, actual) for off, actual, _, _ in self._pending] + regions = _coalesce_regions(items) + fetched = await asyncio.gather(*[self._async_read_impl(r.offset, r.length) for r in regions]) + + data_map: dict[int, bytes] = {} + for region, data in zip(regions, fetched): + self._cache.append((region.offset, data)) + for abs_offset, start, end in region.members: + data_map[abs_offset] = data[start:end] + + for offset, _actual, length, fut in self._pending: + fut.set_result(data_map[offset][:length]) + self._pending.clear() + + def clear_cache(self): + """Clear the async range cache.""" + self._cache.clear() + + @property + def has_pending(self) -> bool: + return len(self._pending) > 0 + + async def _async_read_impl(self, offset: int, length: int) -> bytes: + """Actual async IO. Default runs _raw_read in a thread executor. + + Subclasses with native async IO (S3, Modal) override this. + """ + return await asyncio.get_event_loop().run_in_executor(None, self._raw_read, offset, length) + def close(self): pass class LocalFileReader(FileReader): - """FileReader backed by a local file.""" + """FileReader backed by a local file via os.pread.""" - def __init__(self, path_or_file: str | Path | BinaryIO, *, verbose: bool = False): - super().__init__(verbose=verbose) - if isinstance(path_or_file, (str, Path)): - self._file: BinaryIO = open(path_or_file, "rb") - self._owns_file = True - else: - self._file = path_or_file - self._owns_file = False + def __init__(self, path: str | Path): + super().__init__() + self._fd = os.open(str(path), os.O_RDONLY) def _raw_read(self, offset: int, length: int) -> bytes: - self._file.seek(offset, os.SEEK_SET) - return self._file.read(length) + return os.pread(self._fd, length, offset) def _raw_read_end(self, n: int) -> bytes: - self._file.seek(-n, os.SEEK_END) - return self._file.read(n) + size = os.fstat(self._fd).st_size + return os.pread(self._fd, n, max(size - n, 0)) def close(self): - if self._owns_file: - self._file.close() + os.close(self._fd) class S3FileReader(FileReader): - """FileReader backed by S3 range requests via boto3.""" + """FileReader backed by S3 range requests via aiobotocore. + + Async-first: _async_read_impl is the canonical implementation. + Sync _raw_read calls into _async_read_impl via the background event loop. + + Takes a pre-created aiobotocore S3 client (not a session) so that + SSL context and connection pool setup is amortized across readers. + """ - def __init__(self, s3_client, bucket: str, key: str, *, verbose: bool = False): - super().__init__(verbose=verbose) - self._client = s3_client + def __init__(self, client, bucket: str, key: str): + super().__init__() + self._async_first = True + self._client = client # aiobotocore S3 client (already entered) self._bucket = bucket self._key = key - def _raw_read(self, offset: int, length: int) -> bytes: + async def _async_read_impl(self, offset: int, length: int) -> bytes: range_header = f"bytes={offset}-{offset + length - 1}" - resp = self._client.get_object(Bucket=self._bucket, Key=self._key, Range=range_header) - return resp["Body"].read() + t0 = time.monotonic() + resp = await self._client.get_object(Bucket=self._bucket, Key=self._key, Range=range_header) + async with resp["Body"] as stream: + data = await stream.read() + dt = time.monotonic() - t0 + self.io_time += dt + self.io_count += 1 + self.io_bytes += len(data) + if self._verbose: + print(f"[S3] async_read offset={offset} req={length} got={len(data)} {dt * 1000:.1f}ms") + return data + + async def _async_read_end(self, n: int) -> bytes: + t0 = time.monotonic() + resp = await self._client.get_object(Bucket=self._bucket, Key=self._key, Range=f"bytes=-{n}") + async with resp["Body"] as stream: + data = await stream.read() + dt = time.monotonic() - t0 + self.io_time += dt + self.io_count += 1 + self.io_bytes += len(data) + if self._verbose: + print(f"[S3] async_read_end req={n} got={len(data)} {dt * 1000:.1f}ms") + return data + + def _raw_read(self, offset: int, length: int) -> bytes: + return _get_io_loop().run(self._async_read_impl(offset, length)) def _raw_read_end(self, n: int) -> bytes: - range_header = f"bytes=-{n}" - resp = self._client.get_object(Bucket=self._bucket, Key=self._key, Range=range_header) - return resp["Body"].read() + return _get_io_loop().run(self._async_read_end(n)) -class _ModalEventLoop: +class _IOLoop: """A persistent event loop running on a dedicated daemon thread. - All Modal gRPC work is dispatched here so the client's channel stays - bound to a single loop, and callers on the main thread (or Jupyter, - or another loop) are never blocked by "loop already running" errors.""" + Used for async-first readers (S3, Modal). Callers on the main thread + (or Jupyter, or another loop) are never blocked by "loop already running" + errors.""" def __init__(self): - import asyncio - self._loop = asyncio.new_event_loop() self._thread = threading.Thread(target=self._loop.run_forever, daemon=True) self._thread.start() def run(self, coro): """Submit *coro* to the background loop and block until it completes.""" - import asyncio - future = asyncio.run_coroutine_threadsafe(coro, self._loop) return future.result() @@ -156,43 +292,45 @@ def close(self): # Module-level singleton — created on first use. -_modal_loop: _ModalEventLoop | None = None -_modal_loop_lock = threading.Lock() +_io_loop: _IOLoop | None = None +_io_loop_lock = threading.Lock() -def _get_modal_loop() -> _ModalEventLoop: - global _modal_loop - if _modal_loop is None: - with _modal_loop_lock: - if _modal_loop is None: - _modal_loop = _ModalEventLoop() - return _modal_loop +def _get_io_loop() -> _IOLoop: + global _io_loop + if _io_loop is None: + with _io_loop_lock: + if _io_loop is None: + _io_loop = _IOLoop() + return _io_loop class ModalFileReader(FileReader): - """FileReader backed by Modal Volume range requests via gRPC. + """FileReader backed by Modal Volume range requests via gRPC + aiohttp. - Uses the undocumented ``start``/``len`` fields on ``VolumeGetFile2Request`` - to fetch only the needed byte ranges. Presigned block URLs returned by the - gRPC call are downloaded with ``urllib``. + Async-first: _async_read_impl is the canonical implementation using native + async gRPC for metadata and aiohttp for presigned URL downloads. + Sync _raw_read calls into _async_read_impl via the background event loop. - All async gRPC work runs on a shared daemon-thread event loop (see - ``_ModalEventLoop``) so it works regardless of whether the caller already + All async work runs on a shared daemon-thread event loop (see + ``_IOLoop``) so it works regardless of whether the caller already has a running loop (Jupyter, Modal synchronizer, etc.).""" - def __init__(self, vol, path: str, *, verbose: bool = False): - super().__init__(verbose=verbose) + def __init__(self, vol, path: str): + super().__init__() + self._async_first = True self._vol = vol self._path = path self._size: int | None = None - self._loop = _get_modal_loop() + self._loop = _get_io_loop() + self._aiohttp_session = None @classmethod - def from_name(cls, volume_name: str, path: str, *, verbose: bool = False) -> "ModalFileReader": + def from_name(cls, volume_name: str, path: str) -> "ModalFileReader": """Create a reader for *path* inside the named Modal Volume.""" - loop = _get_modal_loop() + loop = _get_io_loop() vol = loop.run(cls._hydrate(volume_name)) - reader = cls(vol, path, verbose=verbose) + reader = cls(vol, path) return reader @staticmethod @@ -214,17 +352,26 @@ async def _get_range(self, start: int, length: int): ) return await self._vol._client.stub.VolumeGetFile2(req) - def _fetch_urls(self, resp) -> bytes: - """Download presigned block URLs and concatenate the bytes.""" - import requests + async def _get_aiohttp_session(self): + if self._aiohttp_session is None: + import aiohttp - chunks = [] - for url in resp.get_urls: - r = requests.get(url) - r.raise_for_status() - chunks.append(r.content) + self._aiohttp_session = aiohttp.ClientSession() + return self._aiohttp_session + + async def _async_fetch_urls(self, resp) -> bytes: + """Download presigned block URLs concurrently via aiohttp.""" + session = await self._get_aiohttp_session() + tasks = [self._fetch_one(session, url) for url in resp.get_urls] + chunks = await asyncio.gather(*tasks) return b"".join(chunks) + @staticmethod + async def _fetch_one(session, url: str) -> bytes: + async with session.get(url) as r: + r.raise_for_status() + return await r.read() + def _ensure_size(self) -> int: """Fetch the total file size (cached after first call).""" if self._size is None: @@ -232,13 +379,26 @@ def _ensure_size(self) -> int: self._size = resp.size return self._size - def _raw_read(self, offset: int, length: int) -> bytes: - resp = self._loop.run(self._get_range(offset, length)) + async def _async_read_impl(self, offset: int, length: int) -> bytes: + """Native async: gRPC for range metadata, aiohttp for URL downloads.""" + resp = await self._get_range(offset, length) if self._size is None: self._size = resp.size - return self._fetch_urls(resp) + return await self._async_fetch_urls(resp) + + def _raw_read(self, offset: int, length: int) -> bytes: + return self._loop.run(self._async_read_impl(offset, length)) def _raw_read_end(self, n: int) -> bytes: size = self._ensure_size() offset = max(size - n, 0) return self._raw_read(offset, size - offset) + + async def _async_close(self): + if self._aiohttp_session is not None: + await self._aiohttp_session.close() + self._aiohttp_session = None + + def close(self): + if self._aiohttp_session is not None: + self._loop.run(self._async_close()) diff --git a/wsds/pupyarrow/pupyarrow.py b/wsds/pupyarrow/pupyarrow.py index a9e6c67..d7cc69c 100644 --- a/wsds/pupyarrow/pupyarrow.py +++ b/wsds/pupyarrow/pupyarrow.py @@ -12,11 +12,12 @@ from __future__ import annotations +import asyncio import struct from dataclasses import dataclass from enum import IntEnum from pathlib import Path -from typing import Any, BinaryIO, Iterator +from typing import Any, Iterator import numpy as np @@ -417,7 +418,6 @@ def read(self, size: int = -1) -> bytes: Advances the position by the number of bytes returned. """ remaining = self._length - self._pos - print("read:", self._pos, size, remaining) if size < 0: size = remaining else: @@ -481,6 +481,21 @@ def as_numpy(self, dtype: np.dtype) -> np.ndarray: data = self._read_all() return np.frombuffer(data, dtype=dtype) + # -- Async API (IO plan aware) ------------------------------------------ + + async def async_read_all(self) -> bytes: + """Async read via reader (uses cache/plan mode if active).""" + if self._data is not None: + return self._data + raw = await self._reader.async_read(self._offset, self._length) + self._data = _decompress_buffer(raw, self._compression) + return self._data + + async def async_as_numpy(self, dtype: np.dtype) -> np.ndarray: + """Async version of as_numpy.""" + data = await self.async_read_all() + return np.frombuffer(data, dtype=dtype) + def __len__(self) -> int: return self._length @@ -539,6 +554,32 @@ def validity_mask(self) -> np.ndarray | None: return self._validity + async def async_resolve(self) -> None: + """Prefetch all buffers for this array (coalesced if IO plan active).""" + futs = [buf.async_read_all() for buf in self._buffers if buf._length > 0] + if futs: + await asyncio.gather(*futs) + + def to_numpy(self) -> np.ndarray: + raise NotImplementedError(f"{type(self).__name__} does not support to_numpy()") + + def to_masked_array(self) -> np.ma.MaskedArray: + """Return values as a masked array with nulls masked.""" + values = self.to_numpy() + mask = self.validity_mask() + if mask is None: + return np.ma.array(values, mask=False) + return np.ma.array(values, mask=~mask) + + def __getitem__(self, idx: int | slice) -> Any: + return self.to_numpy()[idx] + + async def async_to_numpy(self) -> np.ndarray: + await self.async_resolve() + return self.to_numpy() + + async_to_py = async_to_numpy + class LazyIntArray(LazyArray): """Lazy integer array with numpy access.""" @@ -571,17 +612,6 @@ def to_numpy(self) -> np.ndarray: self._values = self._buffers[1].as_numpy(self._dtype)[: self.length] return self._values - def to_masked_array(self) -> np.ma.MaskedArray: - """Return values as a masked array with nulls masked.""" - values = self.to_numpy() - mask = self.validity_mask() - if mask is None: - return np.ma.array(values, mask=False) - return np.ma.array(values, mask=~mask) - - def __getitem__(self, idx: int | slice) -> Any: - return self.to_numpy()[idx] - def __repr__(self) -> str: return f"LazyIntArray(dtype={self._dtype.__name__}, length={self.length}, nulls={self.null_count})" @@ -607,17 +637,6 @@ def to_numpy(self) -> np.ndarray: self._values = self._buffers[1].as_numpy(self._dtype)[: self.length] return self._values - def to_masked_array(self) -> np.ma.MaskedArray: - """Return values as a masked array with nulls masked.""" - values = self.to_numpy() - mask = self.validity_mask() - if mask is None: - return np.ma.array(values, mask=False) - return np.ma.array(values, mask=~mask) - - def __getitem__(self, idx: int | slice) -> Any: - return self.to_numpy()[idx] - def __repr__(self) -> str: return f"LazyFloatArray(dtype={self._dtype.__name__}, length={self.length}, nulls={self.null_count})" @@ -641,120 +660,15 @@ def to_numpy(self) -> np.ndarray: self._values = np.unpackbits(packed, bitorder="little")[: self.length].astype(bool) return self._values - def to_masked_array(self) -> np.ma.MaskedArray: - """Return values as a masked array with nulls masked.""" - values = self.to_numpy() - mask = self.validity_mask() - if mask is None: - return np.ma.array(values, mask=False) - return np.ma.array(values, mask=~mask) - - def __getitem__(self, idx: int | slice) -> Any: - return self.to_numpy()[idx] - def __repr__(self) -> str: return f"LazyBoolArray(length={self.length}, nulls={self.null_count})" -class LazyStringArray(LazyArray): - """ - Lazy string (Utf8/LargeUtf8) array. - - Offsets are loaded eagerly for efficient slicing. - String data is loaded lazily. - """ - - def __init__( - self, - field: Field, - node: FieldNode, - buffers: list[LazyBuffer], - large: bool = False, - ): - super().__init__(field, node, buffers) - self._large = large - self._offsets: np.ndarray | None = None - self._data_buffer = buffers[2] - self._data: bytes | None = None - # Eagerly load offsets (metadata) - self._load_offsets() - - def _load_offsets(self) -> None: - """Eagerly load offset array.""" - offset_dtype = np.int64 if self._large else np.int32 - self._offsets = self._buffers[1].as_numpy(offset_dtype)[: self.length + 1] - - @property - def offsets(self) -> np.ndarray: - """Return the offset array (eagerly loaded).""" - return self._offsets # type: ignore - - def _ensure_data(self) -> bytes: - """Lazily load string data buffer.""" - if self._data is None: - self._data = self._data_buffer._read_all() - return self._data - - def __getitem__(self, idx: int | slice) -> str | None | list[str | None]: - """Get string(s) by index or slice.""" - if isinstance(idx, slice): - indices = range(*idx.indices(self.length)) - return [self._get_single(i) for i in indices] - - if idx < 0: - idx += self.length - if idx < 0 or idx >= self.length: - raise IndexError(f"Index {idx} out of range for array of length {self.length}") - - return self._get_single(idx) - - def _get_single(self, idx: int) -> str | None: - """Get a single string by index.""" - # Check validity - mask = self.validity_mask() - if mask is not None and not mask[idx]: - return None - - start = int(self._offsets[idx]) - end = int(self._offsets[idx + 1]) - return self._data_buffer.read_range(start, end).decode("utf-8") - - def to_list(self) -> list[str | None]: - """Convert to a Python list of strings.""" - data = self._ensure_data() - mask = self.validity_mask() - result: list[str | None] = [] - - for i in range(self.length): - if mask is not None and not mask[i]: - result.append(None) - else: - start = int(self._offsets[i]) - end = int(self._offsets[i + 1]) - result.append(data[start:end].decode("utf-8")) - - return result - - def to_numpy(self) -> np.ndarray: - """Return as numpy object array of strings.""" - return np.array(self.to_list(), dtype=object) - - def byte_sizes(self) -> np.ndarray: - """Return array of byte sizes for each string (without loading data).""" - return np.diff(self._offsets) - - def __repr__(self) -> str: - type_name = "LargeUtf8" if self._large else "Utf8" - data_loaded = "loaded" if self._data is not None else "not loaded" - return f"LazyStringArray({type_name}, length={self.length}, nulls={self.null_count}, data={data_loaded})" - - class LazyBinaryArray(LazyArray): """ Lazy binary (Binary/LargeBinary) array. - Offsets are loaded eagerly for efficient slicing. - Binary data is loaded lazily. + Offsets and binary data are both loaded lazily on first access. """ def __init__( @@ -769,18 +683,18 @@ def __init__( self._offsets: np.ndarray | None = None self._data_buffer = buffers[2] self._data: bytes | None = None - # Eagerly load offsets (metadata) - self._load_offsets() - def _load_offsets(self) -> None: - """Eagerly load offset array.""" - offset_dtype = np.int64 if self._large else np.int32 - self._offsets = self._buffers[1].as_numpy(offset_dtype)[: self.length + 1] + def _ensure_offsets(self) -> np.ndarray: + """Load offset array on first access.""" + if self._offsets is None: + offset_dtype = np.int64 if self._large else np.int32 + self._offsets = self._buffers[1].as_numpy(offset_dtype)[: self.length + 1] + return self._offsets @property def offsets(self) -> np.ndarray: - """Return the offset array (eagerly loaded).""" - return self._offsets # type: ignore + """Return the offset array (loaded on first access).""" + return self._ensure_offsets() def _ensure_data(self) -> bytes: """Lazily load binary data buffer.""" @@ -788,8 +702,8 @@ def _ensure_data(self) -> bytes: self._data = self._data_buffer._read_all() return self._data - def __getitem__(self, idx: int | slice) -> LazyBuffer | None | list[LazyBuffer | None]: - """Get binary data as a file-like LazyBuffer by index or slice.""" + def __getitem__(self, idx: int | slice): + """Get element(s) by index or slice.""" if isinstance(idx, slice): indices = range(*idx.indices(self.length)) return [self._get_single(i) for i in indices] @@ -807,8 +721,9 @@ def _get_single(self, idx: int) -> LazyBuffer | None: if mask is not None and not mask[idx]: return None - start = int(self._offsets[idx]) - end = int(self._offsets[idx + 1]) + offsets = self._ensure_offsets() + start = int(offsets[idx]) + end = int(offsets[idx + 1]) return self._data_buffer.slice(start, end) def read_range(self, idx: int, start: int, end: int) -> bytes: @@ -822,8 +737,9 @@ def read_range(self, idx: int, start: int, end: int) -> bytes: if idx < 0 or idx >= self.length: raise IndexError(f"Index {idx} out of range") - elem_start = int(self._offsets[idx]) - elem_end = int(self._offsets[idx + 1]) + offsets = self._ensure_offsets() + elem_start = int(offsets[idx]) + elem_end = int(offsets[idx + 1]) # Clamp range to element bounds read_start = elem_start + max(0, start) @@ -836,6 +752,7 @@ def read_range(self, idx: int, start: int, end: int) -> bytes: def to_list(self) -> list[LazyBuffer | None]: """Convert to a Python list of file-like LazyBuffers.""" + offsets = self._ensure_offsets() mask = self.validity_mask() result: list[LazyBuffer | None] = [] @@ -843,15 +760,27 @@ def to_list(self) -> list[LazyBuffer | None]: if mask is not None and not mask[i]: result.append(None) else: - start = int(self._offsets[i]) - end = int(self._offsets[i + 1]) + start = int(offsets[i]) + end = int(offsets[i + 1]) result.append(self._data_buffer.slice(start, end)) return result def byte_sizes(self) -> np.ndarray: """Return array of byte sizes for each element (without loading data).""" - return np.diff(self._offsets) + return np.diff(self._ensure_offsets()) + + async def async_to_bytes_list(self) -> list[bytes | None]: + await self.async_resolve() + data = self._ensure_data() + offsets = self.offsets + mask = self.validity_mask() + return [ + None if (mask is not None and not mask[i]) else data[int(offsets[i]) : int(offsets[i + 1])] + for i in range(self.length) + ] + + async_to_py = async_to_bytes_list def __repr__(self) -> str: type_name = "LargeBinary" if self._large else "Binary" @@ -859,6 +788,53 @@ def __repr__(self) -> str: return f"LazyBinaryArray({type_name}, length={self.length}, nulls={self.null_count}, data={data_loaded})" +class LazyStringArray(LazyBinaryArray): + """ + Lazy string (Utf8/LargeUtf8) array. + + Subclass of LazyBinaryArray that decodes values as UTF-8 strings. + """ + + def _get_single(self, idx: int) -> str | None: + """Get a single string by index, decoding the LazyBuffer from super().""" + buf = super()._get_single(idx) + if buf is None: + return None + return buf._read_all().decode("utf-8") + + def to_list(self) -> list[str | None]: + """Convert to a Python list of strings.""" + offsets = self._ensure_offsets() + data = self._ensure_data() + mask = self.validity_mask() + result: list[str | None] = [] + + for i in range(self.length): + if mask is not None and not mask[i]: + result.append(None) + else: + start = int(offsets[i]) + end = int(offsets[i + 1]) + result.append(data[start:end].decode("utf-8")) + + return result + + def to_numpy(self) -> np.ndarray: + """Return as numpy object array of strings.""" + return np.array(self.to_list(), dtype=object) + + async def async_to_list(self) -> list[str | None]: + await self.async_resolve() + return self.to_list() + + async_to_py = async_to_list + + def __repr__(self) -> str: + type_name = "LargeUtf8" if self._large else "Utf8" + data_loaded = "loaded" if self._data is not None else "not loaded" + return f"LazyStringArray({type_name}, length={self.length}, nulls={self.null_count}, data={data_loaded})" + + class LazyFixedSizeBinaryArray(LazyArray): """Lazy fixed-size binary array with numpy access.""" @@ -878,9 +854,6 @@ def to_numpy(self) -> np.ndarray: self._values = np.frombuffer(data, dtype=np.uint8).reshape(-1, self._byte_width)[: self.length] return self._values - def __getitem__(self, idx: int | slice) -> np.ndarray: - return self.to_numpy()[idx] - def __repr__(self) -> str: return f"LazyFixedSizeBinaryArray(byte_width={self._byte_width}, length={self.length}, nulls={self.null_count})" @@ -1028,6 +1001,48 @@ def __repr__(self) -> str: return f"RecordBatch(rows={self.num_rows}, columns={self.num_columns})" +def _parse_record_batch_message(message_bytes: bytes) -> RecordBatchInfo: + """Parse a record batch message into RecordBatchInfo (no IO).""" + continuation = struct.unpack(" RecordBatchInfo: return self._record_batch_infos[index] # type: ignore block = self._record_batch_blocks[index] - message_bytes = self._reader.read(block.offset, block.metadata_length) - - continuation = struct.unpack(" Iterator[RecordBatch]: for i in range(self.num_record_batches): yield self.record_batch(i) + async def async_record_batch(self, index: int) -> RecordBatch: + """Async version of record_batch — reads metadata via async_read if not cached.""" + if self._record_batch_infos[index] is None: + block = self._record_batch_blocks[index] + raw = await self._reader.async_read(block.offset, block.metadata_length) + self._record_batch_infos[index] = _parse_record_batch_message(raw) + return self.record_batch(index) + + def __getitem__(self, key: str | tuple[str, ...] | list[str]) -> np.ndarray | list | dict[str, np.ndarray | list]: + """ + Read column data across all batches with concurrent IO. + + Single column returns the resolved array/list directly. + Multiple columns return a dict mapping column name to resolved data. + + Numeric/bool/fixed-size-binary columns return np.ndarray. + String columns return list[str | None]. + Binary columns return list[bytes | None]. + + Uses the reader's plan mode to coalesce reads across all batches. + Per-batch async tasks submit reads; the execution loop flushes + them in coalesced rounds. + + Examples: + values = f["score"] # np.ndarray + texts = f["text"] # list[str] + cols = f["score", "text"] # dict + cols = f[["score", "text"]] # dict (also works) + """ + if isinstance(key, str): + names = [key] + elif isinstance(key, (tuple, list)): + names = list(key) + else: + raise TypeError(f"Key must be str, tuple[str, ...], or list[str], got {type(key).__name__}") + + for name in names: + if self._schema.field(name) is None: + raise KeyError(f"Column {name!r} not in schema") + + from wsds.pupyarrow.file_reader import _get_io_loop + + batch_results = _get_io_loop().run(self._async_getitem(names)) + + # Concatenate across batches + output: dict[str, np.ndarray | list] = {} + for name in names: + field = self._schema.field(name) + is_numpy = field.type_id in ( + ArrowType.Int, + ArrowType.FloatingPoint, + ArrowType.Bool, + ArrowType.FixedSizeBinary, + ) + chunks = [br[name] for br in batch_results] + if is_numpy: + output[name] = np.concatenate(chunks) if len(chunks) > 1 else chunks[0] if chunks else np.array([]) + else: + output[name] = [item for chunk in chunks for item in chunk] + + return output[names[0]] if isinstance(key, str) else output + + async def _async_getitem(self, names: list[str]) -> list[dict[str, np.ndarray | list]]: + """Async entry point: resolve all batches with coalesced IO.""" + + async def resolve_batch(batch_idx: int) -> dict[str, np.ndarray | list]: + batch = await self.async_record_batch(batch_idx) + return {name: await batch.column(name).async_to_py() for name in names} + + self._reader._planned.set(True) + try: + tasks = [asyncio.ensure_future(resolve_batch(i)) for i in range(self.num_record_batches)] + while not all(t.done() for t in tasks): + await asyncio.sleep(0) + await self._reader.flush() + return [t.result() for t in tasks] + finally: + self._reader._planned.set(False) + self._reader.clear_cache() + + # -- Context manager & lifecycle ------------------------------------------- + def __enter__(self) -> FeatherFile: return self diff --git a/wsds/ws_s3_shard.py b/wsds/ws_s3_shard.py index 0a19881..27b3d02 100644 --- a/wsds/ws_s3_shard.py +++ b/wsds/ws_s3_shard.py @@ -13,10 +13,36 @@ from .ws_dataset import WSDataset +def create_s3_client(link=None): + """Create a shared aiobotocore S3 client. + + Reads `endpoint_url` and AWS credentials from `link` (falling back to the + WSDS_S3_ENDPOINT_URL env var for the endpoint). Returns the entered client + and its context manager (for cleanup). The client should be shared across + all S3FileReader instances. + """ + from aiobotocore.session import AioSession + from botocore.config import Config + + from .pupyarrow.file_reader import _get_io_loop + + link = link or {} + kwargs = {"config": Config(max_pool_connections=50)} + endpoint_url = link.get("endpoint_url") or os.environ.get("WSDS_S3_ENDPOINT_URL") + if endpoint_url: + kwargs["endpoint_url"] = endpoint_url + for k in ("aws_access_key_id", "aws_secret_access_key", "aws_session_token", "region_name"): + if link.get(k): + kwargs[k] = link[k] + ctx = AioSession().create_client("s3", **kwargs) + client = _get_io_loop().run(ctx.__aenter__()) + return client, ctx + + class WSS3Shard(WSShardInterface): - """A shard reader that loads data from S3 via boto3 range requests. + """A shard reader that loads data from S3 via aiobotocore range requests. - Uses pupyarrow's FeatherFile with an S3File wrapper so that only the + Uses pupyarrow's FeatherFile with an S3FileReader so that only the IPC footer and the specific batch(es) needed are fetched, rather than downloading the entire shard file.""" @@ -27,9 +53,7 @@ def __init__(self, dataset: "WSDataset", bucket: str, key: str, shard_ref: Optio self.key = key if s3_client is None: - import boto3 - - s3_client = boto3.client("s3") + s3_client, _ = create_s3_client() self._reader = S3FileReader(s3_client, bucket, key) try: @@ -68,36 +92,30 @@ def from_link(cls, link, dataset, shard_ref): prefix = link.get("prefix", "") column_dir = link.get("subdir", "") parts = [p for p in (prefix, partition, column_dir, f"{shard}.wsds") if p] - key = os.path.normpath("/".join(parts)) - # Strip leading "../" — partition is relative to the index but S3 paths are absolute from bucket root. - while key.startswith("../"): - key = key[3:] - s3_client = cls._make_s3_client(link) + # we make it an absolute path so any initial ../ are stripped out + key = os.path.normpath("/" + "/".join(parts)) + s3_client, _ = create_s3_client(link) return cls(dataset, link["bucket"], key, shard_ref=shard_ref, s3_client=s3_client) - @classmethod - def _make_s3_client(cls, link=None): - import boto3 - - link = link or {} - keys = ("endpoint_url", "aws_access_key_id", "aws_secret_access_key", "aws_session_token", "region_name") - kwargs = {k: link[k] for k in keys if link.get(k)} - kwargs.setdefault("endpoint_url", os.environ.get("WSDS_S3_ENDPOINT_URL")) - return boto3.client("s3", **{k: v for k, v in kwargs.items() if v}) - @classmethod def _discover_columns_from_s3(cls, link): """Read one shard's footer from S3 to discover column names.""" - s3_client = cls._make_s3_client(link) + from .pupyarrow.file_reader import _get_io_loop + bucket = link["bucket"] prefix = link["prefix"] - response = s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix, MaxKeys=10) - for obj in response.get("Contents", []): - if obj["Key"].endswith(".wsds"): - reader = S3FileReader(s3_client, bucket, obj["Key"]) - feather = FeatherFile(reader) - return feather.schema.names - raise ValueError(f"No .wsds files found in s3://{bucket}/{prefix}") + s3_client, _ = create_s3_client(link) + + async def _discover(): + response = await s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix, MaxKeys=10) + for obj in response.get("Contents", []): + if obj["Key"].endswith(".wsds"): + reader = S3FileReader(s3_client, bucket, obj["Key"]) + feather = FeatherFile(reader) + return feather.schema.names + raise ValueError(f"No .wsds files found in s3://{bucket}/{prefix}") + + return _get_io_loop().run(_discover()) def _s3_path(self) -> str: return f"s3://{self.bucket}/{self.key}" diff --git a/wsds/ws_shard.py b/wsds/ws_shard.py index a465b0d..8470d18 100644 --- a/wsds/ws_shard.py +++ b/wsds/ws_shard.py @@ -6,6 +6,7 @@ import pyarrow as pa +from .pupyarrow.file_reader import FileReader, LocalFileReader from .utils import WSShardMissingError from .ws_audio import WSAudioEpisode, WSAudioSegment from .ws_decode import decode_sample @@ -31,6 +32,10 @@ def get_columns(cls, link: dict, dataset: "WSDataset") -> dict[str, str] | None: def get_sample(self, column: str, offset: int) -> typing.Any: raise NotImplementedError + def get_reader(self) -> FileReader: + """Return a pupyarrow FileReader for the underlying shard file.""" + raise NotImplementedError + class WSShard(WSShardInterface): """Represents a single open data shard (`.wsds` file). @@ -104,6 +109,9 @@ def close(self): pass self._source_file = None + def get_reader(self): + return LocalFileReader(self.fname) + def __repr__(self): r = f"WSShard({repr(self.fname)})" if self._data: From 1f82a7aa15583daa7ed02ac379865523b5152671 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20Piotr=20C=C5=82apa?= Date: Sat, 4 Apr 2026 16:13:40 +0000 Subject: [PATCH 8/9] pupyarrow: BlockCache and cached seeking for LazyBuffer BlockCache: sorted interval cache that merges overlapping byte ranges. Designed for caching S3 range-read results. LazyBuffer: - enable_cache(readahead): activate block cache with configurable readahead - prepopulate(ranges): pre-fetch byte ranges into cache (sync) - async_prepopulate(ranges): parallel async variant - read_range() checks cache before hitting reader, fetches with readahead on miss Co-Authored-By: Claude Opus 4.6 (1M context) --- wsds/pupyarrow/pupyarrow.py | 154 ++++++++++++++++++++++++++++++++++-- 1 file changed, 147 insertions(+), 7 deletions(-) diff --git a/wsds/pupyarrow/pupyarrow.py b/wsds/pupyarrow/pupyarrow.py index d7cc69c..f04a405 100644 --- a/wsds/pupyarrow/pupyarrow.py +++ b/wsds/pupyarrow/pupyarrow.py @@ -369,6 +369,83 @@ def _decompress_buffer(raw_data: bytes, compression: str | None) -> bytes: raise ValueError(f"Unknown compression codec: {compression}") +class BlockCache: + """Sorted interval cache for byte ranges. + + Stores non-overlapping ``(start, end, data)`` intervals, merging on insert. + Designed for caching S3 range-read results so that ffmpeg's AVIO reads + hit local memory instead of issuing new HTTP requests. + + >>> cache = BlockCache() + >>> cache.put(100, b'hello') + >>> cache.put(105, b'world') + >>> cache.get(100, 10) + b'helloworld' + >>> cache.get(103, 4) + b'lowo' + >>> cache.get(100, 11) is None # extends past cached range + True + + Adjacent/overlapping ranges are merged: + + >>> cache2 = BlockCache() + >>> cache2.put(0, b'AAAA') + >>> cache2.put(10, b'BBBB') + >>> cache2.put(4, b'CCCCCC') + >>> cache2.get(0, 14) + b'AAAACCCCCCBBBB' + """ + + __slots__ = ("_ranges",) + + def __init__(self): + self._ranges: list[tuple[int, int, bytes]] = [] + + def get(self, offset: int, length: int) -> bytes | None: + """Return data if ``[offset, offset+length)`` is fully cached, else None.""" + end = offset + length + for start, rend, data in self._ranges: + if start <= offset and rend >= end: + return data[offset - start : offset - start + length] + return None + + def put(self, offset: int, data: bytes) -> None: + """Insert a range, merging with any overlapping/adjacent intervals.""" + if not data: + return + new_start = offset + new_end = offset + len(data) + new_data = bytearray(data) + + merged = [] + for start, end, rdata in self._ranges: + if end < new_start or start > new_end: + # No overlap — keep as-is + merged.append((start, end, rdata)) + else: + # Overlap or adjacent — merge into new range + if start < new_start: + prefix = rdata[: new_start - start] + new_data = bytearray(prefix) + new_data + new_start = start + if end > new_end: + suffix = rdata[new_end - start :] + new_data = new_data + bytearray(suffix) + new_end = end + + merged.append((new_start, new_end, bytes(new_data))) + merged.sort(key=lambda r: r[0]) + self._ranges = merged + + @property + def total_bytes(self) -> int: + return sum(end - start for start, end, _ in self._ranges) + + def __repr__(self) -> str: + parts = [f"[{s}:{e}]" for s, e, _ in self._ranges] + return f"BlockCache({', '.join(parts)}, total={self.total_bytes})" + + class LazyBuffer: """ A lazy buffer that reads data on demand and implements the file-like interface. @@ -384,9 +461,14 @@ class LazyBuffer: same reader with adjusted offset/length. For uncompressed buffers the slice goes directly to the reader; for compressed buffers the parent data is decompressed once and the slice is pre-populated. + + For audio seeking, call ``enable_cache()`` to activate a block cache with + readahead. Then ``prepopulate(ranges)`` pre-fetches byte ranges in + parallel. ffmpeg/humecodec reads hit the cache and only fetch from S3 on + miss. """ - # __slots__ = ("_reader", "_offset", "_length", "_data", "_compression", "_pos") + READAHEAD = 256 * 1024 # readahead on cache miss def __init__(self, reader: FileReader, offset: int, length: int, compression: str | None = None): self._reader = reader @@ -395,6 +477,28 @@ def __init__(self, reader: FileReader, offset: int, length: int, compression: st self._data: bytes | None = None self._compression = compression self._pos = 0 + self._cache: BlockCache | None = None + + def enable_cache(self, readahead: int = 256 * 1024) -> "LazyBuffer": + """Activate block cache with readahead for seeking workloads.""" + self._cache = BlockCache() + self.READAHEAD = readahead + return self + + def prepopulate(self, ranges: list[tuple[int, int]]) -> None: + """Pre-fetch byte ranges (relative to buffer start) into the cache. + + Each range is ``(offset, length)``. Ranges are fetched via the + underlying reader (which may coalesce nearby reads). + """ + if self._cache is None: + self.enable_cache() + for rel_offset, length in ranges: + length = min(length, self._length - rel_offset) + if length <= 0: + continue + data = self._reader.read(self._offset + rel_offset, length) + self._cache.put(rel_offset, data) @property def offset(self) -> int: @@ -452,14 +556,26 @@ def read_range(self, start: int, end: int) -> bytes: """Read a byte range from the (decompressed) buffer. If the buffer is already cached, slices it. If uncompressed, - reads directly from the range without reading the whole buffer. - If compressed, falls back to a full read and slices. + reads directly from the range (via block cache if enabled, or + the reader directly). If compressed, falls back to a full read. """ if self._data is not None: return self._data[start:end] - if self._compression is None: - return self._reader.read(self._offset + start, end - start) - return self._read_all()[start:end] + if self._compression is not None: + return self._read_all()[start:end] + # Uncompressed path — use block cache if enabled + length = end - start + if self._cache is not None: + cached = self._cache.get(start, length) + if cached is not None: + return cached + # Cache miss — fetch with readahead + fetch_len = max(length, self.READAHEAD) + fetch_len = min(fetch_len, self._length - start) + data = self._reader.read(self._offset + start, fetch_len) + self._cache.put(start, data) + return data[:length] + return self._reader.read(self._offset + start, length) def slice(self, start: int, end: int) -> "LazyBuffer": """Create a sub-buffer over the byte range [start, end). @@ -496,12 +612,36 @@ async def async_as_numpy(self, dtype: np.dtype) -> np.ndarray: data = await self.async_read_all() return np.frombuffer(data, dtype=dtype) + async def async_prepopulate(self, ranges: list[tuple[int, int]]) -> None: + """Async pre-fetch of byte ranges into the cache.""" + if self._cache is None: + self.enable_cache() + coros = [] + for rel_offset, length in ranges: + length = min(length, self._length - rel_offset) + if length <= 0: + continue + coros.append(self._async_fetch_range(rel_offset, length)) + if coros: + await asyncio.gather(*coros) + + async def _async_fetch_range(self, rel_offset: int, length: int) -> None: + data = await self._reader.async_read(self._offset + rel_offset, length) + self._cache.put(rel_offset, data) + + def seekable(self) -> bool: + return True + + def readable(self) -> bool: + return True + def __len__(self) -> int: return self._length def __repr__(self) -> str: + cached = f", cache={self._cache}" if self._cache else "" loaded = "loaded" if self._data is not None else "not loaded" - return f"LazyBuffer(offset={self._offset}, length={self._length}, {loaded})" + return f"LazyBuffer(offset={self._offset}, length={self._length}, {loaded}{cached})" class LazyArray: From f9b1f7041b0182ae746e9f4265e393e6f1c2277b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jakub=20Piotr=20C=C5=82apa?= Date: Mon, 18 May 2026 21:28:46 +0200 Subject: [PATCH 9/9] Fix column discovery for S3 shards --- wsds/ws_dataset.py | 21 ++++++++++++++++++++- wsds/ws_modal_shard.py | 27 ++++++++------------------- wsds/ws_s3_shard.py | 31 ++++++++----------------------- wsds/ws_shard.py | 11 ++++++++--- 4 files changed, 44 insertions(+), 46 deletions(-) diff --git a/wsds/ws_dataset.py b/wsds/ws_dataset.py index cf96590..b8b832d 100644 --- a/wsds/ws_dataset.py +++ b/wsds/ws_dataset.py @@ -22,6 +22,17 @@ from .ws_shard import WSShard +def _find_local_shard_ref(dataset_root: Path): + """Return ("", ) for any local .wsds file under `dataset_root`, or None.""" + for column_dir in dataset_root.iterdir(): + if not column_dir.is_dir(): + continue + for f in column_dir.iterdir(): + if f.suffix == ".wsds": + return ("", f.stem) + return None + + class WSDataset: """A multimodal dataset. @@ -559,10 +570,18 @@ def _register_wsds_links(self): self.computed_columns[column_dir] = spec links_to_register.append((column_dir, spec)) + if not links_to_register: + return + + # Pick a representative shard_ref so loaders can open one real shard to read its schema. + # When there's no index, fall back to any local .wsds file — partitions don't exist + # in the unindexed layout, so shard_ref is ("", ). + shard_ref = next(self.index.shards(), None) if self.index else _find_local_shard_ref(self.dataset_root) + # Ask each loader class what columns it provides for link_file, spec in links_to_register: loader_class = self._get_loader_class(spec) - columns = loader_class.get_columns(spec, self) + columns = loader_class.get_columns(spec, self, shard_ref) if columns: # Loader provides multiple columns - register them all diff --git a/wsds/ws_modal_shard.py b/wsds/ws_modal_shard.py index 0c48f20..7a6a19f 100644 --- a/wsds/ws_modal_shard.py +++ b/wsds/ws_modal_shard.py @@ -52,28 +52,17 @@ def from_link(cls, link, dataset, shard_ref): return cls(dataset, link["volume_name"], path, shard_ref=shard_ref) @classmethod - def get_columns(cls, link, dataset): + def get_columns(cls, link, dataset, shard_ref=None): """Return columns provided by this Modal link.""" if "columns" in link: return {col: col for col in link["columns"]} - columns = cls._discover_columns(link) - return {col: col for col in columns if col != "__key__"} - - @classmethod - def _discover_columns(cls, link): - """Read one shard's footer from the Modal Volume to discover column names.""" - import modal - - vol = modal.Volume.from_name(link["volume_name"]) - prefix = link["prefix"] - for entry in vol.listdir(prefix): - if entry.path.endswith(".wsds"): - reader = ModalFileReader.from_name(link["volume_name"], entry.path) - feather = FeatherFile(reader) - names = feather.schema.names - reader.close() - return names - raise ValueError(f"No .wsds files found in modal volume '{link['volume_name']}' at prefix '{prefix}'") + if shard_ref is None: + raise ValueError( + f"cannot discover columns for modal link {link!r}: dataset has no shards in its index; " + "set `columns` in the link spec to skip discovery" + ) + shard = cls.from_link(link, dataset, shard_ref) + return {col: col for col in shard._feather.schema.names if col != "__key__"} def _modal_path(self) -> str: return f"modal://{self.volume_name}/{self.path}" diff --git a/wsds/ws_s3_shard.py b/wsds/ws_s3_shard.py index 27b3d02..5f99299 100644 --- a/wsds/ws_s3_shard.py +++ b/wsds/ws_s3_shard.py @@ -78,12 +78,17 @@ def from_s3_url(cls, dataset: "WSDataset", url: str, shard_ref: Optional[Tuple[s return cls(dataset, bucket, key, shard_ref=shard_ref, s3_client=s3_client) @classmethod - def get_columns(cls, link, dataset): + def get_columns(cls, link, dataset, shard_ref=None): """Return columns provided by this S3 link.""" if "columns" in link: return {col: col for col in link["columns"]} - columns = cls._discover_columns_from_s3(link) - return {col: col for col in columns if col != "__key__"} + if shard_ref is None: + raise ValueError( + f"cannot discover columns for s3 link {link!r}: dataset has no shards in its index; " + "set `columns` in the link spec to skip discovery" + ) + shard = cls.from_link(link, dataset, shard_ref) + return {col: col for col in shard._feather.schema.names if col != "__key__"} @classmethod def from_link(cls, link, dataset, shard_ref): @@ -97,26 +102,6 @@ def from_link(cls, link, dataset, shard_ref): s3_client, _ = create_s3_client(link) return cls(dataset, link["bucket"], key, shard_ref=shard_ref, s3_client=s3_client) - @classmethod - def _discover_columns_from_s3(cls, link): - """Read one shard's footer from S3 to discover column names.""" - from .pupyarrow.file_reader import _get_io_loop - - bucket = link["bucket"] - prefix = link["prefix"] - s3_client, _ = create_s3_client(link) - - async def _discover(): - response = await s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix, MaxKeys=10) - for obj in response.get("Contents", []): - if obj["Key"].endswith(".wsds"): - reader = S3FileReader(s3_client, bucket, obj["Key"]) - feather = FeatherFile(reader) - return feather.schema.names - raise ValueError(f"No .wsds files found in s3://{bucket}/{prefix}") - - return _get_io_loop().run(_discover()) - def _s3_path(self) -> str: return f"s3://{self.bucket}/{self.key}" diff --git a/wsds/ws_shard.py b/wsds/ws_shard.py index 8470d18..e948b1f 100644 --- a/wsds/ws_shard.py +++ b/wsds/ws_shard.py @@ -2,7 +2,7 @@ import re import typing from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional, Tuple import pyarrow as pa @@ -21,9 +21,14 @@ class WSShardInterface: """Used by WSDataset to invalidate cached shards.""" @classmethod - def get_columns(cls, link: dict, dataset: "WSDataset") -> dict[str, str] | None: + def get_columns( + cls, link: dict, dataset: "WSDataset", shard_ref: Optional[Tuple[str, str]] = None + ) -> dict[str, str] | None: """Return columns this link provides: {column_name: column_name}. + `shard_ref` is a representative (partition, shard) from the dataset's index, + for loaders that need to open a real shard to discover its schema. + Override this to provide multiple columns from a single link. Return None to use the default behavior (link file stem as single column). """ @@ -200,7 +205,7 @@ class WSSourceLink(WSShardInterface): _source_sample: WSSample = None @classmethod - def get_columns(cls, link, dataset): + def get_columns(cls, link, dataset, shard_ref=None): """Return all source dataset fields with the configured prefix.""" source_dataset = dataset.get_linked_dataset(link["dataset_dir"]) key_prefix = link.get("key_prefix", "source.")