From e032f153fe469779246a9861b50a9bf22456c151 Mon Sep 17 00:00:00 2001 From: Hynek Kydlicek Date: Wed, 2 Sep 2026 14:50:08 +0200 Subject: [PATCH 01/11] perf(lance): coalesce batches for multipart writes Co-Authored-By: OpenAI Codex --- docs/reading-data/lance.md | 30 ++- docs/reading-data/tabular-files.md | 24 ++- docs/writing-data/lance.md | 32 ++++ src/refiner/__init__.py | 10 +- src/refiner/pipeline/__init__.py | 10 +- src/refiner/pipeline/pipeline.py | 41 ++++- src/refiner/pipeline/sinks/__init__.py | 2 + src/refiner/pipeline/sinks/lance.py | 139 +++++++++++++- src/refiner/pipeline/sources/base.py | 10 + src/refiner/pipeline/sources/lance.py | 18 +- src/refiner/pipeline/sources/limited.py | 13 +- .../pipeline/sources/readers/parquet.py | 13 +- tests/pipeline/test_sinks.py | 173 +++++++++++++++++- tests/readers/test_lance_source.py | 66 ++++++- tests/readers/test_limited_source.py | 7 + tests/readers/test_parquet_reader.py | 55 ++++++ 16 files changed, 606 insertions(+), 37 deletions(-) diff --git a/docs/reading-data/lance.md b/docs/reading-data/lance.md index 0887283df..55c8c63e9 100644 --- a/docs/reading-data/lance.md +++ b/docs/reading-data/lance.md @@ -20,15 +20,39 @@ pipeline = mdr.load_lance( "s3://my-bucket/hands.lance", version=42, columns=["image", "frame_id"], - batch_size=128, num_shards=32, max_rows=2_000, ) ``` When `version` is omitted, Refiner resolves the latest version once and pins it -for the pipeline. Column projection is pushed into Lance, and `batch_size` -controls the streamed Arrow batch size. +for the pipeline. Column projection is pushed into Lance. + +Use the pipeline-wide execution block limit when rows contain large media or +other variable-sized values: + +```python +pipeline = mdr.load_lance("s3://my-bucket/hands.lance").with_max_block_rows(128) +``` + +Refiner propagates this limit into Lance's scanner, so the source does not +materialize a larger Arrow batch and then immediately split it. The same limit +continues to apply to blocks produced by downstream transforms. Omit it to use +Lance's internal read-batch default. + +When source rows are lightweight references but downstream transforms expand +them into large media values, tune the two boundaries independently: + +```python +pipeline = mdr.load_lance( + "s3://my-bucket/hands.lance", + read_batch_rows=256, +).with_max_block_rows(8) +``` + +Here Lance scans 256 rows at a time while transforms and sinks receive blocks +of at most eight rows. An explicit `read_batch_rows` value takes precedence over +scanner inheritance from `with_max_block_rows`. ## Limit the number of rows diff --git a/docs/reading-data/tabular-files.md b/docs/reading-data/tabular-files.md index 776eb8eb7..39cb72019 100644 --- a/docs/reading-data/tabular-files.md +++ b/docs/reading-data/tabular-files.md @@ -19,6 +19,29 @@ pipeline = mdr.read_parquet( Parquet is the best general-purpose tabular format when you control the data. It preserves schema information and works well with vectorized transforms. +For rows with large or variable-sized values, set the pipeline-wide execution +block limit: + +```python +pipeline = mdr.read_parquet("/data/episodes/*.parquet").with_max_block_rows(256) +``` + +Refiner passes this limit to the Parquet scanner and also enforces it on +downstream execution blocks. Parquet may still emit smaller batches at file or +row-group boundaries. + +To tune Parquet scanning independently, set `read_batch_rows`: + +```python +pipeline = mdr.read_parquet( + "/data/episodes/*.parquet", + read_batch_rows=4_096, +).with_max_block_rows(256) +``` + +The explicit read setting controls Parquet scanner batches; the pipeline limit +continues to cap downstream execution blocks. + ## JSON and JSONL ```python @@ -58,4 +81,3 @@ See [Schemas and DTypes](../transforms/schemas-and-dtypes.md). - [Vectorized Transforms](../transforms/vectorized-transforms.md) - [Parquet and JSONL Writers](../writing-data/parquet-and-jsonl.md) - diff --git a/docs/writing-data/lance.md b/docs/writing-data/lance.md index 277cf6a68..840e21b19 100644 --- a/docs/writing-data/lance.md +++ b/docs/writing-data/lance.md @@ -64,6 +64,30 @@ the dataset writer uses Lance for fragments and fsspec for attempt handoff metadata. Credential-bearing URIs and URI query parameters are rejected so pipeline plans do not serialize secrets. +Refiner coalesces small pipeline blocks into approximately 256 MiB Arrow +batches before handing them to Lance. This is the default for efficient +object-store writes. For unusual memory budgets, configure both the target and +the per-shard logical Arrow buffer limit: + +```python +pipeline.write_lance_dataset( + "s3://my-bucket/clean.lance", + io=mdr.LanceIOConfig( + upload_concurrency=32, + multipart_part_bytes=5 * 1024 * 1024, + target_batch_bytes=256 * 1024 * 1024, + max_buffered_bytes=1024 * 1024 * 1024, + ), +) +``` + +The buffer limit must be at least four times the target to cover one batch in +Lance, one queued batch, rows being accumulated, and the transient coalesced +batch. A single row larger than the target is necessarily written alone. +The target must also contain at least one part per upload slot. Upload +concurrency and part size are process-wide Lance settings; a retained worker +must be restarted before using different values. + ## Adding columns Use `AddColumns()` for row-preserving enrichment such as model inference: @@ -147,3 +171,11 @@ a mapping to use per-column defaults instead: ```python mode=mdr.AddColumns(fill={"embedding": None, "status": "not_processed"}) ``` + +## Internal Notes + +Lance drains multipart uploads at input-batch boundaries. Byte-aware +coalescing gives each boundary enough data to use the configured upload +concurrency without requiring an upstream Lance change. The buffer calculation +covers one batch in Lance, one or more queued batches, rows being accumulated, +and the transient Arrow concatenation. diff --git a/src/refiner/__init__.py b/src/refiner/__init__.py index ab77f22cd..35cafffe8 100644 --- a/src/refiner/__init__.py +++ b/src/refiner/__init__.py @@ -18,6 +18,7 @@ "AddColumns": "refiner.pipeline.sinks.lance", "Append": "refiner.pipeline.sinks.lance", "Create": "refiner.pipeline.sinks.lance", + "LanceIOConfig": "refiner.pipeline.sinks.lance", "Overwrite": "refiner.pipeline.sinks.lance", "CUDAVersion": "refiner.pipeline", "BlobAssetConfig": "refiner.pipeline.sinks.assets", @@ -62,6 +63,7 @@ "AddColumns", "Append", "Create", + "LanceIOConfig", "Overwrite", "CUDAVersion", "BlobAssetConfig", @@ -139,7 +141,13 @@ def __dir__() -> list[str]: import refiner.text as text import refiner.video as video from refiner.launchers.secrets import Secrets - from refiner.pipeline.sinks.lance import AddColumns, Append, Create, Overwrite + from refiner.pipeline.sinks.lance import ( + AddColumns, + Append, + Create, + LanceIOConfig, + Overwrite, + ) from refiner.io import read_blob from refiner.pipeline import ( CUDAVersion, diff --git a/src/refiner/pipeline/__init__.py b/src/refiner/pipeline/__init__.py index 3cfa74019..41d117eff 100644 --- a/src/refiner/pipeline/__init__.py +++ b/src/refiner/pipeline/__init__.py @@ -7,6 +7,7 @@ "AddColumns": "refiner.pipeline.sinks.lance", "Append": "refiner.pipeline.sinks.lance", "Create": "refiner.pipeline.sinks.lance", + "LanceIOConfig": "refiner.pipeline.sinks.lance", "Overwrite": "refiner.pipeline.sinks.lance", "Row": "refiner.pipeline.data.row", "Shard": "refiner.pipeline.data.shard", @@ -39,6 +40,7 @@ "AddColumns", "Append", "Create", + "LanceIOConfig", "Overwrite", "CUDAVersion", "GPU", @@ -83,7 +85,13 @@ def __dir__() -> list[str]: if TYPE_CHECKING: - from refiner.pipeline.sinks.lance import AddColumns, Append, Create, Overwrite + from refiner.pipeline.sinks.lance import ( + AddColumns, + Append, + Create, + LanceIOConfig, + Overwrite, + ) from refiner.pipeline.data.row import Row from refiner.pipeline.data.shard import Shard from refiner.pipeline.pipeline import ( diff --git a/src/refiner/pipeline/pipeline.py b/src/refiner/pipeline/pipeline.py index c9de581ba..0b8c26ca5 100644 --- a/src/refiner/pipeline/pipeline.py +++ b/src/refiner/pipeline/pipeline.py @@ -43,6 +43,7 @@ Create, JsonlSink, LanceDatasetSink, + LanceIOConfig, LanceSink, LanceWriteConfig, ParquetSink, @@ -114,6 +115,7 @@ class RefinerPipeline: source: BaseSource pipeline_steps: tuple[RefinerStep, ...] _compiled_segments: tuple[Segment, ...] | None + read_batch_rows: int | None max_block_rows: int | None max_vectorized_block_bytes: int | None sink: BaseSink | None @@ -123,6 +125,7 @@ def __init__( source: BaseSource, pipeline_steps: Sequence[RefinerStep] | None = None, *, + read_batch_rows: int | None = None, max_block_rows: int | None = None, max_vectorized_block_bytes: int | None = None, sink: BaseSink | None = None, @@ -132,19 +135,27 @@ def __init__( Args: source: Source that plans shards and emits source rows or blocks. pipeline_steps: Ordered transform steps applied after the source. + read_batch_rows: Optional source-scanner batch size override. When + omitted, cooperative readers inherit ``max_block_rows``. max_block_rows: Optional maximum rows emitted in one execution block. max_vectorized_block_bytes: Optional soft byte cap for vectorized Arrow blocks. Python row blocks are controlled by ``max_block_rows`` instead. sink: Optional writer sink attached by a ``write_*`` method. """ + if read_batch_rows is not None and read_batch_rows <= 0: + raise ValueError("read_batch_rows must be > 0 when provided") if max_block_rows is not None and max_block_rows <= 0: raise ValueError("max_block_rows must be > 0 when provided") if max_vectorized_block_bytes is not None and max_vectorized_block_bytes <= 0: raise ValueError("max_vectorized_block_bytes must be > 0 when provided") - self.source = source + source_batch_rows = ( + read_batch_rows if read_batch_rows is not None else max_block_rows + ) + self.source = source.with_read_batch_rows(source_batch_rows) self.pipeline_steps = tuple(pipeline_steps) if pipeline_steps else () self._compiled_segments = None + self.read_batch_rows = read_batch_rows self.max_block_rows = max_block_rows self.max_vectorized_block_bytes = max_vectorized_block_bytes self.sink = sink @@ -159,6 +170,7 @@ def add_step(self, step: RefinerStep) -> "RefinerPipeline": return self.__class__( self.source, self.pipeline_steps + (step,), + read_batch_rows=self.read_batch_rows, max_block_rows=self.max_block_rows, max_vectorized_block_bytes=self.max_vectorized_block_bytes, sink=self.sink, @@ -192,6 +204,7 @@ def _add_vectorized_op(self, op: VectorizedOp) -> "RefinerPipeline": return self.__class__( self.source, self.pipeline_steps[:-1] + (merged,), + read_batch_rows=self.read_batch_rows, max_block_rows=self.max_block_rows, max_vectorized_block_bytes=self.max_vectorized_block_bytes, sink=self.sink, @@ -207,6 +220,7 @@ def with_max_block_rows(self, max_block_rows: int | None) -> "RefinerPipeline": return self.__class__( self.source, self.pipeline_steps, + read_batch_rows=self.read_batch_rows, max_block_rows=max_block_rows, max_vectorized_block_bytes=self.max_vectorized_block_bytes, sink=self.sink, @@ -225,6 +239,7 @@ def with_max_vectorized_block_bytes( return self.__class__( self.source, self.pipeline_steps, + read_batch_rows=self.read_batch_rows, max_block_rows=self.max_block_rows, max_vectorized_block_bytes=max_vectorized_block_bytes, sink=self.sink, @@ -240,6 +255,7 @@ def with_sink(self, sink: BaseSink | None) -> "RefinerPipeline": return self.__class__( self.source, self.pipeline_steps, + read_batch_rows=self.read_batch_rows, max_block_rows=self.max_block_rows, max_vectorized_block_bytes=self.max_vectorized_block_bytes, sink=sink, @@ -764,6 +780,7 @@ def write_lance_dataset( mode: LanceWriteConfig = Create(), columns: Sequence[str] | None = None, assets: AssetWriteConfig | None = None, + io: LanceIOConfig = LanceIOConfig(), ) -> "RefinerPipeline": """Attach a distributed Lance dataset writer or schema-evolution sink. @@ -772,6 +789,11 @@ def write_lance_dataset( omitted rows receive the configured scalar or per-column fill value. The legacy ``mode="add_columns"`` remains strict and requires exactly one output for every source row. + + Input blocks are coalesced into byte-sized Arrow batches before Lance + writes them. This keeps multipart uploads busy even when upstream emits + small blocks. Pass ``LanceIOConfig`` to control multipart concurrency, + part size, batch size, and the per-shard logical memory bound. """ source_uri: str | None = None source_version: int | None = None @@ -799,6 +821,7 @@ def write_lance_dataset( source_uri=source_uri, source_version=source_version, assets=assets, + io=io, ) ) @@ -1491,7 +1514,7 @@ def read_parquet( target_shard_bytes: int = DEFAULT_TARGET_SHARD_BYTES, num_shards: int | None = None, max_rows: int | None = None, - arrow_batch_size: int = 65536, + read_batch_rows: int | None = None, columns_to_read: Sequence[str] | None = None, filter: Expr | None = None, split_row_groups: bool = False, @@ -1505,6 +1528,8 @@ def read_parquet( to row groups or row ranges at read time. `filter` uses Arrow expressions for row-group pruning plus row-level filtering during reads. `max_rows` applies after that reader-level filter and before pipeline transforms. + ``read_batch_rows`` overrides the Parquet scanner batch size independently + of downstream execution block sizing. """ max_rows = _validated_max_rows(max_rows) return RefinerPipeline( @@ -1516,7 +1541,6 @@ def read_parquet( recursive=recursive, target_shard_bytes=target_shard_bytes, num_shards=num_shards, - arrow_batch_size=_bounded_reader_window(arrow_batch_size, max_rows), columns_to_read=columns_to_read, filter=filter, split_row_groups=split_row_groups, @@ -1524,7 +1548,8 @@ def read_parquet( dtypes=dtypes, ), max_rows, - ) + ), + read_batch_rows=read_batch_rows, ) @@ -1533,15 +1558,17 @@ def load_lance( *, version: int | str | None = None, columns: Sequence[str] | None = None, - batch_size: int = 65_536, num_shards: int | None = None, max_rows: int | None = None, + read_batch_rows: int | None = None, ) -> RefinerPipeline: """Create a pipeline over one immutable version of a Lance dataset. ``num_shards`` groups whole, adjacent Lance fragments into the requested number of scheduling shards without splitting individual fragments. ``max_rows`` reads at most that many leading rows from the pinned version. + ``read_batch_rows`` overrides Lance's scanner batch size independently of + downstream execution block sizing. """ max_rows = _validated_max_rows(max_rows) return RefinerPipeline( @@ -1550,11 +1577,11 @@ def load_lance( input, version=version, columns=columns, - batch_size=_bounded_reader_window(batch_size, max_rows), num_shards=num_shards, ), max_rows, - ) + ), + read_batch_rows=read_batch_rows, ) diff --git a/src/refiner/pipeline/sinks/__init__.py b/src/refiner/pipeline/sinks/__init__.py index 31d4b602c..79c818237 100644 --- a/src/refiner/pipeline/sinks/__init__.py +++ b/src/refiner/pipeline/sinks/__init__.py @@ -5,6 +5,7 @@ Append, Create, LanceDatasetSink, + LanceIOConfig, LanceWriteConfig, LanceWriteMode, Overwrite, @@ -21,6 +22,7 @@ "Create", "FileCleanupReducerSink", "LanceDatasetSink", + "LanceIOConfig", "LanceSink", "LanceWriteMode", "LanceWriteConfig", diff --git a/src/refiner/pipeline/sinks/lance.py b/src/refiner/pipeline/sinks/lance.py index 4b4413ce4..f02749331 100644 --- a/src/refiner/pipeline/sinks/lance.py +++ b/src/refiner/pipeline/sinks/lance.py @@ -9,6 +9,7 @@ import queue as queue_module import re import tempfile +import threading from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Any, Literal, get_args @@ -95,6 +96,90 @@ class AddColumns: _LANCE_ROW_ADDRESS_FRAGMENT_SHIFT = 32 _LANCE_ROW_ADDRESS_POSITION_MASK = (1 << _LANCE_ROW_ADDRESS_FRAGMENT_SHIFT) - 1 _ADD_COLUMNS_FILL_BATCH_ROWS = 65_536 +_MIN_MULTIPART_PART_BYTES = 5 * 1024 * 1024 +_MAX_MULTIPART_PART_BYTES = 5 * 1024 * 1024 * 1024 - 1 +_LANCE_PROCESS_CONFIG_LOCK = threading.Lock() +_LANCE_PROCESS_UPLOAD_CONFIG: tuple[int, int] | None = None + + +@dataclass(frozen=True, slots=True) +class LanceIOConfig: + """Performance and memory controls for Lance dataset writers.""" + + upload_concurrency: int = 32 + multipart_part_bytes: int = 5 * 1024 * 1024 + target_batch_bytes: int = 256 * 1024 * 1024 + max_buffered_bytes: int = 1024 * 1024 * 1024 + + def __post_init__(self) -> None: + if self.upload_concurrency <= 0: + raise ValueError("upload_concurrency must be greater than zero") + if not ( + _MIN_MULTIPART_PART_BYTES + <= self.multipart_part_bytes + <= _MAX_MULTIPART_PART_BYTES + ): + raise ValueError("multipart_part_bytes must be between 5 MiB and 5 GiB") + if self.target_batch_bytes <= 0: + raise ValueError("target_batch_bytes must be greater than zero") + required_batch_bytes = self.upload_concurrency * self.multipart_part_bytes + if self.target_batch_bytes < required_batch_bytes: + raise ValueError( + "target_batch_bytes must be at least upload_concurrency times " + "multipart_part_bytes" + ) + if self.max_buffered_bytes < 4 * self.target_batch_bytes: + raise ValueError( + "max_buffered_bytes must be at least four times target_batch_bytes" + ) + + def to_plan(self) -> dict[str, int]: + return { + "upload_concurrency": self.upload_concurrency, + "multipart_part_bytes": self.multipart_part_bytes, + "target_batch_bytes": self.target_batch_bytes, + "max_buffered_bytes": self.max_buffered_bytes, + } + + +def _configure_lance_process(config: LanceIOConfig) -> None: + """Set Lance's process-global writer controls before its first upload.""" + global _LANCE_PROCESS_UPLOAD_CONFIG + requested = (config.upload_concurrency, config.multipart_part_bytes) + with _LANCE_PROCESS_CONFIG_LOCK: + if _LANCE_PROCESS_UPLOAD_CONFIG is not None: + if _LANCE_PROCESS_UPLOAD_CONFIG != requested: + raise RuntimeError( + "This worker process already initialized Lance with " + f"upload_concurrency={_LANCE_PROCESS_UPLOAD_CONFIG[0]} and " + f"multipart_part_bytes={_LANCE_PROCESS_UPLOAD_CONFIG[1]}; " + "start a fresh worker to use different Lance I/O settings" + ) + return + os.environ["LANCE_UPLOAD_CONCURRENCY"] = str(config.upload_concurrency) + os.environ["LANCE_INITIAL_UPLOAD_SIZE"] = str(config.multipart_part_bytes) + _LANCE_PROCESS_UPLOAD_CONFIG = requested + logger.info( + "Configured Lance uploads concurrency={} part_bytes={}", + config.upload_concurrency, + config.multipart_part_bytes, + ) + + +def _split_batch_by_bytes( + batch: pa.RecordBatch, + target_bytes: int, +) -> list[pa.RecordBatch]: + """Split a batch until each multi-row slice fits the target byte size.""" + if batch.num_rows <= 1 or batch.nbytes <= target_bytes: + return [batch] + split_at = max(1, min(batch.num_rows - 1, batch.num_rows // 2)) + return [ + *_split_batch_by_bytes(batch.slice(0, split_at), target_bytes), + *_split_batch_by_bytes( + batch.slice(split_at, batch.num_rows - split_at), target_bytes + ), + ] def _import_lance() -> Any: @@ -485,12 +570,20 @@ def __init__( dataset_uri: str, schema: pa.Schema, mode: LanceWriteModeName, + io_config: LanceIOConfig, ) -> None: + _configure_lance_process(io_config) self.dataset_uri = dataset_uri self.schema = schema self.mode = mode + self.io_config = io_config + self._pending_batches: list[pa.RecordBatch] = [] + self._pending_bytes = 0 + queue_slots = max( + 1, io_config.max_buffered_bytes // io_config.target_batch_bytes - 3 + ) self.queue: queue_module.Queue[pa.RecordBatch | object] = queue_module.Queue( - maxsize=8 + maxsize=queue_slots ) self.closed = False self.task_future = _LANCE_WRITER_POOL.submit(self._run) @@ -548,14 +641,38 @@ def put_batches(self, batches: list[pa.RecordBatch]) -> None: if self.closed: raise RuntimeError("Cannot write to a closed Lance shard writer.") for batch in batches: - if self._spool_writer is not None: - self._spool_writer.write_batch(batch) - else: - self._put(batch) + for split_batch in _split_batch_by_bytes( + batch, self.io_config.target_batch_bytes + ): + if ( + self._pending_batches + and self._pending_bytes + split_batch.nbytes + > self.io_config.target_batch_bytes + ): + self._flush_pending() + self._pending_batches.append(split_batch) + self._pending_bytes += split_batch.nbytes + if self._pending_bytes >= self.io_config.target_batch_bytes: + self._flush_pending() + + def _flush_pending(self) -> None: + if not self._pending_batches: + return + if len(self._pending_batches) == 1: + batch = self._pending_batches[0] + else: + batch = pa.concat_batches(self._pending_batches) + self._pending_batches = [] + self._pending_bytes = 0 + if self._spool_writer is not None: + self._spool_writer.write_batch(batch) + else: + self._put(batch) def finish(self) -> list[str]: if not self.closed: self.closed = True + self._flush_pending() if self._spool_writer is None: self._put(_QUEUE_CLOSED) if self._spool_writer is not None: @@ -583,7 +700,9 @@ def __init__( output: DataFolder, fill_missing: bool = False, fill: object | Mapping[str, object] = None, + io_config: LanceIOConfig = LanceIOConfig(), ) -> None: + _configure_lance_process(io_config) self.fragment = fragment self.fragment_id = fragment_id self.num_rows = num_rows @@ -807,6 +926,7 @@ def __init__( source_uri: str | None = None, source_version: int | None = None, assets: AssetWriteConfig | None = None, + io: LanceIOConfig = LanceIOConfig(), ) -> None: mode, fill_missing, fill = _normalize_write_mode(mode) if mode == "add_columns" and not columns: @@ -846,6 +966,7 @@ def __init__( self.source_uri = source_uri self.source_version = source_version self.assets = assets + self.io = io if mode == "add_columns" and self.output.abs_path() != source_uri: raise ValueError("add_columns must write back to the loaded Lance dataset") if isinstance(assets, FileAssetConfig): @@ -1027,6 +1148,7 @@ def _write_add_columns_block(self, shard_id: str, block: Block) -> None: output=self.output, fill_missing=self.fill_missing, fill=self.fill, + io_config=self.io, ) writers[fragment_id] = writer positions = pc.call_function( @@ -1072,6 +1194,7 @@ def write_shard_block(self, shard_id: str, block: Block) -> None: dataset_uri=self._dataset_uri(), schema=table.schema, mode=self.mode, + io_config=self.io, ) self._writers_by_shard[shard_id] = writer writer.put_batches(table.to_batches()) @@ -1242,6 +1365,7 @@ def describe(self) -> tuple[str, str, dict[str, object]]: args["fill_missing"] = True if self.assets is not None: args["assets"] = asset_config_to_plan(self.assets) + args["io"] = self.io.to_plan() return ("write_lance_dataset", "writer", args) def build_reducer(self) -> BaseSink | None: @@ -1253,6 +1377,7 @@ def build_reducer(self) -> BaseSink | None: columns=self.columns, fill_missing=self.fill_missing, fill=self.fill, + io=self.io, ) @@ -1267,6 +1392,7 @@ def __init__( columns: Sequence[str] | None = None, fill_missing: bool = False, fill: object | Mapping[str, object] = None, + io: LanceIOConfig = LanceIOConfig(), ) -> None: _validate_write_mode(mode) self.output = DataFolder.resolve(output) @@ -1277,6 +1403,7 @@ def __init__( self.columns = tuple(columns) if columns is not None else None self.fill_missing = fill_missing self.fill = fill + self.io = io _validate_fill_mapping(fill, columns) self._managed_path_pattern = _compile_output_path_patterns( _METADATA_FILENAME_TEMPLATE @@ -1305,6 +1432,7 @@ def describe(self) -> tuple[str, str, dict[str, object]]: args["assets_subdir"] = self.assets_subdir if self.fill_missing: args["fill_missing"] = True + args["io"] = self.io.to_plan() return ("write_lance_dataset_commit", "writer", args) def write_shard_block(self, shard_id: str, block: Block) -> None: @@ -1577,6 +1705,7 @@ def _fill_add_columns_fragments( output=self.output, fill_missing=True, fill=self.fill, + io_config=self.io, ) updated_fragment, next_schema = active_writer.finish() active_writer = None diff --git a/src/refiner/pipeline/sources/base.py b/src/refiner/pipeline/sources/base.py index 54ba443c9..c5022d32c 100644 --- a/src/refiner/pipeline/sources/base.py +++ b/src/refiner/pipeline/sources/base.py @@ -45,6 +45,16 @@ def read(self) -> Iterator[SourceUnit]: for shard in self.list_shards(): yield from self.iter_shard_units(shard) + def with_read_batch_rows(self, max_rows: int | None) -> "BaseSource": + """Return a source configured for the pipeline's execution block limit. + + Sources that stream native record batches can override this hook to avoid + materializing batches larger than the execution engine will accept. + Other sources may ignore the hint and rely on the engine's hard output + block limit. + """ + return self + @property def schema(self) -> pa.Schema | None: return None diff --git a/src/refiner/pipeline/sources/lance.py b/src/refiner/pipeline/sources/lance.py index 3d151f145..ea79d453c 100644 --- a/src/refiner/pipeline/sources/lance.py +++ b/src/refiner/pipeline/sources/lance.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import Iterator, Sequence +from copy import copy import posixpath from typing import Any @@ -25,6 +26,7 @@ _LANCE_ROW_ADDRESS_COLUMN = "_rowaddr" _LANCE_BLOB_METADATA_KEY = b"lance-encoding:blob" +_DEFAULT_READ_BATCH_ROWS = 65_536 def _import_lance() -> Any: @@ -67,11 +69,8 @@ def __init__( *, version: int | str | None = None, columns: Sequence[str] | None = None, - batch_size: int = 65_536, num_shards: int | None = None, ) -> None: - if batch_size <= 0: - raise ValueError("batch_size must be > 0") if columns is not None and len(set(columns)) != len(columns): raise ValueError("Lance columns must be unique") validate_explicit_num_shards(num_shards) @@ -85,7 +84,7 @@ def __init__( self.dataset_uri = self.input.abs_path() validate_lance_uri(self.dataset_uri) self.columns = tuple(columns) if columns is not None else None - self.batch_size = int(batch_size) + self._read_batch_rows = _DEFAULT_READ_BATCH_ROWS self.num_shards = num_shards self._dataset_cache: Any | None = None @@ -137,6 +136,13 @@ def __init__( ) self._schema = projected_schema + def with_read_batch_rows(self, max_rows: int | None) -> "LanceSource": + source = copy(self) + source._read_batch_rows = ( + _DEFAULT_READ_BATCH_ROWS if max_rows is None else max_rows + ) + return source + @property def schema(self) -> pa.Schema: return self._schema @@ -263,7 +269,7 @@ def read_shard(self, shard: Shard) -> Iterator[SourceUnit]: rows_read = 0 for batch in fragment.to_batches( columns=list(self.columns) if self.columns is not None else None, - batch_size=self.batch_size, + batch_size=self._read_batch_rows, with_row_address=True, blob_handling="blobs_descriptions", ): @@ -291,7 +297,7 @@ def describe(self) -> dict[str, object]: "path": self.dataset_uri, "version": self.version, "columns": list(self.columns) if self.columns is not None else None, - "batch_size": self.batch_size, + "read_batch_rows": self._read_batch_rows, "num_shards": self.num_shards, } diff --git a/src/refiner/pipeline/sources/limited.py b/src/refiner/pipeline/sources/limited.py index 3af51333b..23c1450a8 100644 --- a/src/refiner/pipeline/sources/limited.py +++ b/src/refiner/pipeline/sources/limited.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import Iterator +from copy import copy from typing import Any import pyarrow as pa @@ -25,11 +26,21 @@ class LimitedSource(BaseSource): def __init__(self, source: BaseSource, *, max_rows: int) -> None: if max_rows < 0: raise ValueError("max_rows must be >= 0") - self.source = source self.max_rows = int(max_rows) + self.source = source.with_read_batch_rows(max(1, self.max_rows)) self.name = source.name self._source_shards: tuple[Shard, ...] | None = None + def with_read_batch_rows(self, max_rows: int | None) -> "LimitedSource": + read_batch_rows = self.max_rows + if max_rows is not None: + read_batch_rows = min(read_batch_rows, max_rows) + source = copy(self) + source.source = self.source.with_read_batch_rows( + max(1, read_batch_rows), + ) + return source + @property def schema(self) -> pa.Schema | None: return self.source.schema diff --git a/src/refiner/pipeline/sources/readers/parquet.py b/src/refiner/pipeline/sources/readers/parquet.py index 67908975d..c98a71bb5 100644 --- a/src/refiner/pipeline/sources/readers/parquet.py +++ b/src/refiner/pipeline/sources/readers/parquet.py @@ -2,6 +2,7 @@ import bisect from collections.abc import Iterator, Mapping, Sequence +from copy import copy from dataclasses import dataclass from typing import Any, Optional @@ -96,7 +97,8 @@ def __init__( file_path_column=file_path_column, dtypes=dtypes, ) - self.arrow_batch_size = int(arrow_batch_size) + self._default_arrow_batch_size = int(arrow_batch_size) + self.arrow_batch_size = self._default_arrow_batch_size self.split_row_groups = split_row_groups ## filter @@ -129,6 +131,7 @@ def __init__( "columns_to_read cannot include the synthetic file_path_column; " "omit it from columns_to_read and let the reader append it" ) + # the columns we will actually load into memory. requested+needed for filtering self._scan_columns: list[str] | None = ( None @@ -141,6 +144,13 @@ def __init__( self._open_fragment_file: DataFile | None = None self._open_fragment: ds.ParquetFileFragment | None = None + def with_read_batch_rows(self, max_rows: int | None) -> "ParquetReader": + reader = copy(self) + reader.arrow_batch_size = ( + self._default_arrow_batch_size if max_rows is None else max_rows + ) + return reader + def _get_parquet_file(self, source_file: DataFile) -> pq.ParquetFile: """Get or open a cached ParquetFile for the current path (single-open-file policy).""" fh, opened_new = self._get_file_handle(source_file, mode="rb") @@ -172,6 +182,7 @@ def describe(self) -> dict[str, Any]: "columns_to_read": list(self.columns_to_read) if self.columns_to_read is not None else None, + "read_batch_rows": self.arrow_batch_size, "split_row_groups": self.split_row_groups, "filter": self.filter.to_code() if self.filter is not None else None, "dtypes": list(self.dtypes) if self.dtypes else None, diff --git a/tests/pipeline/test_sinks.py b/tests/pipeline/test_sinks.py index d6ca1ab3a..789f3e18b 100644 --- a/tests/pipeline/test_sinks.py +++ b/tests/pipeline/test_sinks.py @@ -2,6 +2,7 @@ import json import threading +from concurrent.futures import ThreadPoolExecutor from collections.abc import Mapping from typing import cast @@ -11,6 +12,7 @@ import pytest from fsspec.implementations.memory import MemoryFileSystem +import refiner.pipeline.sinks.lance as lance_sink_module from refiner import AddColumns, Append, Create, Overwrite, col, read_blob from refiner.io import DataFolder from refiner.pipeline.data import datatype @@ -28,7 +30,9 @@ from refiner.pipeline.sinks.lance import ( LanceDatasetCommitReducerSink, LanceDatasetSink, + LanceIOConfig, _StreamingAddColumnsWriter, + _StreamingShardWriter, _cast_to_planned_schema, _job_token, _schema_to_base64, @@ -59,6 +63,170 @@ def test_lance_write_mode_dataclasses_normalize(mode, expected, tmp_path) -> Non assert LanceDatasetSink(tmp_path / "output.lance", mode=mode).mode == expected +def test_lance_dataset_sink_uses_byte_aware_batching_defaults(tmp_path) -> None: + sink = LanceDatasetSink(tmp_path / "output.lance") + + assert sink.io == LanceIOConfig() + assert sink.describe()[2]["io"] == { + "upload_concurrency": 32, + "multipart_part_bytes": 5 * 1024 * 1024, + "target_batch_bytes": 256 * 1024 * 1024, + "max_buffered_bytes": 1024 * 1024 * 1024, + } + + +def test_lance_io_config_sets_process_upload_controls_once(monkeypatch) -> None: + monkeypatch.setattr(lance_sink_module, "_LANCE_PROCESS_UPLOAD_CONFIG", None) + monkeypatch.delenv("LANCE_UPLOAD_CONCURRENCY", raising=False) + monkeypatch.delenv("LANCE_INITIAL_UPLOAD_SIZE", raising=False) + config = LanceIOConfig() + + lance_sink_module._configure_lance_process(config) + lance_sink_module._configure_lance_process(config) + + assert lance_sink_module._LANCE_PROCESS_UPLOAD_CONFIG == ( + config.upload_concurrency, + config.multipart_part_bytes, + ) + assert lance_sink_module.os.environ["LANCE_UPLOAD_CONCURRENCY"] == "32" + assert lance_sink_module.os.environ["LANCE_INITIAL_UPLOAD_SIZE"] == str( + 5 * 1024 * 1024 + ) + with pytest.raises(RuntimeError, match="start a fresh worker"): + lance_sink_module._configure_lance_process( + LanceIOConfig( + upload_concurrency=8, + multipart_part_bytes=16 * 1024 * 1024, + target_batch_bytes=128 * 1024 * 1024, + max_buffered_bytes=512 * 1024 * 1024, + ) + ) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"upload_concurrency": 0}, "upload_concurrency"), + ({"multipart_part_bytes": 1024}, "between 5 MiB and 5 GiB"), + ({"target_batch_bytes": 0}, "target_batch_bytes"), + ( + {"upload_concurrency": 2, "target_batch_bytes": 8 * 1024 * 1024}, + "upload_concurrency times", + ), + ({"max_buffered_bytes": 1024}, "four times"), + ], +) +def test_lance_io_config_validates_limits(kwargs, message) -> None: + with pytest.raises(ValueError, match=message): + LanceIOConfig(**kwargs) + + +def _capture_lance_writer_batches( + monkeypatch, + batches: list[pa.RecordBatch], + *, + target_batch_bytes: int = 5 * 1024 * 1024, +) -> tuple[list[pa.RecordBatch], int]: + captured: list[pa.RecordBatch] = [] + writer_pool = ThreadPoolExecutor(max_workers=1) + monkeypatch.setattr(lance_sink_module, "_LANCE_WRITER_POOL", writer_pool) + monkeypatch.setattr(lance_sink_module, "_LANCE_PROCESS_UPLOAD_CONFIG", None) + + def capture_batches( + _writer: _StreamingShardWriter, reader: pa.RecordBatchReader + ) -> list[str]: + captured.extend(reader) + return [] + + monkeypatch.setattr(_StreamingShardWriter, "_write_reader", capture_batches) + writer = _StreamingShardWriter( + dataset_uri="unused", + schema=batches[0].schema, + mode="create", + io_config=LanceIOConfig( + upload_concurrency=1, + multipart_part_bytes=5 * 1024 * 1024, + target_batch_bytes=target_batch_bytes, + max_buffered_bytes=4 * target_batch_bytes, + ), + ) + try: + writer.put_batches(batches) + queue_maxsize = writer.queue.maxsize + writer.finish() + return captured, queue_maxsize + finally: + writer_pool.shutdown(wait=True) + + +def test_lance_streaming_writer_coalesces_and_splits_batches_by_bytes( + monkeypatch, +) -> None: + schema = pa.schema([pa.field("payload", pa.binary())]) + source = [ + pa.record_batch([pa.array([bytes([index]) * 1024 * 1024])], schema=schema) + for index in range(12) + ] + + captured, queue_maxsize = _capture_lance_writer_batches(monkeypatch, source) + + assert queue_maxsize == 1 + assert [ + value[0] for batch in captured for value in batch["payload"].to_pylist() + ] == [*range(12)] + assert len(captured) == 3 + assert all(batch.nbytes <= 5 * 1024 * 1024 for batch in captured) + + captured, _ = _capture_lance_writer_batches( + monkeypatch, + [ + pa.record_batch( + [pa.array([bytes([index]) * 20 for index in range(8)])], + schema=schema, + ) + ], + ) + + assert len(captured) == 1 + assert all(batch.nbytes <= 5 * 1024 * 1024 for batch in captured) + assert [ + value[0] for batch in captured for value in batch["payload"].to_pylist() + ] == [*range(8)] + + captured, _ = _capture_lance_writer_batches( + monkeypatch, + [ + pa.record_batch( + [pa.array([bytes([index]) * 1024 * 1024 for index in range(8)])], + schema=schema, + ) + ], + ) + + assert len(captured) == 2 + assert all(batch.nbytes <= 5 * 1024 * 1024 for batch in captured) + assert [ + value[0] for batch in captured for value in batch["payload"].to_pylist() + ] == [*range(8)] + + +def test_lance_streaming_writer_writes_oversized_single_row_alone(monkeypatch) -> None: + schema = pa.schema([pa.field("payload", pa.binary())]) + captured, _ = _capture_lance_writer_batches( + monkeypatch, + [ + pa.record_batch( + [pa.array([b"x" * (6 * 1024 * 1024), b"small"])], + schema=schema, + ) + ], + ) + + assert [batch.num_rows for batch in captured] == [1, 1] + assert captured[0].nbytes > 5 * 1024 * 1024 + assert captured[1].nbytes <= 5 * 1024 * 1024 + + class _PartialMissingStream: def __init__(self) -> None: self._reads = 0 @@ -417,7 +585,6 @@ def test_lance_add_columns_packs_assets_as_generic_blob_references(tmp_path) -> dataset_uri, version=base.version, columns=["x"], - batch_size=1, ) .map( lambda row: {"image": bytes([int(row["x"])])}, @@ -525,7 +692,6 @@ def test_launch_local_adds_lance_columns_without_rewriting_base_files(tmp_path) dataset_uri, version=base.version, columns=["x"], - batch_size=1, num_shards=1, ) .map( @@ -574,7 +740,6 @@ def test_lance_add_columns_handles_more_fragments_than_io_threads(tmp_path) -> N load_lance( dataset_uri, version=base.version, - batch_size=1, num_shards=1, ) .map(lambda row: {"y": int(row["x"]) * 10}, dtypes={"y": datatype.int64()}) @@ -728,7 +893,7 @@ def reverse_batch(rows: list[Row]): yield row.update({"y": int(row["x"]) * 10}) pipeline = ( - load_lance(dataset_uri, version=base.version, columns=["x"], batch_size=3) + load_lance(dataset_uri, version=base.version, columns=["x"]) .batch_map( reverse_batch, batch_size=3, diff --git a/tests/readers/test_lance_source.py b/tests/readers/test_lance_source.py index 76cc20c2d..5546b07fa 100644 --- a/tests/readers/test_lance_source.py +++ b/tests/readers/test_lance_source.py @@ -55,7 +55,6 @@ def test_load_lance_pins_version_and_shards_by_fragment(tmp_path) -> None: dataset_uri, version=version_one, columns=["x"], - batch_size=1, ) shards = pipeline.list_shards() @@ -70,6 +69,60 @@ def test_load_lance_pins_version_and_shards_by_fragment(tmp_path) -> None: assert [int(row["x"]) for row in pipeline.iter_rows()] == [1, 2, 3] +def test_max_block_rows_bounds_lance_scanner_batches(tmp_path) -> None: + lance = pytest.importorskip("lance") + dataset_uri = tmp_path / "scanner-batches.lance" + lance.write_dataset(pa.table({"x": list(range(9))}), str(dataset_uri)) + + default_pipeline = load_lance(dataset_uri) + bounded_pipeline = default_pipeline.with_max_block_rows(2) + restored_pipeline = bounded_pipeline.with_max_block_rows(None) + overridden_pipeline = ( + load_lance(dataset_uri, read_batch_rows=4).with_max_block_rows(2).select("x") + ) + enlarged_pipeline = load_lance(dataset_uri, read_batch_rows=131_072) + + assert isinstance(default_pipeline.source, LanceSource) + assert isinstance(bounded_pipeline.source, LanceSource) + assert isinstance(restored_pipeline.source, LanceSource) + assert isinstance(overridden_pipeline.source, LanceSource) + assert isinstance(enlarged_pipeline.source, LanceSource) + assert default_pipeline.source._read_batch_rows == 65_536 + assert bounded_pipeline.source._read_batch_rows == 2 + assert restored_pipeline.source._read_batch_rows == 65_536 + assert overridden_pipeline.source._read_batch_rows == 4 + assert enlarged_pipeline.source._read_batch_rows == 131_072 + assert overridden_pipeline.read_batch_rows == 4 + assert overridden_pipeline.max_block_rows == 2 + + units = list(bounded_pipeline.source.read_shard(bounded_pipeline.list_shards()[0])) + tabular_units = [] + for unit in units: + assert isinstance(unit, Tabular) + tabular_units.append(unit) + assert [unit.num_rows for unit in tabular_units] == [2, 2, 2, 2, 1] + + overridden_units = list( + overridden_pipeline.source.read_shard(overridden_pipeline.list_shards()[0]) + ) + assert [ + unit.num_rows for unit in overridden_units if isinstance(unit, Tabular) + ] == [ + 4, + 4, + 1, + ] + + +def test_load_lance_rejects_invalid_read_batch_rows(tmp_path) -> None: + lance = pytest.importorskip("lance") + dataset_uri = tmp_path / "invalid-read-batch.lance" + lance.write_dataset(pa.table({"x": [1]}), str(dataset_uri)) + + with pytest.raises(ValueError, match="read_batch_rows must be > 0"): + load_lance(dataset_uri, read_batch_rows=0) + + def test_load_lance_max_rows_uses_limited_source_and_slices_final_batch( tmp_path, ) -> None: @@ -81,14 +134,14 @@ def test_load_lance_max_rows_uses_limited_source_and_slices_final_batch( max_rows_per_file=2, ) - pipeline = load_lance(dataset_uri, batch_size=2, max_rows=3) + pipeline = load_lance(dataset_uri, max_rows=3) shards = pipeline.list_shards() assert len(shards) == 1 assert isinstance(shards[0].descriptor, ShardGroupDescriptor) assert isinstance(pipeline.source, LimitedSource) assert isinstance(pipeline.source.source, LanceSource) - assert pipeline.source.source.batch_size == 2 + assert pipeline.source.source._read_batch_rows == 3 assert [int(row["x"]) for row in pipeline.iter_rows()] == [0, 1, 2] assert pipeline.source.describe()["max_rows"] == 3 @@ -155,7 +208,7 @@ def test_load_lance_normalizes_classic_blobs_to_references(tmp_path) -> None: max_rows_per_file=1, ) - pipeline = load_lance(dataset_uri, batch_size=1) + pipeline = load_lance(dataset_uri) output_schema = pipeline.output_schema() rows = list(pipeline.iter_rows()) @@ -221,7 +274,6 @@ def test_load_lance_limits_leading_rows_across_fragments(tmp_path) -> None: pipeline = load_lance( dataset_uri, - batch_size=2, num_shards=2, max_rows=5, ) @@ -266,7 +318,7 @@ def test_limited_lance_source_adds_columns_with_null_fill(tmp_path) -> None: ) ( - load_lance(dataset_uri, version=base.version, max_rows=4, batch_size=2) + load_lance(dataset_uri, version=base.version, max_rows=4) .map(lambda row: {"y": int(row["x"]) * 10}, dtypes={"y": datatype.int64()}) .write_lance_dataset( dataset_uri, @@ -295,7 +347,7 @@ def test_load_lance_uses_physical_row_addresses_as_source_row_ids(tmp_path) -> N max_rows_per_file=2, ) - pipeline = load_lance(dataset_uri, batch_size=1) + pipeline = load_lance(dataset_uri) addresses_by_shard = [] for shard in pipeline.list_shards(): addresses = [] diff --git a/tests/readers/test_limited_source.py b/tests/readers/test_limited_source.py index f0f111f63..7c66bc07e 100644 --- a/tests/readers/test_limited_source.py +++ b/tests/readers/test_limited_source.py @@ -221,6 +221,13 @@ def test_builtin_reader_functions_expose_max_rows(reader) -> None: assert "max_rows" in inspect.signature(reader).parameters +def test_lance_and_parquet_use_pipeline_block_sizing_api() -> None: + assert "batch_size" not in inspect.signature(load_lance).parameters + assert "arrow_batch_size" not in inspect.signature(read_parquet).parameters + assert "read_batch_rows" in inspect.signature(load_lance).parameters + assert "read_batch_rows" in inspect.signature(read_parquet).parameters + + def test_parquet_max_rows_is_global_and_applies_after_reader_filter( tmp_path: Path, ) -> None: diff --git a/tests/readers/test_parquet_reader.py b/tests/readers/test_parquet_reader.py index 7846e6adc..a1e7582bb 100644 --- a/tests/readers/test_parquet_reader.py +++ b/tests/readers/test_parquet_reader.py @@ -1,6 +1,8 @@ import pyarrow as pa import pyarrow.parquet as pq +import pytest +from refiner import read_parquet from refiner.pipeline.data.tabular import Tabular from refiner.pipeline.data import datatype from refiner.pipeline.expressions import col @@ -89,6 +91,59 @@ def _rows_from_shard_units(units): yield unit +def test_max_block_rows_bounds_parquet_scanner_batches(tmp_path): + p = _write_parquet(tmp_path) + + default_pipeline = read_parquet(p, file_path_column=None) + bounded_pipeline = default_pipeline.with_max_block_rows(3) + restored_pipeline = bounded_pipeline.with_max_block_rows(None) + overridden_pipeline = read_parquet( + p, + file_path_column=None, + read_batch_rows=7, + ).with_max_block_rows(3) + enlarged_pipeline = read_parquet( + p, + file_path_column=None, + read_batch_rows=131_072, + ) + + assert isinstance(default_pipeline.source, ParquetReader) + assert isinstance(bounded_pipeline.source, ParquetReader) + assert isinstance(restored_pipeline.source, ParquetReader) + assert isinstance(overridden_pipeline.source, ParquetReader) + assert isinstance(enlarged_pipeline.source, ParquetReader) + assert default_pipeline.source.arrow_batch_size == 65_536 + assert bounded_pipeline.source.arrow_batch_size == 3 + assert restored_pipeline.source.arrow_batch_size == 65_536 + assert overridden_pipeline.source.arrow_batch_size == 7 + assert enlarged_pipeline.source.arrow_batch_size == 131_072 + assert overridden_pipeline.read_batch_rows == 7 + assert overridden_pipeline.max_block_rows == 3 + + units = list(bounded_pipeline.source.read_shard(bounded_pipeline.list_shards()[0])) + tabular_units = [] + for unit in units: + assert isinstance(unit, Tabular) + assert unit.num_rows <= 3 + tabular_units.append(unit) + assert sum(unit.num_rows for unit in tabular_units) == 50 + + overridden_units = list( + overridden_pipeline.source.read_shard(overridden_pipeline.list_shards()[0]) + ) + assert all( + isinstance(unit, Tabular) and unit.num_rows <= 7 for unit in overridden_units + ) + + +def test_read_parquet_rejects_invalid_read_batch_rows(tmp_path): + p = _write_parquet(tmp_path) + + with pytest.raises(ValueError, match="read_batch_rows must be > 0"): + read_parquet(p, read_batch_rows=0) + + def test_parquet_reads_all_rows(tmp_path): p = _write_parquet(tmp_path) r = ParquetReader(str(p), target_shard_bytes=200) From 4b487b58a3cd17af93986eb4963ed6ec1fe3f33e Mon Sep 17 00:00:00 2001 From: Hynek Kydlicek Date: Wed, 2 Sep 2026 15:09:21 +0200 Subject: [PATCH 02/11] fix(lance): bound add-columns writer buffering Co-Authored-By: OpenAI Codex --- src/refiner/pipeline/sinks/lance.py | 89 +++++++++++++++++++---------- tests/pipeline/test_sinks.py | 57 ++++++++++++++++-- 2 files changed, 110 insertions(+), 36 deletions(-) diff --git a/src/refiner/pipeline/sinks/lance.py b/src/refiner/pipeline/sinks/lance.py index f02749331..a7cfbf195 100644 --- a/src/refiner/pipeline/sinks/lance.py +++ b/src/refiner/pipeline/sinks/lance.py @@ -182,6 +182,47 @@ def _split_batch_by_bytes( ] +class _ByteAwareBatcher: + """Coalesce ordered Arrow batches without exceeding a target byte size.""" + + def __init__(self, target_bytes: int) -> None: + self.target_bytes = target_bytes + self._pending: list[pa.RecordBatch] = [] + self._pending_bytes = 0 + + def add(self, batch: pa.RecordBatch) -> list[pa.RecordBatch]: + ready: list[pa.RecordBatch] = [] + for split_batch in _split_batch_by_bytes(batch, self.target_bytes): + if ( + self._pending + and self._pending_bytes + split_batch.nbytes > self.target_bytes + ): + ready.append(self.flush()) + self._pending.append(split_batch) + self._pending_bytes += split_batch.nbytes + if self._pending_bytes >= self.target_bytes: + ready.append(self.flush()) + return ready + + def flush(self) -> pa.RecordBatch: + if not self._pending: + raise RuntimeError("Cannot flush an empty Lance batch buffer") + if len(self._pending) == 1: + batch = self._pending[0] + else: + batch = pa.concat_batches(self._pending) + self._pending = [] + self._pending_bytes = 0 + return batch + + def finish(self) -> list[pa.RecordBatch]: + return [] if not self._pending else [self.flush()] + + +def _lance_queue_slots(config: LanceIOConfig) -> int: + return max(1, config.max_buffered_bytes // config.target_batch_bytes - 3) + + def _import_lance() -> Any: check_required_dependencies( "write_lance_dataset", [("lance", "pylance")], dist="lance" @@ -577,13 +618,9 @@ def __init__( self.schema = schema self.mode = mode self.io_config = io_config - self._pending_batches: list[pa.RecordBatch] = [] - self._pending_bytes = 0 - queue_slots = max( - 1, io_config.max_buffered_bytes // io_config.target_batch_bytes - 3 - ) + self._batcher = _ByteAwareBatcher(io_config.target_batch_bytes) self.queue: queue_module.Queue[pa.RecordBatch | object] = queue_module.Queue( - maxsize=queue_slots + maxsize=_lance_queue_slots(io_config) ) self.closed = False self.task_future = _LANCE_WRITER_POOL.submit(self._run) @@ -641,29 +678,10 @@ def put_batches(self, batches: list[pa.RecordBatch]) -> None: if self.closed: raise RuntimeError("Cannot write to a closed Lance shard writer.") for batch in batches: - for split_batch in _split_batch_by_bytes( - batch, self.io_config.target_batch_bytes - ): - if ( - self._pending_batches - and self._pending_bytes + split_batch.nbytes - > self.io_config.target_batch_bytes - ): - self._flush_pending() - self._pending_batches.append(split_batch) - self._pending_bytes += split_batch.nbytes - if self._pending_bytes >= self.io_config.target_batch_bytes: - self._flush_pending() - - def _flush_pending(self) -> None: - if not self._pending_batches: - return - if len(self._pending_batches) == 1: - batch = self._pending_batches[0] - else: - batch = pa.concat_batches(self._pending_batches) - self._pending_batches = [] - self._pending_bytes = 0 + for ready_batch in self._batcher.add(batch): + self._write_batch(ready_batch) + + def _write_batch(self, batch: pa.RecordBatch) -> None: if self._spool_writer is not None: self._spool_writer.write_batch(batch) else: @@ -672,7 +690,8 @@ def _flush_pending(self) -> None: def finish(self) -> list[str]: if not self.closed: self.closed = True - self._flush_pending() + for batch in self._batcher.finish(): + self._write_batch(batch) if self._spool_writer is None: self._put(_QUEUE_CLOSED) if self._spool_writer is not None: @@ -710,9 +729,11 @@ def __init__( self.output = output self.fill_missing = fill_missing self.fill = fill + self.io_config = io_config + self._batcher = _ByteAwareBatcher(io_config.target_batch_bytes) self.base_json = _json_dumps(fragment.metadata.to_json()) self.queue: queue_module.Queue[pa.RecordBatch | object] = queue_module.Queue( - maxsize=8 + maxsize=_lance_queue_slots(io_config) ) self.closed = False self.next_position = 0 @@ -753,6 +774,10 @@ def _raise_if_failed(self) -> None: raise RuntimeError("Lance add-columns writer failed") from error def _put_batch(self, batch: pa.RecordBatch) -> None: + for ready_batch in self._batcher.add(batch): + self._write_batch(ready_batch) + + def _write_batch(self, batch: pa.RecordBatch) -> None: if self._spool_writer is not None: self._spool_writer.write_batch(batch) return @@ -849,6 +874,8 @@ def _close_input(self) -> None: if self.closed: return self.closed = True + for batch in self._batcher.finish(): + self._write_batch(batch) if self._spool_writer is not None: assert self._spool_sink is not None self._spool_writer.close() diff --git a/tests/pipeline/test_sinks.py b/tests/pipeline/test_sinks.py index 789f3e18b..8c8986ffd 100644 --- a/tests/pipeline/test_sinks.py +++ b/tests/pipeline/test_sinks.py @@ -918,8 +918,7 @@ def reverse_batch(rows: list[Row]): } -def test_lance_add_columns_streams_contiguous_reordered_batches(tmp_path) -> None: - first_batch_consumed = threading.Event() +def test_lance_add_columns_orders_contiguous_reordered_batches(tmp_path) -> None: consumed: list[int] = [] class _Metadata: @@ -933,7 +932,6 @@ def merge_columns(self, reader, *, reader_schema): assert reader_schema == pa.schema([("y", pa.int64())]) for batch in reader: consumed.extend(batch.column("y").to_pylist()) - first_batch_consumed.set() return "updated", "schema" writer = _StreamingAddColumnsWriter( @@ -947,17 +945,66 @@ def merge_columns(self, reader, *, reader_schema): pa.chunked_array([[2, 3]], type=pa.uint64()), pa.table({"y": [30, 40]}), ) - assert not first_batch_consumed.is_set() + assert consumed == [] writer.put( pa.chunked_array([[1, 0]], type=pa.uint64()), pa.table({"y": [20, 10]}), ) - assert first_batch_consumed.wait(timeout=2) + assert consumed == [] assert writer.finish() == ("updated", "schema") assert consumed == [10, 20, 30, 40] +def test_lance_add_columns_honors_byte_aware_buffer_limits( + tmp_path, + monkeypatch, +) -> None: + monkeypatch.setattr(lance_sink_module, "_LANCE_PROCESS_UPLOAD_CONFIG", None) + captured: list[pa.RecordBatch] = [] + + class _Metadata: + def to_json(self): + return {"files": [], "physical_rows": 12} + + class _Fragment: + metadata = _Metadata() + + def merge_columns(self, reader, *, reader_schema): + assert reader_schema == pa.schema([("payload", pa.binary())]) + captured.extend(reader) + return "updated", "schema" + + target_bytes = 5 * 1024 * 1024 + writer = _StreamingAddColumnsWriter( + fragment=_Fragment(), + fragment_id=7, + num_rows=12, + schema=pa.schema([("payload", pa.binary())]), + output=LanceDatasetSink(tmp_path / "bounded-columns.lance").output, + io_config=LanceIOConfig( + upload_concurrency=1, + multipart_part_bytes=target_bytes, + target_batch_bytes=target_bytes, + max_buffered_bytes=4 * target_bytes, + ), + ) + + assert writer.queue.maxsize == 1 + for index in range(12): + writer.put( + pa.chunked_array([[index]], type=pa.uint64()), + pa.table({"payload": [bytes([index]) * 1024 * 1024]}), + ) + + assert writer.finish() == ("updated", "schema") + assert len(captured) == 3 + assert all(batch.nbytes <= target_bytes for batch in captured) + assert [ + value[0] for batch in captured for value in batch["payload"].to_pylist() + ] == [*range(12)] + + def test_lance_add_columns_preserves_existing_field_ids(tmp_path) -> None: lance = pytest.importorskip("lance") dataset_uri = tmp_path / "field-ids.lance" From d3572892907aab8459a8ac338d65e2fb22abb338 Mon Sep 17 00:00:00 2001 From: Hynek Kydlicek Date: Wed, 2 Sep 2026 15:32:34 +0200 Subject: [PATCH 03/11] fix(lance): apply byte-weighted writer backpressure Co-Authored-By: OpenAI Codex --- src/refiner/pipeline/sinks/lance.py | 85 +++++++++++++++++++++++------ tests/pipeline/test_sinks.py | 57 ++++++++++++++++++- 2 files changed, 123 insertions(+), 19 deletions(-) diff --git a/src/refiner/pipeline/sinks/lance.py b/src/refiner/pipeline/sinks/lance.py index a7cfbf195..4fee873f6 100644 --- a/src/refiner/pipeline/sinks/lance.py +++ b/src/refiner/pipeline/sinks/lance.py @@ -219,8 +219,69 @@ def finish(self) -> list[pa.RecordBatch]: return [] if not self._pending else [self.flush()] -def _lance_queue_slots(config: LanceIOConfig) -> int: - return max(1, config.max_buffered_bytes // config.target_batch_bytes - 3) +@dataclass(frozen=True, slots=True) +class _QueuedBatch: + batch: pa.RecordBatch + bytes: int + + +class _ByteBoundedBatchQueue: + """Keep queued and Lance-consumed batches within one byte budget.""" + + def __init__(self, config: LanceIOConfig) -> None: + self.max_bytes = config.max_buffered_bytes - 2 * config.target_batch_bytes + self.maxsize = max( + 1, + config.max_buffered_bytes // config.target_batch_bytes - 3, + ) + self._bytes_in_flight = 0 + self._condition = threading.Condition() + self._queue: queue_module.Queue[_QueuedBatch | object] = queue_module.Queue() + + def put( + self, + item: pa.RecordBatch | object, + *, + timeout: float | None = None, + ) -> None: + if item is _QUEUE_CLOSED: + self._queue.put(item) + return + assert isinstance(item, pa.RecordBatch) + item_bytes = item.nbytes + with self._condition: + fits = ( + self._bytes_in_flight == 0 + if item_bytes > self.max_bytes + else self._bytes_in_flight + item_bytes <= self.max_bytes + ) + if not fits: + self._condition.wait(timeout) + raise queue_module.Full + self._bytes_in_flight += item_bytes + self._queue.put(_QueuedBatch(item, item_bytes)) + + def iter_batches(self): + previous_bytes = 0 + try: + while True: + if previous_bytes: + self._release(previous_bytes) + previous_bytes = 0 + item = self._queue.get() + if item is _QUEUE_CLOSED: + return + assert isinstance(item, _QueuedBatch) + previous_bytes = item.bytes + yield item.batch + finally: + if previous_bytes: + self._release(previous_bytes) + + def _release(self, released_bytes: int) -> None: + with self._condition: + self._bytes_in_flight -= released_bytes + self._condition.notify_all() def _import_lance() -> Any: @@ -619,9 +680,7 @@ def __init__( self.mode = mode self.io_config = io_config self._batcher = _ByteAwareBatcher(io_config.target_batch_bytes) - self.queue: queue_module.Queue[pa.RecordBatch | object] = queue_module.Queue( - maxsize=_lance_queue_slots(io_config) - ) + self.queue = _ByteBoundedBatchQueue(io_config) self.closed = False self.task_future = _LANCE_WRITER_POOL.submit(self._run) self._spool_path: str | None = None @@ -635,11 +694,7 @@ def __init__( self._spool_writer = pa.ipc.new_stream(self._spool_sink, self.schema) def _iter_batches(self): - while True: - item = self.queue.get() - if item is _QUEUE_CLOSED: - return - yield item + yield from self.queue.iter_batches() def _write_reader(self, reader: pa.RecordBatchReader) -> list[str]: lance = _import_lance() @@ -732,9 +787,7 @@ def __init__( self.io_config = io_config self._batcher = _ByteAwareBatcher(io_config.target_batch_bytes) self.base_json = _json_dumps(fragment.metadata.to_json()) - self.queue: queue_module.Queue[pa.RecordBatch | object] = queue_module.Queue( - maxsize=_lance_queue_slots(io_config) - ) + self.queue = _ByteBoundedBatchQueue(io_config) self.closed = False self.next_position = 0 self.pending: dict[int, tuple[int, pa.Table]] = {} @@ -752,11 +805,7 @@ def __init__( self._spool_writer = pa.ipc.new_stream(self._spool_sink, self.schema) def _iter_batches(self): - while True: - item = self.queue.get() - if item is _QUEUE_CLOSED: - return - yield item + yield from self.queue.iter_batches() def _merge_reader(self, reader: pa.RecordBatchReader) -> tuple[Any, Any]: return self.fragment.merge_columns(reader, reader_schema=self.schema) diff --git a/tests/pipeline/test_sinks.py b/tests/pipeline/test_sinks.py index 8c8986ffd..f7d91a693 100644 --- a/tests/pipeline/test_sinks.py +++ b/tests/pipeline/test_sinks.py @@ -2,7 +2,7 @@ import json import threading -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import ThreadPoolExecutor, TimeoutError from collections.abc import Mapping from typing import cast @@ -227,6 +227,61 @@ def test_lance_streaming_writer_writes_oversized_single_row_alone(monkeypatch) - assert captured[1].nbytes <= 5 * 1024 * 1024 +def test_lance_streaming_writer_allows_only_one_oversized_batch_in_flight( + monkeypatch, +) -> None: + first_batch_consumed = threading.Event() + release_first_batch = threading.Event() + writer_pool = ThreadPoolExecutor(max_workers=1) + monkeypatch.setattr(lance_sink_module, "_LANCE_WRITER_POOL", writer_pool) + monkeypatch.setattr(lance_sink_module, "_LANCE_PROCESS_UPLOAD_CONFIG", None) + + def consume_batches( + _writer: _StreamingShardWriter, + reader: pa.RecordBatchReader, + ) -> list[str]: + iterator = iter(reader) + next(iterator) + first_batch_consumed.set() + assert release_first_batch.wait(timeout=2) + next(iterator) + with pytest.raises(StopIteration): + next(iterator) + return [] + + monkeypatch.setattr(_StreamingShardWriter, "_write_reader", consume_batches) + schema = pa.schema([pa.field("payload", pa.binary())]) + config = LanceIOConfig( + upload_concurrency=1, + multipart_part_bytes=5 * 1024 * 1024, + target_batch_bytes=5 * 1024 * 1024, + max_buffered_bytes=20 * 1024 * 1024, + ) + writer = _StreamingShardWriter( + dataset_uri="unused", + schema=schema, + mode="create", + io_config=config, + ) + oversized = pa.record_batch([pa.array([b"x" * (15 * 1024 * 1024)])], schema=schema) + + try: + writer.put_batches([oversized]) + assert first_batch_consumed.wait(timeout=2) + + with ThreadPoolExecutor(max_workers=1) as producer: + future = producer.submit(writer.put_batches, [oversized]) + with pytest.raises(TimeoutError): + future.result(timeout=0.1) + release_first_batch.set() + future.result(timeout=2) + + assert writer.finish() == [] + finally: + release_first_batch.set() + writer_pool.shutdown(wait=True) + + class _PartialMissingStream: def __init__(self) -> None: self._reads = 0 From 11da617254c04e0b34c585e93cec7d5233be184d Mon Sep 17 00:00:00 2001 From: Hynek Kydlicek Date: Wed, 2 Sep 2026 15:49:48 +0200 Subject: [PATCH 04/11] fix(lance): share add-columns shard buffer budget Co-Authored-By: OpenAI Codex --- src/refiner/pipeline/sinks/lance.py | 64 ++++++++++++++++++------- tests/pipeline/test_sinks.py | 72 ++++++++++++++++++++++++++--- 2 files changed, 112 insertions(+), 24 deletions(-) diff --git a/src/refiner/pipeline/sinks/lance.py b/src/refiner/pipeline/sinks/lance.py index 4fee873f6..f8a243af9 100644 --- a/src/refiner/pipeline/sinks/lance.py +++ b/src/refiner/pipeline/sinks/lance.py @@ -225,17 +225,47 @@ class _QueuedBatch: bytes: int +class _ByteBudget: + """A byte-weighted, oversized-item-aware budget shared by writer queues.""" + + def __init__(self, max_bytes: int) -> None: + self.max_bytes = max_bytes + self._bytes_in_flight = 0 + self._condition = threading.Condition() + + def acquire(self, item_bytes: int, timeout: float | None) -> None: + with self._condition: + fits = ( + self._bytes_in_flight == 0 + if item_bytes > self.max_bytes + else self._bytes_in_flight + item_bytes <= self.max_bytes + ) + if not fits: + self._condition.wait(timeout) + raise queue_module.Full + self._bytes_in_flight += item_bytes + + def release(self, released_bytes: int) -> None: + with self._condition: + self._bytes_in_flight -= released_bytes + self._condition.notify_all() + + class _ByteBoundedBatchQueue: """Keep queued and Lance-consumed batches within one byte budget.""" - def __init__(self, config: LanceIOConfig) -> None: + def __init__( + self, + config: LanceIOConfig, + *, + budget: _ByteBudget | None = None, + ) -> None: self.max_bytes = config.max_buffered_bytes - 2 * config.target_batch_bytes self.maxsize = max( 1, config.max_buffered_bytes // config.target_batch_bytes - 3, ) - self._bytes_in_flight = 0 - self._condition = threading.Condition() + self._budget = budget or _ByteBudget(self.max_bytes) self._queue: queue_module.Queue[_QueuedBatch | object] = queue_module.Queue() def put( @@ -249,16 +279,7 @@ def put( return assert isinstance(item, pa.RecordBatch) item_bytes = item.nbytes - with self._condition: - fits = ( - self._bytes_in_flight == 0 - if item_bytes > self.max_bytes - else self._bytes_in_flight + item_bytes <= self.max_bytes - ) - if not fits: - self._condition.wait(timeout) - raise queue_module.Full - self._bytes_in_flight += item_bytes + self._budget.acquire(item_bytes, timeout) self._queue.put(_QueuedBatch(item, item_bytes)) def iter_batches(self): @@ -279,9 +300,7 @@ def iter_batches(self): self._release(previous_bytes) def _release(self, released_bytes: int) -> None: - with self._condition: - self._bytes_in_flight -= released_bytes - self._condition.notify_all() + self._budget.release(released_bytes) def _import_lance() -> Any: @@ -775,6 +794,7 @@ def __init__( fill_missing: bool = False, fill: object | Mapping[str, object] = None, io_config: LanceIOConfig = LanceIOConfig(), + byte_budget: _ByteBudget | None = None, ) -> None: _configure_lance_process(io_config) self.fragment = fragment @@ -787,7 +807,7 @@ def __init__( self.io_config = io_config self._batcher = _ByteAwareBatcher(io_config.target_batch_bytes) self.base_json = _json_dumps(fragment.metadata.to_json()) - self.queue = _ByteBoundedBatchQueue(io_config) + self.queue = _ByteBoundedBatchQueue(io_config, budget=byte_budget) self.closed = False self.next_position = 0 self.pending: dict[int, tuple[int, pa.Table]] = {} @@ -918,6 +938,8 @@ def put(self, positions: pa.ChunkedArray, output: pa.Table) -> None: ) run_start = index self._emit_ready() + for batch in self._batcher.finish(): + self._write_batch(batch) def _close_input(self) -> None: if self.closed: @@ -1066,6 +1088,7 @@ def __init__( self._add_columns_writers_by_shard: dict[ str, dict[int, _StreamingAddColumnsWriter] ] = {} + self._add_columns_budget_by_shard: dict[str, _ByteBudget] = {} self._add_columns_schema: pa.Schema | None = None self._planned_output_schema: pa.Schema | None = None self._declared_output_columns: frozenset[str] = frozenset() @@ -1201,6 +1224,10 @@ def _write_add_columns_block(self, shard_id: str, block: Block) -> None: [row_addresses, pa.scalar(_LANCE_ROW_ADDRESS_FRAGMENT_SHIFT, pa.uint64())], ) writers = self._add_columns_writers_by_shard.setdefault(shard_id, {}) + byte_budget = self._add_columns_budget_by_shard.setdefault( + shard_id, + _ByteBudget(self.io.max_buffered_bytes - 2 * self.io.target_batch_bytes), + ) unique_fragment_ids = pc.call_function("unique", [fragment_ids]) for fragment_id_raw in unique_fragment_ids.to_pylist(): fragment_id = int(fragment_id_raw) @@ -1225,6 +1252,7 @@ def _write_add_columns_block(self, shard_id: str, block: Block) -> None: fill_missing=self.fill_missing, fill=self.fill, io_config=self.io, + byte_budget=byte_budget, ) writers[fragment_id] = writer positions = pc.call_function( @@ -1348,6 +1376,7 @@ def on_shard_complete(self, shard_id: str) -> None: def _complete_add_columns_shard(self, shard_id: str) -> None: writers = self._add_columns_writers_by_shard.pop(shard_id, None) + self._add_columns_budget_by_shard.pop(shard_id, None) if not writers: self._write_empty_sidecar(shard_id) return @@ -1425,6 +1454,7 @@ def close(self) -> None: if first_error is None: first_error = err self._add_columns_writers_by_shard.clear() + self._add_columns_budget_by_shard.clear() if first_error is not None: raise first_error diff --git a/tests/pipeline/test_sinks.py b/tests/pipeline/test_sinks.py index f7d91a693..12c7295a7 100644 --- a/tests/pipeline/test_sinks.py +++ b/tests/pipeline/test_sinks.py @@ -974,6 +974,7 @@ def reverse_batch(rows: list[Row]): def test_lance_add_columns_orders_contiguous_reordered_batches(tmp_path) -> None: + first_batch_consumed = threading.Event() consumed: list[int] = [] class _Metadata: @@ -987,6 +988,7 @@ def merge_columns(self, reader, *, reader_schema): assert reader_schema == pa.schema([("y", pa.int64())]) for batch in reader: consumed.extend(batch.column("y").to_pylist()) + first_batch_consumed.set() return "updated", "schema" writer = _StreamingAddColumnsWriter( @@ -1006,7 +1008,7 @@ def merge_columns(self, reader, *, reader_schema): pa.chunked_array([[1, 0]], type=pa.uint64()), pa.table({"y": [20, 10]}), ) - assert consumed == [] + assert first_batch_consumed.wait(timeout=2) assert writer.finish() == ("updated", "schema") assert consumed == [10, 20, 30, 40] @@ -1046,20 +1048,76 @@ def merge_columns(self, reader, *, reader_schema): ) assert writer.queue.maxsize == 1 - for index in range(12): - writer.put( - pa.chunked_array([[index]], type=pa.uint64()), - pa.table({"payload": [bytes([index]) * 1024 * 1024]}), - ) + writer.put( + pa.chunked_array([list(range(12))], type=pa.uint64()), + pa.table( + { + "payload": pa.chunked_array( + [pa.array([bytes([index]) * 1024 * 1024]) for index in range(12)] + ) + } + ), + ) assert writer.finish() == ("updated", "schema") - assert len(captured) == 3 + assert len(captured) == 4 assert all(batch.nbytes <= target_bytes for batch in captured) assert [ value[0] for batch in captured for value in batch["payload"].to_pylist() ] == [*range(12)] +def test_lance_add_columns_shares_buffer_budget_across_fragment_writers( + tmp_path, + monkeypatch, +) -> None: + budgets = [] + + class _Fragment: + num_deletions = 0 + physical_rows = 1 + + class _Dataset: + def get_fragment(self, _fragment_id): + return _Fragment() + + class _Writer: + def __init__(self, **kwargs): + budgets.append(kwargs["byte_budget"]) + + def put(self, _positions, _output): + return None + + monkeypatch.setattr(lance_sink_module, "_StreamingAddColumnsWriter", _Writer) + dataset_uri = tmp_path / "shared-budget.lance" + sink = LanceDatasetSink( + dataset_uri, + mode="add_columns", + columns=["y"], + source_uri=str(dataset_uri), + source_version=1, + ) + sink._source_dataset_cache = _Dataset() + + sink.write_shard_block( + "shard", + Tabular( + pa.table( + { + "y": [10, 20], + SOURCE_ROW_ID_COLUMN: pa.array( + [0, 1 << 32], + type=pa.uint64(), + ), + } + ) + ), + ) + + assert len(budgets) == 2 + assert budgets[0] is budgets[1] + + def test_lance_add_columns_preserves_existing_field_ids(tmp_path) -> None: lance = pytest.importorskip("lance") dataset_uri = tmp_path / "field-ids.lance" From e53e4dbb8fd3d76f38d3509331ff923f7352dcc4 Mon Sep 17 00:00:00 2001 From: Hynek Kydlicek Date: Wed, 2 Sep 2026 16:02:56 +0200 Subject: [PATCH 05/11] fix(lance): coordinate add-columns shard batching Co-Authored-By: OpenAI Codex --- src/refiner/pipeline/sinks/lance.py | 72 ++++++++++++++++++++++++----- tests/pipeline/test_sinks.py | 23 ++++----- 2 files changed, 68 insertions(+), 27 deletions(-) diff --git a/src/refiner/pipeline/sinks/lance.py b/src/refiner/pipeline/sinks/lance.py index f8a243af9..94cd3ddd1 100644 --- a/src/refiner/pipeline/sinks/lance.py +++ b/src/refiner/pipeline/sinks/lance.py @@ -218,6 +218,10 @@ def flush(self) -> pa.RecordBatch: def finish(self) -> list[pa.RecordBatch]: return [] if not self._pending else [self.flush()] + @property + def pending_bytes(self) -> int: + return self._pending_bytes + @dataclass(frozen=True, slots=True) class _QueuedBatch: @@ -303,6 +307,39 @@ def _release(self, released_bytes: int) -> None: self._budget.release(released_bytes) +class _AddColumnsBatchCoordinator: + """Bound pending batches and transport bytes across one add-columns shard.""" + + def __init__(self, config: LanceIOConfig) -> None: + self.target_bytes = config.target_batch_bytes + self.byte_budget = _ByteBudget( + config.max_buffered_bytes - 2 * config.target_batch_bytes + ) + self._writers: list[_StreamingAddColumnsWriter] = [] + + def register(self, writer: _StreamingAddColumnsWriter) -> None: + self._writers.append(writer) + + def prepare(self, active_writer: _StreamingAddColumnsWriter, incoming: int) -> None: + while self._pending_bytes() + incoming > self.target_bytes: + candidate = next( + ( + writer + for writer in self._writers + if writer is not active_writer and writer._batcher.pending_bytes > 0 + ), + None, + ) + if candidate is None and active_writer._batcher.pending_bytes > 0: + candidate = active_writer + if candidate is None: + return + candidate._flush_buffered_batches() + + def _pending_bytes(self) -> int: + return sum(writer._batcher.pending_bytes for writer in self._writers) + + def _import_lance() -> Any: check_required_dependencies( "write_lance_dataset", [("lance", "pylance")], dist="lance" @@ -794,7 +831,7 @@ def __init__( fill_missing: bool = False, fill: object | Mapping[str, object] = None, io_config: LanceIOConfig = LanceIOConfig(), - byte_budget: _ByteBudget | None = None, + batch_coordinator: _AddColumnsBatchCoordinator | None = None, ) -> None: _configure_lance_process(io_config) self.fragment = fragment @@ -806,8 +843,15 @@ def __init__( self.fill = fill self.io_config = io_config self._batcher = _ByteAwareBatcher(io_config.target_batch_bytes) + self._batch_coordinator = batch_coordinator or _AddColumnsBatchCoordinator( + io_config + ) + self._batch_coordinator.register(self) self.base_json = _json_dumps(fragment.metadata.to_json()) - self.queue = _ByteBoundedBatchQueue(io_config, budget=byte_budget) + self.queue = _ByteBoundedBatchQueue( + io_config, + budget=self._batch_coordinator.byte_budget, + ) self.closed = False self.next_position = 0 self.pending: dict[int, tuple[int, pa.Table]] = {} @@ -843,9 +887,14 @@ def _raise_if_failed(self) -> None: raise RuntimeError("Lance add-columns writer failed") from error def _put_batch(self, batch: pa.RecordBatch) -> None: + self._batch_coordinator.prepare(self, batch.nbytes) for ready_batch in self._batcher.add(batch): self._write_batch(ready_batch) + def _flush_buffered_batches(self) -> None: + for batch in self._batcher.finish(): + self._write_batch(batch) + def _write_batch(self, batch: pa.RecordBatch) -> None: if self._spool_writer is not None: self._spool_writer.write_batch(batch) @@ -938,15 +987,12 @@ def put(self, positions: pa.ChunkedArray, output: pa.Table) -> None: ) run_start = index self._emit_ready() - for batch in self._batcher.finish(): - self._write_batch(batch) def _close_input(self) -> None: if self.closed: return self.closed = True - for batch in self._batcher.finish(): - self._write_batch(batch) + self._flush_buffered_batches() if self._spool_writer is not None: assert self._spool_sink is not None self._spool_writer.close() @@ -1088,7 +1134,9 @@ def __init__( self._add_columns_writers_by_shard: dict[ str, dict[int, _StreamingAddColumnsWriter] ] = {} - self._add_columns_budget_by_shard: dict[str, _ByteBudget] = {} + self._add_columns_coordinator_by_shard: dict[ + str, _AddColumnsBatchCoordinator + ] = {} self._add_columns_schema: pa.Schema | None = None self._planned_output_schema: pa.Schema | None = None self._declared_output_columns: frozenset[str] = frozenset() @@ -1224,9 +1272,9 @@ def _write_add_columns_block(self, shard_id: str, block: Block) -> None: [row_addresses, pa.scalar(_LANCE_ROW_ADDRESS_FRAGMENT_SHIFT, pa.uint64())], ) writers = self._add_columns_writers_by_shard.setdefault(shard_id, {}) - byte_budget = self._add_columns_budget_by_shard.setdefault( + batch_coordinator = self._add_columns_coordinator_by_shard.setdefault( shard_id, - _ByteBudget(self.io.max_buffered_bytes - 2 * self.io.target_batch_bytes), + _AddColumnsBatchCoordinator(self.io), ) unique_fragment_ids = pc.call_function("unique", [fragment_ids]) for fragment_id_raw in unique_fragment_ids.to_pylist(): @@ -1252,7 +1300,7 @@ def _write_add_columns_block(self, shard_id: str, block: Block) -> None: fill_missing=self.fill_missing, fill=self.fill, io_config=self.io, - byte_budget=byte_budget, + batch_coordinator=batch_coordinator, ) writers[fragment_id] = writer positions = pc.call_function( @@ -1376,7 +1424,7 @@ def on_shard_complete(self, shard_id: str) -> None: def _complete_add_columns_shard(self, shard_id: str) -> None: writers = self._add_columns_writers_by_shard.pop(shard_id, None) - self._add_columns_budget_by_shard.pop(shard_id, None) + self._add_columns_coordinator_by_shard.pop(shard_id, None) if not writers: self._write_empty_sidecar(shard_id) return @@ -1454,7 +1502,7 @@ def close(self) -> None: if first_error is None: first_error = err self._add_columns_writers_by_shard.clear() - self._add_columns_budget_by_shard.clear() + self._add_columns_coordinator_by_shard.clear() if first_error is not None: raise first_error diff --git a/tests/pipeline/test_sinks.py b/tests/pipeline/test_sinks.py index 12c7295a7..0070abac1 100644 --- a/tests/pipeline/test_sinks.py +++ b/tests/pipeline/test_sinks.py @@ -974,7 +974,6 @@ def reverse_batch(rows: list[Row]): def test_lance_add_columns_orders_contiguous_reordered_batches(tmp_path) -> None: - first_batch_consumed = threading.Event() consumed: list[int] = [] class _Metadata: @@ -988,7 +987,6 @@ def merge_columns(self, reader, *, reader_schema): assert reader_schema == pa.schema([("y", pa.int64())]) for batch in reader: consumed.extend(batch.column("y").to_pylist()) - first_batch_consumed.set() return "updated", "schema" writer = _StreamingAddColumnsWriter( @@ -1008,7 +1006,7 @@ def merge_columns(self, reader, *, reader_schema): pa.chunked_array([[1, 0]], type=pa.uint64()), pa.table({"y": [20, 10]}), ) - assert first_batch_consumed.wait(timeout=2) + assert consumed == [] assert writer.finish() == ("updated", "schema") assert consumed == [10, 20, 30, 40] @@ -1048,19 +1046,14 @@ def merge_columns(self, reader, *, reader_schema): ) assert writer.queue.maxsize == 1 - writer.put( - pa.chunked_array([list(range(12))], type=pa.uint64()), - pa.table( - { - "payload": pa.chunked_array( - [pa.array([bytes([index]) * 1024 * 1024]) for index in range(12)] - ) - } - ), - ) + for index in range(12): + writer.put( + pa.chunked_array([[index]], type=pa.uint64()), + pa.table({"payload": [bytes([index]) * 1024 * 1024]}), + ) assert writer.finish() == ("updated", "schema") - assert len(captured) == 4 + assert len(captured) == 3 assert all(batch.nbytes <= target_bytes for batch in captured) assert [ value[0] for batch in captured for value in batch["payload"].to_pylist() @@ -1083,7 +1076,7 @@ def get_fragment(self, _fragment_id): class _Writer: def __init__(self, **kwargs): - budgets.append(kwargs["byte_budget"]) + budgets.append(kwargs["batch_coordinator"]) def put(self, _positions, _output): return None From 8bf8a0deccfb2e6a9eee9a53f640cac3db8346d0 Mon Sep 17 00:00:00 2001 From: Hynek Kydlicek Date: Wed, 2 Sep 2026 16:18:46 +0200 Subject: [PATCH 06/11] fix(io): isolate reader caches and unordered buffers Co-Authored-By: OpenAI Codex --- src/refiner/pipeline/sinks/lance.py | 55 ++++++++++++++++--- .../pipeline/sources/readers/parquet.py | 6 ++ tests/pipeline/test_sinks.py | 4 ++ tests/readers/test_parquet_reader.py | 31 +++++++++++ 4 files changed, 87 insertions(+), 9 deletions(-) diff --git a/src/refiner/pipeline/sinks/lance.py b/src/refiner/pipeline/sinks/lance.py index 94cd3ddd1..0a2ced8b3 100644 --- a/src/refiner/pipeline/sinks/lance.py +++ b/src/refiner/pipeline/sinks/lance.py @@ -854,7 +854,7 @@ def __init__( ) self.closed = False self.next_position = 0 - self.pending: dict[int, tuple[int, pa.Table]] = {} + self.pending: dict[int, tuple[int, str]] = {} self.task_future = _LANCE_WRITER_POOL.submit(self._run) self._spool_path: str | None = None self._spool_sink: pa.OSFile | None = None @@ -909,11 +909,40 @@ def _write_batch(self, batch: pa.RecordBatch) -> None: def _emit_ready(self) -> None: while self.next_position in self.pending: - end, table = self.pending.pop(self.next_position) - for batch in table.to_batches(): - self._put_batch(batch) + end, path = self.pending.pop(self.next_position) + try: + with pa.memory_map(path, "r") as source: + table = pa.ipc.open_stream(source).read_all() + for batch in table.to_batches(): + self._put_batch(batch) + finally: + os.unlink(path) self.next_position = end + def _spool_pending_table(self, table: pa.Table) -> str: + tmp = tempfile.NamedTemporaryFile( + prefix="refiner-lance-columns-pending-", + delete=False, + ) + path = tmp.name + tmp.close() + try: + with pa.OSFile(path, "wb") as sink: + with pa.ipc.new_stream(sink, table.schema) as writer: + writer.write_table(table) + except Exception: + os.unlink(path) + raise + return path + + def _cleanup_pending_spools(self) -> None: + for _, path in self.pending.values(): + try: + os.unlink(path) + except FileNotFoundError: + pass + self.pending.clear() + def _emit_fill(self, end: int) -> None: while self.next_position < end: batch_end = min(end, self.next_position + _ADD_COLUMNS_FILL_BATCH_ROWS) @@ -981,12 +1010,18 @@ def put(self, positions: pa.ChunkedArray, output: pa.Table) -> None: f"Lance fragment {self.fragment_id} has invalid or duplicate " "row positions" ) - self.pending[start] = ( - end, - sorted_output.slice(run_start, index - run_start), - ) + run_table = sorted_output.slice(run_start, index - run_start) + if start == self.next_position: + for batch in run_table.to_batches(): + self._put_batch(batch) + self.next_position = end + self._emit_ready() + else: + self.pending[start] = ( + end, + self._spool_pending_table(run_table), + ) run_start = index - self._emit_ready() def _close_input(self) -> None: if self.closed: @@ -1027,6 +1062,7 @@ def finish(self) -> tuple[Any, Any]: complete = self.next_position == self.num_rows and not self.pending self._close_input() if not complete: + self._cleanup_pending_spools() if self._spool_path is None: try: result = self.task_future.result() @@ -1049,6 +1085,7 @@ def finish(self) -> tuple[Any, Any]: def abort(self) -> None: self._close_input() + self._cleanup_pending_spools() if self._spool_path is not None: os.unlink(self._spool_path) return diff --git a/src/refiner/pipeline/sources/readers/parquet.py b/src/refiner/pipeline/sources/readers/parquet.py index c98a71bb5..54151ac4f 100644 --- a/src/refiner/pipeline/sources/readers/parquet.py +++ b/src/refiner/pipeline/sources/readers/parquet.py @@ -149,6 +149,12 @@ def with_read_batch_rows(self, max_rows: int | None) -> "ParquetReader": reader.arrow_batch_size = ( self._default_arrow_batch_size if max_rows is None else max_rows ) + reader._open_file = None + reader._open_fh = None + reader._open_pf = None + reader._open_metadata = None + reader._open_fragment_file = None + reader._open_fragment = None return reader def _get_parquet_file(self, source_file: DataFile) -> pq.ParquetFile: diff --git a/tests/pipeline/test_sinks.py b/tests/pipeline/test_sinks.py index 0070abac1..24e4cf1e5 100644 --- a/tests/pipeline/test_sinks.py +++ b/tests/pipeline/test_sinks.py @@ -1000,6 +1000,8 @@ def merge_columns(self, reader, *, reader_schema): pa.chunked_array([[2, 3]], type=pa.uint64()), pa.table({"y": [30, 40]}), ) + pending_path = writer.pending[2][1] + assert lance_sink_module.os.path.exists(pending_path) assert consumed == [] writer.put( @@ -1007,6 +1009,8 @@ def merge_columns(self, reader, *, reader_schema): pa.table({"y": [20, 10]}), ) assert consumed == [] + assert writer.pending == {} + assert not lance_sink_module.os.path.exists(pending_path) assert writer.finish() == ("updated", "schema") assert consumed == [10, 20, 30, 40] diff --git a/tests/readers/test_parquet_reader.py b/tests/readers/test_parquet_reader.py index a1e7582bb..206a41046 100644 --- a/tests/readers/test_parquet_reader.py +++ b/tests/readers/test_parquet_reader.py @@ -137,6 +137,37 @@ def test_max_block_rows_bounds_parquet_scanner_batches(tmp_path): ) +def test_parquet_reader_copy_does_not_share_open_file_caches(tmp_path): + first_dir = tmp_path / "first" + second_dir = tmp_path / "second" + first_dir.mkdir() + second_dir.mkdir() + first_path = _write_parquet(first_dir) + second_path = _write_parquet(second_dir) + pipeline = read_parquet( + [first_path, second_path], + file_path_column=None, + target_shard_bytes=1, + ) + source = pipeline.source + assert isinstance(source, ParquetReader) + shards = source.list_shards() + + list(source.read_shard(shards[0])) + original_handle = source._open_fh + assert original_handle is not None + + derived = pipeline.with_max_block_rows(3) + assert isinstance(derived.source, ParquetReader) + assert derived.source._open_fh is None + assert derived.source._open_pf is None + assert derived.source._open_fragment is None + + list(derived.source.read_shard(shards[-1])) + assert not original_handle.closed + assert list(source.read_shard(shards[0])) + + def test_read_parquet_rejects_invalid_read_batch_rows(tmp_path): p = _write_parquet(tmp_path) From c5d08b83c904e0be3cbd22de76ec9b8a07803f5d Mon Sep 17 00:00:00 2001 From: Hynek Kydlicek Date: Wed, 2 Sep 2026 16:31:16 +0200 Subject: [PATCH 07/11] fix(lance): release failed writer queue budgets Co-Authored-By: OpenAI Codex --- src/refiner/pipeline/sinks/lance.py | 26 ++++++++++++++++++++++---- tests/pipeline/test_sinks.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/src/refiner/pipeline/sinks/lance.py b/src/refiner/pipeline/sinks/lance.py index 0a2ced8b3..d21e558ef 100644 --- a/src/refiner/pipeline/sinks/lance.py +++ b/src/refiner/pipeline/sinks/lance.py @@ -270,6 +270,8 @@ def __init__( config.max_buffered_bytes // config.target_batch_bytes - 3, ) self._budget = budget or _ByteBudget(self.max_bytes) + self._closed = False + self._state_lock = threading.Lock() self._queue: queue_module.Queue[_QueuedBatch | object] = queue_module.Queue() def put( @@ -279,12 +281,19 @@ def put( timeout: float | None = None, ) -> None: if item is _QUEUE_CLOSED: - self._queue.put(item) + with self._state_lock: + if self._closed: + raise queue_module.Full + self._queue.put(item) return assert isinstance(item, pa.RecordBatch) item_bytes = item.nbytes self._budget.acquire(item_bytes, timeout) - self._queue.put(_QueuedBatch(item, item_bytes)) + with self._state_lock: + if self._closed: + self._budget.release(item_bytes) + raise queue_module.Full + self._queue.put(_QueuedBatch(item, item_bytes)) def iter_batches(self): previous_bytes = 0 @@ -300,8 +309,17 @@ def iter_batches(self): previous_bytes = item.bytes yield item.batch finally: - if previous_bytes: - self._release(previous_bytes) + with self._state_lock: + self._closed = True + if previous_bytes: + self._release(previous_bytes) + while True: + try: + queued = self._queue.get_nowait() + except queue_module.Empty: + break + if isinstance(queued, _QueuedBatch): + self._release(queued.bytes) def _release(self, released_bytes: int) -> None: self._budget.release(released_bytes) diff --git a/tests/pipeline/test_sinks.py b/tests/pipeline/test_sinks.py index 24e4cf1e5..0548f7bb7 100644 --- a/tests/pipeline/test_sinks.py +++ b/tests/pipeline/test_sinks.py @@ -28,6 +28,8 @@ FileAssetConfig, ) from refiner.pipeline.sinks.lance import ( + _ByteBoundedBatchQueue, + _ByteBudget, LanceDatasetCommitReducerSink, LanceDatasetSink, LanceIOConfig, @@ -282,6 +284,32 @@ def consume_batches( writer_pool.shutdown(wait=True) +def test_lance_batch_queue_releases_queued_bytes_when_consumer_exits() -> None: + target_bytes = 5 * 1024 * 1024 + config = LanceIOConfig( + upload_concurrency=1, + multipart_part_bytes=target_bytes, + target_batch_bytes=target_bytes, + max_buffered_bytes=4 * target_bytes, + ) + budget = _ByteBudget(2 * target_bytes) + first_queue = _ByteBoundedBatchQueue(config, budget=budget) + second_queue = _ByteBoundedBatchQueue(config, budget=budget) + schema = pa.schema([pa.field("payload", pa.binary())]) + batch = pa.record_batch([pa.array([b"x" * (4 * 1024 * 1024)])], schema=schema) + + first_queue.put(batch, timeout=0) + first_queue.put(batch, timeout=0) + iterator = first_queue.iter_batches() + next(iterator) + iterator.close() + + second_queue.put(batch, timeout=0) + second_iterator = second_queue.iter_batches() + next(second_iterator) + second_iterator.close() + + class _PartialMissingStream: def __init__(self) -> None: self._reads = 0 From 3e6b640c927191b9a51239742ffeff18442c70c1 Mon Sep 17 00:00:00 2001 From: Hynek Kydlicek Date: Wed, 2 Sep 2026 16:39:24 +0200 Subject: [PATCH 08/11] fix(readers): clamp row limits to scanner windows Co-Authored-By: OpenAI Codex --- src/refiner/pipeline/sources/base.py | 4 ++++ src/refiner/pipeline/sources/lance.py | 5 +++++ src/refiner/pipeline/sources/limited.py | 9 +++------ src/refiner/pipeline/sources/readers/parquet.py | 5 +++++ tests/readers/test_lance_source.py | 12 ++++++++++++ tests/readers/test_parquet_reader.py | 11 +++++++++++ 6 files changed, 40 insertions(+), 6 deletions(-) diff --git a/src/refiner/pipeline/sources/base.py b/src/refiner/pipeline/sources/base.py index c5022d32c..c32baef39 100644 --- a/src/refiner/pipeline/sources/base.py +++ b/src/refiner/pipeline/sources/base.py @@ -55,6 +55,10 @@ def with_read_batch_rows(self, max_rows: int | None) -> "BaseSource": """ return self + def with_max_read_batch_rows(self, max_rows: int) -> "BaseSource": + """Return a source whose existing scanner window is no larger than a cap.""" + return self + @property def schema(self) -> pa.Schema | None: return None diff --git a/src/refiner/pipeline/sources/lance.py b/src/refiner/pipeline/sources/lance.py index ea79d453c..15aff42c9 100644 --- a/src/refiner/pipeline/sources/lance.py +++ b/src/refiner/pipeline/sources/lance.py @@ -143,6 +143,11 @@ def with_read_batch_rows(self, max_rows: int | None) -> "LanceSource": ) return source + def with_max_read_batch_rows(self, max_rows: int) -> "LanceSource": + source = copy(self) + source._read_batch_rows = min(source._read_batch_rows, max_rows) + return source + @property def schema(self) -> pa.Schema: return self._schema diff --git a/src/refiner/pipeline/sources/limited.py b/src/refiner/pipeline/sources/limited.py index 23c1450a8..ee67cd39c 100644 --- a/src/refiner/pipeline/sources/limited.py +++ b/src/refiner/pipeline/sources/limited.py @@ -27,18 +27,15 @@ def __init__(self, source: BaseSource, *, max_rows: int) -> None: if max_rows < 0: raise ValueError("max_rows must be >= 0") self.max_rows = int(max_rows) - self.source = source.with_read_batch_rows(max(1, self.max_rows)) + self.source = source.with_max_read_batch_rows(max(1, self.max_rows)) self.name = source.name self._source_shards: tuple[Shard, ...] | None = None def with_read_batch_rows(self, max_rows: int | None) -> "LimitedSource": - read_batch_rows = self.max_rows - if max_rows is not None: - read_batch_rows = min(read_batch_rows, max_rows) source = copy(self) source.source = self.source.with_read_batch_rows( - max(1, read_batch_rows), - ) + max_rows + ).with_max_read_batch_rows(max(1, self.max_rows)) return source @property diff --git a/src/refiner/pipeline/sources/readers/parquet.py b/src/refiner/pipeline/sources/readers/parquet.py index 54151ac4f..b32ce0f76 100644 --- a/src/refiner/pipeline/sources/readers/parquet.py +++ b/src/refiner/pipeline/sources/readers/parquet.py @@ -157,6 +157,11 @@ def with_read_batch_rows(self, max_rows: int | None) -> "ParquetReader": reader._open_fragment = None return reader + def with_max_read_batch_rows(self, max_rows: int) -> "ParquetReader": + reader = self.with_read_batch_rows(self.arrow_batch_size) + reader.arrow_batch_size = min(reader.arrow_batch_size, max_rows) + return reader + def _get_parquet_file(self, source_file: DataFile) -> pq.ParquetFile: """Get or open a cached ParquetFile for the current path (single-open-file policy).""" fh, opened_new = self._get_file_handle(source_file, mode="rb") diff --git a/tests/readers/test_lance_source.py b/tests/readers/test_lance_source.py index 5546b07fa..fe7052b84 100644 --- a/tests/readers/test_lance_source.py +++ b/tests/readers/test_lance_source.py @@ -146,6 +146,18 @@ def test_load_lance_max_rows_uses_limited_source_and_slices_final_batch( assert pipeline.source.describe()["max_rows"] == 3 +def test_load_lance_large_max_rows_preserves_default_scanner_batch(tmp_path) -> None: + lance = pytest.importorskip("lance") + dataset_uri = tmp_path / "large-limit.lance" + lance.write_dataset(pa.table({"x": [1]}), str(dataset_uri)) + + pipeline = load_lance(dataset_uri, max_rows=1_000_000) + + assert isinstance(pipeline.source, LimitedSource) + assert isinstance(pipeline.source.source, LanceSource) + assert pipeline.source.source._read_batch_rows == 65_536 + + def test_load_lance_max_rows_zero_and_negative(tmp_path) -> None: lance = pytest.importorskip("lance") dataset_uri = tmp_path / "zero.lance" diff --git a/tests/readers/test_parquet_reader.py b/tests/readers/test_parquet_reader.py index 206a41046..3414c3b28 100644 --- a/tests/readers/test_parquet_reader.py +++ b/tests/readers/test_parquet_reader.py @@ -6,6 +6,7 @@ from refiner.pipeline.data.tabular import Tabular from refiner.pipeline.data import datatype from refiner.pipeline.expressions import col +from refiner.pipeline.sources.limited import LimitedSource from refiner.pipeline.sources.readers import ParquetReader from refiner.worker.context import set_active_run_context from refiner.worker.metrics.emitter import UserMetricsEmitter @@ -168,6 +169,16 @@ def test_parquet_reader_copy_does_not_share_open_file_caches(tmp_path): assert list(source.read_shard(shards[0])) +def test_parquet_large_max_rows_preserves_default_scanner_batch(tmp_path): + path = _write_parquet(tmp_path) + + pipeline = read_parquet(path, max_rows=1_000_000) + + assert isinstance(pipeline.source, LimitedSource) + assert isinstance(pipeline.source.source, ParquetReader) + assert pipeline.source.source.arrow_batch_size == 65_536 + + def test_read_parquet_rejects_invalid_read_batch_rows(tmp_path): p = _write_parquet(tmp_path) From cdb4791610eb37b4d813afcbaba9f355e133a8fc Mon Sep 17 00:00:00 2001 From: Hynek Kydlicek Date: Wed, 2 Sep 2026 16:48:13 +0200 Subject: [PATCH 09/11] fix(lance): stream reordered spool replay Co-Authored-By: OpenAI Codex --- src/refiner/pipeline/sinks/lance.py | 5 ++--- tests/pipeline/test_sinks.py | 18 +++++++++++++++++- 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/src/refiner/pipeline/sinks/lance.py b/src/refiner/pipeline/sinks/lance.py index d21e558ef..74200d225 100644 --- a/src/refiner/pipeline/sinks/lance.py +++ b/src/refiner/pipeline/sinks/lance.py @@ -930,9 +930,8 @@ def _emit_ready(self) -> None: end, path = self.pending.pop(self.next_position) try: with pa.memory_map(path, "r") as source: - table = pa.ipc.open_stream(source).read_all() - for batch in table.to_batches(): - self._put_batch(batch) + for batch in pa.ipc.open_stream(source): + self._put_batch(batch) finally: os.unlink(path) self.next_position = end diff --git a/tests/pipeline/test_sinks.py b/tests/pipeline/test_sinks.py index 0548f7bb7..5f8e01ce4 100644 --- a/tests/pipeline/test_sinks.py +++ b/tests/pipeline/test_sinks.py @@ -1001,8 +1001,24 @@ def reverse_batch(rows: list[Row]): } -def test_lance_add_columns_orders_contiguous_reordered_batches(tmp_path) -> None: +def test_lance_add_columns_orders_contiguous_reordered_batches( + tmp_path, + monkeypatch, +) -> None: consumed: list[int] = [] + original_open_stream = pa.ipc.open_stream + + class _StreamingOnlyReader: + def __init__(self, source) -> None: + self._reader = original_open_stream(source) + + def __iter__(self): + yield from self._reader + + def read_all(self): + raise AssertionError("Pending Lance batches must be replayed incrementally") + + monkeypatch.setattr(pa.ipc, "open_stream", _StreamingOnlyReader) class _Metadata: def to_json(self): From a371e79848225ee3c61b1b2c45eb1bb3354fdc0d Mon Sep 17 00:00:00 2001 From: Hynek Kydlicek Date: Wed, 2 Sep 2026 17:00:02 +0200 Subject: [PATCH 10/11] fix(lance): stream byte-aware batch splitting Co-Authored-By: OpenAI Codex --- src/refiner/pipeline/sinks/lance.py | 25 +++++++++++-------------- tests/pipeline/test_sinks.py | 20 +++++++++++++++++++- 2 files changed, 30 insertions(+), 15 deletions(-) diff --git a/src/refiner/pipeline/sinks/lance.py b/src/refiner/pipeline/sinks/lance.py index 74200d225..402e724b1 100644 --- a/src/refiner/pipeline/sinks/lance.py +++ b/src/refiner/pipeline/sinks/lance.py @@ -10,7 +10,7 @@ import re import tempfile import threading -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass from typing import Any, Literal, get_args @@ -169,17 +169,16 @@ def _configure_lance_process(config: LanceIOConfig) -> None: def _split_batch_by_bytes( batch: pa.RecordBatch, target_bytes: int, -) -> list[pa.RecordBatch]: +) -> Iterator[pa.RecordBatch]: """Split a batch until each multi-row slice fits the target byte size.""" if batch.num_rows <= 1 or batch.nbytes <= target_bytes: - return [batch] + yield batch + return split_at = max(1, min(batch.num_rows - 1, batch.num_rows // 2)) - return [ - *_split_batch_by_bytes(batch.slice(0, split_at), target_bytes), - *_split_batch_by_bytes( - batch.slice(split_at, batch.num_rows - split_at), target_bytes - ), - ] + yield from _split_batch_by_bytes(batch.slice(0, split_at), target_bytes) + yield from _split_batch_by_bytes( + batch.slice(split_at, batch.num_rows - split_at), target_bytes + ) class _ByteAwareBatcher: @@ -190,19 +189,17 @@ def __init__(self, target_bytes: int) -> None: self._pending: list[pa.RecordBatch] = [] self._pending_bytes = 0 - def add(self, batch: pa.RecordBatch) -> list[pa.RecordBatch]: - ready: list[pa.RecordBatch] = [] + def add(self, batch: pa.RecordBatch) -> Iterator[pa.RecordBatch]: for split_batch in _split_batch_by_bytes(batch, self.target_bytes): if ( self._pending and self._pending_bytes + split_batch.nbytes > self.target_bytes ): - ready.append(self.flush()) + yield self.flush() self._pending.append(split_batch) self._pending_bytes += split_batch.nbytes if self._pending_bytes >= self.target_bytes: - ready.append(self.flush()) - return ready + yield self.flush() def flush(self) -> pa.RecordBatch: if not self._pending: diff --git a/tests/pipeline/test_sinks.py b/tests/pipeline/test_sinks.py index 5f8e01ce4..cd644444d 100644 --- a/tests/pipeline/test_sinks.py +++ b/tests/pipeline/test_sinks.py @@ -188,7 +188,6 @@ def test_lance_streaming_writer_coalesces_and_splits_batches_by_bytes( ) ], ) - assert len(captured) == 1 assert all(batch.nbytes <= 5 * 1024 * 1024 for batch in captured) assert [ @@ -212,6 +211,25 @@ def test_lance_streaming_writer_coalesces_and_splits_batches_by_bytes( ] == [*range(8)] +def test_lance_byte_batcher_consumes_split_batches_incrementally(monkeypatch) -> None: + schema = pa.schema([pa.field("payload", pa.binary())]) + source = pa.record_batch([pa.array([b"a", b"b", b"c"])], schema=schema) + consumed: list[int] = [] + + def tracked_splits(_batch, _target_bytes): + for index in range(source.num_rows): + consumed.append(index) + yield source.slice(index, 1) + + monkeypatch.setattr(lance_sink_module, "_split_batch_by_bytes", tracked_splits) + batcher = lance_sink_module._ByteAwareBatcher(target_bytes=1) + ready = batcher.add(source) + + assert consumed == [] + assert next(ready).to_pydict() == {"payload": [b"a"]} + assert consumed == [0] + + def test_lance_streaming_writer_writes_oversized_single_row_alone(monkeypatch) -> None: schema = pa.schema([pa.field("payload", pa.binary())]) captured, _ = _capture_lance_writer_batches( From 9035b48a51d8a58c9d8a11820acc432fd4a94793 Mon Sep 17 00:00:00 2001 From: Hynek Kydlicek Date: Wed, 2 Sep 2026 17:07:52 +0200 Subject: [PATCH 11/11] fix(lance): accept five gib multipart parts Co-Authored-By: OpenAI Codex --- src/refiner/pipeline/sinks/lance.py | 2 +- tests/pipeline/test_sinks.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/src/refiner/pipeline/sinks/lance.py b/src/refiner/pipeline/sinks/lance.py index 402e724b1..9f0a4ab6a 100644 --- a/src/refiner/pipeline/sinks/lance.py +++ b/src/refiner/pipeline/sinks/lance.py @@ -97,7 +97,7 @@ class AddColumns: _LANCE_ROW_ADDRESS_POSITION_MASK = (1 << _LANCE_ROW_ADDRESS_FRAGMENT_SHIFT) - 1 _ADD_COLUMNS_FILL_BATCH_ROWS = 65_536 _MIN_MULTIPART_PART_BYTES = 5 * 1024 * 1024 -_MAX_MULTIPART_PART_BYTES = 5 * 1024 * 1024 * 1024 - 1 +_MAX_MULTIPART_PART_BYTES = 5 * 1024 * 1024 * 1024 _LANCE_PROCESS_CONFIG_LOCK = threading.Lock() _LANCE_PROCESS_UPLOAD_CONFIG: tuple[int, int] | None = None diff --git a/tests/pipeline/test_sinks.py b/tests/pipeline/test_sinks.py index cd644444d..b974901ae 100644 --- a/tests/pipeline/test_sinks.py +++ b/tests/pipeline/test_sinks.py @@ -123,6 +123,19 @@ def test_lance_io_config_validates_limits(kwargs, message) -> None: LanceIOConfig(**kwargs) +def test_lance_io_config_accepts_five_gib_multipart_endpoint() -> None: + five_gib = 5 * 1024 * 1024 * 1024 + + config = LanceIOConfig( + upload_concurrency=1, + multipart_part_bytes=five_gib, + target_batch_bytes=five_gib, + max_buffered_bytes=4 * five_gib, + ) + + assert config.multipart_part_bytes == five_gib + + def _capture_lance_writer_batches( monkeypatch, batches: list[pa.RecordBatch],