Skip to content

fix: roll fragments on partition-value transitions#463

Merged
hamersaw merged 3 commits into
lance-format:mainfrom
hamersaw:bug/write-multiple-parition-values
Apr 21, 2026
Merged

fix: roll fragments on partition-value transitions#463
hamersaw merged 3 commits into
lance-format:mainfrom
hamersaw:bug/write-multiple-parition-values

Conversation

@hamersaw

Copy link
Copy Markdown
Collaborator

Summary

  • Distributions.clustered + sort on partition columns (from feat: require clustered distribution on write for SPJ #445) does NOT guarantee one partition value per Lance fragment. Hash collisions pack multiple values into one Spark task, and Fragment.create rolls only on maxRowsPerFile/maxBytesPerFile, so fragments can span values — breaking the SPJ contract.
  • Thread partition columns through SparkWriteLanceBatchWriteWriterFactoryLanceDataWriter. In the sorted input stream, detect partition-key transitions and roll a fresh ArrowBatchWriteBuffer + Fragment.create task at each boundary. Streaming semantics preserved — groups still stream through the existing producer/consumer buffer.
  • Adds BasePartitionedWriteTest.testHashCollisionsBreakSinglePartitionPerFragment which forces spark.sql.shuffle.partitions=2 with 20 regions so collisions are guaranteed, and asserts one region per fragment. Also adds the missing PartitionedWriteTest concrete subclass in lance-spark-3.5_2.12 — without it, the abstract tests in feat: require clustered distribution on write for SPJ #445 were never executed.

Test plan

  • mvn -pl lance-spark-3.5_2.12 -am test -Dtest='PartitionedWriteTest' — all 3 tests pass (collision test was failing before fix: Fragment 2 contains 5 distinct regions; expected 1).
  • mvn -pl lance-spark-base_2.12 test -Dtest='LanceDataWriterTest,LanceBatchWriteTest,SparkWriteTest,UpdateColumnsConflictTest' — 14 related tests green.

🤖 Generated with Claude Code

The `lance.partition.columns` write contract in lance-format#445 declares
`Distributions.clustered` plus a sort on the partition columns, intending
that each Lance fragment hold exactly one partition value (SPJ relies on
this). `ClusteredDistribution` only guarantees "same value → same task",
not "one value per task" — hash collisions (distinct values > shuffle
partitions) pack multiple partition values into one task, and
`Fragment.create` rolls only on `maxRowsPerFile`/`maxBytesPerFile`, so a
fragment could still span multiple values.

Thread the partition columns through `SparkWrite` → `LanceBatchWrite` →
`WriterFactory` → `LanceDataWriter`. In the sorted input, detect
partition-key transitions and roll a fresh `ArrowBatchWriteBuffer` +
`Fragment.create` task at each boundary. The existing streaming
producer/consumer buffer is reused per group — groups still stream
through rather than being materialized in memory.

Adds `BasePartitionedWriteTest.testHashCollisionsBreakSinglePartitionPerFragment`
which forces `spark.sql.shuffle.partitions=2` with 20 distinct region
values so collisions are guaranteed, and asserts one region per fragment.
Also adds the missing concrete `PartitionedWriteTest` subclass in
`lance-spark-3.5_2.12` — without it, the abstract base tests from lance-format#445
never actually executed.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
@github-actions github-actions Bot added the bug Something isn't working label Apr 21, 2026
The remaining testHashCollisionsBreakSinglePartitionPerFragment strictly
subsumes it: if one-value-per-fragment holds under forced hash collisions
(2 shuffle partitions, 20 regions), it trivially holds when the default
shuffle cleanly separates values. The removed test never actually
exercised the fragment-rolling logic — it passed only because 200 default
buckets happened to spread 5 regions without collisions.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
Compare partition-column values against the stored key in-place; only clone
on fragment boundaries. Eliminates the per-row Object[] + UTF8String.clone
cost in the writer hot path.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
@hamersaw
hamersaw merged commit e57e13f into lance-format:main Apr 21, 2026
29 of 30 checks passed
@hamersaw
hamersaw deleted the bug/write-multiple-parition-values branch April 21, 2026 03:32
ivscheianu pushed a commit to ivscheianu/lance-spark that referenced this pull request Jul 20, 2026
## Summary
- `Distributions.clustered` + sort on partition columns (from lance-format#445) does
NOT guarantee one partition value per Lance fragment. Hash collisions
pack multiple values into one Spark task, and `Fragment.create` rolls
only on `maxRowsPerFile`/`maxBytesPerFile`, so fragments can span values
— breaking the SPJ contract.
- Thread partition columns through `SparkWrite` → `LanceBatchWrite` →
`WriterFactory` → `LanceDataWriter`. In the sorted input stream, detect
partition-key transitions and roll a fresh `ArrowBatchWriteBuffer` +
`Fragment.create` task at each boundary. Streaming semantics preserved —
groups still stream through the existing producer/consumer buffer.
- Adds
`BasePartitionedWriteTest.testHashCollisionsBreakSinglePartitionPerFragment`
which forces `spark.sql.shuffle.partitions=2` with 20 regions so
collisions are guaranteed, and asserts one region per fragment. Also
adds the missing `PartitionedWriteTest` concrete subclass in
`lance-spark-3.5_2.12` — without it, the abstract tests in lance-format#445 were
never executed.

## Test plan
- [x] `mvn -pl lance-spark-3.5_2.12 -am test
-Dtest='PartitionedWriteTest'` — all 3 tests pass (collision test was
failing before fix: `Fragment 2 contains 5 distinct regions; expected
1`).
- [x] `mvn -pl lance-spark-base_2.12 test
-Dtest='LanceDataWriterTest,LanceBatchWriteTest,SparkWriteTest,UpdateColumnsConflictTest'`
— 14 related tests green.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants