JingsongLi commented on code in PR #8252:
URL: https://github.com/apache/paimon/pull/8252#discussion_r3457410992
##########
paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala:
##########
@@ -215,3 +238,262 @@ case class PaimonStrategy(spark: SparkSession)
SparkShimLoader.shim.classicApi.recacheByPlan(spark, v2Relation)
}
}
+
+case class LateralVectorSearchExec(
+ innerTable: InnerTable,
+ columnName: String,
+ queryVectorExpr: Expression,
+ limit: Int,
+ options: Map[String, String],
+ vectorSearchOutput: Seq[Attribute],
+ projectList: Seq[NamedExpression],
+ projectOutput: Seq[Attribute],
+ child: SparkPlan)
+ extends SparkPlan {
+
+ override def children: Seq[SparkPlan] = Seq(child)
+
+ override def output: Seq[Attribute] = child.output ++ projectOutput
+
+ @transient override lazy val producedAttributes: AttributeSet =
AttributeSet(vectorSearchOutput)
+
+ @transient
+ override lazy val references: AttributeSet = {
+ AttributeSet.fromAttributeSets(expressions.map(_.references)) --
producedAttributes
+ }
+
+ override protected def withNewChildrenInternal(newChildren:
IndexedSeq[SparkPlan]): SparkPlan = {
+ copy(child = newChildren.head)
+ }
+
+ override protected def doExecute(): RDD[InternalRow] = {
+ child.execute().mapPartitions {
+ outerRows =>
+ val strippedQueryExpr = queryVectorExpr.transform {
+ case OuterReference(namedExpression) => namedExpression.toAttribute
+ }
+ val queryVectorProjection =
UnsafeProjection.create(Seq(strippedQueryExpr), child.output)
+ val rightProjection = UnsafeProjection.create(projectList,
vectorSearchOutput)
+ val joinedRow = new JoinedRow
+ lazy val searchContext = createSearchContext(rightProjection)
+ val batchSize = searchContext.batchSize
+
+ outerRows.map(_.copy()).grouped(batchSize).flatMap {
+ outerRowBatch =>
+ val searchBatch = ArrayBuffer[LateralVectorSearchQuery]()
+ outerRowBatch.foreach {
+ outerRow =>
+ toFloatArray(queryVectorProjection(outerRow).get(0,
strippedQueryExpr.dataType))
+ .foreach(
+ queryVector => searchBatch +=
LateralVectorSearchQuery(outerRow, queryVector))
+ }
+
+ if (searchBatch.isEmpty) {
+ Iterator.empty
+ } else {
+ search(searchBatch.toVector, searchContext).map {
+ case (outerRow, rightRow) =>
+ joinedRow(outerRow, rightRow)
+ joinedRow.copy()
+ }
+ }
+ }
+ }
+ }
+
+ private def createSearchContext(rightProjection: UnsafeProjection):
LateralVectorSearchContext = {
+ val rowType = innerTable.rowType()
+ val readFieldNames = vectorSearchOutput
+ .filterNot(_.name == PaimonMetadataColumn.SEARCH_SCORE_COLUMN)
+ .map(_.name)
+ val readFieldNamesWithRowId =
+ if (readFieldNames.contains(SpecialFields.ROW_ID.name())) {
+ readFieldNames
+ } else {
+ readFieldNames :+ SpecialFields.ROW_ID.name()
+ }
+ val rowTypeWithRowId = SpecialFields.rowTypeWithRowId(rowType)
+ val readRowType = rowType.project(readFieldNames.asJava)
+ val readRowTypeWithRowId = SpecialFields.rowTypeWithRowId(readRowType)
+ val readBuilder = innerTable
+ .newReadBuilder()
+ .withReadType(rowTypeWithRowId.project(readFieldNamesWithRowId.asJava))
+ val scoreMetadataColumns =
+ if (vectorSearchOutput.exists(_.name ==
PaimonMetadataColumn.SEARCH_SCORE_COLUMN)) {
+ Seq(PaimonMetadataColumn.SEARCH_SCORE)
+ } else {
+ Seq.empty
+ }
+ val resultRowType =
+ if (scoreMetadataColumns.isEmpty) {
+ readRowTypeWithRowId
+ } else {
+ new RowType(
+ (readRowTypeWithRowId.getFields.asScala ++ scoreMetadataColumns.map(
+ _.toPaimonDataField)).asJava)
+ }
+ val sparkRow = SparkInternalRow.create(resultRowType)
+ val vectorSearchBuilder = innerTable
+ .newBatchVectorSearchBuilder()
+ .withVectorColumn(columnName)
+ .withLimit(limit)
+ .withOptions(options.asJava)
+
+ val vectorPlan = vectorSearchBuilder.newVectorScan().scan()
+ val batchSize =
+ Math.max(1, new
CoreOptions(innerTable.options()).vectorSearchLateralJoinBatchSize())
+
+ LateralVectorSearchContext(
+ readBuilder,
+ vectorSearchBuilder,
+ vectorPlan,
+ scoreMetadataColumns,
+ sparkRow,
+ rowIdOrdinal = resultRowType.getFieldIndex(SpecialFields.ROW_ID.name()),
+ projectionInputOrdinals = vectorSearchOutput.map {
+ attr =>
+ if (attr.name == PaimonMetadataColumn.SEARCH_SCORE_COLUMN) {
+ -1
+ } else {
+ resultRowType.getFieldIndex(attr.name)
+ }
+ },
+ rightProjection,
+ batchSize
+ )
+ }
+
+ private def search(
+ queries: Seq[LateralVectorSearchQuery],
+ context: LateralVectorSearchContext): Iterator[(InternalRow,
InternalRow)] = {
+ val vectors = queries.map(_.queryVector).toArray
+ val globalIndexResults = context.vectorSearchBuilder
+ .withVectors(vectors)
+ .newBatchVectorRead()
+ .readBatch(context.vectorPlan)
+ .asScala
+ .toVector
+ val rowIdToMatches = createRowIdToMatches(queries, globalIndexResults)
+ val batchGlobalIndexResult =
createBatchGlobalIndexResult(globalIndexResults)
+ val scan = context.readBuilder
+ .newScan()
+ .withGlobalIndexResult(batchGlobalIndexResult)
+ .asInstanceOf[InnerTableScan]
+ val read = context.readBuilder.newRead()
+
+ scan.plan().splits().asScala.iterator.flatMap {
+ split =>
+ val reader =
Review Comment:
This reader is only closed when this inner iterator is fully exhausted. If a
downstream operator short-circuits consumption, for example `LIMIT 1`/`take`,
or if the task is interrupted, Spark can stop pulling rows before `hasNext`
returns false, leaving the current `PaimonRecordReaderIterator` and its
underlying `RecordReader` open. Please register a `TaskContext` completion
listener or wrap the returned iterator so the current reader is closed on task
completion/cancellation as well as normal exhaustion.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
