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..541a9fb92 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 @@ -88,7 +88,8 @@ case class AddIndexExec( val btreeBuildMode = IndexUtils.btreeBuildMode(indexType, args) val scalarSegmentIndexType = IndexUtils.scalarSegmentIndexType(method) - val (fragmentWorkloads, canonicalColumns) = { + // Plan and build at one pinned version; the commit opens the live dataset. + val (fragmentWorkloads, canonicalColumns, buildReadOptions) = { val ds = Utils.openDatasetBuilder(readOptions).build() try { val canonical = columns.map { column => @@ -96,10 +97,9 @@ case class AddIndexExec( FieldPathUtils.pathByFieldId(ds.getLanceSchema, field.getId) } ( - ds.getFragments.asScala - .map(fragment => FragmentWorkload(fragment.getId, fragment.metadata().getNumRows)) - .toList, - canonical) + IndexUtils.fragmentWorkloads(ds), + canonical, + IndexUtils.pinVersion(readOptions, ds)) } finally { ds.close() } @@ -163,19 +163,18 @@ case class AddIndexExec( val (nsImpl, nsProps, tableId, initialStorageOpts) = extractNamespaceInfo(lanceDataset, readOptions) - // Range-mode BTree uses preprocessed data from Spark and keeps its dedicated path. if (btreeBuildMode.contains("range")) { val segments = new RangeBasedBTreeIndexJob( this.copy(columns = canonicalColumns), - readOptions, + buildReadOptions, fragmentIds.size, nsImpl, nsProps, tableId, initialStorageOpts).run() - commitIndexSegments(readOptions, canonicalColumns.head, segments) + val indexed = commitIndexSegments(readOptions, canonicalColumns.head, segments) return Seq(new GenericInternalRow(Array[Any]( - fragmentIds.size.toLong, + indexed.toLong, UTF8String.fromString(indexName)))) } @@ -183,7 +182,7 @@ case class AddIndexExec( if (scalarSegmentIndexType.isDefined) { val segmentJob = new ScalarSegmentIndexJob( this.copy(columns = canonicalColumns), - readOptions, + buildReadOptions, fragmentWorkloads, validatedNumSegments, nsImpl, @@ -192,9 +191,9 @@ case class AddIndexExec( initialStorageOpts) val segments = segmentJob.run() // Atomic add+remove via Lance core; see commitIndexSegments - commitIndexSegments(readOptions, canonicalColumns.head, segments) + val indexed = commitIndexSegments(readOptions, canonicalColumns.head, segments) return Seq(new GenericInternalRow(Array[Any]( - fragmentIds.size.toLong, + indexed.toLong, UTF8String.fromString(indexName)))) } @@ -231,16 +230,31 @@ case class AddIndexExec( // Lance core's commitExistingIndexSegments handles atomic replacement: // it finds existing segments whose fragments overlap with incoming ones // and removes them in the same CreateIndex transaction. + // + // Returns the number of fragments the commit actually covers, which is what the command reports. private def commitIndexSegments( readOptions: LanceSparkReadOptions, column: String, - segments: Seq[Index]): Unit = { + segments: Seq[Index]): Int = { val dataset = Utils.openDatasetBuilder(readOptions).build() try { - dataset.commitExistingIndexSegments( + IndexUtils.requireCommittableCoverage( + IndexUtils.liveFragmentIds(dataset), + segments, + indexName) + val committed = dataset.commitExistingIndexSegments( indexName, column, segments.toList.asJava) + // The commit advances this handle to the manifest it wrote, so both the returned metadata and + // the fragment list below describe the committed state rather than the one validated above. + IndexUtils + .establishedCoverage( + segments, + committed.asScala.toSeq, + IndexUtils.liveFragmentIds(dataset), + indexName) + .size } finally { dataset.close() } @@ -295,6 +309,17 @@ class RangeBasedBTreeIndexJob( private val VALUE_COLUMN_NAME = "value" + /** Version `readOptions` is pinned to. Throws if the ref is not a version on main. */ + private def pinnedVersion: Long = { + val ref = readOptions.getRef + if (ref == null || !ref.isMain || !ref.getVersionNumber.isPresent) { + throw new IllegalStateException( + "Range-mode BTree builds need read options pinned to a version on main so the scan and the " + + "segment stamp describe one snapshot; got a ref that names none") + } + ref.getVersionNumber.get.longValue() + } + def run(): Seq[Index] = { if (addIndexExec.columns.size != 1) { throw new UnsupportedOperationException( @@ -316,9 +341,10 @@ class RangeBasedBTreeIndexJob( } val fullTableName = parts.mkString(".") - // Read the indexed column with the row id and fragment id metadata columns. val fragmentColumn = LanceDataset.FRAGMENT_ID_COLUMN.name - val df = session.table(fullTableName) + val df = session.read + .option(LanceSparkReadOptions.CONFIG_VERSION, pinnedVersion) + .table(fullTableName) val selectDf = df.select( df.col(columns.head).as(VALUE_COLUMN_NAME), df.col(LanceDataset.ROW_ID_COLUMN.name), @@ -619,6 +645,89 @@ object IndexUtils extends Logging { .get(method.toLowerCase(Locale.ROOT)) .filter(scalarSegmentIndexTypes.contains) + /** Pins `readOptions` to the version `dataset` is open at. */ + def pinVersion( + readOptions: LanceSparkReadOptions, + dataset: Dataset): LanceSparkReadOptions = + readOptions.withRef(Utils.pinOpenedRef(dataset, readOptions.getRef)) + + /** + * Fragment ids and live row counts of `dataset`, in manifest order. + * + * Reads the primitive fragment-statistics view rather than [[Dataset#getFragments]], which + * materializes a Java object per fragment and per data file. Both commands enumerate fragments on + * the driver, once to plan and once to check the commit, so on a large table that difference is + * the bulk of planning cost. + */ + def fragmentWorkloads(dataset: Dataset): List[FragmentWorkload] = { + val stats = dataset.getFragmentStatistics + val ids = stats.getIds + val rowCounts = stats.getRowCounts + List.tabulate(ids.length)(index => + FragmentWorkload(Integer.valueOf(ids(index)), rowCounts(index))) + } + + /** Fragment ids live in `dataset`. See [[fragmentWorkloads]] for why this avoids getFragments. */ + def liveFragmentIds(dataset: Dataset): Set[Int] = + dataset.getFragmentStatistics.getIds.toSet + + /** Fragment ids the given segments declare coverage of. */ + def declaredCoverage(segments: Seq[Index]): Set[Int] = + segments.iterator + .flatMap(_.fragments().orElse(Collections.emptyList[Integer]()).asScala) + .map(_.intValue) + .toSet + + /** + * Refuses a segment set that declares coverage but intersects no live fragment. Partial loss is + * accepted: fragments can be retired by compaction or deletion during the build, and the segments + * covering what remains are still correct. [[establishedCoverage]] reports what was actually + * committed. + */ + def requireCommittableCoverage( + liveFragmentIds: Set[Int], + segments: Seq[Index], + indexName: String): Unit = { + val declared = declaredCoverage(segments) + if (declared.nonEmpty && declared.intersect(liveFragmentIds).isEmpty) { + throw new IllegalStateException( + s"Index '$indexName' build raced a concurrent operation: every fragment it covers " + + s"(${describeFragmentIds(declared)}) was retired while the build ran, so the segments " + + "would cover nothing. No index change was committed; re-run the command.") + } + } + + /** + * Committed fragment ids from the segments this build produced, intersected with what is still + * live. Only segments whose UUID matches `builtSegments` are counted; pre-existing survivors are + * not this command's coverage. + */ + def establishedCoverage( + builtSegments: Seq[Index], + committedSegments: Seq[Index], + liveFragmentIds: Set[Int], + indexName: String): Set[Int] = { + val builtUuids = builtSegments.map(_.uuid).toSet + val established = + declaredCoverage(committedSegments.filter(segment => builtUuids.contains(segment.uuid))) + .intersect(liveFragmentIds) + val uncovered = declaredCoverage(builtSegments).diff(established) + if (uncovered.nonEmpty) { + logWarning( + s"Index '$indexName' build raced a concurrent operation: fragments " + + s"${describeFragmentIds(uncovered)} are not covered by this commit, because they were " + + "retired or because their indexed field was rewritten while the build ran. The segments " + + "for the remaining fragments are committed; re-run the command to cover them.") + } + established + } + + private def describeFragmentIds(ids: Set[Int]): String = { + val ordered = ids.toSeq.sorted + val shown = ordered.take(10).mkString(", ") + if (ordered.size > 10) s"$shown, ... (${ordered.size} total)" else shown + } + def resolveIndexField( schema: LanceSchema, indexType: IndexType, diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/branch/BaseBranchDDLTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/branch/BaseBranchDDLTest.java index 22d202a21..3892c3aab 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/branch/BaseBranchDDLTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/branch/BaseBranchDDLTest.java @@ -590,6 +590,23 @@ public void testBranchAndVersionOptionsFail() { Assertions.assertTrue(conflictMessages.contains("version")); } + @Test + public void testVersionOptionPinsTheScanToThatVersion() { + DatasetVersions versions = prepareDatasetWithHistory(); + Assertions.assertNotEquals(versions.firstInsertVersion, versions.latestVersion); + + Assertions.assertEquals( + 5, + spark + .read() + .option("version", Long.toString(versions.firstInsertVersion)) + .table(fullTable) + .count(), + "reading with the version option must see only the rows present at that version"); + Assertions.assertEquals( + 10, spark.table(fullTable).count(), "the unpinned read still sees the latest version"); + } + @Test public void testBranchIdentifierRejectsVersionAsOf() { DatasetVersions versions = prepareDatasetWithHistory(); 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..15e56e24b 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 @@ -17,8 +17,11 @@ import org.lance.index.Index; import org.lance.index.IndexCriteria; import org.lance.index.IndexDescription; +import org.lance.index.IndexOptions; +import org.lance.index.IndexParams; import org.lance.index.IndexType; import org.lance.index.OptimizeOptions; +import org.lance.index.scalar.ScalarIndexParams; import org.lance.ipc.FullTextQuery; import org.lance.ipc.LanceScanner; import org.lance.ipc.ScanOptions; @@ -105,6 +108,64 @@ public void tearDown() throws IOException { } } + @Test + public void testSegmentCommitReportsCoverageAsCommitted() { + spark.sql(String.format("create table %s (id int) using lance", fullTable)); + spark.sql(String.format("insert into %s values (0), (1), (2)", fullTable)); + spark.sql(String.format("insert into %s values (3), (4), (5)", fullTable)); + + try (org.lance.Dataset committer = + Utils.openDatasetBuilder(LanceSparkReadOptions.from(tableDir)).build()) { + List fragments = committer.getFragments(); + int coveredFragmentId = fragments.get(fragments.size() - 1).getId(); + + IndexParams indexParams = + IndexParams.builder() + .setScalarIndexParams(ScalarIndexParams.create("zonemap", "{}")) + .build(); + Index built = + committer.createIndex( + IndexOptions.builder(Collections.singletonList("id"), IndexType.ZONEMAP, indexParams) + .withIndexName("idx_committed_coverage") + .replace(true) + .withFragmentIds(Collections.singletonList(coveredFragmentId)) + .build()); + Assertions.assertEquals( + Collections.singletonList(coveredFragmentId), + built.fragments().orElse(Collections.emptyList()), + "the uncommitted segment should declare the fragment it was built for"); + + // Retire every fragment from another handle, leaving the committer on a stale manifest. + spark.sql(String.format("delete from %s where id >= 0", fullTable)); + + List committed = + committer.commitExistingIndexSegments( + "idx_committed_coverage", "id", Collections.singletonList(built)); + + Set ours = Collections.singleton(built.uuid()); + Assertions.assertTrue( + committed.stream().anyMatch(index -> ours.contains(index.uuid())), + "the commit must return the metadata of the segments it was handed"); + + Set liveAfter = + committer.getFragments().stream().map(Fragment::getId).collect(Collectors.toSet()); + Assertions.assertFalse( + liveAfter.contains(coveredFragmentId), + "the committing handle must advance to the manifest the commit wrote"); + + Set established = + committed.stream() + .filter(index -> ours.contains(index.uuid())) + .flatMap(index -> index.fragments().orElse(Collections.emptyList()).stream()) + .filter(liveAfter::contains) + .collect(Collectors.toSet()); + Assertions.assertEquals( + Collections.emptySet(), + established, + "coverage read from the committed state must not count a fragment retired in between"); + } + } + private void prepareDataset() { spark.sql(String.format("create table %s (id int, text string) using lance;", fullTable)); // First insert to create initial fragments diff --git a/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala b/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala index 9644e6f1c..85fab3bd0 100644 --- a/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala +++ b/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala @@ -16,7 +16,9 @@ package org.apache.spark.sql.execution.datasources.v2 import org.apache.spark.sql.catalyst.plans.logical.LanceNamedArgument import org.junit.jupiter.api.Assertions._ import org.junit.jupiter.api.Test -import org.lance.index.IndexType +import org.lance.index.{Index, IndexType} + +import scala.collection.JavaConverters._ /** * Unit tests for [[IndexUtils]] helper methods. @@ -26,6 +28,23 @@ import org.lance.index.IndexType */ class IndexUtilsTest { + /** An index segment carrying only the metadata these helpers read. */ + private def segment( + fragmentIds: Option[Seq[Int]], + indexDetails: Option[Array[Byte]] = None): Index = { + val builder = Index + .builder() + .uuid(java.util.UUID.randomUUID()) + .name("idx_id") + .indexType(IndexType.INVERTED) + fragmentIds.foreach(ids => + builder.fragments(ids.map(java.lang.Integer.valueOf).asJava)) + indexDetails.foreach(builder.indexDetails) + builder.build() + } + + private def coveringSegment(fragmentIds: Int*): Index = segment(Some(fragmentIds)) + private def fragmentWorkloads(rows: Long*): List[FragmentWorkload] = rows.zipWithIndex.map { case (rowCount, fragmentId) => FragmentWorkload(java.lang.Integer.valueOf(fragmentId), rowCount) @@ -257,4 +276,121 @@ class IndexUtilsTest { classOf[ArithmeticException], () => IndexUtils.batchFragments(fragmentWorkloads(Long.MaxValue, 1), Some(1), 1)) } + + @Test + def declaredCoverage_unionsSegmentBitmaps(): Unit = { + assertEquals( + Set(0, 1, 4), + IndexUtils.declaredCoverage(Seq(coveringSegment(0, 1), coveringSegment(4)))) + } + + @Test + def declaredCoverage_treatsAbsentBitmapAsNoCoverage(): Unit = { + assertEquals(Set(2), IndexUtils.declaredCoverage(Seq(segment(None), coveringSegment(2)))) + assertEquals(Set.empty[Int], IndexUtils.declaredCoverage(Seq.empty)) + } + + /** A segment keeping its identity but reporting narrower coverage, as a pruned commit returns. */ + private def prunedTo(built: Index, fragmentIds: Seq[Int]): Index = + Index + .builder() + .uuid(built.uuid) + .name(built.name) + .indexType(built.indexType) + .fragments(fragmentIds.map(java.lang.Integer.valueOf).asJava) + .build() + + @Test + def requireCommittableCoverage_acceptsCoverageThatSurvives(): Unit = { + IndexUtils.requireCommittableCoverage(Set(1, 5), Seq(coveringSegment(0, 1)), "idx_id") + } + + @Test + def requireCommittableCoverage_acceptsSegmentsThatDeclareNothing(): Unit = { + IndexUtils.requireCommittableCoverage(Set(1), Seq(segment(None)), "idx_id") + IndexUtils.requireCommittableCoverage(Set.empty[Int], Seq.empty, "idx_id") + } + + @Test + def requireCommittableCoverage_refusesASetThatWouldCoverNothing(): Unit = { + val error = assertThrows( + classOf[IllegalStateException], + () => IndexUtils.requireCommittableCoverage(Set(5), Seq(coveringSegment(0, 7)), "idx_id")) + + assertTrue(error.getMessage.contains("idx_id"), error.getMessage) + assertTrue(error.getMessage.contains("0, 7"), error.getMessage) + assertTrue(error.getMessage.contains("re-run"), error.getMessage) + } + + @Test + def requireCommittableCoverage_summarizesLargeRetiredSets(): Unit = { + val error = assertThrows( + classOf[IllegalStateException], + () => + IndexUtils.requireCommittableCoverage( + Set.empty[Int], + Seq(coveringSegment(0 to 20: _*)), + "idx_id")) + + assertTrue(error.getMessage.contains("21 total"), error.getMessage) + } + + @Test + def establishedCoverage_reportsWhatTheCommitReturned(): Unit = { + val built = Seq(coveringSegment(0, 1), coveringSegment(2)) + + assertEquals( + Set(0, 1, 2), + IndexUtils.establishedCoverage(built, built, Set(0, 1, 2), "idx_id")) + } + + @Test + def establishedCoverage_followsAPrunedCommitEvenWhileEveryFragmentIsLive(): Unit = { + val built = coveringSegment(0, 1) + + assertEquals( + Set(0), + IndexUtils.establishedCoverage( + Seq(built), + Seq(prunedTo(built, Seq(0))), + Set(0, 1), + "idx_id")) + } + + @Test + def establishedCoverage_excludesFragmentsRetiredByTheCommit(): Unit = { + val built = coveringSegment(0, 1) + + assertEquals( + Set(1), + IndexUtils.establishedCoverage(Seq(built), Seq(built), Set(1, 5), "idx_id")) + } + + @Test + def establishedCoverage_countsOnlyTheSegmentsThisBuildProduced(): Unit = { + val built = coveringSegment(2) + val survivor = coveringSegment(0, 1) + + assertEquals( + Set(2), + IndexUtils.establishedCoverage( + Seq(built), + Seq(survivor, built), + Set(0, 1, 2), + "idx_id")) + } + + @Test + def establishedCoverage_isEmptyWhenNothingOfThisBuildSurvived(): Unit = { + val built = coveringSegment(0) + + assertEquals( + Set.empty[Int], + IndexUtils.establishedCoverage( + Seq(built), + Seq(prunedTo(built, Seq.empty)), + Set(0), + "idx_id")) + } + }