diff --git a/.bumpversion.toml b/.bumpversion.toml index 82e0d1485..2043101ea 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -85,6 +85,11 @@ filename = "lance-spark-4.2_2.13/pom.xml" search = "{current_version}" replace = "{new_version}" +[[tool.bumpversion.files]] +filename = "lance-spark-knn-4.2_2.13/pom.xml" +search = "{current_version}" +replace = "{new_version}" + # Bundle module pom.xml files - parent version [[tool.bumpversion.files]] filename = "lance-spark-bundle-3.4_2.12/pom.xml" diff --git a/lance-spark-knn-4.2_2.13/pom.xml b/lance-spark-knn-4.2_2.13/pom.xml new file mode 100644 index 000000000..dfdcdafad --- /dev/null +++ b/lance-spark-knn-4.2_2.13/pom.xml @@ -0,0 +1,112 @@ + + + 4.0.0 + + + org.lance + lance-spark-root + 0.8.0-beta.1 + ../pom.xml + + + lance-spark-knn-4.2_2.13 + ${project.artifactId} + Indexed nearest-by join on Lance — Spark 4.2 SQL Catalyst integration (SPARK-56395) + jar + + + ${scala213.version} + ${scala213.compat.version} + ${arrow19.version} + ${java17.release} + + + + + org.apache.spark + spark-sql_${scala.compat.version} + ${spark42.version} + provided + + + org.apache.spark + spark-catalyst_${scala.compat.version} + ${spark42.version} + provided + + + + org.lance + lance-spark-base_${scala.compat.version} + ${project.version} + + + org.apache.arrow + arrow-memory-netty-buffer-patch + + + + + + org.lance + lance-spark-4.2_${scala.compat.version} + ${project.version} + test + + + + + + + + net.alchim31.maven + scala-maven-plugin + ${scala-maven-plugin.version} + + + scala-compile-first + process-resources + + compile + + + + scala-test-compile + process-test-resources + + testCompile + + + + + + -feature + -release + ${java.release} + + + + + org.apache.maven.plugins + maven-compiler-plugin + ${maven-compiler-plugin.version} + + ${java.release} + + + + + + + java21 + + 21 + + + + diff --git a/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/internal/LanceProbe.scala b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/internal/LanceProbe.scala new file mode 100644 index 000000000..5dd557f5e --- /dev/null +++ b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/internal/LanceProbe.scala @@ -0,0 +1,700 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.knn.internal + +import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.vector.{BigIntVector, FieldVector, Float4Vector, Float8Vector, UInt8Vector, VectorSchemaRoot} +import org.apache.arrow.vector.ipc.ArrowReader +import org.apache.spark.sql.catalyst.CatalystTypeConverters +import org.apache.spark.sql.types.{DataType, StructField} +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} +import org.lance.Dataset +import org.lance.ipc.{LanceScanner, Query, ScanOptions} +import org.lance.spark.{LanceConstant, LanceRef, LanceRuntime, LanceSparkReadOptions} +import org.lance.spark.utils.Utils +import org.lance.spark.vectorized.LanceArrowColumnVector + +import java.util + +import scala.collection.JavaConverters._ +import scala.collection.mutable + +/** + * Per-task vector-index probe primitive. Opens a Lance dataset once and serves many queries against + * a fixed set of fragments. Two query shapes: + * - [[probe]] returns row references + scores only (no payload); the payload is fetched later via + * [[materialize]]. This split late-materialization is the right shape when the caller OVER-FETCHES + * candidates and trims before materializing — only the survivors are paid for. + * - [[probeRows]] folds the nearest search AND the payload projection into a SINGLE scan. This is + * the right shape when there is NO JVM-side over-fetch (every probed row is kept), so a separate + * materialize scan would just re-fetch the exact rows the search already found. + * + * This is the core primitive Phase 0 of the indexed nearest-by design depends on. Validating its + * cost profile is the first thing to do on a new Lance build: + * - dataset open should be one-time cost + * - per-probe cost should be index traversal + small overhead, not full fragment scan + * - returning top-K row addrs should match Lance's native nearest search recall + * + * Lifecycle: instantiate per task, call `probe(...)` repeatedly, close at end. + * + * @param readOptions Fully resolved Lance read context — dataset URI, storage + * credentials, catalog / cache backends, and the pinned branch / + * version ref. Inside a join this is resolved and pinned once on the + * driver (see [[LanceKnnJoinStage.resolveReadContext]]) so every task + * probes the same snapshot. + * @param initialStorageOptions Driver-side storage options from `namespace.describeTable()`, merged + * into the base storage options at open. May be `null`. + * @param namespaceImpl Namespace implementation type, for reconnecting the namespace on an + * executor where the live handle did not survive serialization. May be + * `null` for a plain URI dataset. + * @param namespaceProperties Namespace connection properties for that reconnection. May be `null`. + * @param fragmentIds Fragments this probe is restricted to. Pass `None` for whole-dataset + * search. + * @param allocator Arrow allocator. Defaults to lance-spark's shared + * `LanceRuntime.allocator()`. + */ +final class LanceProbe( + readOptions: LanceSparkReadOptions, + initialStorageOptions: java.util.Map[String, String], + namespaceImpl: String, + namespaceProperties: java.util.Map[String, String], + fragmentIds: Option[Seq[Int]], + allocator: BufferAllocator = LanceRuntime.allocator()) + extends AutoCloseable { + + /** + * URI + optional pinned-version convenience constructor. Builds read options straight from a bare + * dataset URI and pins `version` on the main branch when present. For callers / tests that have + * only a plain URI and no namespace or storage context. + */ + def this(datasetUri: String, fragmentIds: Option[Seq[Int]], version: Option[Long]) = + this(LanceProbe.readOptionsFor(datasetUri, version), null, null, null, fragmentIds) + + def this(datasetUri: String, fragmentIds: Option[Seq[Int]]) = + this(datasetUri, fragmentIds, None) + + // Open the dataset once. Lance's Java binding caches index metadata against the Dataset handle, + // so reusing it across probes keeps subsequent calls index-warm. Opening through + // `Utils.openDatasetBuilder` (rather than a bare `Dataset.open().uri(...)`) makes the probe honor + // the full resolved read context — storage credentials, catalog / cache backends, the pinned + // branch / version ref, and (on executors, where the live namespace handle is transient) the + // runtime-namespace reconnection — exactly as the connector's own scan path does. + private val dataset: Dataset = openDataset() + + private val javaFragmentIds: Option[util.List[Integer]] = fragmentIds.map { ids => + val javaList = new util.ArrayList[Integer](ids.size) + ids.foreach(i => javaList.add(Integer.valueOf(i))) + javaList: util.List[Integer] + } + + private def openDataset(): Dataset = { + // Apply the connector's exact worker-open namespace policy (`LanceFragmentScanner.create`): + // touch the namespace ONLY when a namespace impl is configured AND executor credential refresh + // is enabled. When refresh is on, either rebuild the namespace client (impls that must run on + // workers) or clear it (impls that must not, so the open falls back to the URI + initial storage + // options). When refresh is OFF we leave the read options' namespace exactly as shipped — this + // is the whole point of `executor_credential_refresh=false`, and forcing a rebuild here would + // turn that policy into a namespace class-load / RPC that can fail before the dataset even opens. + // + // Notably this does NOT call `builder.runtimeNamespace(namespaceImpl, ...)` unconditionally + // (which always loads the namespace impl class) — matching the fragment scanner, which reaches + // the namespace only through the guarded `setNamespace` path above. + if (namespaceImpl != null && readOptions.isExecutorCredentialRefresh()) { + if (LanceRuntime.useNamespaceOnWorkers(namespaceImpl)) { + readOptions.setNamespace( + LanceRuntime.getOrCreateNamespace(namespaceImpl, namespaceProperties)) + } else { + readOptions.setNamespace(null) + } + } + val builder = Utils.openDatasetBuilder(readOptions) + if (initialStorageOptions != null) { + builder.initialStorageOptions(initialStorageOptions) + } + builder.build() + } + + // The dataset's own top-level column names, read once from the open handle. Used by the schema + // eligibility backstop below; cheap (metadata only) and stable for the probe's lifetime. + private lazy val datasetColumnNames: Seq[String] = + dataset.getSchema.getFields.asScala.map(_.getName).toSeq + + /** + * Reject a dataset whose own schema collides with a column the nearest scan injects + * (`_rowid` / `_distance` / `_score`). Every probe route ([[probe]] and [[probeRows]], fused or + * split) runs a `nearest` scan that injects those columns, so such a table cannot be served by ANY + * indexed route: the injected metadata shadows the physical column and it is read out-of-band as + * the ranking score (silently dropped from an all-columns payload) or collides outright when + * projected. There is no fold-vs-split routing that recovers it. The Catalyst rule is expected to + * DECLINE the indexed rewrite for such a table via [[LanceProbe.schemaSupportsNearest]] and fall + * back to the engine's default nearest-by execution; this is the defensive backstop for any caller + * that reached the probe anyway. + */ + private def requireNearestCompatibleSchema(): Unit = { + val collisions = LanceProbe.reservedSchemaColumns(datasetColumnNames) + require( + collisions.isEmpty, + s"Lance dataset schema has column(s) ${collisions.toSeq.sorted.mkString(", ")} whose name(s) " + + "collide with the metadata a nearest scan injects (_rowid, _distance, _score). No indexed " + + "probe can serve this table — the injected metadata shadows the physical column. Decline the " + + "indexed rewrite (see LanceProbe.schemaSupportsNearest) and fall back to default nearest-by " + + "execution.") + } + + /** + * Run a single nearest-neighbor query. Returns up to `k` row references for the configured + * fragments, ordered best-first by `metric`. + * + * Implementation note: lance-spark mandates `prefilter = true` for fragmented vector queries + * (see `LanceFragmentScanner.create`). We mirror that here — Lance's index probe semantics + * require it when fragment scope is restricted. + * + * `vectorColumn` is a per-call argument (not a constructor field) because the same + * `LanceProbe` instance also serves the materialize stage via [[materialize]], which + * doesn't reference any vector column. Keeping it on the call sidesteps the smell of + * passing a placeholder string when constructing for materialize-only use. + * + * `prefilter` is a Lance SQL filter string (DataFusion-flavored). Lance applies it BEFORE the + * vector index lookup when `prefilter = true` (which we always set), so the top-K is computed + * over only the rows matching the filter — exactly what a `Filter(cond, lance) RIGHT JOIN ... + * APPROX NEAREST K` should do. Without prefilter pushdown, a per-fragment vector probe could + * return K rows that are all later filtered out post-join, masking truly-nearest-but-also- + * matching rows further down the index — a recall bug. The translator in + * `IndexedNearestByJoinRule` is responsible for producing only safely-translated SQL; here we + * just hand it through. + */ + def probe( + vectorColumn: String, + query: Array[Float], + k: Int, + metric: Metric, + nprobes: Option[Int] = None, + refineFactor: Option[Int] = None, + ef: Option[Int] = None, + prefilter: Option[String] = None): Seq[ScoredRowRef] = { + require(vectorColumn != null && vectorColumn.nonEmpty, "vectorColumn must be non-empty") + require(query != null && query.length > 0, "Query vector must be non-empty") + require(k > 0, "k must be positive") + requireNearestCompatibleSchema() + + val q = buildNearestQuery(vectorColumn, query, k, metric, nprobes, refineFactor, ef) + + val opts = new ScanOptions.Builder() + .nearest(q) + // EXPERIMENT: drop prefilter(true). The single-machine reference path + // doesn't set it; this LanceProbe call does. Comparing wallclock with + // and without isolates whether the prefilter branch in + // vector_search_source forces a slower index plan than the postfilter + // (default) branch. Re-enable when fragmented probe + prefilter + // pushdown is needed (we know fragmentIds requires prefilter from the + // Lance-side error, but at probeParallelism=1 there are no fragments). + .withRowId(true) + // Project only what we need into the result. The vector column is implied by `nearest`; + // requesting an empty user column list keeps the Arrow batch narrow (just the rowid + + // distance metadata). Materialization fetches payload columns later. + .columns(java.util.Collections.emptyList[String]()) + + if (prefilter.nonEmpty || javaFragmentIds.nonEmpty) { + // Real prefilter or fragment scope is requested — keep prefilter(true) + // so Lance applies the filter / restricts to fragments correctly. + opts.prefilter(true) + } + + prefilter.filter(_.nonEmpty).foreach(opts.filter) + javaFragmentIds.foreach(opts.fragmentIds) + + val scanner: LanceScanner = LanceScanner.create(dataset, opts.build(), allocator) + try { + readScored(scanner.scanBatches()) + } finally { + scanner.close() + } + } + + /** + * Nearest search AND payload projection in a SINGLE scan. Runs the top-`k` search projecting the + * requested payload `projection` columns directly into the result batch, so each hit comes back + * with its row-id, ranking score, AND materialized payload — no second point-fetch scan. + * + * This is the no-overfetch fast path. When the caller does not over-fetch at the JVM level (the + * SQL path: `internalK == k`), every probed row is kept, so the split [[probe]] → trim → + * [[materialize]] would just re-scan the exact rows the search already found — a second Lance scan + * plus a `k`-element `_rowid IN (arrow_cast …)` filter parse, per query, for nothing. Folding the + * two removes both. Rows come back best-first (Lance's native ordering), same as [[probe]]. + * + * Use [[probe]] + [[materialize]] instead when the caller OVER-fetches candidates and trims before + * materializing — there, deferring the payload fetch to only the survivors is the win. `projection` + * / `projectionFields` follow [[materialize]]'s contract: an empty `projection` means "all columns", + * and each projected cell is converted to the Spark EXTERNAL value its declared type expects (see + * [[readRows]]). + */ + def probeRows( + vectorColumn: String, + query: Array[Float], + k: Int, + metric: Metric, + projection: Seq[String], + projectionFields: Seq[StructField], + nprobes: Option[Int] = None, + refineFactor: Option[Int] = None, + ef: Option[Int] = None, + prefilter: Option[String] = None): Seq[MaterializedHit] = { + require(vectorColumn != null && vectorColumn.nonEmpty, "vectorColumn must be non-empty") + require(query != null && query.length > 0, "Query vector must be non-empty") + require(k > 0, "k must be positive") + // The nearest scan injects `_rowid` (via withRowId) and the score columns (via nearest). If the + // dataset's own schema has a column by one of those names the injected metadata shadows it, and + // NO projection shape recovers the physical column: an empty (all-columns) projection reads it + // out-of-band as the ranking score and silently drops it from the payload, while an explicit + // projection of it collides inside the scan. Routing to the split path does not help — its scan + // also injects `_rowid`. So the eligibility is schema-level, not projection-level: reject the + // whole table. The Catalyst rule declines such a table up front via + // [[LanceProbe.schemaSupportsNearest]]; this is the defensive backstop. + requireNearestCompatibleSchema() + + val q = buildNearestQuery(vectorColumn, query, k, metric, nprobes, refineFactor, ef) + + val opts = new ScanOptions.Builder() + .nearest(q) + .withRowId(true) + // Project the payload columns into the nearest scan itself. `_distance` is added by `nearest` + // regardless (Lance includes it even when explicit columns omit it), so the drain still finds a + // score column. An empty projection leaves columns unset → all columns, matching `materialize`. + if (projection.nonEmpty) { + opts.columns(projection.toList.asJava) + } + + if (prefilter.nonEmpty || javaFragmentIds.nonEmpty) { + opts.prefilter(true) + } + prefilter.filter(_.nonEmpty).foreach(opts.filter) + javaFragmentIds.foreach(opts.fragmentIds) + + val scanner: LanceScanner = LanceScanner.create(dataset, opts.build(), allocator) + try { + readScoredRows(scanner.scanBatches(), projectionFields) + } finally { + scanner.close() + } + } + + /** + * Build the Lance nearest-neighbor query. Shared by [[probe]] (refs only) and [[probeRows]] + * (refs + folded payload) so both search the index identically. + */ + private def buildNearestQuery( + vectorColumn: String, + query: Array[Float], + k: Int, + metric: Metric, + nprobes: Option[Int], + refineFactor: Option[Int], + ef: Option[Int]): Query = { + val b = new Query.Builder() + .setColumn(vectorColumn) + .setKey(query) + .setK(k) + .setDistanceType(metric.lanceType) + nprobes.foreach(b.setNprobes(_)) + // refineFactor: IVF-PQ recall knob. Lance fetches `k * refineFactor` approximate + // candidates, then re-ranks them with exact distance and trims to k. Bigger factor = + // better recall, more compute. None leaves Lance's default (= 1, no re-rank). + refineFactor.foreach(b.setRefineFactor(_)) + // ef: HNSW search depth. Higher = better recall, more compute. None leaves Lance's + // index-default. Only meaningful for HNSW indexes; ignored for IVF-PQ. + ef.foreach(b.setEf(_)) + b.build() + } + + /** + * Drain the Arrow stream from a nearest-search scan into `(rowId, score)` pairs. + * + * Expected schema: + * - `_rowid` : UInt8 / BigInt — Lance logical row identifier + * - `_distance` (or score column added by `nearest`) : Float4 / Float8 — ranking value + * + * We resolve columns by name to be encoding-version-agnostic; the underlying primitive type + * (UInt8 vs BigInt for the id, Float4 vs Float8 for score) varies across Arrow / Lance combos + * and we tolerate both. + */ + private def readScored(reader: ArrowReader): Seq[ScoredRowRef] = { + val out = mutable.ArrayBuffer.empty[ScoredRowRef] + try { + while (reader.loadNextBatch()) { + val root = reader.getVectorSchemaRoot + val addrVec: FieldVector = root.getVector(LanceProbe.RowIdColumn) + val scoreVec: FieldVector = resolveScoreVector(root) + + val n = root.getRowCount + var i = 0 + while (i < n) { + out += ScoredRowRef(rowAddrAt(addrVec, i), scoreAt(scoreVec, i)) + i += 1 + } + } + } finally { + reader.close() + } + out.toSeq + } + + /** + * Drain a nearest-search scan that ALSO projected payload columns into `(rowId, score, payload)` + * hits — the folded counterpart of [[readScored]] + [[readRows]]. Row-id and score are read out of + * the `_rowid` / score vectors directly and excluded from the payload map. Every OTHER column the + * scan returned is payload and follows [[readRows]]' contract exactly: a column with a Spark target + * type in `projectionFields` goes through the canonical + * [[org.lance.spark.vectorized.LanceArrowColumnVector]] adapter and back through + * [[CatalystTypeConverters]] to the external value the encoder expects; a projected column WITHOUT + * a supplied type falls back to the generic Arrow conversion ([[LanceProbe.toSparkValue]]) rather + * than being silently dropped. + */ + private def readScoredRows( + reader: ArrowReader, + projectionFields: Seq[StructField]): Seq[MaterializedHit] = { + val schemaByName: Map[String, StructField] = + projectionFields.iterator.map(f => f.name -> f).toMap + val out = mutable.ArrayBuffer.empty[MaterializedHit] + try { + while (reader.loadNextBatch()) { + val root: VectorSchemaRoot = reader.getVectorSchemaRoot + val n = root.getRowCount + val fields = root.getSchema.getFields.asScala.toIndexedSeq + + val addrVec: FieldVector = root.getVector(LanceProbe.RowIdColumn) + val scoreVec: FieldVector = resolveScoreVector(root) + + // `_rowid` and the injected score column are read out-of-band (above) and excluded from the + // payload. Everything else the scan returned IS payload: columns with a Spark target type go + // through the canonical adapter, the rest fall back to the generic Arrow conversion — the + // same split `readRows` / `materialize` apply, so a projected column with no supplied type is + // surfaced (via `toSparkValue`) instead of being silently dropped. + val reserved = Set(LanceProbe.RowIdColumn, scoreVec.getField.getName) + val payloadFields = fields.filterNot(af => reserved.contains(af.getName)) + val mapped = payloadFields.filter(af => schemaByName.contains(af.getName)) + val unmapped = payloadFields.filterNot(af => schemaByName.contains(af.getName)) + val mappedNames: Array[String] = mapped.iterator.map(_.getName).toArray + val mappedTypes: Array[DataType] = + mapped.iterator.map(af => schemaByName(af.getName).dataType).toArray + val mappedConverters: Array[Any => Any] = mapped.iterator + .map(af => + CatalystTypeConverters.createToScalaConverter(schemaByName(af.getName).dataType)) + .toArray + val mappedVectors: Array[ColumnVector] = mapped.iterator + .map(af => + new LanceArrowColumnVector(root.getVector(af.getName), false, schemaByName(af.getName)) + .asInstanceOf[ColumnVector]) + .toArray + // Thin view over the reader-owned Arrow vectors (closeVectorOnClose=false): `reader.close()` + // in the finally frees the buffers once the values below are copied out. + val batch = new ColumnarBatch(mappedVectors, n) + + var i = 0 + while (i < n) { + val rowMap = mutable.LinkedHashMap.empty[String, Any] + val internalRow = batch.getRow(i) + var j = 0 + while (j < mappedNames.length) { + val internal = internalRow.get(j, mappedTypes(j)) + rowMap(mappedNames(j)) = if (internal == null) null else mappedConverters(j)(internal) + j += 1 + } + unmapped.foreach { af => + val v = root.getVector(af.getName) + rowMap(af.getName) = if (v.isNull(i)) null else LanceProbe.toSparkValue(v.getObject(i)) + } + out += MaterializedHit(rowAddrAt(addrVec, i), scoreAt(scoreVec, i), rowMap.toMap) + i += 1 + } + } + } finally { + reader.close() + } + out.toSeq + } + + /** Locate the nearest-search score column by name, tolerant of `_distance` vs `_score`. */ + private def resolveScoreVector(root: VectorSchemaRoot): FieldVector = + LanceProbe.ScoreColumns.iterator + .map(name => Option(root.getVector(name)).orNull) + .find(_ != null) + .getOrElse(throw new IllegalStateException( + "Lance nearest scan did not return a score column. Got: " + + root.getSchema.getFields.asScala.map(_.getName).mkString(", "))) + + /** Read a Lance row id out of an Arrow `_rowid` vector (UInt8 or BigInt, encoding-dependent). */ + private def rowAddrAt(addrVec: FieldVector, i: Int): Long = addrVec match { + case v: UInt8Vector => v.get(i) + case v: BigIntVector => v.get(i) + case other => + throw new IllegalStateException( + s"Unexpected row-address vector type: ${other.getClass.getName}") + } + + /** Read a ranking score out of an Arrow score vector (Float4 or Float8, encoding-dependent). */ + private def scoreAt(scoreVec: FieldVector, i: Int): Float = scoreVec match { + case v: Float4Vector => v.get(i) + case v: Float8Vector => v.get(i).toFloat + case other => + throw new IllegalStateException( + s"Unexpected score vector type: ${other.getClass.getName}") + } + + /** + * Materialize a set of right-side rows by their `_rowaddr`s. Used by the join's materialize + * stage to fetch full payloads after the probe + merge has decided which rows survive. + * + * The row addresses are pushed down as a `_rowaddr IN (...)` filter, which Lance executes via + * its row-address index — the natural point-fetch path. The result is unordered with respect + * to the input list; the caller re-aligns by `_rowaddr`. + * + * @param rowAddrs list of Lance `_rowid` values (parameter name retained for source + * compatibility with callers — semantically these are now row IDs). + * @param projection projected column list. `Seq.empty` means "all columns". + * @param projectionFields Spark target fields (name + dataType) for the projected columns. When + * non-empty, each projected payload cell is converted to the Spark EXTERNAL + * value its declared type expects (see [[readRows]]); when empty, cells fall + * back to a generic Arrow-object conversion. + * @return a sequence of materialized rows, each a `Map[String, Any]` keyed by column name. With + * `projectionFields` supplied, each projected value is already the Spark external + * representation of its target type — the shape the join's `ExpressionEncoder` accepts — + * plus an entry under `LanceProbe.RowIdColumn` (a plain long) so the caller can re-key. + * Building those values into a `Row` / `InternalRow` happens in the join stage. + */ + def materialize( + rowAddrs: Seq[Long], + projection: Seq[String] = Seq.empty, + projectionFields: Seq[StructField] = Seq.empty): Seq[Map[String, Any]] = { + if (rowAddrs.isEmpty) return Seq.empty + + val opts = new ScanOptions.Builder().withRowId(true) + if (projection.nonEmpty) { + opts.columns(projection.toList.asJava) + } + // `_rowid IN (a, b, c)` — Lance lowers this to its row-id lookup path. Same point-fetch + // semantics as `_rowaddr IN (...)` previously used here, but `_rowid` is the universal + // identifier (works on indexed + non-indexed scan paths alike). + // + // Each row ID is rendered as `arrow_cast('', 'UInt64')` for two + // compounding reasons: + // + // 1. Lance row IDs are 64-bit UNSIGNED; storing them as Java signed `long` means + // values >= 2^63 come back negative. `mkString(", ")` would render them as + // negative integer literals and Lance/DataFusion would reject (`Int64(-...) + // cannot convert to UInt64`). + // 2. Even after `Long.toUnsignedString` produces a positive 20-digit decimal, + // DataFusion's SQL parser tries `Int64` first, overflows, then falls back to + // `Float64`. `Float64` loses precision past 2^53 — the literal becomes a + // different number — and DataFusion then can't downcast `Float64` to `UInt64`. + // + // `arrow_cast(string, 'UInt64')` bypasses both: the string literal goes through + // `arrow_cast`'s own coercion, which is precision-preserving for UInt64. + // + // At 100K rows row IDs stay below 2^53 and both layers of the bug are invisible; at + // 1M+ rows they bite. Caught when the DataFrame benchmark hit 1M-row scale. + val rowIdLiterals = rowAddrs.iterator + .map(addr => s"arrow_cast('${java.lang.Long.toUnsignedString(addr)}', 'UInt64')") + .mkString(", ") + opts.filter(s"${LanceProbe.RowIdColumn} IN ($rowIdLiterals)") + javaFragmentIds.foreach(opts.fragmentIds) + + val scanner: LanceScanner = LanceScanner.create(dataset, opts.build(), allocator) + try { + readRows(scanner.scanBatches(), projectionFields) + } finally { + scanner.close() + } + } + + /** + * Drain the materialize scan into row maps, converting each projected payload cell to the Spark + * EXTERNAL value its target type expects — the shape the join's `ExpressionEncoder` (external + * `Row` → `InternalRow`) accepts. + * + * Projected data columns are materialized through the connector's canonical Arrow→Spark adapter + * [[org.lance.spark.vectorized.LanceArrowColumnVector]] — the same vector-and-schema-aware path + * the connector's own reader uses — so every payload type Lance can store round-trips to the + * value Spark produces for that column on an ordinary read: DateType, the various timestamp / time + * units, unsigned integers, decimals, fixed/large binary and varchar, and nested structs / lists / + * maps. A raw Arrow `getObject` would instead hand back e.g. a `LocalDate` for a DateType column, + * which the encoder then rejects. The adapter yields Spark INTERNAL values (days for a date, + * micros for a timestamp, …), so each is run back through [[CatalystTypeConverters]] to the + * external representation the `Row` encoder wants. + * + * Columns absent from `projectionFields` — notably the `_rowid` virtual column, which carries no + * Spark type here — keep the generic Arrow `getObject` fallback; the caller reads `_rowid` only as + * a plain long. + */ + private def readRows( + reader: ArrowReader, + projectionFields: Seq[StructField]): Seq[Map[String, Any]] = { + val schemaByName: Map[String, StructField] = + projectionFields.iterator.map(f => f.name -> f).toMap + val out = mutable.ArrayBuffer.empty[Map[String, Any]] + try { + while (reader.loadNextBatch()) { + val root: VectorSchemaRoot = reader.getVectorSchemaRoot + val n = root.getRowCount + val fields = root.getSchema.getFields.asScala.toIndexedSeq + + // Columns with a Spark target type are converted through the canonical connector adapter; + // the rest (e.g. `_rowid`) keep the raw Arrow-object fallback. + val mapped = fields.filter(af => schemaByName.contains(af.getName)) + val unmapped = fields.filterNot(af => schemaByName.contains(af.getName)) + val mappedNames: Array[String] = mapped.iterator.map(_.getName).toArray + val mappedTypes: Array[DataType] = + mapped.iterator.map(af => schemaByName(af.getName).dataType).toArray + val mappedConverters: Array[Any => Any] = mapped.iterator + .map(af => + CatalystTypeConverters.createToScalaConverter(schemaByName(af.getName).dataType)) + .toArray + val mappedVectors: Array[ColumnVector] = mapped.iterator + .map(af => + new LanceArrowColumnVector(root.getVector(af.getName), false, schemaByName(af.getName)) + .asInstanceOf[ColumnVector]) + .toArray + // The batch is a thin view over the (reader-owned) Arrow vectors; closeVectorOnClose=false + // above means neither the batch nor its column vectors free the underlying buffers — the + // `reader.close()` in the finally does, once the values below have been copied out. + val batch = new ColumnarBatch(mappedVectors, n) + + var i = 0 + while (i < n) { + val rowMap = mutable.LinkedHashMap.empty[String, Any] + val internalRow = batch.getRow(i) + var j = 0 + while (j < mappedNames.length) { + val internal = internalRow.get(j, mappedTypes(j)) + rowMap(mappedNames(j)) = if (internal == null) null else mappedConverters(j)(internal) + j += 1 + } + unmapped.foreach { af => + val v = root.getVector(af.getName) + rowMap(af.getName) = if (v.isNull(i)) null else LanceProbe.toSparkValue(v.getObject(i)) + } + out += rowMap.toMap + i += 1 + } + } + } finally { + reader.close() + } + out.toSeq + } + + override def close(): Unit = dataset.close() +} + +object LanceProbe { + + /** + * Lance row-identity virtual column name. We use `_rowid` rather than `_rowaddr` because + * Lance's INDEXED nearest-search path materializes `_rowid` but not `_rowaddr`, while + * non-indexed scans materialize both. `_rowid` therefore works on every code path that + * calls `probe()` (with or without a vector index built on the column). Sourced from + * `LanceConstant` to keep the literal defined in exactly one place. + */ + val RowIdColumn: String = LanceConstant.ROW_ID + + /** + * Candidate names for the score column in a Lance nearest-search result. Lance's vector indexes + * have used `_distance` historically; tolerate `_score` too in case future versions rename it. + * The lookup is name-based so the consumer is agnostic to where Lance puts the column in its + * output schema. + */ + val ScoreColumns: Seq[String] = Seq("_distance", "_score") + + /** + * Column names Lance's nearest scan injects itself: `_rowid` (from `withRowId`) and the score + * columns (from `nearest`). These are metadata a nearest scan always produces; a right-side + * table column sharing one of these names collides with the injected column. See + * [[reservedSchemaColumns]] / [[schemaSupportsNearest]]. + */ + val ReservedProjectionColumns: Set[String] = ScoreColumns.toSet + RowIdColumn + + /** + * The subset of `schemaColumnNames` that collide with a column the nearest scan injects + * (`_rowid` / `_distance` / `_score`). Empty means the table is nearest-compatible. + */ + def reservedSchemaColumns(schemaColumnNames: Iterable[String]): Set[String] = + schemaColumnNames.iterator.filter(ReservedProjectionColumns.contains).toSet + + /** + * Whether an indexed nearest scan can run against a right-side table with these column names. + * + * Lance's nearest scan ALWAYS injects `_rowid` and the `_distance` / `_score` metadata. If the + * table's own schema already has a column by one of those names, the injected metadata shadows + * it and no projection shape recovers the physical column: the fused scan reads it out-of-band as + * the ranking score and silently drops it from the payload (observed with an empty / all-columns + * projection), and an explicit projection of it collides outright. There is no safe fold-vs-split + * routing around this — every indexed route runs the same nearest scan. The indexed rewrite must + * therefore DECLINE such a table and fall back to the engine's default nearest-by execution. + * + * This is the shared eligibility contract the Catalyst rule consults before intercepting; the + * probe enforces it defensively too (see the schema guard in [[LanceProbe.probe]] / + * [[LanceProbe.probeRows]]). + */ + def schemaSupportsNearest(schemaColumnNames: Iterable[String]): Boolean = + reservedSchemaColumns(schemaColumnNames).isEmpty + + /** + * Build read options from a bare dataset URI, pinning `version` on the main branch when present. + * Backs the URI convenience constructor. + */ + private[knn] def readOptionsFor( + datasetUri: String, + version: Option[Long]): LanceSparkReadOptions = { + val base = LanceSparkReadOptions.from(datasetUri) + version match { + case Some(v) => base.withRef(LanceRef.ofMain(v)) + case None => base + } + } + + /** + * Convert an Arrow-returned cell value into something Spark's encoders accept when stuffed + * into a `Row`. Arrow's `FieldVector.getObject` returns Java types (boxed primitives, + * `JsonStringArrayList` for list cells, `Text` for utf8) which Spark's `RowEncoder` does not + * always understand directly — most painfully, a `java.util.ArrayList` can't satisfy a Spark + * `ArrayType` slot, which expects a `scala.collection.Seq`. + * + * Conversion rules, in order: + * - `java.util.List` → recursively-converted `Seq` + * - `java.util.Map` → recursively-converted Scala `Map` + * - `org.apache.arrow.vector.util.Text` → `String` + * - `Number` boxed primitives → returned as-is (Spark handles them) + * - everything else → returned as-is (caller's responsibility) + * + * Recursive on lists/maps to handle nested types (arrays of structs, etc.) without surprises + * for callers. + */ + def toSparkValue(value: Any): Any = value match { + case null => null + case list: java.util.List[_] => + val out = scala.collection.mutable.ArrayBuffer.empty[Any] + val it = list.iterator + while (it.hasNext) out += toSparkValue(it.next()) + out.toSeq + case map: java.util.Map[_, _] => + val out = scala.collection.mutable.LinkedHashMap.empty[Any, Any] + val it = map.entrySet().iterator + while (it.hasNext) { + val e = it.next() + out(toSparkValue(e.getKey)) = toSparkValue(e.getValue) + } + out.toMap + case t: org.apache.arrow.vector.util.Text => t.toString + case other => other + } +} diff --git a/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/internal/Metric.scala b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/internal/Metric.scala new file mode 100644 index 000000000..2a5ed1726 --- /dev/null +++ b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/internal/Metric.scala @@ -0,0 +1,75 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.knn.internal + +import org.lance.index.DistanceType + +/** + * Vector distance / similarity metric. Mirrors `org.lance.index.DistanceType` but exposed as a + * Scala enumeration so callers don't have to import Lance internals. Each metric fixes the + * "best-first" direction used during merge. + * + * Lance's vector search returns a DISTANCE for every metric — including the similarity-flavored + * ones, which it reports as `1 - cosine_similarity` and `1 - dot_product`. So smaller is better for + * ALL three metrics; there is no larger-is-better case: + * + * - L2: smaller score is better (squared L2 distance) + * - Cosine: smaller score is better (`1 - cosine_similarity`) + * - Dot: smaller score is better (`1 - dot_product`) + */ +sealed trait Metric { + + /** The Lance distance type used when configuring a `Query`. */ + def lanceType: DistanceType + + /** True if smaller scores rank better (distance), false if larger (similarity). */ + def smallerIsBetter: Boolean +} + +object Metric { + + case object L2 extends Metric { + val lanceType: DistanceType = DistanceType.L2 + val smallerIsBetter: Boolean = true + } + + case object Cosine extends Metric { + val lanceType: DistanceType = DistanceType.Cosine + // Lance returns `1 - cosine_similarity`, a distance: smaller is better. + val smallerIsBetter: Boolean = true + } + + case object Dot extends Metric { + val lanceType: DistanceType = DistanceType.Dot + // Lance returns `1 - dot_product`, a distance: smaller is better. + val smallerIsBetter: Boolean = true + } + + /** + * Parse a metric name. Accepts the same set of names Lance accepts plus a few synonyms commonly + * used in Spark vector functions: + * + * - "l2" | "euclidean" → L2 + * - "cosine" → Cosine + * - "dot" | "inner" | "ip" → Dot + */ + def fromName(name: String): Metric = name.trim.toLowerCase match { + case "l2" | "euclidean" => L2 + case "cosine" => Cosine + case "dot" | "inner" | "ip" => Dot + case other => + throw new IllegalArgumentException( + s"Unknown metric '$other'. Expected one of: l2, cosine, dot.") + } +} diff --git a/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/internal/ScoredRowRef.scala b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/internal/ScoredRowRef.scala new file mode 100644 index 000000000..bacd984c9 --- /dev/null +++ b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/internal/ScoredRowRef.scala @@ -0,0 +1,55 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.knn.internal + +/** + * A reference to a single right-side row produced by a vector index probe, along with the ranking + * score. Carries no payload — payloads are fetched in the materialize stage by row address. Pairing + * a tiny ref with a score is the unit of work passing through the shuffle and is what keeps the + * shuffle volume to `O(|L| × tasks × K × ~24B)` instead of `O(|L| × tasks × K × payload_bytes)`. + * + * @param rowAddr Lance row address (`_rowaddr`): packed `(frag_id << 32) | row_in_frag`. Stable + * within a Lance dataset version. + * @param score Distance or similarity returned by Lance's vector search. Smaller-is-better for + * distance metrics (L2), larger-is-better for similarity metrics (cosine/dot). + * Direction is carried out-of-band in the operator config; this struct stays metric- + * agnostic. + */ +final case class ScoredRowRef(rowAddr: Long, score: Float) + +object ScoredRowRef { + + /** Order best-first for distance metrics (smallest score wins). */ + val distanceOrdering: Ordering[ScoredRowRef] = + Ordering.by[ScoredRowRef, Float](_.score) + + /** Order best-first for similarity metrics (largest score wins). */ + val similarityOrdering: Ordering[ScoredRowRef] = + Ordering.by[ScoredRowRef, Float](-_.score) +} + +/** + * A single probe hit WITH its materialized payload — the unit produced by the no-overfetch fast path + * ([[LanceProbe.probeRows]]), which folds the nearest search and the payload projection into one + * scan. Unlike [[ScoredRowRef]] (a payload-free ref bound for a later materialize point-fetch), this + * already carries the right row. + * + * @param rowAddr Lance row id of the hit (kept for de-duplication / re-keying by the join stage). + * @param score Distance or similarity from Lance's vector search; direction is carried out-of-band + * in the operator config, so this stays metric-agnostic (see [[ScoredRowRef.score]]). + * @param row Materialized right-side payload keyed by column name, each value already the Spark + * EXTERNAL representation its target type expects (the join's `ExpressionEncoder` + * accepts it directly). + */ +final case class MaterializedHit(rowAddr: Long, score: Float, row: Map[String, Any]) diff --git a/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/internal/TopKHeap.scala b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/internal/TopKHeap.scala new file mode 100644 index 000000000..7c562517f --- /dev/null +++ b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/internal/TopKHeap.scala @@ -0,0 +1,118 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.knn.internal + +import scala.collection.mutable + +/** + * Bounded top-K heap with metric-aware ordering. Used by the probe stage for map-side combine + * across fragments owned by a single task — keeps the per-left-row state at exactly K entries no + * matter how many fragments contribute, and again on the reduce side to merge contributions from + * different tasks for the same `leftId`. + * + * Semantics: + * - `smallerIsBetter = true` (distance, e.g. L2): retain the K smallest-score entries. + * - `smallerIsBetter = false` (similarity, e.g. cosine): retain the K largest-score entries. + * + * Internally, the heap's *head* holds the worst surviving element so eviction is O(log K). Scala's + * `mutable.PriorityQueue` is a max-heap by the supplied `Ordering`, so the ordering is chosen to + * place "worst surviving" at the top: + * - distance → max-heap on `score` (largest score is worst) + * - similarity → max-heap on `-score` (smallest score is worst) + * + * Not thread-safe. Each left row in a probe stage gets its own heap. + */ +final class TopKHeap(k: Int, smallerIsBetter: Boolean) { + require(k > 0, "k must be positive") + + private val ord: Ordering[ScoredRowRef] = + if (smallerIsBetter) Ordering.by[ScoredRowRef, Float](_.score) + else Ordering.by[ScoredRowRef, Float](-_.score) + + private val heap = new mutable.PriorityQueue[ScoredRowRef]()(ord) + + /** + * Insert `ref` if it would survive the top-K cut. Either grows the heap up to K or evicts the + * current worst-surviving element if `ref` is strictly better than it. + */ + def offer(ref: ScoredRowRef): Unit = { + if (heap.size < k) { + heap.enqueue(ref) + } else { + val worst = heap.head + // Admission must use the SAME total ordering as the heap, not a raw float `<` / `>`. Float + // comparisons involving NaN are always false, so a NaN worst-survivor would never test as + // "beatable" and could never be evicted — a single NaN score would then pin a slot forever. + // `ord` is `java.lang.Float.compare`-based (NaN sorts as the largest Float, i.e. the worst + // for a distance), and `ord.lt(ref, worst)` means exactly "ref ranks strictly better than the + // current worst" in both directions — so a finite score correctly displaces a NaN worst. + val isBetter = ord.lt(ref, worst) + if (isBetter) { + heap.dequeue() + heap.enqueue(ref) + } + } + } + + def offerAll(refs: TraversableOnce[ScoredRowRef]): Unit = refs.foreach(offer) + + /** + * Drain the heap into a best-first sorted Array. After this call the heap is empty. Best-first + * means index 0 is the top-ranked entry (smallest score for distance, largest for similarity). + */ + def drain(): Array[ScoredRowRef] = { + val out = new Array[ScoredRowRef](heap.size) + var i = heap.size - 1 + // PriorityQueue.dequeue returns the worst surviving element first; walking the array in + // reverse places best at index 0. + while (i >= 0) { + out(i) = heap.dequeue() + i -= 1 + } + out + } + + def size: Int = heap.size + def isEmpty: Boolean = heap.isEmpty +} + +object TopKHeap { + + /** + * Convenience: merge several already-sorted (best-first) ref arrays into one top-K array. Used + * by the merge stage as the `reduceByKey` combine function. + */ + def merge( + a: Array[ScoredRowRef], + b: Array[ScoredRowRef], + k: Int, + smallerIsBetter: Boolean): Array[ScoredRowRef] = { + if (a.isEmpty) return takeBest(b, k, smallerIsBetter) + if (b.isEmpty) return takeBest(a, k, smallerIsBetter) + val heap = new TopKHeap(k, smallerIsBetter) + heap.offerAll(a) + heap.offerAll(b) + heap.drain() + } + + private def takeBest( + arr: Array[ScoredRowRef], + k: Int, + smallerIsBetter: Boolean): Array[ScoredRowRef] = { + if (arr.length <= k) return arr + val heap = new TopKHeap(k, smallerIsBetter) + heap.offerAll(arr) + heap.drain() + } +} diff --git a/lance-spark-knn-4.2_2.13/src/test/scala/org/lance/spark/knn/internal/LanceProbeValidationTest.scala b/lance-spark-knn-4.2_2.13/src/test/scala/org/lance/spark/knn/internal/LanceProbeValidationTest.scala new file mode 100644 index 000000000..05bfe4580 --- /dev/null +++ b/lance-spark-knn-4.2_2.13/src/test/scala/org/lance/spark/knn/internal/LanceProbeValidationTest.scala @@ -0,0 +1,459 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.knn.internal + +import org.apache.spark.sql.{Row, RowFactory, SparkSession} +import org.apache.spark.sql.types._ +import org.junit.jupiter.api.{AfterEach, BeforeEach, Test} +import org.junit.jupiter.api.Assertions._ +import org.junit.jupiter.api.io.TempDir +import org.lance.spark.LanceSparkReadOptions + +import java.nio.file.Path +import java.util.{Collections, Random} + +import scala.collection.JavaConverters._ + +/** + * End-to-end validation of [[LanceProbe]] against a real Lance dataset written by Spark. These are + * the day-1 validation tasks the implementation plan calls out: + * + * - Per-probe call should succeed and return Lance's nearest neighbors. + * - Repeated probes against the same `LanceProbe` instance should reuse the open dataset + * handle; the second call should not re-pay the dataset open cost. + * - `fragmentIds` restriction should narrow the search to specified fragments only. + * - Without an explicit vector index the probe falls back to a brute-force per-fragment scan, + * which gives recall = 1.0 — making the no-index path the natural correctness oracle. + * + * These tests do NOT require an actual vector index; that is exercised in the indexed test + * suites which build IVF-PQ via Lance's index DDL. Validating the brute-force path first lets us + * isolate any LanceProbe bugs from index-quality issues. + */ +class LanceProbeValidationTest { + + @TempDir var tempDir: Path = _ + private var spark: SparkSession = _ + + // Small synthetic dataset: 64 vectors, dim 8. Enough to exercise the probe loop without making + // the test slow. + private val NumRows = 64 + private val VectorDim = 8 + private val Seed = 42L + + @BeforeEach def setup(): Unit = { + spark = SparkSession.builder() + .appName("lance-probe-validation") + .master("local[2]") + // Pin the driver to loopback so test JVMs in restricted networks (CI sandboxes, dev + // containers) can bind without scanning the host's interfaces. + .config("spark.driver.bindAddress", "127.0.0.1") + .config("spark.driver.host", "127.0.0.1") + .getOrCreate() + } + + @AfterEach def teardown(): Unit = { + if (spark != null) spark.stop() + } + + /** + * Smoke test: write a dataset, probe it, get K rows back. No correctness assertion beyond + * "result has the right shape" — the brute-force-equivalence test below covers semantics. + */ + @Test def testProbeReturnsKResults(): Unit = { + val datasetUri = writeSyntheticDataset() + val query = randomVector(new Random(7L), VectorDim) + + val probe = new LanceProbe(datasetUri, fragmentIds = None) + try { + val results = probe.probe(vectorColumn = "vec", query, k = 5, metric = Metric.L2) + assertEquals(5, results.size, "probe should return exactly k results") + // Distances must be monotonically non-decreasing for L2 (best-first). + val scores = results.map(_.score) + assertEquals(scores, scores.sorted, "L2 results should be sorted ascending by distance") + // Row addresses are stable u64s; we just sanity-check they aren't all zero. + assertTrue(results.exists(_.rowAddr != 0L), "row addresses should be populated") + } finally probe.close() + } + + /** + * Without a vector index, Lance does an exact per-fragment scan. That makes it a recall = 1.0 + * oracle: the probe result should equal the ground-truth top-K computed in plain Scala. + */ + @Test def testProbeMatchesBruteForceOracle(): Unit = { + val rng = new Random(Seed) + val (rows, vectors) = generateRows(rng, NumRows, VectorDim) + val datasetUri = writeRows(rows) + + val query = randomVector(new Random(123L), VectorDim) + val k = 10 + + val oracle: Seq[(Int, Float)] = vectors.zipWithIndex + .map { case (v, idx) => (idx, l2Distance(query, v)) } + .sortBy(_._2) + .take(k) + + val probe = new LanceProbe(datasetUri, fragmentIds = None) + val actual = + try probe.probe("vec", query, k, Metric.L2) + finally probe.close() + + assertEquals(k, actual.size) + // Compare scores within float tolerance. + val expectedScores = oracle.map(_._2) + val actualScores = actual.map(_.score) + expectedScores.zip(actualScores).foreach { case (expected, actualScore) => + assertEquals( + expected, + actualScore, + 1e-4f, + s"top-K distance mismatch: oracle=$expectedScores actual=$actualScores") + } + } + + /** + * Cosine and Dot must rank best-first the same direction L2 does. Lance returns a DISTANCE for + * every metric (`1 - cosine_similarity`, `1 - dot_product` for the similarity-flavored ones), so + * smaller is better for all three — [[Metric.smallerIsBetter]] must be `true` for each. This + * reproduces the gatekeeper's failure directly: run a real Lance query, then merge the results + * through the size-1 [[TopKHeap]] the join stage uses, keyed by the metric's own direction flag. + * The nearest ref (Lance returns best-first, so `refs.head`) must survive; a wrong flag (treating a + * Lance distance as larger-is-better) would retain the FARTHEST ref instead. + */ + @Test def testMetricFlagsKeepNearestThroughHeap(): Unit = { + val datasetUri = writeSyntheticDataset() + val query = randomVector(new Random(555L), VectorDim) + + val probe = new LanceProbe(datasetUri, fragmentIds = None) + try { + Seq[Metric](Metric.L2, Metric.Cosine, Metric.Dot).foreach { metric => + val refs = probe.probe("vec", query, k = 5, metric) + assertEquals(5, refs.size, s"$metric probe should return k results") + // Lance returns best-first, so refs.head is the true nearest for this metric. + val nearest = refs.head + val heap = new TopKHeap(k = 1, metric.smallerIsBetter) + heap.offerAll(refs) + val survivor = heap.drain() + assertEquals(1, survivor.length, s"$metric: size-1 heap should retain one ref") + assertEquals( + nearest.rowAddr, + survivor.head.rowAddr, + s"$metric: size-1 heap must keep the nearest ref (rowAddr=${nearest.rowAddr}), " + + s"got ${survivor.head.rowAddr} — wrong smallerIsBetter direction?") + assertEquals(nearest.score, survivor.head.score, 1e-6f, s"$metric: kept score mismatch") + } + } finally probe.close() + } + + /** + * A projected payload column WITHOUT a supplied Spark type must be preserved through the generic + * Arrow conversion — the same fallback [[LanceProbe.materialize]] / `readRows` apply — not silently + * dropped. Regression: `projection = Seq("id")` with empty `projectionFields` must return payload + * keys `Set("id")` (the injected `_rowid` / score columns stay out of the payload). + */ + @Test def testProbeRowsPreservesUnmappedProjectedFields(): Unit = { + val datasetUri = writeSyntheticDataset() + val query = randomVector(new Random(321L), VectorDim) + + val probe = new LanceProbe(datasetUri, fragmentIds = None) + try { + val hits = probe.probeRows( + "vec", + query, + k = 5, + Metric.L2, + projection = Seq("id"), + projectionFields = Seq.empty) + assertEquals(5, hits.size, "probeRows should return k hits") + hits.foreach { h => + assertEquals( + Set("id"), + h.row.keySet, + s"unmapped projected field must be preserved (id only, no _rowid/_distance); " + + s"got ${h.row.keySet}") + assertNotNull(h.row("id"), "unmapped id payload must be populated") + } + } finally probe.close() + } + + /** + * The eligibility contract is SCHEMA-level, not projection-level. Lance's nearest scan always + * injects `_rowid` and the `_distance` / `_score` metadata; if the dataset's OWN schema has a + * column by one of those names, the injected metadata shadows it and no probe route recovers the + * physical column. Empirically, a physical `_distance` column writes fine and probes silently drop + * it: an all-columns fused scan reads it out-of-band as the ranking score, so its payload comes + * back without `_distance`. That silent data loss is why the indexed rewrite must DECLINE such a + * table (see [[LanceProbe.schemaSupportsNearest]]) rather than pick a materialization shape. + * + * This regression writes a real dataset WITH a `_distance` column and asserts (a) the schema is + * reported non-nearest-compatible, and (b) BOTH probe entry points — `probe` and the all-columns + * `probeRows(projection = empty)` — fail fast with a clear error naming the offending column, + * instead of returning results that silently omit it. + */ + @Test def testReservedSchemaColumnIsDeclinedNotSilentlyDropped(): Unit = { + val datasetUri = writeDatasetWithReservedColumn() + val query = randomVector(new Random(1L), VectorDim) + + // (a) schema eligibility primitive reports the collision. + assertFalse( + LanceProbe.schemaSupportsNearest(Seq("id", "vec", "_distance")), + "a schema owning a reserved column (_distance) must not be nearest-compatible") + + val probe = new LanceProbe(datasetUri, fragmentIds = None) + try { + // (b) probe() declines with a clear error naming the offending column. + val probeEx = assertThrows( + classOf[IllegalArgumentException], + () => probe.probe("vec", query, k = 5, Metric.L2)) + assertTrue( + String.valueOf(probeEx.getMessage).contains("_distance"), + s"probe guard message should name the offending column '_distance'; got: ${probeEx.getMessage}") + + // (b) all-columns probeRows — the exact path that used to SILENTLY drop the physical + // `_distance` — declines with the same clear error rather than returning a lossy payload. + val probeRowsEx = assertThrows( + classOf[IllegalArgumentException], + () => + probe.probeRows( + "vec", + query, + k = 5, + Metric.L2, + projection = Seq.empty, + projectionFields = Seq.empty)) + assertTrue( + String.valueOf(probeRowsEx.getMessage).contains("_distance"), + s"probeRows guard message should name the offending column '_distance'; " + + s"got: ${probeRowsEx.getMessage}") + } finally probe.close() + } + + /** + * Pure eligibility primitive: [[LanceProbe.schemaSupportsNearest]] is false iff the schema names + * any column the nearest scan injects, and [[LanceProbe.reservedSchemaColumns]] returns exactly + * that colliding set. No backend needed — this is the contract the Catalyst rule consults. + */ + @Test def testSchemaSupportsNearestContract(): Unit = { + assertTrue( + LanceProbe.schemaSupportsNearest(Seq("id", "vec", "payload")), + "a schema with no reserved column names must be nearest-compatible") + assertTrue( + LanceProbe.schemaSupportsNearest(Seq.empty), + "an empty schema must be nearest-compatible") + LanceProbe.ReservedProjectionColumns.foreach { reserved => + assertFalse( + LanceProbe.schemaSupportsNearest(Seq("id", reserved, "vec")), + s"a schema owning reserved column '$reserved' must not be nearest-compatible") + assertEquals( + Set(reserved), + LanceProbe.reservedSchemaColumns(Seq("id", reserved, "vec")), + s"reservedSchemaColumns must report exactly the colliding column '$reserved'") + } + assertEquals( + LanceProbe.ReservedProjectionColumns, + LanceProbe.reservedSchemaColumns(Seq("id") ++ LanceProbe.ReservedProjectionColumns.toSeq), + "reservedSchemaColumns must report every reserved column present") + } + + /** + * The folded fast path ([[LanceProbe.probeRows]]) must be observationally identical to the split + * probe + materialize path: the SAME (rowAddr, score) hits in the SAME order, and the SAME + * materialized payload per row. This is the invariant the SQL join relies on when it single-scans + * (internalK == k) instead of probing then late-materializing. Runs on the brute-force path so the + * search itself is deterministic. + */ + @Test def testProbeRowsMatchesProbeThenMaterialize(): Unit = { + val datasetUri = writeSyntheticDataset() + val query = randomVector(new Random(321L), VectorDim) + val k = 8 + val projection = Seq("id", "vec") + val projectionFields = Seq( + StructField("id", IntegerType, nullable = false), + StructField("vec", ArrayType(FloatType, containsNull = false), nullable = false)) + + val probe = new LanceProbe(datasetUri, fragmentIds = None) + try { + // Split path: search for refs, then late-materialize the payload by _rowid. + val refs = probe.probe("vec", query, k, Metric.L2) + val expectedPayload: Map[Long, Map[String, Any]] = probe + .materialize(refs.map(_.rowAddr), projection, projectionFields) + .map(m => rowAddrOf(m) -> m) + .toMap + + // Folded path: search AND project the payload in one scan. + val hits = probe.probeRows("vec", query, k, Metric.L2, projection, projectionFields) + + assertEquals(k, hits.size, "probeRows should return exactly k hits") + // Same search: identical (rowAddr, score) sequence, in order. + assertEquals( + refs.map(r => (r.rowAddr, r.score)), + hits.map(h => (h.rowAddr, h.score)), + "probeRows hits must match probe refs (rowAddr + score), in order") + // Same payload: each folded hit equals the split materialize's row for that rowAddr, on the + // projected columns. + hits.foreach { h => + val expected = expectedPayload(h.rowAddr) + assertEquals(expected("id"), h.row("id"), s"id mismatch for rowAddr=${h.rowAddr}") + assertEquals( + expected("vec").asInstanceOf[Seq[_]].toList, + h.row("vec").asInstanceOf[Seq[_]].toList, + s"vec mismatch for rowAddr=${h.rowAddr}") + } + } finally probe.close() + } + + /** + * Validate the dataset handle is reused across calls. The exact perf invariant ("second call + * faster than first by some factor") is too brittle for CI, so we only assert that repeated + * probes succeed and don't OOM — i.e., no JNI handle / Arrow buffer leak per call. + */ + @Test def testRepeatedProbesShareDatasetHandle(): Unit = { + val datasetUri = writeSyntheticDataset() + val probe = new LanceProbe(datasetUri, None) + try { + val rng = new Random(99L) + val k = 4 + var i = 0 + while (i < 50) { + val results = probe.probe("vec", randomVector(rng, VectorDim), k, Metric.L2) + assertEquals(k, results.size, s"iteration $i returned wrong size") + i += 1 + } + } finally probe.close() + } + + /** Empty fragment-id list ⇒ no rows match. Confirms the pushdown actually narrows search. */ + @Test def testEmptyFragmentRestrictionReturnsNothing(): Unit = { + val datasetUri = writeSyntheticDataset() + val probe = new LanceProbe(datasetUri, Some(Seq.empty)) + try { + val results = probe.probe("vec", randomVector(new Random(1L), VectorDim), 5, Metric.L2) + assertTrue(results.isEmpty, s"empty fragmentIds should yield no results, got ${results.size}") + } finally probe.close() + } + + /** + * When `executor_credential_refresh = false`, the probe must NOT rebuild (or even load) the + * runtime namespace on the worker — exactly the policy `LanceFragmentScanner.create` applies. + * Regression against the earlier `openDataset` that called `builder.runtimeNamespace(impl, ...)` + * unconditionally, which forced the namespace impl class to load regardless of the refresh flag. + * + * We pass a namespace impl that does not exist on the classpath and point the probe at a + * non-existent dataset URI. The open must fail because the DATASET is missing — not because it + * tried to load the bogus namespace class. Asserting the failure message does not mention the + * namespace class proves the namespace path was skipped. + */ + @Test def testExecutorCredentialRefreshFalseSkipsNamespaceRebuild(): Unit = { + val missingUri = tempDir.resolve("does_not_exist").toString + val readOptions = LanceSparkReadOptions + .builder() + .datasetUri(missingUri) + .executorCredentialRefresh(false) + .build() + val ex = assertThrows( + classOf[RuntimeException], + () => + new LanceProbe( + readOptions, + null, + "example.namespace.MustNotBeLoaded", + Collections.emptyMap[String, String](), + None)) + assertFalse( + String.valueOf(ex.getMessage).contains("MustNotBeLoaded"), + s"namespace impl must not be loaded when executor credential refresh is off; got: " + + ex.getMessage) + } + + // -- helpers ------------------------------------------------------------------------------ + + /** Write a fresh dataset and return its file:// URI. */ + private def writeSyntheticDataset(): String = { + val rng = new Random(Seed) + val (rows, _) = generateRows(rng, NumRows, VectorDim) + writeRows(rows) + } + + private def writeRows(rows: Seq[Row]): String = { + val schema = new StructType(Array( + StructField("id", IntegerType, nullable = false), + StructField( + "vec", + ArrayType(FloatType, containsNull = false), + nullable = false, + new MetadataBuilder().putLong("arrow.fixed-size-list.size", VectorDim.toLong).build()))) + val df = spark.createDataFrame(rows.asJava, schema) + + val outDir = tempDir.resolve(s"probe_test_${System.nanoTime()}").toString + df.write.format("lance").save(outDir) + outDir + } + + /** + * Write a dataset whose OWN schema carries a `_distance` column — a name Lance's nearest scan + * injects. Used to prove the schema eligibility backstop declines such a table. The `_distance` + * values are arbitrary payload; the point is that the physical column exists in the stored schema. + */ + private def writeDatasetWithReservedColumn(): String = { + val rng = new Random(Seed) + val (baseRows, _) = generateRows(rng, NumRows, VectorDim) + val rows = baseRows.zipWithIndex.map { case (r, idx) => + RowFactory.create(r.get(0), r.get(1), java.lang.Float.valueOf(idx.toFloat)) + } + val schema = new StructType(Array( + StructField("id", IntegerType, nullable = false), + StructField( + "vec", + ArrayType(FloatType, containsNull = false), + nullable = false, + new MetadataBuilder().putLong("arrow.fixed-size-list.size", VectorDim.toLong).build()), + StructField("_distance", FloatType, nullable = false))) + val df = spark.createDataFrame(rows.asJava, schema) + val outDir = tempDir.resolve(s"reserved_col_test_${System.nanoTime()}").toString + df.write.format("lance").save(outDir) + outDir + } + + private def generateRows(rng: Random, n: Int, dim: Int): (Seq[Row], Seq[Array[Float]]) = { + val vectors = (0 until n).map(_ => randomVector(rng, dim)) + val rows = vectors.zipWithIndex.map { case (v, idx) => + RowFactory.create(Integer.valueOf(idx), v) + } + (rows, vectors) + } + + /** Read the `_rowid` key out of a materialized row map (a boxed / stringy long). */ + private def rowAddrOf(m: Map[String, Any]): Long = m(LanceProbe.RowIdColumn) match { + case l: java.lang.Long => l.longValue() + case l: Long => l + case other => other.toString.toLong + } + + private def randomVector(rng: Random, dim: Int): Array[Float] = { + val v = new Array[Float](dim) + var i = 0 + while (i < dim) { v(i) = rng.nextFloat(); i += 1 } + v + } + + private def l2Distance(a: Array[Float], b: Array[Float]): Float = { + var s = 0.0f + var i = 0 + while (i < a.length) { + val d = a(i) - b(i) + s += d * d + i += 1 + } + s + } +} diff --git a/lance-spark-knn-4.2_2.13/src/test/scala/org/lance/spark/knn/internal/TopKHeapTest.scala b/lance-spark-knn-4.2_2.13/src/test/scala/org/lance/spark/knn/internal/TopKHeapTest.scala new file mode 100644 index 000000000..33481b4f8 --- /dev/null +++ b/lance-spark-knn-4.2_2.13/src/test/scala/org/lance/spark/knn/internal/TopKHeapTest.scala @@ -0,0 +1,108 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.knn.internal + +import org.junit.jupiter.api.Assertions._ +import org.junit.jupiter.api.Test + +/** + * Unit tests for [[TopKHeap]]. The heap's correctness is the foundation of the merge stage — + * any off-by-one or wrong-direction ordering would silently corrupt top-K results. We test both + * metric directions explicitly. + */ +class TopKHeapTest { + + private def ref(addr: Long, score: Float): ScoredRowRef = ScoredRowRef(addr, score) + + /** Distance metric: smaller score is better. Top-K must hold the K smallest. */ + @Test def testDistanceKeepsKSmallest(): Unit = { + val heap = new TopKHeap(k = 3, smallerIsBetter = true) + Seq(5.0f, 1.0f, 4.0f, 2.0f, 8.0f, 0.5f).zipWithIndex.foreach { case (s, i) => + heap.offer(ref(i.toLong, s)) + } + val out = heap.drain() + val scores = out.map(_.score).toSeq + assertEquals(Seq(0.5f, 1.0f, 2.0f), scores, "distance heap should retain three smallest") + } + + /** Similarity metric: larger score is better. Top-K must hold the K largest. */ + @Test def testSimilarityKeepsKLargest(): Unit = { + val heap = new TopKHeap(k = 3, smallerIsBetter = false) + Seq(5.0f, 1.0f, 4.0f, 2.0f, 8.0f, 0.5f).zipWithIndex.foreach { case (s, i) => + heap.offer(ref(i.toLong, s)) + } + val out = heap.drain() + val scores = out.map(_.score).toSeq + assertEquals(Seq(8.0f, 5.0f, 4.0f), scores, "similarity heap should retain three largest") + } + + /** Drain order is best-first regardless of insertion order. */ + @Test def testDrainOrderIsBestFirst(): Unit = { + val heap = new TopKHeap(k = 4, smallerIsBetter = true) + heap.offerAll(Seq(ref(1, 9f), ref(2, 1f), ref(3, 5f), ref(4, 3f), ref(5, 2f))) + val drained = heap.drain() + val scores = drained.map(_.score).toSeq + assertEquals(Seq(1f, 2f, 3f, 5f), scores) + assertTrue(heap.isEmpty, "drain should leave the heap empty") + } + + /** Heap with fewer than K elements drains them all in best-first order. */ + @Test def testFewerThanKReturnsAll(): Unit = { + val heap = new TopKHeap(k = 10, smallerIsBetter = true) + heap.offerAll(Seq(ref(1, 3f), ref(2, 1f), ref(3, 2f))) + assertEquals(Seq(1f, 2f, 3f), heap.drain().map(_.score).toSeq) + } + + /** A worse-than-current-worst candidate is rejected. */ + @Test def testRejectsWorseCandidate(): Unit = { + val heap = new TopKHeap(k = 2, smallerIsBetter = true) + heap.offer(ref(1, 1f)) + heap.offer(ref(2, 2f)) + heap.offer(ref(3, 5f)) // worse than existing 2 → rejected + val drained = heap.drain() + assertEquals(Seq(1f, 2f), drained.map(_.score).toSeq) + assertEquals(Seq(1L, 2L), drained.map(_.rowAddr).toSeq) + } + + /** + * A NaN score must not pin a heap slot forever. Admission derives from the heap's total ordering + * (`java.lang.Float.compare`, which sorts NaN as the largest Float), not a raw float `<` — with + * raw `<`, `1.0f < NaN` is false, so a NaN worst-survivor would test as un-beatable and never be + * evicted. Regression for exactly that: offer NaN then a finite score into a size-1 distance heap; + * the finite score must win. + */ + @Test def testNaNScoreIsEvictable(): Unit = { + val heap = new TopKHeap(k = 1, smallerIsBetter = true) + heap.offer(ref(1, Float.NaN)) + heap.offer(ref(2, 1.0f)) + val drained = heap.drain() + assertEquals(Seq(2L), drained.map(_.rowAddr).toSeq, "finite score must evict the NaN worst") + assertEquals(1.0f, drained.head.score, "size-1 distance heap must retain the finite minimum") + } + + /** `merge` combines two pre-sorted arrays preserving top-K. */ + @Test def testMergeCombinesTwoArrays(): Unit = { + val a = Array(ref(1, 1f), ref(2, 3f), ref(3, 5f)) + val b = Array(ref(4, 2f), ref(5, 4f), ref(6, 6f)) + val merged = TopKHeap.merge(a, b, k = 4, smallerIsBetter = true) + assertEquals(Seq(1f, 2f, 3f, 4f), merged.map(_.score).toSeq) + } + + /** Merging with one empty input is a noop modulo trim to K. */ + @Test def testMergeWithEmpty(): Unit = { + val a = Array(ref(1, 1f), ref(2, 2f), ref(3, 3f)) + val merged = TopKHeap.merge(a, Array.empty[ScoredRowRef], k = 2, smallerIsBetter = true) + assertEquals(Seq(1f, 2f), merged.map(_.score).toSeq) + } +} diff --git a/pom.xml b/pom.xml index 488425ef2..5e400b7cf 100644 --- a/pom.xml +++ b/pom.xml @@ -149,6 +149,7 @@ lance-spark-bundle-4.1_2.13 lance-spark-4.2_2.13 lance-spark-bundle-4.2_2.13 + lance-spark-knn-4.2_2.13