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..180f1a506 --- /dev/null +++ b/lance-spark-knn-4.2_2.13/pom.xml @@ -0,0 +1,111 @@ + + + 4.0.0 + + + org.lance + lance-spark-root + 0.7.1 + ../pom.xml + + + lance-spark-knn-4.2_2.13 + ${project.artifactId} + Catalyst integration for the indexed nearest-by join on Lance (Spark 4.2 SQL, 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-knn_${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/catalyst/IndexedNearestByJoinRule.scala b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/catalyst/IndexedNearestByJoinRule.scala new file mode 100644 index 000000000..655e1ae12 --- /dev/null +++ b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/catalyst/IndexedNearestByJoinRule.scala @@ -0,0 +1,444 @@ +/* + * 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.catalyst + +import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeReference, AttributeSet, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, In, IsNotNull, IsNull, LessThan, LessThanOrEqual, Literal, Not, Or, VectorCosineSimilarity, VectorInnerProduct, VectorL2Distance} +import org.apache.spark.sql.catalyst.plans.{JoinType, LeftOuter, NearestByDirection, NearestByDistance, NearestBySimilarity} +import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan, NearestByJoin, Project, SubqueryAlias} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.connector.catalog.Table +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation +import org.apache.spark.sql.types.{BooleanType, ByteType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, StructField, StructType} +import org.apache.spark.unsafe.types.UTF8String +import org.lance.spark.knn.internal.{LanceKnnJoinStage, Metric} + +/** + * Catalyst rule that rewrites a Spark [[NearestByJoin]] (`approx = true`) over a Lance scan with + * a recognized vector-distance ranking expression into a single [[LanceKnnJoinLogicalPlan]], + * wrapped in a top-level [[Project]] that restores `NearestByJoin.output` exactly. The paired + * [[LanceKnnJoinStrategy]] then lowers that node to [[LanceKnnJoinExec]], which drives the same + * no-shuffle `LanceKnnJoinStage.runPartition` the DataFrame API path uses. + * + * == Why this rule must be a `postHocResolutionRule`, not an optimizer rule == + * + * Spark's built-in [[org.apache.spark.sql.catalyst.optimizer.RewriteNearestByJoin]] rule runs in + * the optimizer's `FinishAnalysis` batch — the very first batch. `injectOptimizerRule` adds + * rules to `operatorOptimizationBatch`, which runs AFTER `FinishAnalysis`. By the time an + * injected optimizer rule fires, the `NearestByJoin` operator has already been replaced with the + * cross-product + `MaxMinByK` rewrite, and we have nothing to pattern-match. + * + * `injectPostHocResolutionRule` runs after analysis but before any optimizer batch — it is the + * only injection point that sees the unrewritten `NearestByJoin`. The same constraint applies to + * any future engine wanting to substitute a different physical strategy for `NearestByJoin`. + * + * == Pattern match == + * + * The rule fires on the conjunction of: + * - `NearestByJoin(_, right, joinType, approx = true, k, rankingExpression, direction)` + * - `right` resolves to a Lance DSv2 relation (immediate or under a `SubqueryAlias`) + * - `rankingExpression` is one of three recognized vector functions, AND its direction matches + * the direction declared on `NearestByJoin`: + * + * | Spark expression | direction | metric | + * |---------------------------------|------------------------|----------------| + * | `VectorL2Distance(L, R)` | `NearestByDistance` | `Metric.L2` | + * | `VectorCosineSimilarity(L, R)` | `NearestBySimilarity` | `Metric.Cosine`| + * | `VectorInnerProduct(L, R)` | `NearestBySimilarity` | `Metric.Dot` | + * + * Any other shape is left alone — Spark's default cross-product rewrite handles it. + * + * The two arguments of the ranking function must each resolve to an [[Attribute]] from one + * specific side of the join. Mixed-side compounds (e.g. `l2_distance(left.vec, left.vec)`) and + * derived expressions (e.g. `l2_distance(left.vec, slice(right.vec, ...))`) are out of scope and + * fall through to the cross-product rewrite. + * + * == Lance scan detection == + * + * Class-name match: `getClass.getName.contains("Lance")`. The probe / materialize path drives + * Lance's Java API directly, so the indexed-path executor is Lance-specific by construction — + * there's no general "any vector-capable backend" extension point here. URI extracted from the + * standard `path` / `datasetUri` option. The rule is opt-in via + * `spark.lance.knn.indexedNearestByJoin.enabled`, so a false positive can only fire when the user + * explicitly enabled the feature against a non-Lance backend; the runtime probe would surface the + * mismatch. + * + * == Prefilter pushdown == + * + * If the right side is a `Filter(cond, lance)` (a `WHERE` clause on the indexed table), the + * rule translates the predicate to a Lance SQL filter string and threads it through to the + * probe. Lance applies the filter BEFORE the index lookup (we always pass `prefilter = true`), + * so the top-K is computed over only the rows matching the filter — the only correct behavior + * for `right WHERE p APPROX NEAREST K`. + * + * Translation is conservative: it handles binary comparisons (=, !=, <, <=, >, >=), `IN`, + * `IS [NOT] NULL`, `AND`/`OR`/`NOT` over right-side attributes vs. literals. Anything else + * (UDFs, subqueries, computed expressions) means the rule REFUSES the rewrite and returns the + * original `NearestByJoin`, falling through to Spark's brute-force cross-product. Refusal — not + * "push what we can, drop the rest" — because dropping a residual would silently change result + * semantics. The job becomes slow rather than wrong. + * + * Filter pushdown into the V2 relation does NOT happen at this point: this rule runs as a + * `postHocResolutionRule` (before the optimizer), so the right side is still the freshly + * analyzed `Filter` over `DataSourceV2Relation` — the V2 `SupportsPushDownFilters` step has not + * yet run. After we rewrite, the right side is absorbed into our plan, so V2 pushdown never + * gets a chance to drop the filter on the floor. + */ +object IndexedNearestByJoinRule extends Rule[LogicalPlan] { + + /** Configuration key that gates the rule. Off by default to keep the rule opt-in for now. */ + val EnabledConfKey: String = "spark.lance.knn.indexedNearestByJoin.enabled" + + /** + * IVF cluster count to visit per query. Higher = better recall, more compute. Default + * (None) leaves Lance's index-default (typically 1). + */ + val NprobesConfKey: String = "spark.lance.knn.nprobes" + + /** + * IVF-PQ refine factor — Lance fetches `K * refineFactor` PQ candidates and re-ranks them + * with exact distance using full vectors. Highest-leverage recall knob for IVF-PQ. Default + * (None) leaves Lance's index-default (= 1, no re-rank). + */ + val RefineFactorConfKey: String = "spark.lance.knn.refineFactor" + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!conf.getConfString(EnabledConfKey, "false").toBoolean) { + return plan + } + val nprobes = optInt(NprobesConfKey) + val refineFactor = optInt(RefineFactorConfKey) + plan.transformDown { + case j @ NearestByJoin(left, right, joinType, true, k, rankingExpr, direction) => + rewriteIfApplicable( + j, + left, + right, + joinType, + k, + rankingExpr, + direction, + nprobes, + refineFactor).getOrElse(j) + } + } + + private def optInt(key: String): Option[Int] = + Option(conf.getConfString(key, null)).map(_.toInt) + + /** + * Rewrite `NearestByJoin` into a single [[LanceKnnJoinLogicalPlan]] carrying the + * [[LanceKnnJoinStage.Conf]] the executor runs per partition — the same stage the DataFrame + * API path drives: + * + * {{{ + * Project(j.output, drop __score) + * +- LanceKnnJoinLogicalPlan output = left ++ right ++ __score + * +- left + * }}} + * + * We add a top-level `Project` because `NearestByJoin.output` is `left ++ right` (no + * score), but the join node emits `left ++ right ++ __score`. The Project slices the trailing + * score attribute — Catalyst's ColumnPruning won't interfere because + * `LanceKnnJoinLogicalPlan` overrides `references = child.outputSet`. + */ + private def rewriteIfApplicable( + j: NearestByJoin, + left: LogicalPlan, + right: LogicalPlan, + joinType: JoinType, + k: Int, + rankingExpr: Expression, + direction: NearestByDirection, + nprobes: Option[Int], + refineFactor: Option[Int]): Option[LogicalPlan] = { + for { + (metric, leftVecAttr, rightVecCol) <- recognizeRanking(rankingExpr, direction, left, right) + lance <- unwrapLanceScan(right) + } yield { + val leftVecIdx = left.output.indexWhere(_.exprId == leftVecAttr.exprId) + require(leftVecIdx >= 0, s"left vector attr not found in left.output: $leftVecAttr") + + val rightFields: Seq[StructField] = + lance.output.map(a => StructField(a.name, a.dataType, nullable = true)) + val rightProjection: Seq[String] = lance.output.map(_.name) + + val stageConf = LanceKnnJoinStage.Conf( + datasetUri = lance.uri, + version = lance.version, + vectorColumn = rightVecCol, + metric = metric, + k = k, + internalK = k, // no overfetch on the SQL path + nprobes = nprobes, + refineFactor = refineFactor, + ef = None, + prefilter = lance.prefilter, + leftVecIdx = leftVecIdx, + rightProjection = rightProjection, + rightFields = rightFields, + leftFieldCount = left.output.size, + outerJoin = joinType == LeftOuter, + smallerIsBetter = metric.smallerIsBetter) + + // The join node emits left ++ right ++ __score. The SQL output is j.output (= left ++ right, + // no score). Set finalOutput = j.output :+ scoreAttr so the node's output is stable; the + // top-level Project drops __score. + // + // `NearestByJoin.output` widens every left+right attribute to `nullable = true` — a contract + // the base Spark rewrite also honors. `finalSchema` feeds the `ExpressionEncoder` in + // `LanceKnnJoinExec.doExecute`; if we left left fields at raw `nullable = false` while the + // logical output declares them nullable, the encoder's binary layout would drift from what + // downstream consumers expect. Widen left here to keep the encoder consistent. + val leftSchemaStruct = StructType( + left.output.map(a => StructField(a.name, a.dataType, a.nullable))) + val scoreAttr = AttributeReference("__score", FloatType, nullable = true)() + val finalSchema = StructType( + leftSchemaStruct.fields.map(_.copy(nullable = true)) ++ + rightFields.map(f => f.copy(nullable = true)) :+ + StructField("__score", FloatType, nullable = true)) + val finalOutput: Seq[Attribute] = j.output :+ scoreAttr + + val node = LanceKnnJoinLogicalPlan( + child = left, + stageConf = stageConf, + leftSchema = leftSchemaStruct, + finalSchema = finalSchema, + finalOutput = finalOutput) + + // Top-level Project drops the __score attr so the plan's external output matches + // NearestByJoin.output exactly. + Project(j.output, node) + } + } + + /** Lance scan info extracted from a DSv2 relation, optionally with a translated prefilter. */ + final private case class LanceScanInfo( + uri: String, + version: Option[Long], + output: Seq[Attribute], + prefilter: Option[String]) + + private def unwrapLanceScan(plan: LogicalPlan): Option[LanceScanInfo] = plan match { + case SubqueryAlias(_, child) => unwrapLanceScan(child) + case v: org.apache.spark.sql.catalyst.plans.logical.View => + // SQL `createOrReplaceTempView` + `spark.sql(... FROM ...)` wraps the underlying + // DataSourceV2Relation in a `View`. Unwrap to find the actual relation underneath. + unwrapLanceScan(v.children.head) + case Filter(cond, child) => + // Right-side `WHERE` clause. Recurse first so we have the relation's output to validate + // attribute references against, then translate the predicate. If translation fails we + // bail entirely (return None, no rewrite) — pushing only PART of a `WHERE` would silently + // change query semantics. The user's filter must be pushed in full or not at all. + unwrapLanceScan(child).flatMap { info => + translateFilter(cond, AttributeSet(info.output)).map { sql => + val combined = info.prefilter match { + case Some(prev) => Some(s"($prev) AND ($sql)") + case None => Some(sql) + } + info.copy(prefilter = combined) + } + } + case Project(projectList, child) if isPassthroughProject(projectList, child) => + // `SELECT * FROM lance` analyzes to `Project(, lance)` — a pass-through + // that preserves attrs and exprIds. Unwrap it. Non-pass-through Projects (renames, drops, + // computed columns) would change the schema we rely on for `j.output` mapping, so we + // refuse those by falling through to the default `_ => None` case. + unwrapLanceScan(child) + case rel: DataSourceV2Relation if isLanceTable(rel.table) => + // The probe / materialize path drives Lance's Java API directly, so this rule is + // Lance-specific by construction — there's no plug-in point for a non-Lance backend + // here. We detect Lance via class-name match and pull the URI from the standard `path` + // / `datasetUri` option. If neither is present we fall through (returning None lets + // Spark's brute-force rewrite handle the query). + val opts = rel.options + val uri = Option(opts.get("path")).orElse(Option(opts.get("datasetUri"))) + uri.map { u => + LanceScanInfo( + uri = u, + version = Option(opts.get("version")).map(_.toLong), + output = rel.output, + prefilter = None) + } + case _ => None + } + + /** + * Translate a Spark `Filter` predicate into a Lance SQL filter string. Returns `None` if any + * sub-expression isn't supported — refusal, not partial pushdown. + * + * Supported shapes (over right-side attributes only): + * - `attr literal` and `literal attr` for `=`, `!=`, `<`, `<=`, `>`, `>=` + * - `attr IS NULL` / `attr IS NOT NULL` + * - `attr IN (lit, lit, ...)` (the IN list must be all foldable literals) + * - `AND` / `OR` over supported children + * - `NOT` over supported child + * + * Anything else — UDFs, joins, subqueries, expressions on both sides referencing the LEFT + * input, computed sub-expressions on the right (e.g. `year(ts) = 2025`) — returns `None`. + * Lance's SQL dialect is DataFusion-flavored; the constructs above all parse identically + * there, so we don't need to translate operator names beyond literal serialization. + */ + private[catalyst] def translateFilter( + expr: Expression, + rightAttrs: AttributeSet): Option[String] = expr match { + case And(l, r) => + for { + a <- translateFilter(l, rightAttrs) + b <- translateFilter(r, rightAttrs) + } yield s"($a) AND ($b)" + case Or(l, r) => + for { + a <- translateFilter(l, rightAttrs) + b <- translateFilter(r, rightAttrs) + } yield s"($a) OR ($b)" + case Not(EqualTo(l, r)) => + // Render `NOT (a = b)` as `(a != b)` so it's the natural Lance form. + binaryOp(l, r, rightAttrs, "!=") + case Not(child) => + translateFilter(child, rightAttrs).map(s => s"NOT ($s)") + case IsNull(c) => + asRightColumn(c, rightAttrs).map(name => s"$name IS NULL") + case IsNotNull(c) => + asRightColumn(c, rightAttrs).map(name => s"$name IS NOT NULL") + case EqualTo(l, r) => binaryOp(l, r, rightAttrs, "=") + case GreaterThan(l, r) => binaryOp(l, r, rightAttrs, ">") + case GreaterThanOrEqual(l, r) => binaryOp(l, r, rightAttrs, ">=") + case LessThan(l, r) => binaryOp(l, r, rightAttrs, "<") + case LessThanOrEqual(l, r) => binaryOp(l, r, rightAttrs, "<=") + case In(value, list) if list.nonEmpty => + for { + col <- asRightColumn(value, rightAttrs) + lits <- list.foldLeft(Option(Vector.empty[String])) { (accOpt, e) => + accOpt.flatMap(acc => asLiteral(e).map(acc :+ _)) + } + } yield s"$col IN (${lits.mkString(", ")})" + case _ => None + } + + private def binaryOp( + l: Expression, + r: Expression, + rightAttrs: AttributeSet, + op: String): Option[String] = { + // attr literal — the natural shape + val attrLit = for { + col <- asRightColumn(l, rightAttrs) + lit <- asLiteral(r) + } yield s"$col $op $lit" + // literal attr — flip when the parser/optimizer emitted args in this order. Renders + // as `lit op col`, which DataFusion also accepts. + attrLit.orElse { + for { + col <- asRightColumn(r, rightAttrs) + lit <- asLiteral(l) + } yield s"$lit $op $col" + } + } + + private def asRightColumn(e: Expression, rightAttrs: AttributeSet): Option[String] = e match { + case a: Attribute if rightAttrs.contains(a) => Some(a.name) + case _ => None + } + + /** + * Render a Spark literal as a Lance SQL literal. Dispatch is by `dataType`, NOT by the boxed + * value class — Catalyst stores e.g. `Literal(0, DateType)` with the value as a plain `Int`, + * so a value-class match would silently let a date literal through as the integer "0", a + * recall-corrupting mistranslation. + * + * Supports nulls, booleans, numeric primitives, and strings (with `'`-escaped quoting). Bails + * on dates, timestamps, decimals, binary, arrays, structs — those have non-trivial cross- + * dialect renderings and we'd rather refuse pushdown than risk a wrong filter. + */ + private def asLiteral(e: Expression): Option[String] = e match { + case Literal(null, _) => Some("NULL") + case Literal(v, BooleanType) => Some(v.toString) + case Literal(v, ByteType) => Some(v.toString) + case Literal(v, ShortType) => Some(v.toString) + case Literal(v, IntegerType) => Some(v.toString) + case Literal(v, LongType) => Some(v.toString) + case Literal(v, FloatType) => Some(v.toString) + case Literal(v, DoubleType) => Some(v.toString) + case Literal(v: UTF8String, StringType) => Some(quoteString(v.toString)) + case Literal(v: String, StringType) => Some(quoteString(v)) + case _ => None + } + + private def quoteString(s: String): String = "'" + s.replace("'", "''") + "'" + + /** + * True iff the Project is the canonical `SELECT *` form: same number of outputs as the child, + * each entry a bare `AttributeReference` whose `exprId` matches the child's output in order. + * Any aliasing, reordering, dropping, or computed column — return false and refuse to + * unwrap, since those change the schema we'd surface as the join's right-side output. + */ + private def isPassthroughProject( + projectList: Seq[org.apache.spark.sql.catalyst.expressions.NamedExpression], + child: LogicalPlan): Boolean = { + val childOut = child.output + if (projectList.size != childOut.size) return false + projectList.zip(childOut).forall { + case (a: Attribute, c) => a.exprId == c.exprId + case _ => false + } + } + + private def isLanceTable(table: Table): Boolean = { + val cls = table.getClass.getName + // Loose by design — the rule is opt-in via spark.lance.knn.indexedNearestByJoin.enabled, so + // a false positive here would only fire when the user explicitly turned the feature on + // against a non-Lance backend, and the runtime probe would surface the mismatch. + cls.contains("Lance") || cls.contains("lance") + } + + /** + * Recognize `rankingExpr` as one of the supported vector-distance functions, AND verify the + * declared `direction` on `NearestByJoin` matches the function's natural ordering. + * + * Returns `(metric, leftVecAttr, rightVecColName)` on success. + */ + private def recognizeRanking( + rankingExpr: Expression, + direction: NearestByDirection, + left: LogicalPlan, + right: LogicalPlan): Option[(Metric, Attribute, String)] = { + val (metric, lhs, rhs) = rankingExpr match { + case VectorL2Distance(l, r) if direction == NearestByDistance => (Metric.L2, l, r) + case VectorCosineSimilarity(l, r) if direction == NearestBySimilarity => (Metric.Cosine, l, r) + case VectorInnerProduct(l, r) if direction == NearestBySimilarity => (Metric.Dot, l, r) + case _ => return None + } + // Each argument must be a bare attribute from one side of the join. + (asAttr(lhs), asAttr(rhs)) match { + case (Some(la), Some(ra)) => + val leftAttrIds = left.outputSet + val rightAttrIds = right.outputSet + if (leftAttrIds.contains(la) && rightAttrIds.contains(ra)) { + Some((metric, la, ra.name)) + } else if (leftAttrIds.contains(ra) && rightAttrIds.contains(la)) { + // Argument order swapped — still valid for symmetric metrics. All three of L2/Cosine/Dot + // are symmetric so we don't have to retain the original orientation. + Some((metric, ra, la.name)) + } else { + None + } + case _ => None + } + } + + private def asAttr(e: Expression): Option[Attribute] = e match { + case a: Attribute => Some(a) + case _ => None + } +} diff --git a/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/catalyst/LanceKnnJoinExec.scala b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/catalyst/LanceKnnJoinExec.scala new file mode 100644 index 000000000..ae62196d3 --- /dev/null +++ b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/catalyst/LanceKnnJoinExec.scala @@ -0,0 +1,78 @@ +/* + * 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.catalyst + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet} +import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.types.StructType +import org.lance.spark.knn.internal.LanceKnnJoinStage + +/** + * Physical operator for the indexed nearest-by join. The whole join is one no-shuffle + * `mapPartitions` over the left input — [[requiredChildDistribution]] is intentionally NOT + * overridden, so Catalyst inserts NO `Exchange` above this node. Each task: + * + * 1. decodes the child's `RDD[InternalRow]` (the left input) into typed `Row`s, + * 2. drives [[LanceKnnJoinStage.runPartition]], which opens R's index once, probes + trims + + * late-materializes per left row (see that object's doc for why a shuffle/merge pipeline only + * adds cost), and + * 3. re-encodes the assembled `left ++ right ++ __score` rows back to `RDD[InternalRow]`. + * + * The decode/encode uses `ExpressionEncoder`; `.copy()` on both sides because Spark reuses the + * `InternalRow` buffer across iterations of the upstream/downstream operators. + */ +case class LanceKnnJoinExec( + override val child: SparkPlan, + stageConf: LanceKnnJoinStage.Conf, + leftSchema: StructType, + finalSchema: StructType, + finalOutput: Seq[Attribute]) + extends UnaryExecNode { + + override def output: Seq[Attribute] = finalOutput + + override def nodeName: String = "LanceKnnJoin" + + // The right-side + score attrs in `output` are synthesised per row from the probe results; + // they do not appear in `child.output`. Declare them produced so Spark's `missingInput` check + // (and the `!` marker in tree-string output) doesn't flag this node. + override def producedAttributes: AttributeSet = AttributeSet(output) -- child.outputSet + + override protected def doExecute(): RDD[InternalRow] = { + val childRdd = child.execute() + val leftSchemaCaptured = leftSchema + val finalSchemaCaptured = finalSchema + val confCaptured = stageConf + + // Encoders are created on the driver and captured into the closure (ExpressionEncoder is + // serializable). The deserializer/serializer instances are NOT thread-safe, so build them + // per partition inside `mapPartitions`. + val leftEnc = ExpressionEncoder(leftSchemaCaptured).resolveAndBind() + val finalEnc = ExpressionEncoder(finalSchemaCaptured).resolveAndBind() + + childRdd.mapPartitions { iter => + val deser = leftEnc.createDeserializer() + val ser = finalEnc.createSerializer() + val leftRows: Iterator[Row] = iter.map(ir => deser(ir.copy())) + LanceKnnJoinStage.runPartition(leftRows, confCaptured).map(row => ser(row).copy()) + } + } + + override protected def withNewChildInternal(newChild: SparkPlan): LanceKnnJoinExec = + copy(child = newChild) +} diff --git a/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/catalyst/LanceKnnJoinLogicalPlan.scala b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/catalyst/LanceKnnJoinLogicalPlan.scala new file mode 100644 index 000000000..6348ffc48 --- /dev/null +++ b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/catalyst/LanceKnnJoinLogicalPlan.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.catalyst + +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeSet} +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, UnaryNode} +import org.apache.spark.sql.types.StructType +import org.lance.spark.knn.internal.LanceKnnJoinStage + +/** + * The single logical node the SQL rewrite ([[IndexedNearestByJoinRule]]) emits for an indexed + * `APPROX NEAREST K` join. Its one child is the LEFT input; the right (Lance) side is captured in + * `stageConf` (URI + version + probe parameters), NOT as a plan child — the join runs entirely + * inside one `mapPartitions` over the left rows, with NO shuffle. The matching + * [[LanceKnnJoinStrategy]] lowers this to [[LanceKnnJoinExec]], which drives the same + * [[LanceKnnJoinStage.runPartition]] the DataFrame API path uses. + * + * `output` is `left ++ right ++ __score` — the right-side and score attributes are synthesised + * here from the probe results, so `producedAttributes = output -- child.outputSet` marks them as + * introduced by this node (Catalyst's `missingInput` check would otherwise flag them). + * + * The `references = child.outputSet` override is load-bearing: the matching exec decodes the WHOLE + * left row per partition to feed the probe, so no left column can be pruned. Without this override + * Catalyst's `ColumnPruning` sees a downstream consumer that references only a subset (or nothing, + * e.g. `count(*)`) and wraps the child in a narrowing `Project`, which would change the row shape + * the executor's left-side encoder expects. + */ +case class LanceKnnJoinLogicalPlan( + override val child: LogicalPlan, + stageConf: LanceKnnJoinStage.Conf, + leftSchema: StructType, + finalSchema: StructType, + finalOutput: Seq[Attribute]) + extends UnaryNode { + + override def output: Seq[Attribute] = finalOutput + + override def producedAttributes: AttributeSet = AttributeSet(output) -- child.outputSet + + override lazy val references: AttributeSet = child.outputSet + + override protected def withNewChildInternal(newChild: LogicalPlan): LanceKnnJoinLogicalPlan = + copy(child = newChild) +} diff --git a/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/catalyst/LanceKnnJoinStrategy.scala b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/catalyst/LanceKnnJoinStrategy.scala new file mode 100644 index 000000000..23eb85b1c --- /dev/null +++ b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/catalyst/LanceKnnJoinStrategy.scala @@ -0,0 +1,36 @@ +/* + * 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.catalyst + +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.execution.{SparkPlan, SparkStrategy} + +/** + * Lowers [[LanceKnnJoinLogicalPlan]] to [[LanceKnnJoinExec]]. Registered via + * `SparkSessionExtensions.injectPlannerStrategy` in + * [[org.lance.spark.knn.extensions.LanceKnnSparkSessionExtensions]]. `planLater(p.child)` defers + * planning of the left input to the rest of the planner. + */ +object LanceKnnJoinStrategy extends SparkStrategy { + override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match { + case p: LanceKnnJoinLogicalPlan => + LanceKnnJoinExec( + planLater(p.child), + p.stageConf, + p.leftSchema, + p.finalSchema, + p.finalOutput) :: Nil + case _ => Nil + } +} diff --git a/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/extensions/LanceKnnSparkSessionExtensions.scala b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/extensions/LanceKnnSparkSessionExtensions.scala new file mode 100644 index 000000000..8a8bdd766 --- /dev/null +++ b/lance-spark-knn-4.2_2.13/src/main/scala/org/lance/spark/knn/extensions/LanceKnnSparkSessionExtensions.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.extensions + +import org.apache.spark.sql.SparkSessionExtensions +import org.lance.spark.knn.catalyst.{IndexedNearestByJoinRule, LanceKnnJoinStrategy} + +/** + * Registers the Catalyst integration for the indexed nearest-by join (the SQL + * `APPROX NEAREST K BY DISTANCE ...` syntax added in Spark 4.2 by SPARK-56395). + * + * Wire this into a SparkSession with: + * + * {{{ + * SparkSession.builder() + * .config("spark.sql.extensions", + * "org.lance.spark.knn.extensions.LanceKnnSparkSessionExtensions") + * .config("spark.lance.knn.indexedNearestByJoin.enabled", "true") + * ... + * }}} + * + * The `enabled` flag gates the rule itself — see [[IndexedNearestByJoinRule.EnabledConfKey]]. Off + * by default to keep the integration opt-in. + * + * == Injection point: postHocResolutionRule, NOT optimizerRule == + * + * Spark's `RewriteNearestByJoin` runs in `FinishAnalysis`, which precedes the + * `operatorOptimizationBatch` that `injectOptimizerRule` adds rules to. By the time an injected + * optimizer rule fires, the `NearestByJoin` operator has already been replaced with the + * cross-product + `MaxMinByK` rewrite. `injectPostHocResolutionRule` runs after analysis but + * before any optimizer batch — this is the only injection point that sees the unrewritten + * `NearestByJoin`. See [[IndexedNearestByJoinRule]]'s class doc for the full rationale. + * + * Coexistence: this extension does not replace `LanceSparkSessionExtensions` from the connector + * modules; both can be wired together in a comma-separated `spark.sql.extensions` value. + */ +class LanceKnnSparkSessionExtensions extends (SparkSessionExtensions => Unit) { + override def apply(extensions: SparkSessionExtensions): Unit = { + extensions.injectPostHocResolutionRule(_ => IndexedNearestByJoinRule) + // Lowers the single `LanceKnnJoinLogicalPlan` the rule emits to `LanceKnnJoinExec` — the same + // no-shuffle `LanceKnnJoinStage.runPartition` the DataFrame API path drives. + extensions.injectPlannerStrategy(_ => LanceKnnJoinStrategy) + } +} diff --git a/lance-spark-knn-4.2_2.13/src/test/scala/org/lance/spark/knn/catalyst/IndexedNearestByJoinE2ETest.scala b/lance-spark-knn-4.2_2.13/src/test/scala/org/lance/spark/knn/catalyst/IndexedNearestByJoinE2ETest.scala new file mode 100644 index 000000000..b1f7b0c7d --- /dev/null +++ b/lance-spark-knn-4.2_2.13/src/test/scala/org/lance/spark/knn/catalyst/IndexedNearestByJoinE2ETest.scala @@ -0,0 +1,322 @@ +/* + * 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.catalyst + +import org.apache.spark.sql.{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 java.nio.file.Path +import java.util.Random + +import scala.collection.JavaConverters._ + +/** + * End-to-end SQL test for the Catalyst integration. Drives the full path: + * + * ANTLR parser ─▶ Analyzer ─▶ IndexedNearestByJoinRule (our postHoc) ─▶ + * Optimizer ─▶ LanceKnnJoinStrategy ─▶ LanceKnnJoinExec ─▶ + * Lance native per-row probe + late materialize ─▶ Rows + * + * Requires Spark 4.2 (the release where `NearestByJoin` exists, added by SPARK-56395) AND the + * `lance-spark-4.2_2.13` connector built against the same Spark version. + * + * Coverage: + * - SQL `APPROX NEAREST k BY DISTANCE vector_l2_distance(...)` parses, the rule rewrites to a + * single `LanceKnnJoinLogicalPlan`, the strategy lowers it to `LanceKnnJoinExec`, which + * executes against a real Lance dataset; results match the brute-force oracle. With no vector + * index built, Lance does an exact per-fragment scan, so any disagreement is a bug. + * - Right-side `WHERE` round-trips through the prefilter pushdown. + * - With the gating config disabled, the same SQL falls through to Spark's + * `RewriteNearestByJoin` (cross-product + `MaxMinByK`) — proves the rule's opt-in behavior. + */ +class IndexedNearestByJoinE2ETest { + + @TempDir var tempDir: Path = _ + private var spark: SparkSession = _ + + private val Dim = 16 + private val NumRight = 64 + private val NumLeft = 8 + private val Seed = 4242L + + @BeforeEach def setup(): Unit = { + spark = SparkSession.builder() + .appName("indexed-nearest-by-join-e2e") + .master("local[2]") + .config("spark.driver.bindAddress", "127.0.0.1") + .config("spark.driver.host", "127.0.0.1") + .config( + "spark.sql.extensions", + "org.lance.spark.knn.extensions.LanceKnnSparkSessionExtensions") + .config("spark.sql.crossJoin.enabled", "true") + .getOrCreate() + spark.sparkContext.setLogLevel("WARN") + } + + @AfterEach def teardown(): Unit = if (spark != null) spark.stop() + + /** + * Full SQL path with the rule enabled. The physical plan must contain the `LanceKnnJoin` exec + * AND the result must match the brute-force oracle on every left row. + */ + @Test def testSqlApproxNearestRoutesThroughIndexedPathAndMatchesOracle(): Unit = { + spark.conf.set(IndexedNearestByJoinRule.EnabledConfKey, "true") + + val (leftRows, leftVectors, leftIds) = generateLeft(NumLeft, Dim, Seed) + val (rightRows, rightVectors, rightIds) = generateRight(NumRight, Dim, Seed + 1) + val rightUri = writeRightLance(rightRows) + + spark.createDataFrame(leftRows.asJava, leftSchema()).createOrReplaceTempView("queries") + spark.read.format("lance").load(rightUri).createOrReplaceTempView("docs") + + val k = 5 + val sql = + s"""SELECT q.lid, d.rid + |FROM queries q INNER JOIN docs d + |APPROX NEAREST $k BY DISTANCE vector_l2_distance(q.lvec, d.rvec)""".stripMargin + val df = spark.sql(sql) + + // Plan-shape: confirm the rule fired (logical node present, AQE-independent) AND the strategy + // lowered it to the `LanceKnnJoin` physical exec. + val optimized = df.queryExecution.optimizedPlan + val joinLogicals = optimized.collect { case p: LanceKnnJoinLogicalPlan => p } + assertTrue( + joinLogicals.nonEmpty, + s"expected LanceKnnJoinLogicalPlan in optimized plan; got:\n$optimized") + val executed = df.queryExecution.executedPlan + val tree = executed.treeString + assertTrue(tree.contains("LanceKnnJoin"), s"expected LanceKnnJoin exec in tree:\n$tree") + + // Correctness: oracle equivalence. + val rows = df.collect() + assertEquals(NumLeft * k, rows.length, "expected k results per left row") + val byLid = rows.groupBy(_.getAs[Int]("lid")) + leftIds.zip(leftVectors).foreach { case (lid, lvec) => + val oracle = rightVectors.indices + .map(i => (rightIds(i), l2(lvec, rightVectors(i)))) + .sortBy(_._2) + .take(k) + .map(_._1) + .toSet + val actual = byLid(lid).map(_.getAs[Int]("rid")).toSet + assertEquals( + oracle, + actual, + s"top-K mismatch for lid=$lid (rule on, brute-force oracle)") + } + } + + /** + * Right-side `WHERE` clause must round-trip through the prefilter pushdown — Lance computes + * top-K only over rows matching the filter, so the result must equal the brute-force oracle + * computed AFTER applying the same filter. If the rule pushed the filter wrong (or dropped + * it), this test would diverge from the oracle. + * + * Two right-side rows in this test share each `category` value, so a `WHERE category = 'A'` + * shrinks the candidate pool meaningfully without zeroing it out. + */ + @Test def testSqlWherePushdownMatchesFilteredOracle(): Unit = { + spark.conf.set(IndexedNearestByJoinRule.EnabledConfKey, "true") + + val (leftRows, leftVectors, leftIds) = generateLeft(NumLeft, Dim, Seed + 200) + val (rightRows, rightVectors, rightIds, rightCategories) = + generateRightWithCategories(NumRight, Dim, Seed + 201) + val rightUri = writeRightLanceWithCategories(rightRows) + + spark.createDataFrame(leftRows.asJava, leftSchema()).createOrReplaceTempView("queries") + spark.read.format("lance").load(rightUri).createOrReplaceTempView("docs") + + val k = 4 + val targetCat = "A" + val sql = + s"""SELECT q.lid, d.rid + |FROM queries q INNER JOIN (SELECT * FROM docs WHERE category = '$targetCat') d + |APPROX NEAREST $k BY DISTANCE vector_l2_distance(q.lvec, d.rvec)""".stripMargin + val df = spark.sql(sql) + + val optimized = df.queryExecution.optimizedPlan + val joinLogicals = optimized.collect { case p: LanceKnnJoinLogicalPlan => p } + assertTrue( + joinLogicals.nonEmpty, + s"expected LanceKnnJoinLogicalPlan; optimized plan was:\n$optimized") + val prefilter = joinLogicals.head.stageConf.prefilter + assertTrue( + prefilter.exists(_.contains(s"'$targetCat'")), + s"expected prefilter to carry category='$targetCat'; got: $prefilter") + + // Oracle: brute-force top-K computed AFTER applying the same filter on the right side. + val filteredIdxs = rightCategories.indices.filter(rightCategories(_) == targetCat) + val rows = df.collect() + val byLid = rows.groupBy(_.getAs[Int]("lid")) + leftIds.zip(leftVectors).foreach { case (lid, lvec) => + val oracle = filteredIdxs + .map(i => (rightIds(i), l2(lvec, rightVectors(i)))) + .sortBy(_._2) + .take(k) + .map(_._1) + .toSet + assertTrue( + oracle.nonEmpty, + s"oracle is empty for lid=$lid — test setup didn't produce filterable rows") + val actual = byLid(lid).map(_.getAs[Int]("rid")).toSet + assertEquals( + oracle, + actual, + s"top-K mismatch under WHERE pushdown for lid=$lid (filtered brute-force oracle)") + } + } + + /** + * With the gating config disabled, the SAME SQL falls through to Spark's + * `RewriteNearestByJoin` (cross-product + `MaxMinByK`). The plan contains NO + * `LanceKnnJoinLogicalPlan` and (importantly) results still match the oracle. This proves + * the rule's opt-in behavior at the SQL level: turning it off doesn't break correctness. + */ + @Test def testSqlFallsThroughToBruteForceWhenRuleDisabled(): Unit = { + spark.conf.set(IndexedNearestByJoinRule.EnabledConfKey, "false") + + val (leftRows, leftVectors, leftIds) = generateLeft(NumLeft, Dim, Seed + 100) + val (rightRows, rightVectors, rightIds) = generateRight(NumRight, Dim, Seed + 101) + val rightUri = writeRightLance(rightRows) + + spark.createDataFrame(leftRows.asJava, leftSchema()).createOrReplaceTempView("queries") + spark.read.format("lance").load(rightUri).createOrReplaceTempView("docs") + + val k = 4 + val df = spark.sql( + s"""SELECT q.lid, d.rid + |FROM queries q INNER JOIN docs d + |APPROX NEAREST $k BY DISTANCE vector_l2_distance(q.lvec, d.rvec)""".stripMargin) + + val optimized = df.queryExecution.optimizedPlan + val joinLogicals = optimized.collect { case p: LanceKnnJoinLogicalPlan => p } + assertTrue( + joinLogicals.isEmpty, + s"rule disabled — expected NO LanceKnnJoinLogicalPlan; got:\n$optimized") + + val rows = df.collect() + assertEquals(NumLeft * k, rows.length) + val byLid = rows.groupBy(_.getAs[Int]("lid")) + leftIds.zip(leftVectors).foreach { case (lid, lvec) => + val oracle = rightVectors.indices + .map(i => (rightIds(i), l2(lvec, rightVectors(i)))) + .sortBy(_._2) + .take(k) + .map(_._1) + .toSet + val actual = byLid(lid).map(_.getAs[Int]("rid")).toSet + assertEquals(oracle, actual, s"top-K mismatch for lid=$lid (rule off, brute-force fallback)") + } + } + + // -- helpers ------------------------------------------------------------------------------ + + private def leftSchema(): StructType = new StructType(Array( + StructField("lid", IntegerType, nullable = false), + StructField( + "lvec", + ArrayType(FloatType, containsNull = false), + nullable = false, + new MetadataBuilder().putLong("arrow.fixed-size-list.size", Dim.toLong).build()))) + + private def rightSchema(): StructType = new StructType(Array( + StructField("rid", IntegerType, nullable = false), + StructField( + "rvec", + ArrayType(FloatType, containsNull = false), + nullable = false, + new MetadataBuilder().putLong("arrow.fixed-size-list.size", Dim.toLong).build()))) + + private def generateLeft( + n: Int, + dim: Int, + seed: Long): (Seq[org.apache.spark.sql.Row], Array[Array[Float]], Array[Int]) = { + val rng = new Random(seed) + val vectors = (0 until n).map(_ => randomVector(rng, dim)).toArray + val ids = (0 until n).toArray + val rows = ids.zip(vectors).map { case (id, v) => RowFactory.create(Integer.valueOf(id), v) } + (rows.toSeq, vectors, ids) + } + + private def generateRight( + n: Int, + dim: Int, + seed: Long): (Seq[org.apache.spark.sql.Row], Array[Array[Float]], Array[Int]) = { + val rng = new Random(seed) + val vectors = (0 until n).map(_ => randomVector(rng, dim)).toArray + val ids = (0 until n).map(_ + 1000).toArray + val rows = ids.zip(vectors).map { case (id, v) => RowFactory.create(Integer.valueOf(id), v) } + (rows.toSeq, vectors, ids) + } + + private def writeRightLance(rows: Seq[org.apache.spark.sql.Row]): String = { + val df = spark.createDataFrame(rows.asJava, rightSchema()) + val out = tempDir.resolve(s"right_${System.nanoTime()}").toString + df.write.format("lance").save(out) + out + } + + private def rightSchemaWithCategories(): StructType = new StructType(Array( + StructField("rid", IntegerType, nullable = false), + StructField("category", StringType, nullable = false), + StructField( + "rvec", + ArrayType(FloatType, containsNull = false), + nullable = false, + new MetadataBuilder().putLong("arrow.fixed-size-list.size", Dim.toLong).build()))) + + /** + * Like `generateRight`, but every row also carries a category drawn from a small alphabet so + * the e2e WHERE-pushdown test has a non-trivial filter to apply. + */ + private def generateRightWithCategories(n: Int, dim: Int, seed: Long): ( + Seq[org.apache.spark.sql.Row], + Array[Array[Float]], + Array[Int], + Array[String]) = { + val rng = new Random(seed) + val vectors = (0 until n).map(_ => randomVector(rng, dim)).toArray + val ids = (0 until n).map(_ + 2000).toArray + val alphabet = Array("A", "B", "C", "D") + val categories = (0 until n).map(i => alphabet(i % alphabet.length)).toArray + val rows = ids.zip(vectors).zip(categories).map { case ((id, v), cat) => + RowFactory.create(Integer.valueOf(id), cat, v) + } + (rows.toSeq, vectors, ids, categories) + } + + private def writeRightLanceWithCategories(rows: Seq[org.apache.spark.sql.Row]): String = { + val df = spark.createDataFrame(rows.asJava, rightSchemaWithCategories()) + val out = tempDir.resolve(s"right_cat_${System.nanoTime()}").toString + df.write.format("lance").save(out) + out + } + + 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 l2(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/catalyst/IndexedNearestByJoinRuleTest.scala b/lance-spark-knn-4.2_2.13/src/test/scala/org/lance/spark/knn/catalyst/IndexedNearestByJoinRuleTest.scala new file mode 100644 index 000000000..463d499d2 --- /dev/null +++ b/lance-spark-knn-4.2_2.13/src/test/scala/org/lance/spark/knn/catalyst/IndexedNearestByJoinRuleTest.scala @@ -0,0 +1,454 @@ +/* + * 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.catalyst + +import org.apache.spark.sql.{RowFactory, SparkSession} +import org.apache.spark.sql.catalyst.expressions.{Add, And, Attribute, AttributeSet, EqualTo, Expression, GreaterThan, In, IsNotNull, IsNull, LessThanOrEqual, Literal, Not, Or, VectorCosineSimilarity, VectorInnerProduct, VectorL2Distance} +import org.apache.spark.sql.catalyst.plans.Inner +import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan, NearestByJoin, Project, SubqueryAlias} +import org.apache.spark.sql.catalyst.plans.{NearestByDistance, NearestBySimilarity} +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation +import org.apache.spark.sql.types._ +import org.apache.spark.unsafe.types.UTF8String +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.knn.internal.Metric + +import java.nio.file.Path + +import scala.collection.JavaConverters._ + +/** + * Unit tests for [[IndexedNearestByJoinRule]]. The rule's responsibility is purely Catalyst-side + * pattern-matching — we don't need a Lance backend to exercise it. Each test constructs a small + * resolved plan and runs the rule, asserting either a rewrite to + * `Project(..., LanceKnnJoinLogicalPlan(left, ...))` or a no-op fallthrough. + * + * Coverage: + * - Happy path: VectorL2Distance + NearestByDistance over a Lance DSv2 relation rewrites. + * - Direction mismatch (e.g. L2 distance with NearestBySimilarity) does NOT rewrite. + * - EXACT (`approx = false`) does NOT rewrite — Spark's brute-force keeps owning that path. + * - Non-Lance right side does NOT rewrite (duck-type check via class name). + * - Disabled by default — fires only when the gating config is set. + * - Prefilter pushdown: right-side `WHERE` translates to a Lance SQL filter string, or refuses + * the rewrite entirely when the predicate can't be pushed in full. + * + * The rule's runtime behavior beyond the rewrite (probe execution against real Lance) is covered + * by the oracle tests in lance-spark-knn_2.12 and the e2e test in this module. + */ +class IndexedNearestByJoinRuleTest { + + @TempDir var tempDir: Path = _ + private var spark: SparkSession = _ + + @BeforeEach def setup(): Unit = { + spark = SparkSession.builder() + .appName("indexed-nearest-by-join-rule-test") + .master("local[2]") + .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() + + /** L2 + NearestByDistance + Lance scan + enabled config → rewrite. */ + @Test def testL2RewritesToIndexedPlan(): Unit = { + spark.conf.set(IndexedNearestByJoinRule.EnabledConfKey, "true") + val (left, leftVec, right, rightVec) = buildPlans(metricFunction = "l2") + val join = NearestByJoin( + left = left, + right = right, + joinType = Inner, + approx = true, + numResults = 5, + rankingExpression = VectorL2Distance(leftVec, rightVec), + direction = NearestByDistance) + val rewritten = IndexedNearestByJoinRule(join) + val plan = expectRewritten(rewritten) + assertEquals(Metric.L2, plan.metric) + assertEquals(5, plan.k) + assertEquals(rightVec.name, plan.rightVecCol) + assertEquals(leftVec.exprId, plan.leftVecAttr.exprId) + } + + /** Cosine similarity + NearestBySimilarity → rewrite. */ + @Test def testCosineRewrites(): Unit = { + spark.conf.set(IndexedNearestByJoinRule.EnabledConfKey, "true") + val (left, leftVec, right, rightVec) = buildPlans(metricFunction = "cosine") + val join = NearestByJoin( + left, + right, + Inner, + approx = true, + numResults = 3, + rankingExpression = VectorCosineSimilarity(leftVec, rightVec), + direction = NearestBySimilarity) + val rewritten = IndexedNearestByJoinRule(join) + assertEquals(Metric.Cosine, expectRewritten(rewritten).metric) + } + + /** Inner product + NearestBySimilarity → rewrite as Dot. */ + @Test def testDotRewrites(): Unit = { + spark.conf.set(IndexedNearestByJoinRule.EnabledConfKey, "true") + val (left, leftVec, right, rightVec) = buildPlans(metricFunction = "dot") + val join = NearestByJoin( + left, + right, + Inner, + approx = true, + numResults = 4, + rankingExpression = VectorInnerProduct(leftVec, rightVec), + direction = NearestBySimilarity) + val rewritten = IndexedNearestByJoinRule(join) + assertEquals(Metric.Dot, expectRewritten(rewritten).metric) + } + + /** L2 distance with NearestBySimilarity is inconsistent — rule should NOT fire. */ + @Test def testDirectionMismatchDoesNotRewrite(): Unit = { + spark.conf.set(IndexedNearestByJoinRule.EnabledConfKey, "true") + val (left, leftVec, right, rightVec) = buildPlans(metricFunction = "l2") + val join = NearestByJoin( + left, + right, + Inner, + approx = true, + numResults = 5, + rankingExpression = VectorL2Distance(leftVec, rightVec), + direction = NearestBySimilarity) + val rewritten = IndexedNearestByJoinRule(join) + assertSame(join, rewritten, "rule should not fire on direction/metric mismatch") + } + + /** EXACT mode (approx = false) is owned by Spark's brute-force rewrite. */ + @Test def testExactModeDoesNotRewrite(): Unit = { + spark.conf.set(IndexedNearestByJoinRule.EnabledConfKey, "true") + val (left, leftVec, right, rightVec) = buildPlans(metricFunction = "l2") + val join = NearestByJoin( + left, + right, + Inner, + approx = false, + numResults = 5, + rankingExpression = VectorL2Distance(leftVec, rightVec), + direction = NearestByDistance) + val rewritten = IndexedNearestByJoinRule(join) + assertSame(join, rewritten, "EXACT queries must not be intercepted") + } + + /** Disabled flag (default) → no rewrite even when otherwise applicable. */ + @Test def testDisabledByDefault(): Unit = { + spark.conf.unset(IndexedNearestByJoinRule.EnabledConfKey) + val (left, leftVec, right, rightVec) = buildPlans(metricFunction = "l2") + val join = NearestByJoin( + left, + right, + Inner, + approx = true, + numResults = 5, + rankingExpression = VectorL2Distance(leftVec, rightVec), + direction = NearestByDistance) + val rewritten = IndexedNearestByJoinRule(join) + assertSame(join, rewritten, "rule must be opt-in") + } + + /** Non-Lance right side (regular DataFrame as Project, no DSv2 relation) → no rewrite. */ + @Test def testNonLanceRightDoesNotRewrite(): Unit = { + spark.conf.set(IndexedNearestByJoinRule.EnabledConfKey, "true") + val left = trivialPlan("lid", "lvec") + val right = trivialPlan("rid", "rvec") + val leftVec = left.output.find(_.name == "lvec").get + val rightVec = right.output.find(_.name == "rvec").get + val join = NearestByJoin( + left, + right, + Inner, + approx = true, + numResults = 5, + rankingExpression = VectorL2Distance(leftVec, rightVec), + direction = NearestByDistance) + val rewritten = IndexedNearestByJoinRule(join) + assertSame(join, rewritten, "non-Lance right must fall through") + } + + /** Right side wrapped in SubqueryAlias still rewrites — alias unwrapping happens in the rule. */ + @Test def testSubqueryAliasOnRightStillRewrites(): Unit = { + spark.conf.set(IndexedNearestByJoinRule.EnabledConfKey, "true") + val (left, leftVec, right, rightVec) = buildPlans(metricFunction = "l2") + val aliased = SubqueryAlias("d", right) + val join = NearestByJoin( + left, + aliased, + Inner, + approx = true, + numResults = 5, + rankingExpression = VectorL2Distance(leftVec, rightVec), + direction = NearestByDistance) + val rewritten = IndexedNearestByJoinRule(join) + // Rule emits `Project(j.output, LanceKnnJoinLogicalPlan(left, ...))`. Asserting on the top + // Project wrapping the join node is enough for the "did the rule fire" check. + assertTrue( + rewritten.isInstanceOf[Project] && + rewritten.asInstanceOf[Project].child.isInstanceOf[LanceKnnJoinLogicalPlan], + s"expected Project(..., LanceKnnJoinLogicalPlan(...)), got: " + + s"${rewritten.getClass.getSimpleName}") + } + + // -- prefilter pushdown ------------------------------------------------------------------- + + /** + * Right side wrapped in `Filter(simple predicate)` rewrites AND the predicate lands on the + * indexed plan as a Lance SQL filter string. The filter must be pushed in full (not dropped) + * for the result to be semantically equivalent to the original plan. + */ + @Test def testFilterOverLancePushesAsPrefilter(): Unit = { + spark.conf.set(IndexedNearestByJoinRule.EnabledConfKey, "true") + val (left, leftVec, right, rightVec) = buildPlans(metricFunction = "l2") + val category = right.output.find(_.name == "category").get + val bucket = right.output.find(_.name == "bucket").get + val cond = And( + EqualTo(category, Literal(UTF8String.fromString("A"), StringType)), + GreaterThan(bucket, Literal(5, IntegerType))) + val filtered = Filter(cond, right) + val join = NearestByJoin( + left, + filtered, + Inner, + approx = true, + numResults = 5, + rankingExpression = VectorL2Distance(leftVec, rightVec), + direction = NearestByDistance) + val rewritten = IndexedNearestByJoinRule(join) + val plan = expectRewritten(rewritten) + assertTrue(plan.prefilter.isDefined, "prefilter should be populated") + val sql = plan.prefilter.get + assertTrue(sql.contains("category"), s"prefilter missing column ref: $sql") + assertTrue(sql.contains("'A'"), s"prefilter missing string literal: $sql") + assertTrue(sql.contains("bucket"), s"prefilter missing column ref: $sql") + assertTrue(sql.contains("> 5"), s"prefilter missing numeric comparison: $sql") + assertTrue(sql.contains("AND"), s"prefilter missing conjunction: $sql") + } + + /** + * Predicate touches a left-side attribute — translator can't safely render that as a Lance + * SQL string (Lance only sees the right table's columns). Rule must REFUSE the rewrite, not + * drop the predicate. We verify the original `NearestByJoin` is returned unchanged. + */ + @Test def testPredicateReferencingLeftAttrRefusesRewrite(): Unit = { + spark.conf.set(IndexedNearestByJoinRule.EnabledConfKey, "true") + val (left, leftVec, right, rightVec) = buildPlans(metricFunction = "l2") + val lid = left.output.find(_.name == "lid").get + val cond = EqualTo(lid, Literal(0, IntegerType)) + val filtered = Filter(cond, right) + val join = NearestByJoin( + left, + filtered, + Inner, + approx = true, + numResults = 5, + rankingExpression = VectorL2Distance(leftVec, rightVec), + direction = NearestByDistance) + val rewritten = IndexedNearestByJoinRule(join) + assertSame( + join, + rewritten, + "predicate touching left side must refuse pushdown — not partial-push") + } + + /** + * Predicate is a computed expression (e.g. `bucket + 1 = 6`), not a bare attr-vs-literal + * comparison. Translator returns None, rule refuses. + */ + @Test def testComputedPredicateRefusesRewrite(): Unit = { + spark.conf.set(IndexedNearestByJoinRule.EnabledConfKey, "true") + val (left, leftVec, right, rightVec) = buildPlans(metricFunction = "l2") + val bucket = right.output.find(_.name == "bucket").get + val cond = EqualTo(Add(bucket, Literal(1, IntegerType)), Literal(6, IntegerType)) + val filtered = Filter(cond, right) + val join = NearestByJoin( + left, + filtered, + Inner, + approx = true, + numResults = 5, + rankingExpression = VectorL2Distance(leftVec, rightVec), + direction = NearestByDistance) + val rewritten = IndexedNearestByJoinRule(join) + assertSame(join, rewritten, "computed expression must refuse pushdown") + } + + /** Filter wrapped in SubqueryAlias still pushes — order of unwrap shouldn't matter. */ + @Test def testFilterUnderSubqueryAliasPushes(): Unit = { + spark.conf.set(IndexedNearestByJoinRule.EnabledConfKey, "true") + val (left, leftVec, right, rightVec) = buildPlans(metricFunction = "l2") + val category = right.output.find(_.name == "category").get + val cond = EqualTo(category, Literal(UTF8String.fromString("X"), StringType)) + val plan = SubqueryAlias("d", Filter(cond, right)) + val join = NearestByJoin( + left, + plan, + Inner, + approx = true, + numResults = 3, + rankingExpression = VectorL2Distance(leftVec, rightVec), + direction = NearestByDistance) + val rewritten = IndexedNearestByJoinRule(join) + val p = expectRewritten(rewritten) + assertTrue(p.prefilter.isDefined, s"prefilter should be set; got ${p.prefilter}") + } + + // -- predicate translator unit tests ----------------------------------------------------- + + /** + * Direct unit tests on `translateFilter` to lock in the supported shapes. Uses a synthetic + * AttributeSet so we don't need a logical plan. + */ + @Test def testTranslatorHandlesSupportedShapes(): Unit = { + val rid = makeAttr("rid", IntegerType) + val category = makeAttr("category", StringType) + val bucket = makeAttr("bucket", IntegerType) + val attrs = AttributeSet(Seq(rid, category, bucket)) + + val cases: Seq[(Expression, String)] = Seq( + EqualTo(category, lit("A")) -> "category = 'A'", + Not(EqualTo(category, lit("A"))) -> "category != 'A'", + GreaterThan(bucket, lit(5)) -> "bucket > 5", + LessThanOrEqual(bucket, lit(5)) -> "bucket <= 5", + IsNull(category) -> "category IS NULL", + IsNotNull(category) -> "category IS NOT NULL", + In(bucket, Seq(lit(1), lit(2), lit(3))) -> "bucket IN (1, 2, 3)", + And(EqualTo(category, lit("A")), GreaterThan(bucket, lit(5))) -> + "(category = 'A') AND (bucket > 5)", + Or(EqualTo(category, lit("A")), EqualTo(category, lit("B"))) -> + "(category = 'A') OR (category = 'B')", + // String-literal escape — single quotes inside the value get doubled. + EqualTo(category, lit("O'Brien")) -> "category = 'O''Brien'", + // literal-on-left flip + EqualTo(lit(5), bucket) -> "5 = bucket") + cases.foreach { case (expr, expected) => + val got = IndexedNearestByJoinRule.translateFilter(expr, attrs) + assertEquals(Some(expected), got, s"translation mismatch for: $expr") + } + } + + /** Translator must return None for unsupported expressions so the rule refuses pushdown. */ + @Test def testTranslatorRefusesUnsupportedShapes(): Unit = { + val rid = makeAttr("rid", IntegerType) + val ts = makeAttr("ts", DateType) // date literals not in our supported set + val attrs = AttributeSet(Seq(rid, ts)) + + val rejected: Seq[Expression] = Seq( + // Two attributes — no literal — translator can't render `attr op attr` safely (Lance can, + // but we don't promise it; refuse to keep the rule conservative). + EqualTo(rid, makeAttr("rid2", IntegerType)), + // Foreign attribute (not in `attrs`) — translator must reject. + EqualTo(makeAttr("foreign", IntegerType), lit(1)), + // Empty IN list. + In(rid, Seq.empty), + // Date literal — out of supported types. + EqualTo(ts, Literal(0, DateType))) + rejected.foreach { e => + assertEquals( + None, + IndexedNearestByJoinRule.translateFilter(e, attrs), + s"expected refusal for: $e") + } + } + + // -- helpers ------------------------------------------------------------------------------ + + /** + * Construct a left-side regular plan and a right-side that resembles a Lance DSv2 scan via the + * duck-type check. Avoids the need for a real Lance reader. + */ + private def buildPlans(metricFunction: String) + : (LogicalPlan, Attribute, LogicalPlan, Attribute) = { + val left = trivialPlan("lid", "lvec") + val rightLance = lanceLikeDsv2Relation() + val leftVec = left.output.find(_.name == "lvec").get + val rightVec = rightLance.output.find(_.name == "rvec").get + (left, leftVec, rightLance, rightVec) + } + + private def trivialPlan(idCol: String, vecCol: String): LogicalPlan = { + val schema = new StructType(Array( + StructField(idCol, IntegerType, nullable = false), + StructField(vecCol, ArrayType(FloatType, containsNull = false), nullable = false))) + val rows = (0 until 4).map(i => RowFactory.create(Integer.valueOf(i), Array.fill(8)(0.0f))) + spark.createDataFrame(rows.asJava, schema).queryExecution.analyzed + } + + /** + * Build a `DataSourceV2Relation` whose `table.getClass.getName.contains("Lance")` so the + * rule's duck-type check accepts it. We don't actually run any I/O. Includes a `category` + * (string) and `bucket` (int) column so prefilter-pushdown tests can build realistic + * filter predicates without needing to extend the schema separately. + */ + private def lanceLikeDsv2Relation(): LogicalPlan = { + val schema = new StructType(Array( + StructField("rid", IntegerType, nullable = false), + StructField("category", StringType, nullable = true), + StructField("bucket", IntegerType, nullable = true), + StructField("rvec", ArrayType(FloatType, containsNull = false), nullable = false))) + val table = new FakeLanceTable(schema) + val opts = new java.util.HashMap[String, String]() + opts.put("path", tempDir.resolve("fake_lance").toString) + val cims = new org.apache.spark.sql.util.CaseInsensitiveStringMap(opts) + DataSourceV2Relation.create(table, None, None, cims) + } + + /** + * Extract an assertion-friendly summary of the rule's rewrite output. The rule produces + * `Project(j.output, LanceKnnJoinLogicalPlan(left, stageConf, ...))`; this helper pulls out the + * fields the test cases want to check straight off `stageConf`. + */ + private case class RewriteSummary( + metric: Metric, + k: Int, + rightVecCol: String, + leftVecAttr: Attribute, + prefilter: Option[String]) + + private def expectRewritten(plan: LogicalPlan): RewriteSummary = plan match { + case Project(_, node: LanceKnnJoinLogicalPlan) => + val conf = node.stageConf + RewriteSummary( + metric = conf.metric, + k = conf.k, + rightVecCol = conf.vectorColumn, + leftVecAttr = node.child.output(conf.leftVecIdx), + prefilter = conf.prefilter) + case other => + fail(s"expected Project(LanceKnnJoinLogicalPlan(...)), got: $other"); ??? + } + + private def makeAttr(name: String, dt: DataType): Attribute = + org.apache.spark.sql.catalyst.expressions.AttributeReference(name, dt, nullable = true)() + + private def lit(v: Int): Literal = Literal(v, IntegerType) + private def lit(s: String): Literal = Literal(UTF8String.fromString(s), StringType) +} + +/** + * Stub Table whose class name ends with "Lance" so the rule's duck-type check accepts it. No I/O + * — the rule only reads schema and options. Lives in the test source tree. + */ +class FakeLanceTable(_schema: StructType) extends org.apache.spark.sql.connector.catalog.Table { + override def name(): String = "fake_lance" + override def schema(): StructType = _schema + override def capabilities() + : java.util.Set[org.apache.spark.sql.connector.catalog.TableCapability] = + java.util.Collections.emptySet() +} diff --git a/lance-spark-knn_2.12/pom.xml b/lance-spark-knn_2.12/pom.xml new file mode 100644 index 000000000..9dc20cf66 --- /dev/null +++ b/lance-spark-knn_2.12/pom.xml @@ -0,0 +1,73 @@ + + + 4.0.0 + + + org.lance + lance-spark-root + 0.7.1 + ../pom.xml + + + lance-spark-knn_2.12 + ${project.artifactId} + Indexed nearest-neighbor join for Lance datasets in Spark + jar + + + + org.lance + lance-spark-base_2.12 + ${project.version} + + + org.apache.spark + spark-sql_${scala.compat.version} + provided + + + + org.lance + lance-spark-3.5_2.12 + ${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 + + + + + + diff --git a/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/IndexedNearestJoin.scala b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/IndexedNearestJoin.scala new file mode 100644 index 000000000..1aee5714b --- /dev/null +++ b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/IndexedNearestJoin.scala @@ -0,0 +1,132 @@ +/* + * 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 + +import org.apache.spark.sql.{DataFrame, Row} +import org.apache.spark.sql.types._ +import org.lance.spark.knn.internal.{LanceKnnJoinStage, Metric} + +/** + * Public entry point for the indexed nearest-by join over a Lance dataset. + * + * The join runs entirely inside one `mapPartitions` over the left DataFrame, with NO shuffle: each + * partition opens R's index once and, per left row, runs a native `LanceProbe.probe(...)` (a + * complete distributed top-K search in Lance's own threads), trims the overfetched candidates to + * `k`, and late-materializes the surviving right rows by `_rowid`. See + * [[org.lance.spark.knn.internal.LanceKnnJoinStage]] for why a probe → shuffle → merge → + * materialize pipeline only adds cost over letting the single native call do the work. + * + * The result is assembled through public Spark APIs (`left.rdd.mapPartitions` + + * `SparkSession.createDataFrame`) rather than a custom Catalyst node; the SQL module wires the same + * `LanceKnnJoinStage.runPartition` into a physical operator for the `NEAREST BY` rewrite. + */ +object IndexedNearestJoin { + + /** + * Run an approximate-nearest-neighbor join. + * + * @param left left DataFrame; one query vector per row in `leftVecCol` + * @param rightLanceUri Lance dataset URI for the right side + * @param leftVecCol name of the vector column in `left`. Must be `ArrayType[Float]`. + * @param rightVecCol name of the indexed (or to-be-searched) vector column on the right + * @param k top-K rows per left row + * @param metric distance/similarity metric: "l2" / "cosine" / "dot" (and synonyms) + * @param rightProjection columns to materialize from the right side. Defaults to all data + * columns. The score column is added separately. + * @param outerJoin when true, left rows with zero matches are preserved with NULL right- + * side columns. Defaults to false (inner-join semantics). + * @param scoreCol name of the synthesized score column added to the output. Defaults to + * `__score`. + * @param overfetch ratio of internal candidates to `k`. Lance fetches `k × overfetch` + * candidates natively and this stage trims to `k`. Defaults to 1. + * @param nprobes optional override of Lance's `nprobes` for IVF-PQ indexes + * @param version optional Lance version pin; if unset, latest version is used + * @param refineFactor IVF-PQ recall knob. When set, Lance fetches `k * refineFactor` + * approximate candidates, re-ranks them with exact distance, and trims + * back to k. Higher = better recall, more compute. `None` leaves Lance's + * default (= 1, no re-rank). Ignored for non-IVF-PQ indexes / unindexed. + * @param ef HNSW search depth. Higher = better recall, more compute. `None` leaves + * Lance's default (the index's build-time `ef_construction` value). + * Ignored for non-HNSW indexes / unindexed. + */ + def apply( + left: DataFrame, + rightLanceUri: String, + leftVecCol: String, + rightVecCol: String, + k: Int, + metric: String = "l2", + rightProjection: Option[Seq[String]] = None, + outerJoin: Boolean = false, + scoreCol: String = "__score", + overfetch: Int = 1, + nprobes: Option[Int] = None, + version: Option[Long] = None, + refineFactor: Option[Int] = None, + ef: Option[Int] = None): DataFrame = { + + require(k > 0, "k must be positive") + require(overfetch >= 1, "overfetch must be >= 1") + + val spark = left.sparkSession + val parsedMetric = Metric.fromName(metric) + val internalK = k * overfetch + + // Snapshot right-side schema on the driver before any executor work happens. + val rightSchema: StructType = { + val reader = spark.read.format("lance") + version.foreach(v => reader.option("version", v.toString)) + val raw = reader.load(rightLanceUri) + val pruned = rightProjection match { + case Some(cols) if cols.nonEmpty => raw.select(cols.head, cols.tail: _*) + case _ => raw + } + pruned.schema + } + + val outputSchema = buildOutputSchema(left.schema, rightSchema, scoreCol) + val rightProjectionCols: Seq[String] = + rightProjection.getOrElse(rightSchema.fieldNames.toSeq) + + val conf = LanceKnnJoinStage.Conf( + datasetUri = rightLanceUri, + version = version, + vectorColumn = rightVecCol, + metric = parsedMetric, + k = k, + internalK = internalK, + nprobes = nprobes, + refineFactor = refineFactor, + ef = ef, + prefilter = None, + leftVecIdx = left.schema.fieldIndex(leftVecCol), + rightProjection = rightProjectionCols, + rightFields = rightSchema.fields.toSeq, + leftFieldCount = left.schema.fields.length, + outerJoin = outerJoin, + smallerIsBetter = parsedMetric.smallerIsBetter) + + val rowRdd = left.rdd.mapPartitions(iter => LanceKnnJoinStage.runPartition(iter, conf)) + spark.createDataFrame(rowRdd, outputSchema) + } + + private def buildOutputSchema( + left: StructType, + right: StructType, + scoreCol: String): StructType = { + val rightNullable = right.fields.map(f => f.copy(nullable = true)) + val score = StructField(scoreCol, FloatType, nullable = true) + StructType(left.fields ++ rightNullable :+ score) + } +} diff --git a/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/LanceKnnImplicits.scala b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/LanceKnnImplicits.scala new file mode 100644 index 000000000..306a3d482 --- /dev/null +++ b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/LanceKnnImplicits.scala @@ -0,0 +1,155 @@ +/* + * 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 + +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation +import org.lance.spark.knn.internal.LanceKnnSizeGate + +/** + * Idiomatic DataFrame extension for the indexed nearest-K join. The SQL syntax + * (`APPROX NEAREST K BY DISTANCE ...`) requires Spark 4.2+ because that's where the + * `NearestByJoin` operator landed. The DataFrame API path here works on every Spark version + * the lance-spark connector supports (3.5, 4.0, 4.1, 4.2+) — it just calls Lance's Java probe + * API directly through `IndexedNearestJoin.apply`, no Catalyst rule, no SQL. + * + * Usage: + * {{{ + * import org.lance.spark.knn.LanceKnnImplicits._ + * + * val docs = spark.read.format("lance").load("/path/to/lance/dataset") + * val joined = queries.kNearestJoin( + * right = docs, + * leftVecCol = "qvec", + * rightVecCol = "vec", + * k = 10, + * metric = "l2") + * }}} + * + * The right DataFrame must be a Lance scan (`spark.read.format("lance").load(uri)`): the extension + * extracts the underlying dataset URI from the right-side analyzed plan and runs the native probe + * pipeline against it directly. A non-Lance right (parquet, delta, in-memory, arbitrary subplan) + * is rejected with a clear error — build a Lance dataset (with a vector index) for R first. + * + * == Why an extension method, not a builder == + * + * Builder-style APIs (`new KNearestJoin(...).build()`) are heavier syntactically than what + * users want for what should be a one-line call. The extension method makes the verb + * (`kNearestJoin`) hang off the left DataFrame the same way `join` does, so users discover it + * via IDE autocomplete and can reach for it without learning a new pattern. + */ +object LanceKnnImplicits { + + implicit class LanceKnnDataFrameOps(val df: DataFrame) extends AnyVal { + + /** + * Approximate top-K nearest-neighbor join against a Lance-backed right DataFrame. + * + * @param right right DataFrame — must be a Lance scan + * @param leftVecCol name of the vector column on `this` (left) + * @param rightVecCol name of the vector column on `right` + * @param k number of nearest neighbors per left row + * @param metric distance / similarity metric: "l2" | "cosine" | "dot" + * @param rightProjection columns to materialize from `right`. `None` = all of R's columns. + * @param outerJoin left-outer mode: emit a left row even if zero neighbors found + * @param scoreCol name of the appended score column (default `__score`) + * @param overfetch multiplier on `k` during the probe before final trim + * @param nprobes IVF cluster count to visit per query (None = Lance default) + * @param refineFactor IVF-PQ exact-distance re-rank factor (None = no re-rank) + * @param ef HNSW search depth (None = Lance default; only meaningful for + * HNSW indexes) + */ + def kNearestJoin( + right: DataFrame, + leftVecCol: String, + rightVecCol: String, + k: Int, + metric: String = "l2", + rightProjection: Option[Seq[String]] = None, + outerJoin: Boolean = false, + scoreCol: String = "__score", + overfetch: Int = 1, + nprobes: Option[Int] = None, + refineFactor: Option[Int] = None, + ef: Option[Int] = None): DataFrame = { + // R must be a Lance scan — the native probe runs against the underlying dataset URI. + val (uri, version) = LanceKnnImplicits.extractLanceUri(right).getOrElse { + throw new IllegalArgumentException( + "kNearestJoin requires the right DataFrame to be a Lance scan " + + "(spark.read.format(\"lance\").load(uri)); got a non-Lance relation. Build a Lance " + + "dataset with a vector index on the search column for R first.") + } + // Plan-time size gate: estimate the per-executor resident footprint of opening R's + // whole-dataset index and, per `spark.lance.knn.sizeGate.mode`, fail fast / warn / skip. + // Gated here (the user-facing entry point) rather than in `IndexedNearestJoin.apply` so + // internal callers/benchmarks that drive the lower-level API on large R stay ungated. + LanceKnnSizeGate.check(df.sparkSession, uri, version, rightVecCol) + IndexedNearestJoin( + left = df, + rightLanceUri = uri, + leftVecCol = leftVecCol, + rightVecCol = rightVecCol, + k = k, + metric = metric, + rightProjection = rightProjection, + outerJoin = outerJoin, + scoreCol = scoreCol, + overfetch = overfetch, + nprobes = nprobes, + version = version, + refineFactor = refineFactor, + ef = ef) + } + } + + /** + * Walk a DataFrame's analyzed plan looking for a `LanceTable`-backed + * `DataSourceV2Relation`. Skips through wrappers that don't change the underlying + * relation: `SubqueryAlias`, `View`, `Project`, `Filter`. Returns + * `Some((uri, optional version))` pulled from the relation's options when a Lance scan + * is found, or `None` otherwise. + * + * Lance detection mirrors `IndexedNearestByJoinRule.isLanceTable` — + * class-name match (`getClass.getName.contains("Lance")`) — to keep the user-facing + * extension working without a hard dependency on the connector's internal types. The + * extension only needs to be able to spot a Lance relation; it doesn't operate on it + * directly. + * + * Public for tests. + */ + private[knn] def extractLanceUri(df: DataFrame): Option[(String, Option[Long])] = { + findLanceRelation(df.queryExecution.analyzed).flatMap { rel => + val opts = rel.options + val uri = Option(opts.get("path")).orElse(Option(opts.get("datasetUri"))) + uri.map { u => + val version = Option(opts.get("version")).map(_.toLong) + (u, version) + } + } + } + + private def findLanceRelation(plan: LogicalPlan): Option[DataSourceV2Relation] = plan match { + case rel: DataSourceV2Relation if isLanceTable(rel) => Some(rel) + case other => + // Iterator.find avoids 2.13's `nextOption()` so this stays Scala 2.12-compatible. + val it = other.children.iterator.map(findLanceRelation).filter(_.isDefined) + if (it.hasNext) it.next() else None + } + + private def isLanceTable(rel: DataSourceV2Relation): Boolean = { + val cls = rel.table.getClass.getName + cls.contains("Lance") || cls.contains("lance") + } +} diff --git a/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/LanceKnnJoinStage.scala b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/LanceKnnJoinStage.scala new file mode 100644 index 000000000..ed53fc1f5 --- /dev/null +++ b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/LanceKnnJoinStage.scala @@ -0,0 +1,208 @@ +/* + * 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 +import org.apache.spark.sql.types.StructField + +import scala.collection.mutable + +/** + * The whole indexed nearest-by join, done per Spark partition with NO shuffle. + * + * A single native `LanceProbe.probe(...)` call is already a complete distributed search: Lance + * probes the IVF index and scans the candidate fragments across its own threads, heap-merges in + * process, and returns the final top-K for one query. Handing that orchestration to Spark (a + * probe → shuffle → merge → materialize pipeline) only adds a shuffle round-trip, a redundant + * merge stage, a second materialize scan, and `M × N_frag × K` refs crossing the Rust→JVM + * boundary. So this stage keeps everything local: + * + * {{{ + * left.rdd.mapPartitions { rows => + * val probe = new LanceProbe(uri, fragmentIds = None, version) // whole-index, once per task + * rows.flatMap { leftRow => + * val hits = probe.probe(query(leftRow), internalK, ...) // native top-K search + * val topK = trimToK(hits) // overfetch → K + * val payloads = probe.materialize(topK.map(_.rowAddr)) // late point-fetch by _rowid + * topK.map(ref => assembleRow(leftRow, payloads(ref), ref.score)) + * } + * } + * }}} + * + * No `requiredChildDistribution`, no Exchange. Each task opens R's whole index (`fragmentIds = + * None`) — Lance does the cross-fragment merge internally — so per-executor resident memory grows + * with `|R|`, which is what `LanceKnnSizeGate` gates on the DataFrame path. + * + * Both the DataFrame API ([[org.lance.spark.knn.IndexedNearestJoin]]) and the SQL Catalyst node + * (`LanceKnnJoinExec` in the 4.2 module) drive this same `runPartition`, so probe/trim/materialize + * semantics stay defined in exactly one place. + */ +object LanceKnnJoinStage { + + /** + * Everything a probe task needs, shipped from the driver. `internalK` is the overfetch count + * handed to Lance (`k × overfetch`); `k` is the final per-left-row cut applied after the native + * search. `leftVecIdx` is the position of the query-vector column in the left row. + */ + final case class Conf( + datasetUri: String, + version: Option[Long], + vectorColumn: String, + metric: Metric, + k: Int, + internalK: Int, + nprobes: Option[Int], + refineFactor: Option[Int], + ef: Option[Int], + prefilter: Option[String], + leftVecIdx: Int, + rightProjection: Seq[String], + rightFields: Seq[StructField], + leftFieldCount: Int, + outerJoin: Boolean, + smallerIsBetter: Boolean) + extends Serializable + + /** + * Run the join for one partition of left rows. Opens the probe once, probes + materializes per + * row, and returns assembled join rows. Materializing into an `ArrayBuffer` before returning the + * iterator is deliberate: Spark pulls from `mapPartitions` lazily, so a bare lazy iterator would + * let the consumer outlive the `try`/`finally` and read from a closed probe handle. + */ + def runPartition(leftRows: Iterator[Row], conf: Conf): Iterator[Row] = { + if (leftRows.isEmpty) return Iterator.empty + + val probe = new LanceProbe(conf.datasetUri, fragmentIds = None, version = conf.version) + val out = mutable.ArrayBuffer.empty[Row] + try { + leftRows.foreach { leftRow => + val q = extractVector(leftRow, conf.leftVecIdx) + if (q == null) { + // Null query vector: nothing to search. Emit a null-right row only for an outer join. + if (conf.outerJoin) { + out += assembleRow(leftRow, conf.leftFieldCount, conf.rightFields, null, null) + } + } else { + // Overfetch `internalK` candidates natively, then trim to the final `k`. Lance already + // returns them best-first, so when it hands back no more than `k` we keep them as-is and + // skip the heap entirely. + val refs = probe + .probe( + conf.vectorColumn, + q, + conf.internalK, + conf.metric, + conf.nprobes, + conf.refineFactor, + conf.ef, + conf.prefilter) + .toArray + val trimmed = + if (refs.length <= conf.k) refs + else { + val heap = new TopKHeap(conf.k, conf.smallerIsBetter) + heap.offerAll(refs) + heap.drain() + } + + if (trimmed.isEmpty) { + if (conf.outerJoin) { + out += assembleRow(leftRow, conf.leftFieldCount, conf.rightFields, null, null) + } + } else { + // Late materialization: point-fetch the surviving right rows by `_rowid`. Building the + // `rowAddr -> row` map collapses any duplicate rowAddr to one payload; the loop below + // still emits one output row per surviving ref. + val materialized: Map[Long, Map[String, Any]] = probe + .materialize(trimmed.iterator.map(_.rowAddr).toSeq, conf.rightProjection) + .map(m => extractRowAddr(m) -> m) + .toMap + trimmed.foreach { ref => + val rightMap = materialized.getOrElse(ref.rowAddr, null) + out += assembleRow( + leftRow, + conf.leftFieldCount, + conf.rightFields, + rightMap, + ref.score) + } + } + } + } + } finally probe.close() + out.iterator + } + + /** + * Pull a query vector out of a Spark `Row`'s ArrayType column. The Scala 2.13 `Seq` gotcha is + * real: `Row.get` on `ArrayType` returns `mutable.ArraySeq`, which `case s: Seq[_]` only matches + * against the root `scala.collection.Seq` trait (the default `Seq` alias is `immutable.Seq` on + * 2.13). + */ + private[knn] def extractVector(row: Row, idx: Int): Array[Float] = { + if (row.isNullAt(idx)) return null + row.get(idx) match { + case s: scala.collection.Seq[_] => + s.iterator.map { + case f: java.lang.Float => f.floatValue() + case f: Float => f + case d: java.lang.Double => d.doubleValue().toFloat + case d: Double => d.toFloat + case other => + throw new IllegalStateException( + s"Unsupported vector element type: ${other.getClass.getName}") + }.toArray + case arr: Array[Float] => arr + case arr: Array[java.lang.Float] => arr.map(_.floatValue()) + case other => + throw new IllegalStateException( + s"Unsupported vector column representation: ${other.getClass.getName}") + } + } + + /** Read the `_rowid` key out of a materialized row map (tolerating boxed / stringy longs). */ + private def extractRowAddr(m: Map[String, Any]): Long = + m.get(LanceProbe.RowIdColumn) match { + case Some(l: java.lang.Long) => l.longValue() + case Some(l: Long) => l + case Some(other) => other.toString.toLong + case None => + throw new IllegalStateException( + s"Materialized row missing ${LanceProbe.RowIdColumn}; " + + s"got keys: ${m.keys.mkString(", ")}") + } + + /** + * Assemble one output row: `left fields ++ right fields ++ score`. A null `rightValues` (outer + * join with no hit) fills the right side with nulls. + */ + private def assembleRow( + leftRow: Row, + leftFieldCount: Int, + rightFields: Seq[StructField], + rightValues: Map[String, Any], + score: Any): Row = { + val arr = new Array[Any](leftFieldCount + rightFields.size + 1) + var i = 0 + while (i < leftFieldCount) { arr(i) = leftRow.get(i); i += 1 } + var j = 0 + while (j < rightFields.size) { + arr(leftFieldCount + j) = + if (rightValues == null) null else rightValues.getOrElse(rightFields(j).name, null) + j += 1 + } + arr(leftFieldCount + rightFields.size) = score + Row.fromSeq(arr.toSeq) + } +} diff --git a/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/LanceKnnSizeGate.scala b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/LanceKnnSizeGate.scala new file mode 100644 index 000000000..de63ffdee --- /dev/null +++ b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/LanceKnnSizeGate.scala @@ -0,0 +1,290 @@ +/* + * 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.network.util.{ByteUnit, JavaUtils} +import org.apache.spark.sql.SparkSession +import org.lance.{Dataset, ReadOptions} +import org.lance.spark.LanceRuntime + +import scala.collection.JavaConverters._ + +/** + * Plan-time size gate for the indexed native KNN probe. + * + * Each probe task opens R's whole-dataset index (`fragmentIds = None`), so per-executor resident + * memory grows with `|R|`, not with query count. We model that resident footprint as + * + * {{{ + * residentFootprint(|R|) ≈ a + b · |R| # per executor + * apply native probe iff residentFootprint(|R|) ≤ SAFETY · executorRAM + * }}} + * + * and compare it to a safe fraction of the executor container budget + * (`spark.executor.memory` + `spark.executor.memoryOverhead`). The constants were measured on a + * standalone Spark cluster (synthetic uniform `dim=128` vectors, IVF-PQ, 4 concurrent probes per + * executor): `a ≈ 0.70 GiB`, `b ≈ 64 MiB per 1M rows` at `dim=128`. The per-row slope is scaled + * from that single `dim=128` point (see [[perRowBytesForDim]]); that extrapolation is UNVALIDATED + * at other dimensions — the resident PQ footprint tracks the index's `num_sub_vectors`, not the raw + * vector dimension — so treat the number as advisory. See sezruby/lance-spark#11 "Memory & size + * gate" for the full write-up. + * + * This is a soft, latency-oriented guardrail rather than a hard-OOM safety check: the measured + * `VmHWM` growth is dominated by reclaimable memory-mapped index pages, so past the budget the + * native path degrades in latency (index pages fault in and out) rather than crashing. And the + * native probe is the ONLY path that scales here — the default cross-product rewrite is strictly + * worse and would OOM at the same `|R|` — so the gate never reroutes anywhere; it only surfaces the + * working-set cliff. That is why the default mode is `warn`, not a hard failure. + * + * == Modes (`spark.lance.knn.sizeGate.mode`) == + * - `warn` (default): log a warning with the numbers and proceed with the native probe. The + * estimate is advisory and the native path degrades rather than crashes, so the default never + * blocks a job that would otherwise complete. + * - `error`: throw when the estimate exceeds the budget — opt-in fail-fast for callers who would + * rather not start an under-provisioned job at all. + * - `off`: skip the check entirely (no dataset open). + * + * == Overrides == + * - `spark.lance.knn.sizeGate.safetyFraction` (default 0.7) + * - `spark.lance.knn.sizeGate.baselineBytes` (default 0.70 GiB — the `a` term) + * - `spark.lance.knn.sizeGate.bytesPerRow` (default extrapolated from dim — see caveat above) + * - `spark.lance.knn.sizeGate.executorRamBytes` (default `spark.executor.memory` + overhead) + * + * The two arithmetic helpers ([[budgetBytes]], [[estimateFor]]) are pure and deployment-agnostic so + * they can be unit-tested without a Spark session or a live Lance dataset. + */ +object LanceKnnSizeGate { + + private val LOG = org.slf4j.LoggerFactory.getLogger("org.lance.spark.knn.sizeGate") + + // ---- configuration keys ------------------------------------------------------------------ + val ModeConfKey: String = "spark.lance.knn.sizeGate.mode" + val SafetyFractionConfKey: String = "spark.lance.knn.sizeGate.safetyFraction" + val BaselineBytesConfKey: String = "spark.lance.knn.sizeGate.baselineBytes" + val BytesPerRowConfKey: String = "spark.lance.knn.sizeGate.bytesPerRow" + val ExecutorRamBytesConfKey: String = "spark.lance.knn.sizeGate.executorRamBytes" + + // ---- measured constants (sezruby/lance-spark#11 "Memory & size gate") -------------------- + private val GiB: Long = 1L << 30 + private val MiB: Long = 1L << 20 + + /** `a` — baseline per-executor footprint (JVM + native runtime + per-task buffers). */ + val DefaultBaselineBytes: Long = (0.70 * GiB).toLong + + /** Reference dimension the per-row slope `b` was measured at. */ + val ReferenceDim: Int = 128 + + /** `b` at the reference dim: 64 MiB per 1M rows ≈ 67.1 bytes/row. */ + val BytesPerRowAtReferenceDim: Double = 64.0 * MiB / 1e6 + + val DefaultSafetyFraction: Double = 0.7 + + // Spark's own defaults for the executor container budget. + private val DefaultExecutorMemory: String = "1g" + private val MinOverheadBytes: Long = 384L * MiB + private val DefaultOverheadFactor: Double = 0.1 + + /** + * The result of a size estimate. `numRows`/`dim` come from R; the rest are the resolved + * model inputs. `fits` is the gate decision. + */ + final case class Estimate( + numRows: Long, + dim: Option[Int], + bytesPerRow: Double, + baselineBytes: Long, + footprintBytes: Long, + execRamBytes: Long, + safety: Double, + mode: String) { + def budgetBytes: Long = (safety * execRamBytes).toLong + def fits: Boolean = footprintBytes <= budgetBytes + def footprintGiB: String = f"${footprintBytes.toDouble / GiB}%.1f" + def budgetGiB: String = f"${budgetBytes.toDouble / GiB}%.1f" + def execRamGiB: String = f"${execRamBytes.toDouble / GiB}%.1f" + } + + /** + * Plan-time entry point. Reads config off `spark`, estimates the per-executor footprint for a + * whole-dataset probe of R, and acts per [[ModeConfKey]]. `warn` (default) logs and proceeds, + * `error` throws (opt-in fail fast), `off` skips without opening the dataset. + */ + def check( + spark: SparkSession, + datasetUri: String, + version: Option[Long], + vectorColumn: String): Unit = { + val mode = modeOf(spark) + if (mode == "off") return + val est = estimate(spark, datasetUri, version, vectorColumn, mode) + if (!est.fits) { + val sizing = + s"KNN native probe estimated at ~${est.footprintGiB} GiB/exec vs budget " + + s"${est.safety}×${est.execRamGiB} = ${est.budgetGiB} GiB at |R|=${est.numRows}" + + est.dim.map(d => s" (dim=$d)").getOrElse("") + "." + mode match { + // Opt-in fail-fast: refuse to start an under-provisioned job. + case "error" => + throw new IllegalArgumentException( + sizing + s" Failing fast ($ModeConfKey=error): raise spark.executor.memory, lower " + + s"executor cores, or set $ModeConfKey=warn to proceed anyway.") + // Default `warn` (and any unrecognized value): the estimate is advisory and the native + // path degrades in latency rather than crashing, so never block — just surface it. + case _ => + LOG.warn( + sizing + " Proceeding: the native probe still runs (it degrades in latency rather " + + "than crashing, and it is the only path that scales at this size). Raise " + + s"spark.executor.memory / lower executor cores if latency suffers, or set " + + s"$ModeConfKey=error to fail fast instead.") + } + } + } + + private def modeOf(spark: SparkSession): String = + spark.sessionState.conf.getConfString(ModeConfKey, "warn").trim.toLowerCase + + /** + * Resolve the model inputs from config + R's metadata and compute the [[Estimate]]. The row + * count and vector dimension are read once via the Lance Java API on the driver. + */ + def estimate( + spark: SparkSession, + datasetUri: String, + version: Option[Long], + vectorColumn: String, + mode: String): Estimate = { + val safety = optDouble(spark, SafetyFractionConfKey).getOrElse(DefaultSafetyFraction) + val baselineBytes = optLong(spark, BaselineBytesConfKey).getOrElse(DefaultBaselineBytes) + val execRamBytes = + optLong(spark, ExecutorRamBytesConfKey).getOrElse(resolveExecutorRamBytes(spark)) + val bytesPerRowOverride = optDouble(spark, BytesPerRowConfKey) + + val (numRows, dim) = readRowsAndDim(datasetUri, version, vectorColumn) + if (dim.isEmpty && bytesPerRowOverride.isEmpty) { + LOG.warn( + s"Could not determine the vector dimension for column '$vectorColumn' at $datasetUri; the " + + s"size estimate falls back to the reference dim ($ReferenceDim) and may be inaccurate. " + + s"Set $BytesPerRowConfKey to calibrate the per-row footprint.") + } + estimateFor(numRows, dim, safety, baselineBytes, bytesPerRowOverride, execRamBytes, mode) + } + + /** + * Pure footprint arithmetic — no Spark, no Lance. `bytesPerRowOverride` wins when set; otherwise + * the per-row slope is scaled linearly from the measured `dim=128` point (falling back to the + * reference dim when R's dimension can't be determined). + */ + private[knn] def estimateFor( + numRows: Long, + dim: Option[Int], + safety: Double, + baselineBytes: Long, + bytesPerRowOverride: Option[Double], + execRamBytes: Long, + mode: String): Estimate = { + val bytesPerRow = bytesPerRowOverride.getOrElse(perRowBytesForDim(dim)) + val footprintBytes = baselineBytes + math.round(bytesPerRow * numRows) + Estimate(numRows, dim, bytesPerRow, baselineBytes, footprintBytes, execRamBytes, safety, mode) + } + + /** + * Per-row resident-index bytes, scaled linearly from the measured `dim=128` slope. The linear + * scaling is an UNVALIDATED extrapolation — the resident PQ footprint really tracks the index's + * `num_sub_vectors`, which does not grow one-for-one with the raw vector dimension, so this can + * over-predict at high dim. It is advisory only (the default `warn` mode never blocks); override + * with [[BytesPerRowConfKey]] to calibrate against a measured workload. + */ + private[knn] def perRowBytesForDim(dim: Option[Int]): Double = + BytesPerRowAtReferenceDim * (dim.getOrElse(ReferenceDim).toDouble / ReferenceDim) + + /** + * Executor container budget = `spark.executor.memory` + `spark.executor.memoryOverhead`. RSS + * here is off-heap-dominated, so comparing against `-Xmx` (executor memory) alone would be + * dimensionally wrong. Mirrors Spark's own overhead default: + * `max(384 MiB, overheadFactor × executorMemory)`. + */ + private[knn] def resolveExecutorRamBytes(spark: SparkSession): Long = { + val sparkConf = spark.sparkContext.getConf + val execMem = sparkConf.get("spark.executor.memory", DefaultExecutorMemory) + val overhead = sparkConf.getOption("spark.executor.memoryOverhead") + val overheadFactor = sparkConf.getOption("spark.executor.memoryOverheadFactor").map(_.toDouble) + budgetBytes(execMem, overhead, overheadFactor) + } + + /** + * Pure budget arithmetic. Both `spark.executor.memory` and `spark.executor.memoryOverhead` + * follow Spark's `ByteUnit.MiB` convention (a bare number means MiB); a unit suffix (`g`, `m`, + * …) is honored. When overhead is unset, use `max(384 MiB, factor × executorMemory)`. + */ + private[knn] def budgetBytes( + execMem: String, + overhead: Option[String], + overheadFactor: Option[Double]): Long = { + val execMemBytes = memBytes(execMem) + val overheadBytes = overhead match { + case Some(s) => memBytes(s) + case None => + val factor = overheadFactor.getOrElse(DefaultOverheadFactor) + math.max(MinOverheadBytes, (factor * execMemBytes).toLong) + } + execMemBytes + overheadBytes + } + + /** Parse a Spark memory string (MiB convention: bare number = MiB) to bytes. */ + private def memBytes(str: String): Long = JavaUtils.byteStringAs(str, ByteUnit.MiB) * MiB + + /** Read R's total row count and the vector column's dimension via the Lance Java API. */ + private[knn] def readRowsAndDim( + datasetUri: String, + version: Option[Long], + vectorColumn: String): (Long, Option[Int]) = { + val ds = openDataset(datasetUri, version) + try { + (ds.countRows(), dimensionOf(ds.getSchema, vectorColumn)) + } finally ds.close() + } + + /** Open the Lance dataset at `datasetUri`, optionally pinned to `version`. */ + private def openDataset(datasetUri: String, version: Option[Long]): Dataset = { + val readOpts = { + val b = new ReadOptions.Builder() + version.foreach(v => b.setVersion(v)) + b.build() + } + Dataset + .open() + .uri(datasetUri) + .allocator(LanceRuntime.allocator()) + .readOptions(readOpts) + .build() + } + + /** Vector dimension = the `FixedSizeList` width of the vector field, if the schema declares it. */ + private def dimensionOf( + schema: org.apache.arrow.vector.types.pojo.Schema, + vectorColumn: String): Option[Int] = { + schema.getFields.asScala + .find(_.getName == vectorColumn) + .map(_.getType) + .collect { + case fsl: org.apache.arrow.vector.types.pojo.ArrowType.FixedSizeList => fsl.getListSize + } + } + + private def optDouble(spark: SparkSession, key: String): Option[Double] = + Option(spark.sessionState.conf.getConfString(key, null)).map(_.trim.toDouble) + + private def optLong(spark: SparkSession, key: String): Option[Long] = + Option(spark.sessionState.conf.getConfString(key, null)).map(_.trim.toLong) +} diff --git a/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/LanceProbe.scala b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/LanceProbe.scala new file mode 100644 index 000000000..0ebac1ab5 --- /dev/null +++ b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/LanceProbe.scala @@ -0,0 +1,356 @@ +/* + * 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.lance.{Dataset, ReadOptions} +import org.lance.ipc.{LanceScanner, Query, ScanOptions} +import org.lance.spark.{LanceConstant, LanceRuntime} + +import java.util + +import scala.collection.JavaConverters._ +import scala.collection.mutable + +/** + * Per-task vector-index probe primitive. Opens a Lance dataset once and runs many `probe()` calls + * against a fixed set of fragments. Returns row references + scores only — no payload. Late + * materialization happens elsewhere (`LanceMaterialize`). + * + * 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 datasetUri Lance dataset URI (passed straight to `Dataset.open`). + * @param fragmentIds Fragments this probe is restricted to. Pass `None` for whole-dataset search. + * @param version Optional Lance version to pin. Required when used inside a join, so all + * probe / materialize stages see the same snapshot. + * @param allocator Arrow allocator. Defaults to lance-spark's shared `LanceRuntime.allocator()`. + */ +final class LanceProbe( + datasetUri: String, + fragmentIds: Option[Seq[Int]], + version: Option[Long] = None, + allocator: BufferAllocator = LanceRuntime.allocator()) + extends AutoCloseable { + + // 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. + 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 = { + val readOpts = { + val b = new ReadOptions.Builder() + version.foreach(v => b.setVersion(v)) + b.build() + } + Dataset.open() + .uri(datasetUri) + .allocator(allocator) + .readOptions(readOpts) + .build() + } + + /** + * 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") + + val q = { + 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() + } + + 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() + } + } + + /** + * 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 = LanceProbe.ScoreColumns.iterator + .map(name => Option(root.getVector(name)).orNull) + .find(_ != null) + .getOrElse(throw new IllegalStateException( + s"Lance nearest scan did not return a score column. Got: " + + root.getSchema.getFields.asScala.map(_.getName).mkString(", "))) + + val n = root.getRowCount + var i = 0 + while (i < n) { + val addr = 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}") + } + val score = 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}") + } + out += ScoredRowRef(addr, score) + i += 1 + } + } + } finally { + reader.close() + } + out.toSeq + } + + /** + * 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". + * @return a sequence of materialized rows, each represented as a `Map[String, Any]` for the + * projected columns plus an entry under `LanceProbe.RowIdColumn` so the caller can + * re-key. Returning a Map keeps this primitive Spark-agnostic; conversion to + * `InternalRow` happens in the API layer. + */ + def materialize( + rowAddrs: Seq[Long], + projection: Seq[String] = 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()) + } finally { + scanner.close() + } + } + + private def readRows(reader: ArrowReader): Seq[Map[String, Any]] = { + val out = mutable.ArrayBuffer.empty[Map[String, Any]] + try { + while (reader.loadNextBatch()) { + val root: VectorSchemaRoot = reader.getVectorSchemaRoot + val n = root.getRowCount + var i = 0 + while (i < n) { + val rowMap = mutable.LinkedHashMap.empty[String, Any] + val fields = root.getSchema.getFields.asScala + var f = 0 + while (f < fields.size) { + val name = fields(f).getName + val v = root.getVector(name) + rowMap(name) = if (v.isNull(i)) null else LanceProbe.toSparkValue(v.getObject(i)) + f += 1 + } + 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") + + /** + * 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_2.12/src/main/scala/org/lance/spark/knn/internal/Metric.scala b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/Metric.scala new file mode 100644 index 000000000..798bcedd4 --- /dev/null +++ b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/Metric.scala @@ -0,0 +1,68 @@ +/* + * 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: + * + * - L2: smaller score is better (distance) + * - Cosine / Dot: larger score is better (similarity) + */ +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 + val smallerIsBetter: Boolean = false + } + + case object Dot extends Metric { + val lanceType: DistanceType = DistanceType.Dot + val smallerIsBetter: Boolean = false + } + + /** + * 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_2.12/src/main/scala/org/lance/spark/knn/internal/ScoredRowRef.scala b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/ScoredRowRef.scala new file mode 100644 index 000000000..ac1b194fc --- /dev/null +++ b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/ScoredRowRef.scala @@ -0,0 +1,40 @@ +/* + * 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) +} diff --git a/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/TopKHeap.scala b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/TopKHeap.scala new file mode 100644 index 000000000..673179c8b --- /dev/null +++ b/lance-spark-knn_2.12/src/main/scala/org/lance/spark/knn/internal/TopKHeap.scala @@ -0,0 +1,114 @@ +/* + * 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 + val isBetter = + if (smallerIsBetter) ref.score < worst.score + else ref.score > worst.score + 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_2.12/src/test/scala/org/lance/spark/knn/IndexedNearestJoinIvfPqRecallTest.scala b/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/IndexedNearestJoinIvfPqRecallTest.scala new file mode 100644 index 000000000..8eb5e64c8 --- /dev/null +++ b/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/IndexedNearestJoinIvfPqRecallTest.scala @@ -0,0 +1,324 @@ +/* + * 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 + +import org.apache.spark.sql.{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.knn.internal.LanceVectorIndexBuilder +import org.lance.spark.knn.testutil.ClusteredEmbeddings + +import java.nio.file.Path +import java.util.Random + +import scala.collection.JavaConverters._ + +/** + * Phase 3 — real-recall validation against an IVF-PQ-indexed Lance dataset. + * + * Builds an IVF-PQ vector index via Lance's `Dataset.createIndex` Java binding, then runs + * `IndexedNearestJoin` and measures recall@K vs. the brute-force ground truth. With an index + * Lance returns *approximate* top-K, so recall is < 1.0 — the point of this test is to verify: + * + * 1. The indexed path actually engages (Lance's `useIndex` defaults to true on a Query + * against an indexed column; our `LanceProbe.probe` doesn't override it). + * 2. Recall at the default settings is in a sane range — our small synthetic dataset is + * small enough that recall should be high (most rows survive the IVF cluster cut). + * 3. `refineFactor > 1` improves recall by re-ranking more candidates with exact distance. + * + * Until this test, the 608x / 17.4x benchmark headlines were on a NO-INDEX Lance dataset where + * Lance's brute-force per-fragment scan made everything exact (recall = 1.0). The + * approximate-vs-exact recall trade-off that an indexed connector exposes was unmeasured. This + * test closes that gap. + * + * Setup specifics: 1024 right rows, dim 32, 4 IVF partitions, 8 PQ sub-vectors. The dataset + * is intentionally tiny so the test runs in a few seconds — production-realistic dataset + * sizes would need much larger N. + */ +class IndexedNearestJoinIvfPqRecallTest { + + @TempDir var tempDir: Path = _ + private var spark: SparkSession = _ + + private val Dim = 32 + private val NumRight = 1024 + private val NumLeft = 32 + private val K = 10 + private val Seed = 0xCAFEL + + @BeforeEach def setup(): Unit = { + spark = SparkSession.builder() + .appName("indexed-nearest-ivfpq-recall") + .master("local[2]") + .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() + + /** + * The headline test: build IVF-PQ, run IndexedNearestJoin, measure recall@10 against the + * brute-force oracle. With 1024 rows × 4 IVF partitions, each partition holds ~256 rows; + * a default-`nprobes` query hits ~1 partition, so we expect recall to be lower than 1.0 + * but still substantial. + */ + @Test def testIvfPqRecallReasonableAtDefaults(): Unit = { + val (leftDf, leftIds, leftVecs) = buildLeft() + val (rightUri, rightIds, rightVecs) = writeRight() + LanceVectorIndexBuilder.buildIvfPq( + datasetUri = rightUri, + vectorColumn = "rvec", + numPartitions = 4, + numSubVectors = 8, + numBits = 8) + assertEquals( + 1, + LanceVectorIndexBuilder.listIndexCount(rightUri), + "expected exactly one index after build") + + val joined = IndexedNearestJoin( + left = leftDf, + rightLanceUri = rightUri, + leftVecCol = "lvec", + rightVecCol = "rvec", + k = K, + metric = "l2", + rightProjection = Some(Seq("rid"))) + + val rows = joined.collect() + val recall = computeRecallAtK(rows, leftIds, leftVecs, rightIds, rightVecs, K) + println(s" IVF-PQ recall@$K (no refine, default nprobes): $recall") + // With 1024 rows and 4 IVF partitions, default nprobes = 1, the index returns ~256 + // candidates per query. Recall should be substantially > 0; our acceptance threshold + // is loose because IVF-PQ recall depends on the random data layout — anything < 0.3 + // would suggest a real bug, not an inherent IVF limitation. + assertTrue(recall > 0.3, s"recall@$K=$recall too low; index path probably not engaging") + } + + /** + * Production-realistic distribution: clustered Gaussian mixture, unit-sphere-normalized — + * the geometry of typical sentence-transformer / image-feature embeddings. The benchmark and + * the recall test elsewhere use uniform-random vectors over [0, 1]^d, which is the WORST + * case for IVF (k-means has no natural cluster structure to latch onto). This test exercises + * the indexed path on a more realistic distribution and asserts: + * + * 1. Recall@K on clustered data >= 0.5 at default IVF-PQ settings. If realistic data + * collapsed to coin-flip recall, the indexed path wouldn't be useful in production. + * 2. Both uniform and clustered recall numbers are printed, so a reviewer can see whether + * the realistic case actually helps in practice (it should — see the file's preamble). + * + * Why we don't `assert(clustered >= uniform)`: Lance's IVF training (k-means initialization) + * is non-deterministic across JVM sessions, so on a tiny 1024-row dataset the run-to-run + * noise in either recall number routinely exceeds the structural advantage of realistic + * data. A reliable comparison would need either (a) averaging over many seeds, which is + * slow and fragile in CI, or (b) much larger N where the structural effect dominates noise. + * We chose (c): print both, assert only the realistic-floor invariant. + */ + @Test def testClusteredEmbeddingsRecallSurvives(): Unit = { + val (uniformDf, uniformIds, uniformVecs) = + buildLeftFromVectors(generateUniform(NumLeft, Dim, Seed)) + val (uniformUri, uniformRightIds, uniformRightVecs) = + writeRightFromVectors(generateUniform(NumRight, Dim, Seed + 1)) + LanceVectorIndexBuilder.buildIvfPq(uniformUri, "rvec", numPartitions = 4, numSubVectors = 8) + + val (clusteredDf, clusteredIds, clusteredVecs) = buildLeftFromVectors( + ClusteredEmbeddings.generate(NumLeft, Dim, numClusters = 4, seed = Seed + 2)) + val (clusteredUri, clusteredRightIds, clusteredRightVecs) = writeRightFromVectors( + ClusteredEmbeddings.generate(NumRight, Dim, numClusters = 16, seed = Seed + 3)) + LanceVectorIndexBuilder.buildIvfPq( + clusteredUri, + "rvec", + numPartitions = 4, + numSubVectors = 8) + + val uniformRecall = recallAgainst( + uniformDf, + uniformUri, + uniformIds, + uniformVecs, + uniformRightIds, + uniformRightVecs) + val clusteredRecall = recallAgainst( + clusteredDf, + clusteredUri, + clusteredIds, + clusteredVecs, + clusteredRightIds, + clusteredRightVecs) + println( + s" IVF-PQ recall@$K: uniform=$uniformRecall, clustered=$clusteredRecall " + + "(uniform = IVF worst case; clustered = production-shaped)") + + assertTrue( + clusteredRecall >= 0.5, + s"clustered-data recall@$K=$clusteredRecall is unexpectedly low; " + + "defaults should comfortably exceed 0.5 on production-shaped embeddings — " + + "if this fails, suspect a regression in Lance's index path or in our probe wiring") + } + + /** + * `refineFactor > 1` engages Lance's exact-distance re-rank: fetch `K * refineFactor` + * approximate candidates, re-rank, trim back to K. Strictly improves (or matches) recall + * vs. no refine. We assert the >= relation rather than a strict > so the test isn't flaky + * on tiny datasets where both paths happen to find the same K rows. + */ + @Test def testRefineFactorImprovesRecall(): Unit = { + val (leftDf, leftIds, leftVecs) = buildLeft() + val (rightUri, rightIds, rightVecs) = writeRight() + LanceVectorIndexBuilder.buildIvfPq(rightUri, "rvec", numPartitions = 4) + + val baseline = IndexedNearestJoin( + left = leftDf, + rightLanceUri = rightUri, + leftVecCol = "lvec", + rightVecCol = "rvec", + k = K, + metric = "l2", + rightProjection = Some(Seq("rid"))) + val refined = IndexedNearestJoin( + left = leftDf, + rightLanceUri = rightUri, + leftVecCol = "lvec", + rightVecCol = "rvec", + k = K, + metric = "l2", + rightProjection = Some(Seq("rid")), + refineFactor = Some(8)) + + val recallBaseline = + computeRecallAtK(baseline.collect(), leftIds, leftVecs, rightIds, rightVecs, K) + val recallRefined = + computeRecallAtK(refined.collect(), leftIds, leftVecs, rightIds, rightVecs, K) + println(s" IVF-PQ recall@$K: no refine = $recallBaseline, refineFactor=8 = $recallRefined") + assertTrue( + recallRefined >= recallBaseline, + s"refineFactor should not hurt recall: baseline=$recallBaseline, refined=$recallRefined") + } + + // -- helpers ------------------------------------------------------------------------------ + + private def buildLeft(): (org.apache.spark.sql.DataFrame, Array[Int], Array[Array[Float]]) = + buildLeftFromVectors(generateUniform(NumLeft, Dim, Seed)) + + private def writeRight(): (String, Array[Int], Array[Array[Float]]) = + writeRightFromVectors(generateUniform(NumRight, Dim, Seed + 1)) + + /** Build a left-side DataFrame from a pre-generated vector array. */ + private def buildLeftFromVectors( + vectors: Array[Array[Float]]) + : (org.apache.spark.sql.DataFrame, Array[Int], Array[Array[Float]]) = { + val schema = new StructType(Array( + StructField("lid", IntegerType, nullable = false), + StructField( + "lvec", + ArrayType(FloatType, containsNull = false), + nullable = false, + new MetadataBuilder().putLong("arrow.fixed-size-list.size", Dim.toLong).build()))) + val ids = (0 until vectors.length).toArray + val rows = ids.zip(vectors).map { case (id, v) => + RowFactory.create(Integer.valueOf(id), v) + } + val df = spark.createDataFrame(rows.toSeq.asJava, schema) + (df, ids, vectors) + } + + /** Write a right-side Lance dataset from a pre-generated vector array. */ + private def writeRightFromVectors( + vectors: Array[Array[Float]]): (String, Array[Int], Array[Array[Float]]) = { + val schema = new StructType(Array( + StructField("rid", IntegerType, nullable = false), + StructField( + "rvec", + ArrayType(FloatType, containsNull = false), + nullable = false, + new MetadataBuilder().putLong("arrow.fixed-size-list.size", Dim.toLong).build()))) + val ids = (0 until vectors.length).map(_ + 100000).toArray + val rows = ids.zip(vectors).map { case (id, v) => + RowFactory.create(Integer.valueOf(id), v) + } + val df = spark.createDataFrame(rows.toSeq.asJava, schema) + val out = tempDir.resolve(s"right_${System.nanoTime()}").toString + df.write.format("lance").save(out) + (out, ids, vectors) + } + + /** Run an indexed nearest join against the given right dataset and compute recall@K. */ + private def recallAgainst( + leftDf: org.apache.spark.sql.DataFrame, + rightUri: String, + leftIds: Array[Int], + leftVecs: Array[Array[Float]], + rightIds: Array[Int], + rightVecs: Array[Array[Float]]): Double = { + val joined = IndexedNearestJoin( + left = leftDf, + rightLanceUri = rightUri, + leftVecCol = "lvec", + rightVecCol = "rvec", + k = K, + metric = "l2", + rightProjection = Some(Seq("rid"))) + computeRecallAtK(joined.collect(), leftIds, leftVecs, rightIds, rightVecs, K) + } + + /** Uniform-random vectors over the unit hypercube — the IVF-worst-case data distribution. */ + private def generateUniform(n: Int, dim: Int, seed: Long): Array[Array[Float]] = { + val rng = new Random(seed) + Array.fill(n)(randomVector(rng, dim)) + } + + /** + * Mean recall@K across all left rows: |intersection of indexed top-K with brute-force + * top-K| divided by K. A value of 1.0 means the indexed path returned the same K rows as + * brute force; lower values mean the IVF cluster cut excluded some true neighbors. + */ + private def computeRecallAtK( + joinedRows: Array[org.apache.spark.sql.Row], + leftIds: Array[Int], + leftVecs: Array[Array[Float]], + rightIds: Array[Int], + rightVecs: Array[Array[Float]], + k: Int): Double = { + val byLid = joinedRows.groupBy(_.getAs[Int]("lid")) + val perLidRecall = leftIds.zip(leftVecs).map { case (lid, lvec) => + val oracle = rightVecs.indices + .map(i => (rightIds(i), l2(lvec, rightVecs(i)))) + .sortBy(_._2) + .take(k) + .map(_._1) + .toSet + val actual = byLid.getOrElse(lid, Array.empty).map(_.getAs[Int]("rid")).toSet + val hit = (oracle intersect actual).size.toDouble + hit / k + } + perLidRecall.sum / perLidRecall.length + } + + 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 l2(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_2.12/src/test/scala/org/lance/spark/knn/internal/LanceKnnSizeGateTest.scala b/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/internal/LanceKnnSizeGateTest.scala new file mode 100644 index 000000000..5ddb453ed --- /dev/null +++ b/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/internal/LanceKnnSizeGateTest.scala @@ -0,0 +1,110 @@ +/* + * 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 [[LanceKnnSizeGate]]'s pure arithmetic — the footprint model, the executor-budget + * parsing, and the dimension scaling. The live path (opening a Lance dataset, throwing/warning per + * mode) is exercised by the integration tests; here we pin the constants and the decision math so a + * regression in the published `a + b·|R|` model is caught without a cluster. + */ +class LanceKnnSizeGateTest { + + private val GiB: Long = 1L << 30 + private val MiB: Long = 1L << 20 + + // ---- budget parsing (executor.memory + memoryOverhead) ----------------------------------- + + @Test def testBudgetDefaultOverhead(): Unit = { + // 8g heap + default overhead = max(384 MiB, 0.1 × 8 GiB) = 819.2 MiB. + val expected = 8L * GiB + math.max(384L * MiB, (0.1 * (8L * GiB)).toLong) + assertEquals(expected, LanceKnnSizeGate.budgetBytes("8g", None, None)) + } + + @Test def testBudgetExplicitOverhead(): Unit = { + assertEquals(8L * GiB + 2L * GiB, LanceKnnSizeGate.budgetBytes("8g", Some("2g"), None)) + } + + @Test def testBudgetBareNumberIsMiB(): Unit = { + // Bare numbers follow Spark's ByteUnit.MiB convention. + assertEquals(512L * MiB + 128L * MiB, LanceKnnSizeGate.budgetBytes("512", Some("128"), None)) + } + + @Test def testBudgetSmallHeapClampsOverheadFloor(): Unit = { + // 1g heap: 0.1 × 1 GiB = 102.4 MiB < 384 MiB floor, so overhead clamps to 384 MiB. + assertEquals(1L * GiB + 384L * MiB, LanceKnnSizeGate.budgetBytes("1g", None, None)) + } + + // ---- per-row slope scales linearly with dimension ---------------------------------------- + + @Test def testPerRowBytesReferenceDim(): Unit = { + assertEquals(64.0 * MiB / 1e6, LanceKnnSizeGate.perRowBytesForDim(Some(128)), 1e-9) + } + + @Test def testPerRowBytesScalesWithDim(): Unit = { + val base = LanceKnnSizeGate.perRowBytesForDim(Some(128)) + assertEquals(2.0 * base, LanceKnnSizeGate.perRowBytesForDim(Some(256)), 1e-9) + } + + @Test def testPerRowBytesUnknownDimUsesReference(): Unit = { + assertEquals( + LanceKnnSizeGate.perRowBytesForDim(Some(128)), + LanceKnnSizeGate.perRowBytesForDim(None), + 1e-9) + } + + // ---- footprint model vs. the published threshold table ----------------------------------- + + private def estimateAt(numRows: Long, execRamBytes: Long) = + LanceKnnSizeGate.estimateFor( + numRows = numRows, + dim = Some(128), + safety = LanceKnnSizeGate.DefaultSafetyFraction, + baselineBytes = LanceKnnSizeGate.DefaultBaselineBytes, + bytesPerRowOverride = None, + execRamBytes = execRamBytes, + mode = "error") + + @Test def testThreshold8GiBAt78MFits(): Unit = { + // Issue #11 table: 8 GiB @ SAFETY=0.7 → ~78M rows (dim=128). + assertTrue(estimateAt(78000000L, 8L * GiB).fits, "78M rows should fit an 8 GiB budget") + assertFalse(estimateAt(80000000L, 8L * GiB).fits, "80M rows should exceed an 8 GiB budget") + } + + @Test def testThreshold16GiBAt167MFits(): Unit = { + // Issue #11 table: 16 GiB @ SAFETY=0.7 → ~167M rows (dim=128). + assertTrue(estimateAt(167000000L, 16L * GiB).fits, "167M rows should fit a 16 GiB budget") + assertFalse(estimateAt(170000000L, 16L * GiB).fits, "170M rows should exceed a 16 GiB budget") + } + + @Test def testBaselineDominatesForTinyR(): Unit = { + // Small R: footprint ≈ baseline `a`, comfortably under any real budget. + val est = estimateAt(1000L, 8L * GiB) + assertTrue(est.fits) + assertEquals( + LanceKnnSizeGate.DefaultBaselineBytes + math.round(est.bytesPerRow * 1000L), + est.footprintBytes) + } + + @Test def testGiBFormattingIsHumanReadable(): Unit = { + val est = estimateAt(80000000L, 8L * GiB) + assertEquals("8.0", est.execRamGiB) + assertEquals("5.6", est.budgetGiB) // 0.7 × 8 GiB + // footprint carries the score through the `%.1f` GiB formatter used in the error message. + assertTrue(est.footprintGiB.matches("""\d+\.\d"""), s"unexpected format: ${est.footprintGiB}") + } +} diff --git a/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/internal/LanceProbeValidationTest.scala b/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/internal/LanceProbeValidationTest.scala new file mode 100644 index 000000000..0ae287138 --- /dev/null +++ b/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/internal/LanceProbeValidationTest.scala @@ -0,0 +1,202 @@ +/* + * 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 java.nio.file.Path +import java.util.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") + } + } + + /** + * 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() + } + + // -- 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 + } + + 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) + } + + 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_2.12/src/test/scala/org/lance/spark/knn/internal/LanceVectorIndexBuilder.scala b/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/internal/LanceVectorIndexBuilder.scala new file mode 100644 index 000000000..2a8f319c7 --- /dev/null +++ b/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/internal/LanceVectorIndexBuilder.scala @@ -0,0 +1,105 @@ +/* + * 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.{Dataset, ReadOptions} +import org.lance.index.{IndexOptions, IndexParams, IndexType} +import org.lance.index.vector.VectorIndexParams +import org.lance.spark.LanceRuntime + +import scala.collection.JavaConverters._ + +/** + * Test-only helper to build an IVF-PQ vector index on a Lance dataset via + * `Dataset.createIndex`. Exists so recall tests can construct the indexed scan path + * without writing the Lance Java boilerplate inline. + * + * Lives in `src/test/scala` because the production code path doesn't need to build + * indexes — users build them via Lance's Python / Rust / SQL DDL on their own datasets, + * and we just probe whatever's there. The helper exists for closed-loop recall validation. + */ +object LanceVectorIndexBuilder { + + /** + * Build an IVF-PQ index on `vectorColumn` of the dataset at `datasetUri`. Defaults are + * tuned for tiny test datasets — production users would size these much larger. + * + * @param numPartitions IVF cluster count. Should divide cleanly into the dataset row count. + * For a 4K-row dataset, 4-8 partitions is reasonable. + * @param numSubVectors PQ sub-vector count. Must divide vector dim evenly. + * @param numBits PQ bits per sub-vector. 8 is the standard. + * @param metric distance type. Must match the metric used at probe time. + * @param maxIters KMeans iteration cap during IVF training. 50 is enough for tests. + */ + def buildIvfPq( + datasetUri: String, + vectorColumn: String, + numPartitions: Int = 4, + numSubVectors: Int = 8, + numBits: Int = 8, + metric: Metric = Metric.L2, + maxIters: Int = 50): Unit = { + val dataset = openDataset(datasetUri) + try { + // Arg order in lance-core is (numPartitions, numBits, numSubVectors, distanceType, maxIters) + // — numBits precedes numSubVectors. Both default to 8 here so a swap is silent; pin the + // documented order explicitly. + val vectorParams = + VectorIndexParams.ivfPq(numPartitions, numBits, numSubVectors, metric.lanceType, maxIters) + val indexParams = IndexParams.builder().setVectorIndexParams(vectorParams).build() + val opts = IndexOptions + .builder(java.util.Collections.singletonList(vectorColumn), IndexType.VECTOR, indexParams) + .build() + dataset.createIndex(opts) + } finally dataset.close() + } + + /** + * Build an IVF_FLAT index — IVF clustering without PQ compression. Exact distances within + * visited clusters (no PQ noise), so recall depends purely on `nprobes` coverage. Higher + * memory/disk footprint than IVF-PQ (full vectors stored per cluster) but better recall on + * high-dim or random workloads where PQ compression drops too much information. + */ + def buildIvfFlat( + datasetUri: String, + vectorColumn: String, + numPartitions: Int = 4, + metric: Metric = Metric.L2): Unit = { + val dataset = openDataset(datasetUri) + try { + val vectorParams = VectorIndexParams.ivfFlat(numPartitions, metric.lanceType) + val indexParams = IndexParams.builder().setVectorIndexParams(vectorParams).build() + val opts = IndexOptions + .builder(java.util.Collections.singletonList(vectorColumn), IndexType.VECTOR, indexParams) + .build() + dataset.createIndex(opts) + } finally dataset.close() + } + + private def openDataset(uri: String): Dataset = { + Dataset + .open() + .uri(uri) + .allocator(LanceRuntime.allocator()) + .readOptions(new ReadOptions.Builder().build()) + .build() + } + + /** Number of indexes on the dataset (sanity check after building). */ + def listIndexCount(datasetUri: String): Int = { + val dataset = openDataset(datasetUri) + try dataset.listIndexes.asScala.size + finally dataset.close() + } +} diff --git a/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/internal/TopKHeapTest.scala b/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/internal/TopKHeapTest.scala new file mode 100644 index 000000000..41bd1de7d --- /dev/null +++ b/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/internal/TopKHeapTest.scala @@ -0,0 +1,92 @@ +/* + * 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) + } + + /** `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/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/testutil/ClusteredEmbeddings.scala b/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/testutil/ClusteredEmbeddings.scala new file mode 100644 index 000000000..47ccc3c6d --- /dev/null +++ b/lance-spark-knn_2.12/src/test/scala/org/lance/spark/knn/testutil/ClusteredEmbeddings.scala @@ -0,0 +1,137 @@ +/* + * 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.testutil + +import java.util.Random + +/** + * Generate a clustered Gaussian-mixture embedding sample as a stand-in for real production + * embeddings (SIFT / sentence-transformer / image features). Real embeddings are not uniform + * over the unit hypercube — they cluster around a small number of topic centroids with each + * cluster occupying a relatively narrow region of the space. Uniform-random vectors are the + * worst case for IVF: there's no natural cluster structure for k-means to latch onto, so the + * IVF partitions cover the space arbitrarily and the per-cluster recall is essentially random. + * + * Method: + * 1. Pick `numClusters` cluster centers, each drawn uniformly from the unit hypercube. + * 2. For each row, pick a cluster (round-robin so each cluster gets equal mass) and sample + * a Gaussian centered on it with standard deviation `sigma * cluster_separation`. + * 3. L2-normalize so vectors live on the unit sphere — the natural geometry for cosine / + * inner-product retrieval, and what most production embedding models produce. + * + * The cluster-separation factor is the median pairwise distance between centers; scaling sigma + * by it keeps the cluster radius proportional to inter-cluster spacing regardless of `dim` or + * `numClusters`. With sigma ≈ 0.15 the clusters overlap a little but stay distinguishable — + * a reasonable proxy for production embedding distributions. + * + * The generator is deterministic given the seed so test runs are reproducible. + */ +object ClusteredEmbeddings { + + /** + * Build a clustered-Gaussian-mixture sample. + * + * @param n number of vectors to generate + * @param dim vector dimension + * @param numClusters number of cluster centers (small relative to `n` — typical 16-64) + * @param sigma per-cluster standard deviation, in units of inter-cluster distance. + * 0.05 = tight clusters (high recall floor); 0.5 = loose, near-uniform + * @param seed RNG seed for reproducibility + * @return an array of `n` float vectors of dimension `dim`, L2-normalized + */ + def generate( + n: Int, + dim: Int, + numClusters: Int, + sigma: Double = 0.15, + seed: Long = 0L): Array[Array[Float]] = { + require(n > 0 && dim > 0 && numClusters > 0, "n, dim, numClusters must all be positive") + require(numClusters <= n, "numClusters cannot exceed n") + val rng = new Random(seed) + + // Step 1: cluster centers, uniform on [0, 1]^dim. Stored as Doubles so the noise pass keeps + // numerical headroom — L2 normalization at the end folds back to Float precision. + val centers = Array.fill(numClusters)(Array.fill(dim)(rng.nextDouble())) + + // Step 2: median pairwise distance between centers, used to scale sigma. We don't want sigma + // expressed in absolute distance units — the right notion is "fraction of cluster spacing," + // which keeps clustering tightness behavior stable across (dim, numClusters) settings. + val sep = medianPairwiseDistance(centers) + val scaledSigma = sigma * sep + + // Step 3: sample each row from a Gaussian centered on a round-robin cluster. Round-robin + // (rather than uniformly random cluster choice) gives every cluster the same mass — a more + // controlled benchmark setup than letting some clusters get sparsely populated. + val out = new Array[Array[Float]](n) + var i = 0 + while (i < n) { + val center = centers(i % numClusters) + val v = new Array[Float](dim) + var d = 0 + while (d < dim) { + v(d) = (center(d) + rng.nextGaussian() * scaledSigma).toFloat + d += 1 + } + l2Normalize(v) + out(i) = v + i += 1 + } + out + } + + /** + * Median pairwise L2 distance between centers. We sample up to 1024 random center pairs + * rather than computing all `O(K^2)` of them — for `numClusters = 64` that's 2016 pairs, + * trivial; for larger K we'd otherwise pay cost the rest of the test doesn't need. + */ + private def medianPairwiseDistance(centers: Array[Array[Double]]): Double = { + val k = centers.length + if (k < 2) return 1.0 + val rng = new Random(0L) + val numPairs = math.min(1024, k * (k - 1) / 2) + val dists = new Array[Double](numPairs) + var p = 0 + while (p < numPairs) { + var i = rng.nextInt(k) + var j = rng.nextInt(k) + while (j == i) j = rng.nextInt(k) + dists(p) = euclidean(centers(i), centers(j)) + p += 1 + } + java.util.Arrays.sort(dists) + dists(dists.length / 2) + } + + private def euclidean(a: Array[Double], b: Array[Double]): Double = { + var s = 0.0 + var i = 0 + while (i < a.length) { + val d = a(i) - b(i) + s += d * d + i += 1 + } + math.sqrt(s) + } + + private def l2Normalize(v: Array[Float]): Unit = { + var s = 0.0 + var i = 0 + while (i < v.length) { s += v(i) * v(i); i += 1 } + val norm = math.sqrt(s).toFloat + if (norm > 0f) { + i = 0 + while (i < v.length) { v(i) = v(i) / norm; i += 1 } + } + } +} diff --git a/lance-spark-knn_2.13/pom.xml b/lance-spark-knn_2.13/pom.xml new file mode 100644 index 000000000..ace622e1c --- /dev/null +++ b/lance-spark-knn_2.13/pom.xml @@ -0,0 +1,77 @@ + + + 4.0.0 + + + org.lance + lance-spark-root + 0.7.1 + ../pom.xml + + + lance-spark-knn_2.13 + ${project.artifactId} + Indexed nearest-neighbor join for Lance datasets in Spark + jar + + + ${scala213.version} + ${scala213.compat.version} + + + + + org.lance + lance-spark-base_2.13 + ${project.version} + + + org.apache.spark + spark-sql_${scala.compat.version} + provided + + + + org.lance + lance-spark-3.5_2.13 + ${project.version} + test + + + + + + + ../lance-spark-knn_2.12/src/main/scala + ../lance-spark-knn_2.12/src/test/scala + + + net.alchim31.maven + scala-maven-plugin + ${scala-maven-plugin.version} + + + scala-compile-first + process-resources + + add-source + compile + + + + scala-test-compile + process-test-resources + + testCompile + + + + + + -feature + + + + + + diff --git a/pom.xml b/pom.xml index 572c3ce07..d4f006f06 100644 --- a/pom.xml +++ b/pom.xml @@ -149,6 +149,9 @@ lance-spark-bundle-4.1_2.13 lance-spark-4.2_2.13 lance-spark-bundle-4.2_2.13 + lance-spark-knn_2.12 + lance-spark-knn_2.13 + lance-spark-knn-4.2_2.13