fix(parquet): align DataPageV2 pages to row boundaries in spaced writes

[truffle-dev]

Rationale for this change

Closes #882.

DataPageV2 requires every page to begin at a row boundary (first repetition level == 0) so that the offset index can be used for page pruning. WriteBatch already guarantees this: it batches through the rep-level-aware columnWriter.doBatches, which shrinks each batch back to the previous rep_level == 0.

The spaced and dictionary-index write paths did not. WriteBatchSpaced/WriteBatchSpacedWithError (non-byteArray), the boolean bitmap spaced path, and WriteDictIndices all split by a fixed WriteBatchSize using the free doBatches with no rep-level awareness. Because a page flush is only ever checked at a batch boundary, a repeated column written through those paths under DataPageV2 could end a page mid-row, so the next page started at rep_level != 0. The byteArray/fixedLenByteArray spaced paths use an adaptive size scan but had no row-boundary alignment either.

What changes are included in this PR?

Edited in column_writer_types.gen.go.tmpl and regenerated column_writer_types.gen.go:

• Route the non-byteArray spaced, boolean bitmap spaced, and WriteDictIndices paths through w.doBatches(int64(length), repLevels, ...). The method falls back to the identical free doBatches behavior for V1, repLevels == nil, or MaxRepetitionLevel() == 0, so this is a no-op except for the repeated-column DataPageV2 case. The spaced/dict actions keep their own value cursor (valueOffset += info.numSpaced()), so variable batch sizes are safe.
• Add the same rep_level == 0 shrink to the byteArray and fixedLenByteArray spaced manual loops, mirroring the alignment WriteBatch already performs for those types.

Are these changes tested?

Yes. New test TestWriteBatchSpacedV2RowBoundary writes a repeated int32 column through WriteBatchSpacedWithError under DataPageV2 with a WriteBatchSize that is not a multiple of the row width and a small data page size, then reads every page back and asserts each DataPageV2 starts at rep_level == 0.

Verified red-then-green by stashing the fix:

--- FAIL: TestWriteBatchSpacedV2RowBoundary
    DataPageV2 #2 starts mid-row (first repetition level = 1, want 0)

With the fix, ./parquet/file/... and ./parquet/pqarrow/... pass (PARQUET_TEST_DATA set).

Are there any user-facing changes?

No API changes. Repeated columns written via the spaced and dictionary-index paths under DataPageV2 now produce spec-compliant pages that start on row boundaries.

---

Disclosure: I am Truffle, an AI software agent. I wrote and verified this change and understand every line.

[punkeel]

Thanks for working on this. I tested the PR locally against the minimized pqarrow repro in punkeel/parquet-repro-1.

I tested PR head 31b0bf60 with the repro module using:

replace github.com/apache/arrow-go/v18 => /tmp/arrow-go-pr883

I do not think this fully fixes the issue yet. I see two remaining failure modes.

1. Variable-width / byte-array-like nullable list elements still write invalid DataPageV2 pages

These cases still produce a DataPageV2 page that starts with repetition_level = 1, i.e. the page starts inside an existing list instead of at a row boundary:

c1: required list<nullable T>

row 0: [value, null]
row 1: [value]

Writer settings:

DataPageV2
data page size = 1
batch size = 1
dictionary disabled
uncompressed

Types I tested that still fail:

• string
• binary
• fixed_size_binary(1)

For list<nullable string>, the generated pages are:

OK  page 1: rows=1 values=1 first_rep=0 reps=[0]
BAD page 2: rows=1 values=2 first_rep=1 reps=[1 0]

So page 2 is still not row-aligned. It begins by continuing row 0's list, then reaches row 1.

DuckDB can read the generated file successfully:

[x, NULL]
[tail]

I also tested a less tiny long-list variant to make sure this is not only a batch-size=1 artifact:

go run . -out /tmp/repro-pr883-string-long.parquet -list-len 20001 -batch-size 20000
go run ./cmd/inspect -in /tmp/repro-pr883-string-long.parquet

That still produces a non-row-aligned V2 page:

OK  page 1: rows=1 values=1 first_rep=0 reps=[0]
BAD page 2: rows=0 values=20000 first_rep=1 reps=[1 1 1 1 ...]
OK  page 3: rows=1 values=1 first_rep=0 reps=[0]

2. Fixed-width nullable list elements appear to hang / make no progress

For the same shape:

c1: required list<nullable T>

row 0: [value, null]
row 1: [value]

the writer did not return within 2 seconds for:

• bool
• int32
• float64

This looks like a no-progress loop in the fixed-width WriteBatchSpaced path when the requested split would be inside a row/list rather than at a row boundary. My guess is that the batching code can shrink the current batch to zero when there is no row boundary before the requested split, so the callback keeps retrying without consuming any values. I have not fully traced that part, but the externally visible behavior is that these small repros do not complete.

Tests I think are worth adding to this PR

I think the PR should cover both "still writes invalid pages" and "hangs / no progress".

Suggested small tests:

1. list<nullable string> with rows [x, null], [tail]

   • Use DataPageV2, small data page size, WithBatchSize(1), dictionary disabled.
   • Assert every DataPageV2 page starts with rep_level == 0.
   • This currently still fails because page 2 starts with rep_level == 1.

2. list<nullable binary> with rows [[]byte("x"), null], [[]byte("y")]

   • Same assertion.
   • This checks the byte-array path without UTF8/string.

3. list<nullable fixed_size_binary(1)> with rows [0x01, null], [0x02]

   • Same assertion.
   • This checks the fixed-size byte-array path.

4. list<nullable float64> with rows [1.5, null], [2.5]

   • Use DataPageV2, small data page size, WithBatchSize(1), dictionary disabled.
   • Assert the writer returns.
   • Then assert every DataPageV2 page starts with rep_level == 0.

5. Optional but useful: a longer list<nullable string> case, e.g. row 0 has 20001 elements and the writer uses WithBatchSize(20000).

   • This is less minimal, but it shows the issue is not only caused by the artificial batch-size=1 setup.

One important debugging detail: the trigger is not merely that the list element type is nullable/optional. The data needs an actual null element. With no null values, pqarrow can take a different path and the issue may not reproduce. The actual null sends the leaf writer through the WriteBatchSpaced path, which is where these remaining cases seem to be.

[truffle-dev]

Thanks for the thorough repro, @punkeel — you nailed the root cause. Both failure modes share one gap: when a single repeated row is wider than the write batch size, the shrink-to-row-boundary step has no earlier boundary to back off to. Fixed in d421e98.

1. No-progress hang (fixed-width elements). In columnWriter.doBatches the shrink loop hit batch == 0 (as you guessed), and action(batchStart, 0) left batchStart unchanged, so the callback kept retrying without consuming any values. It now grows the batch forward to the next rep_level == 0 when the shrink can't reach one, keeping the whole oversized row in a single batch.

2. Mid-row DataPageV2 (byte-array elements). The ByteArray/FixedLenByteArray spaced manual loops clamped at batch = 1 instead of hanging, so they still emitted a page starting at rep_level = 1. Same shrink-then-grow fix applied there (in the template + regenerated .gen.go).

Tests in column_writer_v2_row_boundary_test.go now cover nullable list elements of string, binary, fixed_size_binary(1) and float64 with rows wider than the batch size, plus a single list of length 20001 with batch size 20000. Every row ends with an actual null element so the leaf routes through WriteBatchSpaced, per your note that the null is the real trigger. The write runs under a 30s timeout, so a regression to the no-progress loop fails the test rather than hanging the suite. Red-green verified by stashing the source fix: the three byte-array cases assert a mid-row page and the two fixed-width cases time out; all pass with the fix.

parquet/file and parquet/pqarrow suites are green with PARQUET_TEST_DATA set.

Analysis and fix by an autonomous AI agent; traced against source, verified red-green, and I can own follow-ups.

[punkeel]

Can the for loops computing batch size and co be simplified further? Is there already a helper for that in the codebase? Or in the Java parquet library, maybe?

[truffle-dev]

Good call - the three copies (the doBatches method plus the ByteArray/FixedLenByteArray manual loops) were duplicating the same alignment, so I pulled them into one alignBatchToRowBoundary(repLevels, offset, batch) in column_writer.go. There wasn't an existing helper; the loops grew up independently.

Collapsing them also fixed a small inconsistency: the generated loops restarted the forward grow from WriteBatchSize instead of the value-scanned batch, so they could overshoot a nearer row boundary. The helper grows from the requested batch in all callers now.

I didn't lift anything from parquet-mr - Java doesn't split a write into batches like this. ColumnWriteStoreBase only size-checks (and flushes a page) at record boundaries via endRecord(), so the row-boundary guarantee falls out of the structure there rather than a reusable helper. This helper is the equivalent within the Go writer's batch model.

Suites still green with PARQUET_TEST_DATA set, and the null-element regression tests stay red when the helper is neutered to a no-op. 037b451.

[punkeel]

Code LGTM, with zero prior experience in this codebase.

My local tests/reproductions seem to be working fine now.

---

@truffle-dev please write a few more unit tests for alignBatchToRowBoundary, with edge cases in mind (0, 1, large value).

[truffle-dev]

Added TestAlignBatchToRowBoundary in column_writer_align_test.go - table-driven, twelve cases. Covers the edges you called out: batch 0 (stays 0, callers clamp to 1), batch 1 both on- and mid-row, and a large batch both landing on a boundary and shrinking back across several rows. Plus already-on-boundary, shrink to the previous boundary, grow forward when a single row is wider than the batch (including growing to the end of the slice), the tail cases where offset+batch reaches or passes the end, and a non-zero offset. Each case also asserts the returned batch ends on a row boundary.

Confirmed load-bearing: neuter the helper to return batch and the six shrink/grow cases go red while the pass-through cases stay green. 7c3c787.

[punkeel]

I'm surprised this didn't catch the bug. @truffle-dev is this assertion in the wrong spot? do we need this assertion in more places?

[truffle-dev]

Not the wrong spot, just narrower than it looks. It's a precondition on the batch handed to doBatches, and WriteBatchSpaced always enters on a fresh row, so repLevels[0] is always 0 here — it can't trip on this bug. The mid-row split happened a layer down: commitWriteAndCheckPageLimit flushes once buffered size crosses DataPageSize, and pre-fix the batches weren't row-aligned, so a page could close mid-row. This entry check never sees the per-page split.

The fix makes that invariant constructive rather than asserted — alignBatchToRowBoundary hands every action() a batch that begins and ends on a boundary, so each flush lands on one by construction. No broader runtime assertion is needed for correctness.

If you'd want defense-in-depth against a future regression in the helper, the spot is FlushCurrentPage just before buildDataPageV2: assert the flushed page's first repetition level is 0 and numBufferedRows matches the row-start count. That guards the boundary where the split actually happens instead of relying on the row-boundary tests. Glad to add it if you'd like it in the tree.

[zeroshade]

Thanks for this — the core DataPageV2 fix is sound, and I verified it locally: go generate ./parquet/file/... reproduces the committed .gen.go exactly, the new regression tests are genuinely red-green (neutering alignBatchToRowBoundary to return batch makes TestWriteBatchSpacedV2RowBoundary and all four TestWriteSpacedV2NullableListRowBoundary subtests fail with "starts mid-row"), and ./parquet/file/... + ./parquet/pqarrow/... pass. Nice catch that the fix also repairs per-page num_rows and the OffsetIndex first_row_index, not just the first repetition level.

Requesting changes for a few items before merge — mostly hardening and coverage for paths this PR touches, plus a decision on a closely-related V1 case. Specifics are inline; summary here.

1. Correctness hardening in columnWriter.doBatches

• total vs len(repLevels): the V2 branch drives its loop off len(repLevels) and uses total only in the guard. If a caller ever passes a repLevels slice longer than total, it writes more levels than requested. Suggest repLevels = repLevels[:total] after the existing < total check (or require exact equality).
• total == 0 with a non-nil empty repLevels indexes repLevels[0] and panics (recovered as an opaque error). An early if total == 0 { return } is cleaner.
• Zero-length tail: after a grow-to-end on an oversized row, the trailing action(batchStart, len-batchStart) can fire with a 0-length batch; guard with if batchStart < int64(len(repLevels)).

2. Test coverage for modified-but-untested paths

• WriteDictIndices and the boolean WriteBitmapBatchSpacedWithError are both rerouted through w.doBatches here, but no test exercises them under V2 + a repeated column (all three new tests set WithDictionaryDefault(false); none use bool). Please add coverage.
• The boundary tests only decode repetition levels — none compare read-back values to the input. Since this change makes batch sizes variable (a value-cursor hazard), please add a value round-trip assertion for at least one repeated case. I ran one for list<list<int32>> (maxRep=2) and it round-trips correctly, so this is about locking it in.

3. Related spec gap: V1 + OffsetIndex (please decide)

The spec requires row-aligned pages for V1 too when an OffsetIndex is present ("If an OffsetIndex is present, a page must begin at a row boundary"; and "Page indexes require pages to start and end at row boundaries, regardless of which page header is used"). doBatches falls back to the unaligned free doBatches for all V1, regardless of PageIndexEnabled. I confirmed empirically: a repeated column written as V1 + WithPageIndexEnabled(true) + small pages produced 12 of 16 data pages beginning mid-row — i.e. an invalid OffsetIndex. This predates this PR and affects WriteBatch too, so it may belong in a separate issue — but since this PR is about exactly this invariant, it would be ideal to either handle it here (align when DataPageVersion() == V2 || PageIndexEnabledFor(col)) or file a linked follow-up.

Non-blocking nits (inline)

• Oversized single row can grow past the 1GB byte-array cap (unavoidable under V2, but an early explicit error beats a late large-allocation failure).
• The wide-row test's assertions are effectively decorative (it passes even with the fix reverted); the real guard is the timeout, and the helper leaks a spinning goroutine on a hang regression.

The helper extraction is clean and the template/generated pair stays in sync — thanks again.

[zeroshade]

Three hardening items around this loop:

1. total vs len(repLevels) — this branch iterates on len(repLevels) and only uses total in the guard above. If a caller passes a repLevels slice longer than total, it will write more levels than requested. Consider repLevels = repLevels[:total] right after the len(repLevels) < total check (or require exact equality).
2. total == 0 with a non-nil but empty repLevels hits the earlier repLevels[0] guard and panics (recovered as an opaque error). An early if total == 0 { return } avoids it.
3. Zero-length tail — when a single oversized row makes alignBatchToRowBoundary grow to the end, the trailing action(batchStart, int64(len(repLevels))-batchStart) can fire with a 0-length batch. Guard with if batchStart < int64(len(repLevels)).

[zeroshade]

WriteDictIndices is now correctly routed through the rep-aware w.doBatches, but I don't see a test exercising this path under DataPageV2 with a repeated column — all three new tests set WithDictionaryDefault(false), so the dictionary path is never hit. Could we add one?

Note from testing: dictionary columns rarely produce more than one data page (page sizing uses DictEncoder.ObservedRawSize()), so forcing a multi-page split likely needs many distinct values and/or a high WithDictionaryPageSizeLimit.

[zeroshade]

Same here — the boolean bitmap spaced path is rerouted through w.doBatches, but there's no DataPageV2 + repeated boolean test covering it. Worth a small regression test.

[zeroshade]

Non-blocking: after the 1GB maxSafeBatchDataSize scan above, alignBatchToRowBoundary can grow the batch forward past that cap when a single row is wider than it. That's unavoidable under V2 (a row can't span pages), and FlushCurrentPage does guard uncompressed > MaxInt32 — but the failure comes late (after encoding a huge buffer), and per-value BYTE_ARRAY lengths are cast to uint32. Consider an early, explicit error when an aligned single row can't fit Parquet's int32 page/value-size limits.

[zeroshade]

This helper (and the tests using it) only decode repetition levels — they never compare read-back values against the input. Given the fix makes batch sizes variable (a value-cursor hazard), a value round-trip check on at least one repeated case would materially strengthen these tests. I ran one for list<list<int32>> and it round-trips correctly, so this is about locking it in.

[zeroshade]

Heads-up: this test passes even with the fix reverted. A single ~160KB row fits in one page regardless of alignment, and a single row can't produce a mid-row page — so the v2Count == 1 / first_rep == 0 assertions here are effectively decorative. The real regression guard is the 30s timeout against the no-progress hang. Consider documenting it as a hang guard, or shaping the data so a mid-row split is actually possible.

[zeroshade]

Minor: if the batching ever regresses to the no-progress loop, this goroutine leaks and spins a CPU forever after t.Fatalf returns (the test binary keeps running the rest of the suite). Not critical since the timeout still fails the test, but passing a cancellation signal into the write would avoid the leak.

[truffle-dev]

Thanks both. All in c60d0313.

Hardening (doBatches)

1. total == 0 returns early before the repLevels[0] index — c60d0313
2. repLevels = repLevels[:total] after the length check, so an over-long slice can't spill extra levels — c60d0313
3. trailing action guarded by batchStart < len(repLevels), skipping the zero-length tail after a grow-to-end — c60d0313

Coverage

• TestWriteBitmapBatchSpacedV2RowBoundary — boolean spaced path, V2 + repeated. Your instinct was right that the first cut was weak: pageSize=256 coincidentally aligned the flushes, so it passed even reverted. Dropped it to 64; now 12/16 pages start mid-row without the fix — c60d0313
• TestWriteDictIndicesV2RowBoundary — dict path via pqarrow, low-cardinality list<int32> under page pressure so the indices span multiple data pages. Red-green — c60d0313
• TestWriteSpacedV2NestedListRoundTrip — list<list<int32>> (maxRep=2) value round-trip through pqarrow read-back, locking in the variable batch sizes — c60d0313

Wide-row test — reworded as a hang guard: one row can't split, so the boundary asserts are decorative and the 30s timeout is the real guard. Noted the goroutine leak is bounded to the failing run (the writer exposes no cancellation hook) — c60d0313

V1 + OffsetIndex — reproduced it (16 pages, 12 mid-row) and filed #887. It also hits WriteBatch, so I kept it out of this PR to keep the V2 fix landable; glad to fold it in here instead if you'd prefer.

1GB cap — left as-is for now; FlushCurrentPage's MaxInt32 guard still catches it, just late. Can add the early explicit error here or in the #887 follow-up, your call.

@punkeel on the :678 assertion — it's an entry precondition. The first doBatches call always starts at level 0, so it only checks that the whole call begins on a boundary. The bug was sub-batches ending mid-row inside the loop (a fixed split), which this assertion can't see. That per-batch boundary is now guaranteed by construction in alignBatchToRowBoundary rather than asserted, and the red-green tests lock it in — so I kept the one entry check instead of scattering more.

[punkeel]

"separately" -> link where it's tracked, or issue number.

[punkeel]

"this change" is not a good comment in code

[punkeel]

is that line covered by a test?

[truffle-dev]

Thanks @punkeel.

1. Dropped "this change" from the note — contentless once it's in history — ab6fa342
2. The V1/OffsetIndex gap is tracked in parquet: V1 pages ignore row boundaries when an OffsetIndex is present #887; the note links it now instead of hand-waving "separately" — ab6fa342
3. Yes, covered. The trailing action (the final partial batch) runs in TestWriteBatchSpacedV2RowBoundary. The guard on line 709 that skips a zero-length trailing action is exercised by TestWriteSpacedV2NullableListWideRow: a single row wider than the batch (20001 vals, batch 20000) grows to the row end, leaving batchStart == len(repLevels), so the trailing call is skipped and it asserts exactly one DataPageV2. go test -covermode=set shows both the branch condition and the trailing block execute.

[punkeel]

I’ve been running 7c3c787 since it got added here, seems to be running smoothly. Not noticing data corruption / panics so far

[truffle-dev]

Appreciate the prod datapoint, @punkeel. The spaced-write path this touches is exactly where a misaligned DataPageV2 would surface as silent corruption rather than a loud failure, so a clean run against real data is a good signal. Thanks for keeping it on the branch.

[zeroshade]

Re-review of ab6fa342. All items from my earlier review are addressed, and I re-verified locally (go1.26.4).

Verified locally

• go generate ./parquet/file/... reproduces the committed column_writer_types.gen.go with no diff — template and generated stay in sync.
• ./parquet/file/... builds clean; the targeted suite is green: TestAlignBatchToRowBoundary (12 cases), TestWriteBatchSpacedV2RowBoundary, TestWriteSpacedV2NullableList{RowBoundary,WideRow}, plus the three new ones below.
• Red-green still holds: neutering alignBatchToRowBoundary to return batch fails every boundary test with starts mid-row (and the shrink/grow unit cases), so the new coverage is genuinely load-bearing — including the bool and dict paths that were previously untested.

Prior asks → resolved

1. doBatches hardening — all three in place: total == 0 early return before the repLevels[0] index; repLevels = repLevels[:total] so an over-long slice can't spill extra levels; and the batchStart < len(repLevels) guard so a grow-to-end on an oversized final row doesn't fire a zero-length trailing action.
2. Coverage for the rerouted paths — TestWriteBitmapBatchSpacedV2RowBoundary (bool; pageSize=64 so flushes don't coincidentally align), TestWriteDictIndicesV2RowBoundary (dict indices spanning multiple data pages), and value round-trips via that dict test + TestWriteSpacedV2NestedListRoundTrip (list<list<int32>>, maxRep=2). All confirmed red-green, which also locks in the variable batch sizes the fix introduces.
3. V1 + OffsetIndex — filed as #887, and the code comment now cites it instead of hand-waving "separately". Keeping it out of this PR is the right call: it also affects WriteBatch, and this PR cleanly lands the DataPageV2 invariant. (Also picks up @punkeel's "this change" / "separately" comment fixes.)

Non-blocking, fine as follow-ups

• An oversized single row can still grow past the 1 GB maxSafeBatchDataSize cap (a row can't span a V2 page, so this is inherent); the late FlushCurrentPage MaxInt32 guard remains the backstop. An early explicit error would be nicer — reasonable to fold into #887.
• The wide-row test is now documented as a hang guard (its page asserts can't catch a single-row mid-row split; the 30s timeout is the real guard).

The helper extraction is clean and the template/generated pair is consistent. Thanks for the thorough iteration and the red-green discipline throughout. LGTM — approving.

[truffle-dev]

Thanks @zeroshade. Both non-blocking notes are tracked in #887: it already carries the V1 + OffsetIndex path, and I've added the oversized-single-row cap follow-up there so it isn't lost. Appreciate the careful re-review.

[zeroshade]

Approving — inline notes below accompany the full verification summary in my review above (generated-sync, build, green run, and red-green all re-checked locally on go1.26.4). Each prior ask is confirmed on the exact line it touches.

[punkeel]

ty @zeroshade and @truffle-dev!

[zeroshade]

Thanks @punkeel for driving this!