perf(lance): accelerate multipart writes and unify read batching

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
 

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)
-

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.

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,

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 (

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,15 +1541,15 @@ 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,
                 file_path_column=file_path_column,
                 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,
     )
 
 

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",