diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala index 085090996..e0bb09697 100755 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala @@ -335,6 +335,7 @@ class RangeBasedBTreeIndexJob( val indexBuilder = RangeBTreeIndexBuilder( encode(readOptions), + addIndexExec.indexName, columns, zoneSize, nsImpl, @@ -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. @@ -374,6 +376,7 @@ class RangeBasedBTreeIndexJob( */ case class RangeBTreeIndexBuilder( encodedReadOptions: String, + indexName: String, columns: List[String], zoneSize: Option[Long], namespaceImpl: Option[String], @@ -437,10 +440,7 @@ 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. + // replace is for Lance's name check, and the driver commit still publishes. val btreeParamsBuilder = BTreeIndexParams.builder() if (zoneSize.isDefined) { btreeParamsBuilder.zoneSize(zoneSize.get) @@ -451,7 +451,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() @@ -498,6 +499,7 @@ class ScalarSegmentIndexJob( val tasks = fragmentBatches.map { batch => ScalarSegmentIndexTask( encodedReadOptions, + addIndexExec.indexName, columns, addIndexExec.method, argsJson, @@ -523,6 +525,7 @@ final private[v2] case class FragmentWorkload(fragmentId: Integer, numRows: Long */ case class ScalarSegmentIndexTask( encodedReadOptions: String, + indexName: String, columns: List[String], method: String, argsJson: String, @@ -545,8 +548,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) diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java index 711292373..5d6dea51b 100755 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java @@ -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 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 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) { @@ -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");