feat: require clustered distribution on write for SPJ#445
Conversation
There was a problem hiding this comment.
@beinan thanks for the PR! I think this is in the right direction, Can we add some tests (unit + integration) to verify? And then just one concern (I'm hoping the tests will validate):
I don't think this does exactly what you're expecting. This PR imposes global sort and hash-based distribution of data on partition columns during write. Presumably this is so we have a dataset that adheres to the partitioned read path, namely one partition column variant per fragment.
However, this can have multiple partition key combinations sent to the same Spark partition (just a hash distribution) and without the writer being told to create a new fragment per partition column combination we would just get partition keys sorted within the same Lance fragment.
| private final java.util.Map<String, String> tableProperties; | ||
|
|
||
| static final String TABLE_OPT_PARTITION_COLUMNS = "lance.partition.columns"; | ||
| static final String TABLE_OPT_PARTITION_COLUMNS = LanceConstant.TABLE_OPT_PARTITION_COLUMNS; |
There was a problem hiding this comment.
I don't think we need to define this constant twice. Lets just use LanceConstant.TABLE_OPT_PARTITION_COLUMNS directly everywhere.
There was a problem hiding this comment.
Done — removed the duplicate constant. Now using LanceConstant.TABLE_OPT_PARTITION_COLUMNS directly in LanceScanBuilder.
|
Thanks for the detailed review! You're right that However, the combination of clustered distribution + sort ordering ensures that within each task, rows are grouped contiguously by partition value. Since each I've added an integration test ( The key insight: even though hash collisions can send multiple partition values to the same task, the sort ordering within the task means the writer sees values in order (e.g., all "east" rows, then all "north" rows), and lance-core's fragment creation naturally produces one fragment per contiguous block. |
jackye1995
left a comment
There was a problem hiding this comment.
looks like there is a merge conflict, once resolved I can merge it!
…lumns is set When a Lance table has the `lance.partition.columns` table property, SparkWrite now implements RequiresDistributionAndOrdering to tell Spark to automatically cluster (partition) and sort the input data by those columns before writing. This ensures each Lance fragment contains exactly one distinct value for the partition column(s), which is the prerequisite for Storage-Partitioned Joins (SPJ) on the read path. Previously users had to manually call repartitionByRange() before writing to achieve SPJ-compatible fragment layout. With this change, just setting the table property is sufficient — Spark handles the shuffle automatically. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
…ration tests - Remove duplicate TABLE_OPT_PARTITION_COLUMNS in LanceScanBuilder; use LanceConstant.TABLE_OPT_PARTITION_COLUMNS directly everywhere - Add unit tests for requiredDistribution() and requiredOrdering() with and without lance.partition.columns table property - Add BasePartitionedWriteTest integration test verifying that writes with partition columns produce fragments with one distinct value each Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
d5c36e4 to
e6a0d41
Compare
## Summary - `Distributions.clustered` + sort on partition columns (from #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 #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]>
Self-review with the lance-spark prompt at repos/spark_review_prompt flagged seven findings; this commit addresses all of them. Important: 1. LanceProbe.RowAddressColumn duplicated the literal "_rowaddr" already defined in LanceConstant.ROW_ADDRESS. lance-spark convention (PR lance-format#445 in the corpus) is one definition. Now references the constant. 2. Inline FQNs in IndexedNearestByJoinRule.scala violated AGENTS.md's "use top-level imports" rule. Three references — JoinType, StructField, Table — moved to top-level imports. 3. Duck-type fallback's `Table.name()` URI fallback was unsafe: the dataset name returned by LanceDataset is generally not a URI openable by Dataset.open(). Drop it; rule returns None when no `path` / `datasetUri` option exists, leaving Spark's brute-force rewrite to handle the query. Substantive: 4. Stale documentation: - PHASE_PROGRESS.md "Phase 1 known limitations (carried forward UNCHANGED from Phase 0)" — qualified with "when probeParallelism = 1" since Phase 1.5 ships an opt-in non-degenerate path. - PHASE_PROGRESS.md "Phase 3 — hardening (NOT STARTED)" was contradicted by the Phase 3 done section earlier in the same file; dropped, the IMPL_PLAN.md "What's left" table is the canonical backlog source. - LanceProbeStage class doc said "exercised by future tests once fragment-grouping is wired in" — Phase 1.5 wired it in. - LanceFragments class doc said "Skew-aware splitting (...) is a Phase 3 concern" — the same file ships enumerateGroupsByRowCount. - Metric.scala scaladoc listed "L2 / Hamming" but no Hamming case exists in the enum. 5. LanceMaterializeStage's `rowAddr -> materialized row` map silently collapsed duplicate refs. Added a comment documenting that this is intentional: TopKHeap.merge does not dedupe across contributions, so duplicate refs mean the same right row appeared as the K-th nearest along multiple fragment-group paths — emitting one output row per ref is correct behavior. All 28 + 10 tests still pass. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
## Summary - Add `unenforced_clustering_key` metadata to the Lance schema format, mirroring the existing `unenforced_primary_key` pattern - Clustering keys hint at the physical ordering of data within a table, enabling query engine optimizations such as storage-partitioned joins (SPJ) - Unlike primary keys, clustering key fields may be nullable Changes across all layers: - **Protobuf**: `unenforced_clustering_key` (bool) + `unenforced_clustering_key_position` (uint32) fields 14-15 - **Rust core**: field struct, constants, Arrow metadata parsing, schema method - **Protobuf serialization**: round-trip support with backward compat - **Java JNI + LanceField**: constructor args and getters - **Python bindings + type stubs**: `is_unenforced_clustering_key()` / `unenforced_clustering_key_position()` - **Format docs**: clustering key metadata section ## Motivation This was discussed in the lance-spark SPJ PR (lance-format/lance-spark#445). Rather than using custom table properties, embedding clustering key info in the schema metadata follows the established pattern and avoids migration issues. ## Test plan - [x] `cargo check -p lance-core -p lance-file` passes - [x] `cargo test -p lance-core -p lance-file` passes (all tests including existing primary key tests) - [x] `cargo clippy -p lance-core -p lance-file --tests -- -D warnings` clean 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
## Summary Adds `RequiresDistributionAndOrdering` to `SparkWrite` so that when `lance.partition.columns` is set, Spark clusters and sorts the write data by the partition column(s). This ensures each Lance fragment contains exactly one distinct partition value — the prerequisite for Storage-Partitioned Joins (SPJ) on the read path. ### Changes - `SparkWrite` implements `RequiresDistributionAndOrdering` — returns `Distributions.clustered()` and ascending `SortOrder` for partition columns - `LanceConstant.TABLE_OPT_PARTITION_COLUMNS` centralizes the table property key - `LanceScanBuilder` uses `LanceConstant` directly (no duplicate constant) - Thread `tableProperties` through `SparkWriteBuilder` → `SparkWrite` ### Tests - Unit tests for `requiredDistribution()` / `requiredOrdering()` with and without partition columns - `BasePartitionedWriteTest` integration test verifying that fragments written with partition columns each contain exactly one distinct partition value 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
## 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]>
Summary
Adds
RequiresDistributionAndOrderingtoSparkWriteso that whenlance.partition.columnsis set, Spark clusters and sorts the write data by the partition column(s). This ensures each Lance fragment contains exactly one distinct partition value — the prerequisite for Storage-Partitioned Joins (SPJ) on the read path.Changes
SparkWriteimplementsRequiresDistributionAndOrdering— returnsDistributions.clustered()and ascendingSortOrderfor partition columnsLanceConstant.TABLE_OPT_PARTITION_COLUMNScentralizes the table property keyLanceScanBuilderusesLanceConstantdirectly (no duplicate constant)tablePropertiesthroughSparkWriteBuilder→SparkWriteTests
requiredDistribution()/requiredOrdering()with and without partition columnsBasePartitionedWriteTestintegration test verifying that fragments written with partition columns each contain exactly one distinct partition value🤖 Generated with Claude Code