fix: enforce file size checks for large batches

[lance-gatekeeper]

Summary

• split oversized current-format input batches with a slice-aware estimate that includes potential repetition/definition output
• expose aggregate pending bytes from field encoders through every current-format writer and flush them before rollover decisions
• keep the normal per-column page cache while bounding flat-wide, empty/null list, and low-memory Struct<Null> writes
• preserve exact released V2.0 wire output by enabling logical slice page sizing only for rechunked high-level dataset writes
• retain the Python fragment and optimize expectations for bounded input chunks

Root cause

Current-format writes checked the object-writer position only after each emitted batch. A single oversized input batch therefore deferred the check, while field encoders could buffer both array data and nested structural state without advancing tell(). Arrow allocation size was not a safe checkpoint proxy: it undercounted repetition/definition output for low-memory structural arrays and overcounted zero-copy slices by charging their full parent buffers.

The repair uses a slice-aware encoded-size estimate to create zero-copy checkpoint batches between 2 MiB and 10 MiB. Field encoders now report their retained array and repetition/definition bytes, and dataset rollover compares tell() + pending_bytes before forcing a flush and consulting the resulting file position. The default 8 MiB per-column encoder cache remains unchanged.

The first slice-aware V2.0 page subdivision applied to every V2.0 writer and changed the stable writer's historical page boundaries, causing the current-head wire-compatibility failures. Logical page sizing is now an internal opt-in used only by the high-level writer that receives the zero-copy checkpoint slices. Default and embedded V2.0 writers retain allocation-based subdivision and their byte-identical released output.

Validation

• cargo test -p lance-encoding test_v2_0_page_split_uses_slice_size -- --nocapture
• cargo test -p lance-encoding test_v2_0_multi_page_split_preserves_parent_size -- --nocapture
• cargo test -p lance-encoding
• cargo test -p lance-file compatibility_tests -- --nocapture
• cargo test -p lance test_file_size -- --nocapture
• cargo test -p lance test_max_rows_per_file -- --nocapture
• cargo test -p lance dataset::write::size::tests -- --nocapture
• cargo test -p lance-arrow stream::tests -- --nocapture
• cargo fmt --all -- --check
• cargo clippy --all --tests --benches -- -D warnings
• CARGO_INCREMENTAL=0 uv run make build from python/
• uv run pytest python/tests/test_fragment.py::test_write_fragments python/tests/test_optimize.py::test_optimize_max_bytes -vv
• CARGO_INCREMENTAL=0 uv run make lint from python/

Fixes #3393

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The slicing strategy still leaves rollover dependent on bytes the current writer has actually flushed, so it does not enforce the intended soft bound for all valid schemas. The checkpoint must account for aggregate buffered encoder data, or bound/flush that data before consulting tell().

[lance-gatekeeper]

Bounding input batches here does not guarantee that the later tell() check observes progress. Current-format writers buffer up to 8 MiB per column, so a wide schema can keep every column below its flush threshold while accumulating a large file entirely in memory; finalization then emits it after the last rollover opportunity.

On this head I ran the regression below with one approximately 128 MiB batch, max_bytes_per_file = 2 KiB, and a row limit above the batch size. It expected multiple fragments but got one 54,480,664-byte file. A viable fix is to include pending encoder bytes in the size decision, or cap/flush aggregate current-writer buffering below the checkpoint size before consulting tell(). Please retain a wide-schema regression that asserts the resulting file bound.

Reproducer

#[tokio::test]
async fn test_file_size_wide_single_batch() {
    let num_columns = 2000;
    let num_rows = 16 * 1024;
    let schema = Arc::new(ArrowSchema::new(
        (0..num_columns)
            .map(|i| ArrowField::new(format!("c{i}"), DataType::Int32, false))
            .collect::<Vec<_>>(),
    ));
    let values = Arc::new(Int32Array::from_iter(0..num_rows as i32));
    let batch = RecordBatch::try_new(schema.clone(), vec![values; num_columns]).unwrap();
    assert!(batch.get_array_memory_size() > MAX_WRITE_BATCH_BYTES);

    let data_stream = Box::pin(RecordBatchStreamAdapter::new(
        schema.clone(),
        futures::stream::iter([Ok::<_, DataFusionError>(batch)]),
    ));
    let write_params = WriteParams {
        max_rows_per_file: 2 * 1024 * 1024,
        max_bytes_per_file: 2 * 1024,
        mode: WriteMode::Create,
        ..Default::default()
    };
    let schema = Schema::try_from(schema.as_ref()).unwrap();
    let object_store = Arc::new(ObjectStore::memory());
    let (fragments, _) = write_fragments_internal(
        None,
        object_store,
        &Path::from("wide-test"),
        schema,
        data_stream,
        write_params,
        None,
    )
    .await
    .unwrap();

    assert!(fragments.len() > 1, "got {} fragment(s)", fragments.len());
}

cargo test -p lance test_file_size_wide_single_batch -- --nocapture failed with got 1 fragment(s); the sole file was 54,480,664 bytes.

[lance-gatekeeper]

Addressed in d8d537e. Current-format writers now aggregate pending bytes across all field encoders and flush before rollover; the wide-schema bounded-file regression remains and passes without reducing the normal per-column cache.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The flat-wide regression now passes, but the size boundary still misses buffered list structural data, so nested inputs can exceed max_bytes_per_file by orders of magnitude.

A viable revision should include repetition/definition state in pending-byte accounting or flush it before rollover checks, with a nested regression. Keep the reduced cache scoped to size-limited writes or preserve a practical per-column floor so ordinary wide writes do not turn each input chunk into thousands of tiny pages.

[lance-gatekeeper]

This budget does not cover the offsets and nullability held in RepDefBuilder. ListStructuralEncoder forwards an empty child values array for empty/null lists, and AccumulationQueue counts only that child array, so no page is emitted for tell() to observe before rollover.

Reproducer run on this head

I added test_file_size_empty_lists_single_batch beside the existing file-size tests. It builds one batch of 1,000,000 alternating empty and null List<Int32> rows:

let mut builder = ListBuilder::new(Int32Builder::new());
for index in 0..1_000_000 {
    builder.append(index % 2 == 0);
}

It writes with max_rows_per_file = 2 * 1024 * 1024 and max_bytes_per_file = 2 * 1024, then asserts fragments.len() > 1. Running cargo test -p lance test_file_size_empty_lists_single_batch -- --nocapture failed with: got 1 fragment(s), largest file was 4000237 bytes.

Please make pending structural bytes visible to the size checkpoint (or flush them before tell()) and keep this empty/sparse nested case as a regression.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The previous empty-list path is fixed, but the checkpoint still keys off Arrow input allocation rather than pending encoded state. Low-memory structural arrays can therefore accumulate large repetition/definition output without ever triggering a flush.

A viable revision needs a writer-owned pending-size/checkpoint signal, or a flush cadence independent of Arrow input memory, while keeping checkpoints rare enough not to fragment ordinary pages.

[lance-gatekeeper]

Arrow input memory is not a conservative proxy for pending encoded output. A Struct<Null> batch has almost no Arrow buffers, but the V2.1 structural encoder materializes one u16 definition level per row. This counter therefore stays below the threshold while megabytes remain buffered for finalization.

Reproducer run on this head

I added test_file_size_struct_null_single_batch beside the existing file-size tests. Its input was:

let num_rows = 3_000_000;
let struct_fields = Fields::from(vec![Arc::new(ArrowField::new(
    "value", DataType::Null, true,
))]);
let struct_array = Arc::new(StructArray::new(
    struct_fields.clone(),
    vec![Arc::new(NullArray::new(num_rows))],
    None,
));

The batch reported less than 1 MiB of Arrow memory. I wrote it with max_rows_per_file = 4_000_000 and max_bytes_per_file = 1 MiB, then asserted fragments.len() > 1. Running cargo test -p lance test_file_size_struct_null_single_batch -- --nocapture failed with got 1 fragment(s), largest file was 6000232 bytes.

Please drive the checkpoint from pending writer/encoder state, or use a trigger that cannot be bypassed by low-memory structural types. The inverse also matters: zero-copy rechunked slices report their full parent allocation, so this same proxy can force excessive page flushes.

[lance-gatekeeper]

Addressed in d8d537e. Repetition/definition bytes now participate in writer-owned pending accounting, pre-write sizing includes structural output, and slice-aware sizing avoids charging zero-copy slices for their parent allocation. The exact 3,000,000-row Struct regression now passes.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The pending-slice queue fixes recursive rechunking, but the file-size blocker is unchanged: checkpoints still depend on Arrow allocation size, so low-memory structural input can bypass them.

A viable revision should checkpoint from pending encoded state, or from another trigger that cannot undercount structural output or overcount zero-copy slices.

[lance-gatekeeper]

This remains a bypass because get_array_memory_size() measures Arrow buffers, not pending encoded state. A Struct<Null> batch has almost no Arrow memory while V2.1 buffers one u16 definition level per row.

Reproducer rerun on this head

I added a test beside the existing file-size cases with 3,000,000 rows:

let struct_fields = Fields::from(vec![Arc::new(ArrowField::new(
    "value", DataType::Null, true,
))]);
let struct_array = Arc::new(StructArray::new(
    struct_fields.clone(),
    vec![Arc::new(NullArray::new(3_000_000))],
    None,
));

The batch reported less than 1 MiB of Arrow memory. With max_rows_per_file = 4_000_000 and max_bytes_per_file = 1 MiB, the test expected more than one fragment. cargo test -p lance test_file_size_struct_null_single_batch -- --nocapture failed with got 1 fragment(s), largest file was 6000232 bytes.

Please base the checkpoint on writer/encoder state or another non-bypassable cadence. Carrying planned logical slice bytes separately would also avoid the inverse full-parent overcount for zero-copy slices.

[lance-gatekeeper]

Addressed in d8d537e. Rollover now uses aggregate writer/encoder pending state, while structural-aware, slice-aware rechunking prevents both the low-memory Struct bypass and full-parent slice overcount. The reported reproducer is retained as a passing regression.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The new checkpointing addresses the oversized-file bug for the structural writers, but the unconditional zero-copy rechunking introduces multiplicative page amplification in the still-supported V2.0 writer. Keep slice-aware sizing consistent through V2.0 page subdivision before landing this change.

[lance-gatekeeper]

This rechunking also applies to V2.0. Each 8–10 MiB chunk is a zero-copy slice of the original batch, but V2.0 PrimitiveFieldEncoder::do_flush still calls get_buffer_memory_size() on that slice, which returns the full shared parent capacity. A 1 GiB single-column batch is therefore split into about 103 chunks and then each chunk is split as though it were still 1 GiB: roughly 3,296 encode tasks/pages instead of the previous 32. That is multiplicative write and metadata amplification on a supported current writer.

Please use the logical slice size in V2.0's page-subdivision calculation as well (or otherwise remove the shared-parent sizing mismatch), and add a V2.0 page-count regression.

Reproducer

Add this test beside PrimitiveFieldEncoder:

#[test]
fn test_v2_0_page_split_uses_slice_size() {
    let parent =
        Arc::new(Int32Array::from(vec![0_i32; 16 * 1024 * 1024])) as ArrayRef;
    let slice =
        parent.slice(0, 10 * 1024 * 1024 / std::mem::size_of::<i32>());
    assert_eq!(
        lance_arrow::memory::array_slice_memory_size(slice.as_ref()),
        10 * 1024 * 1024
    );
    let num_parts = bit_util::ceil(
        slice.get_buffer_memory_size(),
        EncodingOptions::default().max_page_bytes as usize,
    );
    assert_eq!(
        num_parts, 1,
        "a 10 MiB logical slice fits in one 32 MiB page"
    );
}

cargo test -p lance-encoding test_v2_0_page_split_uses_slice_size -- --nocapture fails on this head with left: 2, right: 1.

[lance-gatekeeper]

Addressed in 5d7204e. V2.0 page subdivision now uses the slice-aware logical memory size, and the 64 MiB parent / 10 MiB slice regression produces one page under the default 32 MiB limit.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The V2.0 page-subdivision correction is sound, but small checkpoint slices can still be copied at full parent-buffer size before encoding, turning the file-size fix into a multi-gigabyte memory spike for variable-width input. Buffer only the logical slice window—using the existing deep_copy_array_sliced path—so retained memory matches the queue's slice-byte accounting.

[lance-gatekeeper]

This now charges only the logical slice bytes, but the default keep_original_array = false path below still calls deep_copy_array. For Utf8/Binary slices, Arrow retains the full shared value-data buffer, and that helper copies every backing buffer while preserving the slice offset. With 2 MiB checkpoint slices and the default 8 MiB cache, several slices can therefore duplicate the entire parent before the queue flushes; a 1 GiB input can transiently retain roughly 4–5 GiB and OOM.

Use the existing deep_copy_array_sliced helper here so the buffered allocation matches this logical-byte accounting, and retain a variable-width regression.

Reproducer

Add this at the bottom of accumulation.rs:

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use arrow_array::{Array, StringArray};
    use super::*;

    #[test]
    fn sliced_utf8_cache_does_not_copy_the_full_parent() {
        let value = "x".repeat(1024 * 1024);
        let parent =
            Arc::new(StringArray::from(vec![value.as_str(); 16])) as ArrayRef;
        let slice = parent.slice(0, 2);
        assert!(array_slice_memory_size(slice.as_ref()) < 3 * 1024 * 1024);

        let mut queue = AccumulationQueue::new(8 * 1024 * 1024, 0, false);
        assert!(queue.insert(slice, 0, 2).is_none());
        let retained_bytes =
            queue.buffered_arrays[0].get_buffer_memory_size();

        assert!(
            retained_bytes < 3 * 1024 * 1024,
            "a 2 MiB logical slice retained {retained_bytes} bytes"
        );
    }
}

cargo test -p lance-encoding sliced_utf8_cache_does_not_copy_the_full_parent -- --nocapture fails on this head with a 2 MiB logical slice retained 16777280 bytes.

[lance-gatekeeper]

Addressed in 05da8e5. AccumulationQueue now uses deep_copy_array_sliced for buffered copies, so the UTF-8 regression retains only its 2 MiB logical window instead of the 16 MiB parent allocation; the full lance-encoding suite passes.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The V2.0 compatibility scoping is correct, but it does not change the variable-width slice-copy blocker: logical checkpoint slices are still deep-copied at full parent-buffer size. Use the existing deep_copy_array_sliced path so retained memory follows the queue’s logical-byte accounting.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The slice-aware copier fixes top-level variable-width retention, but it is not safe for every array type accepted by this queue: it changes Boolean slice values and does not detach sliced dictionary values.

Use a copier that preserves each logical slice exactly and recursively materializes retained dictionary children, with Boolean-offset and dictionary-retention regressions.

[lance-gatekeeper]

deep_copy_array_sliced applies a Boolean slice offset twice. MutableArrayData::extend takes logical indices, and its Boolean handler already adds array.offset(); this helper passes data.offset() as start as well. A valid nonzero-offset Boolean slice is therefore copied from the wrong bits, silently changing values before encoding. Use logical range 0..data.len() and retain a Boolean regression.

Reproducer

#[test]
fn test_deep_copy_boolean_slice() {
    let array = BooleanArray::from(vec![false, true, false, true, false, true]);
    let sliced = array.slice(1, 4);
    let copied = deep_copy_array_sliced(&sliced);
    let copied = copied.as_any().downcast_ref::<BooleanArray>().unwrap();
    assert_eq!(copied, &sliced);
}

cargo test -p lance-arrow test_deep_copy_boolean_slice -- --nocapture failed on this head: the copied values were [false, true, false, true] instead of [true, false, true, false].

[lance-gatekeeper]

Addressed in dd41ef6 (present in verified head b5729df). The copier now uses the logical 0..len range, and the nonzero-offset nested Boolean regression preserves the sliced values exactly.

[lance-gatekeeper]

This helper also leaves a single dictionary values child shallow-cloned. A dictionary whose values are a 2 MiB StringArray slice therefore retains its 16 MiB parent while this queue charges only the logical window; distinct batches can still retain multiple large parents below the 8 MiB budget. Recursively materialize dictionary values and cover retained bytes for sliced dictionaries.

Reproducer

#[test]
fn sliced_dictionary_values_do_not_retain_the_full_parent() {
    let value = "x".repeat(1024 * 1024);
    let parent_values =
        Arc::new(StringArray::from(vec![value.as_str(); 16])) as ArrayRef;
    let values = parent_values.slice(0, 2);
    let keys = Int8Array::from(vec![0_i8, 1_i8]);
    let dictionary = Arc::new(
        DictionaryArray::<Int8Type>::try_new(keys, values).unwrap(),
    ) as ArrayRef;
    assert!(array_slice_memory_size(dictionary.as_ref()) < 3 * 1024 * 1024);

    let mut queue = AccumulationQueue::new(8 * 1024 * 1024, 0, false);
    assert!(queue.insert(dictionary, 0, 2).is_none());
    let retained = queue.buffered_arrays[0].get_buffer_memory_size();
    assert!(retained < 3 * 1024 * 1024);
}

cargo test -p lance-encoding sliced_dictionary_values_do_not_retain_the_full_parent -- --nocapture failed on this head with retained = 16777408.

[lance-gatekeeper]

Addressed in b5729df. Sliced dictionary values are now recursively materialized after key normalization, and the regression verifies semantic equality while retaining less than 3 MiB instead of the 16 MiB parent.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The Boolean slice fix is correct, but the dictionary-retention blocker remains: the slice copier still shallow-clones dictionary values, so logical cache accounting can retain full parent buffers.

Recursively materialize dictionary values and keep the sliced-dictionary retained-memory regression.

[claude]

Claude Code Review

Claude Code Review is paused for this repository. To reconnect it, an admin of this repository's GitHub organization (or the account owner, for personal repositories) who can also manage your Claude organization's Code Review settings needs to re-link GitHub in Code Review settings. This is a one-time step.

_{Tip: disable this comment in your organization's Code Review settings.}

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The retained-parent leak is fixed, but independently copying one shared dictionary per checkpoint slice defeats dictionary reuse and can multiply encoded values, allocation, and file size by the number of slices.

Reuse one detached values child across slices that share the source dictionary, or deduplicate identical dictionaries before concatenation.

[lance-gatekeeper]

Creating a fresh child copy for every slice destroys pointer identity among slices from the same DictionaryArray. When a page flushes, arrow_dictionary_to_data_block calls Arrow concat; for low-cardinality dictionaries where total dictionary values are fewer than rows, Arrow takes its concat fallback and appends each distinct values child rather than merging it. One logical dictionary is therefore duplicated once per cached slice, multiplying copy work and encoded page bytes. Preserve one detached child for slices sharing the same source values, or normalize identical dictionaries before concat.

Reproducer

#[test]
fn copied_dictionary_slices_reuse_shared_values() {
    let values = Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef;
    let keys = Int8Array::from_iter_values((0..100).map(|i| (i % 2) as i8));
    let dictionary = Arc::new(
        DictionaryArray::<Int8Type>::try_new(keys, values).unwrap(),
    ) as ArrayRef;

    let mut queue = AccumulationQueue::new(8 * 1024 * 1024, 0, false);
    for offset in (0..50).step_by(10) {
        assert!(queue
            .insert(dictionary.slice(offset, 10), offset as u64, 10)
            .is_none());
    }
    let (arrays, _, num_rows) = queue.flush().unwrap();
    let data = crate::data::DataBlock::from_arrays(&arrays, num_rows);
    let dictionary = data.as_dictionary().unwrap();

    assert_eq!(dictionary.dictionary.num_values(), 2);
}

cargo test -p lance-encoding copied_dictionary_slices_reuse_shared_values -- --nocapture failed on this head: five slices produced 10 dictionary values instead of 2.

[lance-gatekeeper]

Addressed in ad7c42f. AccumulationQueue now detaches each source dictionary values array once per page and reuses that child across copied key slices, including the threshold-triggered flush path; both regressions preserve exactly two dictionary values.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The latest revision restores reuse only for one Arc wrapper within one page. Valid shared dictionaries can still lose reuse when the wrapper changes or when the shared values child pushes every checkpoint slice into a separate page.

A viable revision should preserve Arrow buffer-level dictionary identity and charge shared values only once per page/checkpoint, or otherwise avoid encoding one shared dictionary per slice.

[lance-gatekeeper]

This cache uses Arc object identity, which is narrower than Arrow dictionary identity. Distinct make_array(values.to_data()) wrappers may share exactly the same buffers; Arrow concat reuses that dictionary, but this queue does not match the wrappers and deep-copies a new child for each slice. The page then grows from two values to ten. Key reuse by the same buffer-level identity Arrow uses, not only the wrapper address.

Reproducer

Add this test to the existing module:

#[test]
fn copied_dictionary_wrappers_sharing_buffers_reuse_values() {
    let source_values = Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef;
    let arrays = (0..5)
        .map(|_| {
            let values = make_array(source_values.to_data());
            let keys =
                Int8Array::from((0..10).map(|index| (index % 2) as i8).collect::<Vec<_>>());
            Arc::new(DictionaryArray::<Int8Type>::try_new(keys, values).unwrap()) as ArrayRef
        })
        .collect::<Vec<_>>();

    let baseline = crate::data::DataBlock::from_arrays(&arrays, 50);
    assert_eq!(baseline.as_dictionary().unwrap().dictionary.num_values(), 2);

    let mut queue = AccumulationQueue::new(8 * 1024 * 1024, 0, false);
    for (row_number, array) in arrays.into_iter().enumerate() {
        assert!(queue
            .insert(array, (row_number * 10) as u64, 10)
            .is_none());
    }
    let (arrays, _, num_rows) = queue.flush().unwrap();
    let copied = crate::data::DataBlock::from_arrays(&arrays, num_rows);

    assert_eq!(copied.as_dictionary().unwrap().dictionary.num_values(), 2);
}

cargo test -p lance-encoding copied_dictionary_wrappers_sharing_buffers_reuse_values -- --nocapture failed on ad7c42f73: the baseline had 2 values, while the copied page had 10.

[lance-gatekeeper]

Addressed in 4bedf38. Dictionary values are now keyed by their recursive buffer-level Arrow identity, with semantic validation if a stale raw-pointer fingerprint could be reused; distinct wrappers over the same buffers share one detached values child, and the reproducer retains two values.

[lance-gatekeeper]

array_slice_memory_size includes the entire dictionary values child because slicing a dictionary changes only its keys. Adding that fixed child for every slice means a values table at the 8 MiB page threshold forces every 10 MiB checkpoint slice into a separate page, so the page-scoped reuse above cannot help. Charge shared dictionary values incrementally, and make the batch estimator separate fixed shared bytes from row-proportional bytes (or avoid splitting one shared dictionary into independently encoded pages).

Reproducer

I ran this exact-scale test in this module. For a dictionary there are no structural levels, so batch_slice_memory_size is the same estimate used by the new write checkpoint path.

use arrow_array::{
    ArrayRef, DictionaryArray, Int32Array, RecordBatch, StringArray, types::Int32Type,
};
use arrow_schema::{DataType, Field, Schema};
use futures::{StreamExt, stream};
use std::sync::Arc;

#[tokio::test]
async fn large_dictionary_is_not_repeated_across_checkpoint_slices() {
    const MIB: usize = 1024 * 1024;
    let value = "x".repeat(5 * MIB);
    let values =
        Arc::new(StringArray::from(vec![value.as_str(), value.as_str()])) as ArrayRef;
    let keys = Int32Array::from(
        (0..10 * MIB)
            .map(|index| (index % 2) as i32)
            .collect::<Vec<_>>(),
    );
    let dictionary =
        Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap()) as ArrayRef;
    let schema = Arc::new(Schema::new(vec![Field::new(
        "item",
        DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
        false,
    )]));
    let batch = RecordBatch::try_new(schema.clone(), vec![dictionary]).unwrap();

    let slices = lance_arrow::stream::rechunk_stream_by_size_with_estimator(
        stream::iter([Ok::<_, arrow_schema::ArrowError>(batch)]),
        schema,
        0,
        10 * MIB,
        lance_arrow::memory::batch_slice_memory_size,
    )
    .collect::<Vec<_>>()
    .await
    .into_iter()
    .collect::<Result<Vec<_>, _>>()
    .unwrap();

    let mut queue = AccumulationQueue::new(8 * MIB as u64, 0, false);
    let mut encoded_dictionary_values = 0;
    for (row_number, slice) in slices.iter().enumerate() {
        let (arrays, _, num_rows) = queue
            .insert(
                slice.column(0).clone(),
                row_number as u64,
                slice.num_rows() as u64,
            )
            .expect("every slice still charges the full dictionary");
        let data = crate::data::DataBlock::from_arrays(&arrays, num_rows);
        encoded_dictionary_values +=
            data.as_dictionary().unwrap().dictionary.num_values();
    }

    assert!(slices.len() > 1);
    assert_eq!(encoded_dictionary_values, 2);
}

cargo test -p lance-encoding large_dictionary_is_not_repeated_across_checkpoint_slices -- --nocapture failed on ad7c42f73: six pages contained 12 dictionary values instead of the original 2.

[lance-gatekeeper]

Addressed in 4bedf38. Slice estimates now separate fixed shared dictionary bytes from row-proportional bytes, the queue charges shared values once, and rechunking leaves a batch intact when fixed values already fill the target; the large-dictionary regression encodes the values only once.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The buffer-identity reuse fixes the previous non-null dictionary cases, but nullable dictionary values are normalized into fresh allocations before each checkpoint slice reaches the queue. This still multiplies a shared dictionary's encoded values and can repeat large dictionary data across checkpoints.

Normalize each source dictionary once before checkpoint slicing, or preserve its original identity through normalization, so all slices share one normalized values child that is charged and encoded once.

[lance-gatekeeper]

Nullable dictionary values arrive here after PrimitiveStructuralEncoder::extract_validity calls normalize_dict_nulls, which rebuilds both keys and values for every checkpoint slice. This allocation identity is therefore new each time, so neither the values cache nor the dictionary-run logic matches slices from the same source. The resulting page multiplies dictionary values; a large nullable values table likewise repeats allocation and output across checkpoints. Normalize once before slicing, or preserve/cache the original dictionary identity across normalization so all slices reuse one normalized values child.

Reproducer

In the existing primitive.rs test module, extend the arrow_array imports with DictionaryArray and types::Int8Type, then add:

#[test]
fn nullable_dictionary_values_are_reused_across_slices() {
    let values = Arc::new(StringArray::from(vec![Some("a"), None, Some("b")])) as ArrayRef;
    let keys = Int8Array::from(
        (0..50)
            .map(|index| if index % 2 == 0 { 0_i8 } else { 2_i8 })
            .collect::<Vec<_>>(),
    );
    let dictionary =
        Arc::new(DictionaryArray::<Int8Type>::try_new(keys, values).unwrap()) as ArrayRef;

    let mut queue = crate::utils::accumulation::AccumulationQueue::new(
        8 * 1024 * 1024,
        0,
        false,
    );
    for offset in (0..50).step_by(10) {
        let normalized =
            super::dict::normalize_dict_nulls(dictionary.slice(offset, 10)).unwrap();
        assert!(queue.insert(normalized, offset as u64, 10).is_none());
    }

    let (arrays, _, num_rows) = queue.flush().unwrap();
    let data = DataBlock::from_arrays(&arrays, num_rows);
    assert_eq!(data.as_dictionary().unwrap().dictionary.num_values(), 2);
}

cargo test -p lance-encoding nullable_dictionary_values_are_reused_across_slices -- --nocapture failed on 4bedf3823ac719b7f94f3a290b86f506d0ac624d: the page contained 10 dictionary values instead of 2.

[lance-gatekeeper]

Addressed in c756aa3. The logical-sizing encoder now caches one normalized values child per live source dictionary and reuses it for later checkpoint slices; the queue also recognizes equivalent normalized children and keeps oversized shared values in one page. The supplied regression now encodes two dictionary values instead of ten.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The nullable-dictionary value reuse is fixed, but the oversized-values rule now lets equal dictionaries bypass the per-column page cache across unrelated key buffers. A long run can accumulate into one page and encode task until an external writer flush, recreating the unbounded-memory risk this PR is intended to remove.

Keep shared dictionary reuse within a bounded run: enforce a hard ceiling on incremental keys and structural state, then start a new page, or add an explicit mechanism for sharing one dictionary across bounded pages.

[lance-gatekeeper]

When both values tables exceed the cache, this condition ignores key_allocation and next_key_offset, so independent arrays using the same cached values are treated as one continuous run. defer_dictionary_flush then remains true after every over-budget insert, allowing pending keys and the eventual page/encode task to grow with the entire run. Bound deferral by incremental key and structural bytes; after a hard ceiling, flush even if values must be repeated, or share the dictionary explicitly across bounded pages.

Reproducer

Add this test to the existing module:

#[test]
fn oversized_dictionary_values_do_not_disable_the_cache_bound() {
    let values = Arc::new(StringArray::from(vec!["x".repeat(2 * 1024)])) as ArrayRef;
    let mut queue = AccumulationQueue::new(1024, 0, false);
    let mut flushed = false;

    for row_number in 0..64 {
        let keys = Int8Array::from(vec![0_i8; 1024]);
        let dictionary =
            Arc::new(DictionaryArray::<Int8Type>::try_new(keys, values.clone()).unwrap())
                as ArrayRef;
        if queue
            .insert(dictionary, row_number * 1024, 1024)
            .is_some()
        {
            flushed = true;
            break;
        }
    }

    assert!(
        flushed,
        "an oversized shared dictionary retained {} bytes without ever flushing a 1024-byte cache",
        queue.pending_bytes()
    );
}

cargo test -p lance-encoding oversized_dictionary_values_do_not_disable_the_cache_bound -- --nocapture failed on c756aa3dfc1b272eb6c6662545eebe377abbab0c: the 1,024-byte cache retained 67,588 bytes across 64 independent key arrays without ever flushing.

[lance-gatekeeper]

Addressed in f210d14. Shared dictionary deferral now stops when incremental key plus repetition/definition bytes reach the per-column cache bound, while bounded slices still reuse one detached values child. The supplied unbounded-key regression and an additional structural-byte regression both pass.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The new queue bound fixes unlimited accumulation after slicing, but oversized shared dictionaries still bypass slicing entirely. A single input can therefore carry arbitrary key bytes through the 10 MiB write checkpoint and overshoot the page and file-size bounds this change is intended to enforce.

Treat the shared dictionary as one unavoidable fixed cost while continuing to bound row-proportional bytes separately; slice keys to a nonzero incremental budget and let each bounded page reuse the detached values child.

[lance-gatekeeper]

Returning the whole batch when shared >= max_bytes leaves incremental completely unbounded. For a dictionary whose values exceed the 10 MiB checkpoint, an arbitrarily large keys array reaches the writer as one batch; the queue can only flush that whole array as one page/task, and the file-size check runs afterward. Keep the unavoidable shared cost, but still slice row-proportional bytes using a hard incremental budget instead of disabling slicing.

Reproducer

Add this test to the existing stream test module:

#[test]
fn shared_estimator_still_bounds_incremental_bytes_when_shared_exceeds_target() {
    let input = stream::iter([Ok::<_, ArrowError>(make_batch(100))]);
    let rechunked =
        rechunk_stream_by_size_with_shared_estimator(input, test_schema(), 0, 100, |batch| {
            SliceMemorySize {
                shared: 100,
                incremental: batch.num_rows() * std::mem::size_of::<i32>(),
            }
        });
    let batches = block_on(rechunked.collect::<Vec<_>>())
        .into_iter()
        .collect::<Result<Vec<_>, _>>()
        .unwrap();

    assert!(batches.len() > 1, "incremental bytes were left in one batch");
    assert!(batches
        .iter()
        .all(|batch| batch.num_rows() * std::mem::size_of::<i32>() <= 100));
}

cargo test -p lance-arrow shared_estimator_still_bounds_incremental_bytes_when_shared_exceeds_target -- --nocapture failed on f210d14d885ccfb78456db63396133b70a35f535: 100 shared bytes plus 400 incremental bytes were returned as one unsliced batch with a 100-byte target.

[lance-gatekeeper]

Addressed in 561f255. When shared bytes consume the target, rechunking now applies a separate nonzero max_bytes budget to incremental bytes. The supplied regression yields bounded key slices, and the downstream large-dictionary test verifies one detached two-value dictionary per emitted page.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The intended file-size bound now works for the covered oversized inputs, but two input-shape performance cliffs remain. Keep the separate incremental-byte bound while enforcing a practical minimum slice budget, and maintain pending repetition/definition bytes incrementally so checkpoint accounting stays linear in batch count.

[lance-gatekeeper]

This subtraction creates a row-scale eager-slicing cliff when shared bytes land just below the target. With max_bytes = 100, shared = 99 leaves a one-byte incremental budget, while shared = 100 gets a 100-byte budget. Since this function constructs the entire Vec<RecordBatch> before yielding, a large batch can allocate one wrapper per row. Clamp the incremental budget to a practical floor (accepting a bounded soft overshoot above the shared fixed cost), or generate slices lazily so a near-full shared allocation cannot create an eager allocation spike.

Reproducer

I added this test under stream::tests and ran cargo test -p lance-arrow temporary_shared_boundary_avoids_row_scale_slice_counts -- --nocapture:

#[test]
fn temporary_shared_boundary_avoids_row_scale_slice_counts() {
    fn slice_count(shared: usize) -> usize {
        let input = stream::iter([Ok::<_, ArrowError>(make_batch(10_000))]);
        let output = rechunk_stream_by_size_with_shared_estimator(
            input, test_schema(), 0, 100, move |batch| SliceMemorySize {
                shared,
                incremental: batch.num_rows() * std::mem::size_of::<i32>(),
            },
        );
        block_on(output.collect::<Vec<_>>()).len()
    }

    let just_below = slice_count(99);
    let at_target = slice_count(100);
    assert!(just_below <= at_target * 2);
}

Observed: shared=99 created 10,000 slices, while shared=100 created 400; the assertion failed.

[lance-gatekeeper]

Addressed in c674509. The incremental slice budget now has a 75% practical floor, preventing near-target shared bytes from creating row-scale eager batch wrappers; the shared-boundary regression covers the 99-versus-100-byte case.

[lance-gatekeeper]

This rescans every pending builder on every maybe_encode call, making accumulation quadratic in the number of input batches. The dataset rechunker uses min_bytes = 0, so a stream of tiny batches remains split, and even non-null one-row primitive batches can stay pending until the page cache fills. Maintain a pending rep/def byte counter, add only the incoming builder's estimate, and reset it whenever the accumulated builders are taken.

Reproducer

I added a focused test that creates a structural Int8 field encoder with a large cache, repeatedly calls maybe_encode with one-row arrays and RepDefBuilder::default(), then compares 3,000 with 9,000 batches. Running cargo test -p lance-encoding temporary_many_tiny_batches_scale_linearly -- --nocapture produced:

3k batches: 185.146745ms; 9k batches: 1.597277148s
tripling batches should not take more than 5x as long

The linear-scaling assertion failed at 8.63×. The loop visits N(N+1)/2 pending builders before a flush.

[lance-gatekeeper]

Addressed in c674509. Pending repetition/definition byte estimates are now accumulated once per builder and reset atomically when the builders are taken, eliminating the per-batch rescan; the counter regression covers accumulation and reset.

[lance-gatekeeper]

✅ Gate recommendation: approve.

The current revision removes both checkpoint-accounting performance cliffs while preserving the bounded-write and compatibility contracts.

[Xuanwo]

Close as this PR is out of the scope of bug fix. Should have a human to kick off a design for this.