Skip to content

feat: require clustered distribution on write for SPJ#445

Merged
jackye1995 merged 2 commits into
lance-format:mainfrom
beinan:feat/write-required-distribution
Apr 20, 2026
Merged

feat: require clustered distribution on write for SPJ#445
jackye1995 merged 2 commits into
lance-format:mainfrom
beinan:feat/write-required-distribution

Conversation

@beinan

@beinan beinan commented Apr 16, 2026

Copy link
Copy Markdown
Contributor

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 SparkWriteBuilderSparkWrite

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

@hamersaw hamersaw left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@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;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think we need to define this constant twice. Lets just use LanceConstant.TABLE_OPT_PARTITION_COLUMNS directly everywhere.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done — removed the duplicate constant. Now using LanceConstant.TABLE_OPT_PARTITION_COLUMNS directly in LanceScanBuilder.

@beinan

beinan commented Apr 20, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for the detailed review! You're right that Distributions.clustered() uses hash partitioning, so multiple distinct values can land in the same Spark task if they hash to the same bucket.

However, the combination of clustered distribution + sort ordering ensures that within each task, rows are grouped contiguously by partition value. Since each LanceDataWriter (one per task) creates fragments via Fragment.create(), and lance-core respects the sorted order when writing, the result is that each fragment contains exactly one partition value — as long as fragment size limits don't split mid-value.

I've added an integration test (BasePartitionedWriteTest.testWriteWithPartitionColumnClustersData) that verifies this end-to-end: it writes 50 rows across 5 regions in a single INSERT and asserts that each resulting fragment has exactly 1 distinct region value using _fragid.

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 jackye1995 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks like there is a merge conflict, once resolved I can merge it!

beinan and others added 2 commits April 20, 2026 18:17
…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]>
@beinan
beinan force-pushed the feat/write-required-distribution branch from d5c36e4 to e6a0d41 Compare April 20, 2026 18:21
@jackye1995
jackye1995 merged commit 24a4139 into lance-format:main Apr 20, 2026
16 checks passed
hamersaw added a commit that referenced this pull request Apr 21, 2026
## 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]>
sezruby added a commit to sezruby/lance-spark that referenced this pull request May 7, 2026
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]>
jackye1995 pushed a commit to lance-format/lance that referenced this pull request May 8, 2026
## 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]>
ivscheianu pushed a commit to ivscheianu/lance-spark that referenced this pull request Jul 20, 2026
## 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]>
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

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants