Skip to content

fix: name the index on uncommitted segment builds #793

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
geruh merged 3 commits into from
Sep 3, 2026
Merged

fix: name the index on uncommitted segment builds #793

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d392e75
fix: name the index on uncommitted segment builds
ivscheianu Aug 26, 2026
8d61777
docs: trim comments to match code, not PR history
ivscheianu Sep 2, 2026
3bfdfe6
chore: update comments
ivscheianu Sep 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 16 additions & 6 deletions ...base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -335,6 +335,7 @@ class RangeBasedBTreeIndexJob(

val indexBuilder = RangeBTreeIndexBuilder(
encode(readOptions),
addIndexExec.indexName,
columns,
zoneSize,
nsImpl,
Expand Down Expand Up @@ -364,6 +365,7 @@ class RangeBasedBTreeIndexJob(
* This class is serialized and sent to executors to build the index for a specific range of data.
*
* @param encodedReadOptions Serialized configuration for Lance dataset access.
* @param indexName Name of the logical index the segment will belong to.
* @param columns The names of the columns to be indexed.
* @param zoneSize Optional size of zones within the B-tree index.
* @param namespaceImpl Optional implementation class for namespace operations, used for credential vending.
Expand All @@ -374,6 +376,7 @@ class RangeBasedBTreeIndexJob(
*/
case class RangeBTreeIndexBuilder(
encodedReadOptions: String,
indexName: String,
columns: List[String],
zoneSize: Option[Long],
namespaceImpl: Option[String],
Expand Down Expand Up @@ -438,9 +441,10 @@ case class RangeBTreeIndexBuilder(
Data.exportArrayStream(allocator, reader, stream)

// Build an uncommitted BTree segment for this fragment group from the
// pre-sorted data. No index name or UUID is set: Lance generates the
// segment UUID, and the fragment ids declare the segment's coverage so
// the per-partition segments stay disjoint.
// pre-sorted data. No UUID is set: Lance generates the segment UUID, and

@geruh geruh Sep 2, 2026

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.

nit: can we clean this comment up while we are here we don't have to leave the story there. wdyt just something like:

// replace is for Lance's name check, and the driver commit still publishes.

// the fragment ids declare the segment's coverage so the per-partition
// segments stay disjoint. Named and replace-flagged to suppress Lance's
// collision pre-check against the existing index.
val btreeParamsBuilder = BTreeIndexParams.builder()
if (zoneSize.isDefined) {
btreeParamsBuilder.zoneSize(zoneSize.get)
Expand All @@ -451,7 +455,8 @@ case class RangeBTreeIndexBuilder(

val indexOptions = IndexOptions
.builder(columns.asJava, IndexType.BTREE, indexParams)
.replace(false)
.withIndexName(indexName)
.replace(true)
.withFragmentIds(fragmentIds.toList.asJava)
.withPreprocessedData(stream)
.build()
Expand Down Expand Up @@ -498,6 +503,7 @@ class ScalarSegmentIndexJob(
val tasks = fragmentBatches.map { batch =>
ScalarSegmentIndexTask(
encodedReadOptions,
addIndexExec.indexName,
columns,
addIndexExec.method,
argsJson,
Expand All @@ -519,10 +525,13 @@ class ScalarSegmentIndexJob(
final private[v2] case class FragmentWorkload(fragmentId: Integer, numRows: Long)

/**
* A task to create a scalar index segment on a batch of fragments.
* A task to create a scalar index segment on a batch of fragments. Named after the logical index

@geruh geruh Sep 2, 2026

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.

nit: we can drop the change here. the comment is enough for me to follow

* and replace-flagged to suppress Lance's collision pre-check against the existing index; actual
* replacement is handled by the driver's {@code commitExistingIndexSegments} transaction.
*/
case class ScalarSegmentIndexTask(
encodedReadOptions: String,
indexName: String,
columns: List[String],
method: String,
argsJson: String,
Expand All @@ -545,8 +554,9 @@ case class ScalarSegmentIndexTask(

val indexOptions = IndexOptions
.builder(java.util.Arrays.asList(columns: _*), indexType, params)
.withIndexName(indexName)
.withFragmentIds(fragmentIds.asJava)
.replace(false)
.replace(true)
.build()

val dataset = Utils.openDatasetBuilder(readOptions)
Expand Down
48 changes: 46 additions & 2 deletions lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -670,6 +670,50 @@ public void testRepeatedCreateZonemapIndexReplacesExistingSegments() {
}
}

@ParameterizedTest(name = "{0}")
@MethodSource("segmentBuildIndexMethods")
public void testRepeatedCreateIndexUnderLanceDefaultName(
String caseName, String method, String options) {
prepareDataset();

String sql =
String.format(
"alter table %s create index id_idx using %s (id) %s", fullTable, method, options);

spark.sql(sql);
checkIndex("id_idx");
spark.sql(sql);
checkIndex("id_idx");

org.lance.Dataset lanceDataset = org.lance.Dataset.open().uri(tableDir).build();
try {
int fragmentCount = lanceDataset.getFragments().size();
int coveredFragments =
lanceDataset.getIndexes().stream()
.filter(index -> "id_idx".equals(index.name()))
.map(index -> index.fragments().orElse(Collections.emptyList()).size())
.mapToInt(Integer::intValue)
.sum();
Assertions.assertEquals(
fragmentCount,
coveredFragments,
"Expected the recreated " + caseName + " segments to cover all fragments exactly once");
} finally {
lanceDataset.close();
}

// The index has to answer queries after the replacement, not merely exist in the manifest.
Dataset<Row> query = spark.sql(String.format("select * from %s where id=15", fullTable));
Assertions.assertEquals(1L, query.count());
Assertions.assertEquals("text_15", query.collectAsList().get(0).getString(1));
}

private static Stream<Arguments> segmentBuildIndexMethods() {
return Stream.of(
Arguments.of("zonemap", "zonemap", ""),
Arguments.of("btree-range", "btree", "with (build_mode = 'range')"));
}

@ParameterizedTest(name = "{0}")
@MethodSource("singleColumnIndexMethods")
public void testIndexesRejectMultipleColumns(String method, IndexType indexType) {
Expand Down Expand Up @@ -994,8 +1038,8 @@ public void testRepeatedCreateBTreeRangeIndex() {
firstRunUuids.size(),
"Expected one disjoint range segment per fragment on first create");

// Re-create with the same name: exercises replace(false) on the segment builds plus
// atomic replacement at commit time. The old segments must be replaced, not duplicated.
// Re-create with the same name: exercises the named segment builds plus atomic replacement at
// commit time. The old segments must be replaced, not duplicated.
spark.sql(sql);
checkIndex("test_range_repeat");

Expand Down
Loading