HeartSaVioR commented on code in PR #59910: URL: https://github.com/apache/spark/pull/59910#discussion_r4419820080
########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PrepareExternalLookupJoin.scala: ########## @@ -0,0 +1,233 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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.apache.spark.sql.execution.datasources.v2 + +import java.util.{Collections, IdentityHashMap} + +import scala.collection.mutable + +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, EqualTo, Expression, + Literal, PredicateHelper} +import org.apache.spark.sql.catalyst.plans.LeftOuter +import org.apache.spark.sql.catalyst.plans.logical.{ExternalLookupJoin, Filter, Join, LocalRelation, + LogicalPlan, PreparedExternalLookupJoin, Project} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.EXTERNAL_LOOKUP_JOIN +import org.apache.spark.sql.catalyst.util.CharVarcharUtils +import org.apache.spark.sql.connector.catalog.{SupportsRead, TableCapability} +import org.apache.spark.sql.connector.expressions.{Expression => V2Expression, FieldReference} +import org.apache.spark.sql.connector.expressions.filter.{Predicate => V2Predicate} +import org.apache.spark.sql.connector.read.{LookupParameter, SupportsExternalLookup} +import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{BooleanType, DataType} + +/** Finalizes an optimized external lookup side immediately before ordinary V2 scan preparation. */ +object PrepareExternalLookupJoin extends Rule[LogicalPlan] with PredicateHelper + with QueryErrorsBase { + override def apply(plan: LogicalPlan): LogicalPlan = { + val prepared = plan.transformUpWithPruning(_.containsPattern(EXTERNAL_LOOKUP_JOIN)) { + case join: ExternalLookupJoin if provablyEmpty(join) => emptyJoinResult(join) + case join: ExternalLookupJoin => prepare(join) + } + // Prepared lookup joins no longer contain their lookup relation, so a lookup-only table that + // is still read here cannot run. Pre-order visits a join before the relation it reads. + prepared.foreach { + // Classic Dataset.join(right, usingColumns) first analyzes Join(left, right, joinType, None), + // including CheckAnalysis, to disambiguate a self-join, and only then builds the USING join. + // That plan cannot be told apart from a join without a condition, so analysis keeps it, and + // a missing condition is reported here, on the final plan. + case Join(_, right, _, None, _) if readsLookupOnlyTable(right) => + throw QueryCompilationErrors.unsupportedExternalLookupJoinError("MISSING_JOIN_CONDITION") + case node => + ExternalLookupRelations.lookupOnlyTable(node).foreach { table => + throw QueryCompilationErrors.unsupportedBatchReadError(table) + } + } + prepared + } + + private def readsLookupOnlyTable(plan: LogicalPlan): Boolean = plan match { + case node @ (_: Project | _: Filter) => readsLookupOnlyTable(node.children.head) + case node => ExternalLookupRelations.lookupOnlyTable(node).isDefined + } + + // The optimizer folds `WHERE false` on the lookup side into an empty relation, and a condition + // such as `s.id = r.id AND false` into `false`. Such a join needs no connector call. + private def provablyEmpty(join: ExternalLookupJoin): Boolean = + join.right.maxRows.contains(0L) || (join.condition match { + case Some(Literal(value, BooleanType)) => value == null || value == false + case _ => false + }) + + private def emptyJoinResult(join: ExternalLookupJoin): LogicalPlan = join.joinType match { + case LeftOuter => + val nullLookupColumns = join.right.output.map { attr => + Alias(Literal(null, attr.dataType), attr.name)(attr.exprId) + } + Project(join.left.output ++ nullLookupColumns, join.left) + case _ => LocalRelation(join.output, Seq.empty, join.isStreaming) + } + + private def prepare(join: ExternalLookupJoin): PreparedExternalLookupJoin = { + val relation = extractLookupPlan(join.right).getOrElse { + throw QueryCompilationErrors.unsupportedExternalLookupJoinError("LOOKUP_PLAN") + } + val tableName = relation.table.name() + val supportsRead = relation.table match { + case table: SupportsRead + if table.capabilities().contains(TableCapability.EXTERNAL_LOOKUP) => table + case _ => + throw QueryCompilationErrors.unsupportedExternalLookupJoinError( + "TABLE", Map("table" -> toSQLId(tableName))) + } + val scanBuilder = supportsRead.newScanBuilder(relation.options) + val lookupBuilder = scanBuilder match { + case builder: SupportsExternalLookup => builder + case _ => + throw QueryCompilationErrors.unsupportedExternalLookupJoinError( + "SCAN_BUILDER", Map("table" -> toSQLId(tableName))) + } + + val conjuncts = join.condition.toSeq.flatMap(splitConjunctivePredicates) + val offeredCandidates = mutable.ArrayBuffer.empty[OfferedLookupPredicate] + val residual = mutable.ArrayBuffer.empty[Expression] + conjuncts.foreach { + case equality @ EqualTo(leftExpr, rightAttr: Attribute) + if leftExpr.references.subsetOf(join.left.outputSet) && + join.right.outputSet.contains(rightAttr) => + offeredCandidates += OfferedLookupPredicate(leftExpr, rightAttr, equality, "=", false) + case equality @ EqualTo(rightAttr: Attribute, leftExpr) + if leftExpr.references.subsetOf(join.left.outputSet) && + join.right.outputSet.contains(rightAttr) => + offeredCandidates += OfferedLookupPredicate(leftExpr, rightAttr, equality, "=", false) + case other => residual += other + } + + val offeredPredicates = offeredCandidates.map { candidate => + new V2Predicate(candidate.predicateName, Array[V2Expression]( + FieldReference.column(candidate.projectedField.name), + new LookupParameter(candidate.parameter.dataType))) + }.toArray + // A copy, so a connector that reorders the array in place cannot break the identity mapping. + val (selectedPredicates, unselectedPredicates) = selectedCandidates(tableName, + offeredPredicates, offeredCandidates, + lookupBuilder.selectLookupPredicates(offeredPredicates.clone())) + residual ++= unselectedPredicates.map(_.condition) + + val scan = scanBuilder.build() + val resolver = SQLConf.get.resolver + // Fields match the way ordinary V2 scans match them (`PushDownUtils.toOutputAttrs`): CHAR and + // VARCHAR are read as STRING, and nullability, including nested nullability, does not count. + val readSchema = CharVarcharUtils.replaceCharVarcharWithStringInSchema(scan.readSchema()) + val lookupInput = readSchema.fields.map { field => + relation.output.find { attr => Review Comment: > [!NOTE] > **DISCLAIMER:** This comment was written by an LLM, Claude Opus 5.5 (`claude-opus-5-5`, 1M context), posting on behalf of @HeartSaVioR. Fixed in #59909 ([73b62b7c559](https://github.com/apache/spark/pull/59909/commits/73b62b7c559b90caf5ac9088417679591118c5d5)) and [01060c7f72e](https://github.com/apache/spark/pull/59910/commits/01060c7f72e4e7256c177939e0e2306e463dfae0). In #59909, `lookupInput` takes the read schema's types and nullability and keeps the relation's exprIds. `lookupOutput` keeps the analyzed attributes in analyzed order, and a read schema that misses one of them fails with `MISSING_SCAN_FIELD`. In this PR, `ExternalLookupJoinExec` reads the connector's value rows through `lookupInput`. Its lookup columns keep their analyzed order and take the read schema's types and nullability, and both the residual and the output projection are bound to the read-schema fields. New tests: - A read schema in another column order returns the right columns. - For the `nested` table (`s.a INT NOT NULL`, read as nullable), the null that the connector returns comes back as null: `to_json(s)` is `{}`. With the table's types it came back as `{"a":0}`. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PrepareExternalLookupJoin.scala: ########## @@ -0,0 +1,233 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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.apache.spark.sql.execution.datasources.v2 + +import java.util.{Collections, IdentityHashMap} + +import scala.collection.mutable + +import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, EqualTo, Expression, + Literal, PredicateHelper} +import org.apache.spark.sql.catalyst.plans.LeftOuter +import org.apache.spark.sql.catalyst.plans.logical.{ExternalLookupJoin, Filter, Join, LocalRelation, + LogicalPlan, PreparedExternalLookupJoin, Project} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.catalyst.trees.TreePattern.EXTERNAL_LOOKUP_JOIN +import org.apache.spark.sql.catalyst.util.CharVarcharUtils +import org.apache.spark.sql.connector.catalog.{SupportsRead, TableCapability} +import org.apache.spark.sql.connector.expressions.{Expression => V2Expression, FieldReference} +import org.apache.spark.sql.connector.expressions.filter.{Predicate => V2Predicate} +import org.apache.spark.sql.connector.read.{LookupParameter, SupportsExternalLookup} +import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{BooleanType, DataType} + +/** Finalizes an optimized external lookup side immediately before ordinary V2 scan preparation. */ +object PrepareExternalLookupJoin extends Rule[LogicalPlan] with PredicateHelper + with QueryErrorsBase { + override def apply(plan: LogicalPlan): LogicalPlan = { + val prepared = plan.transformUpWithPruning(_.containsPattern(EXTERNAL_LOOKUP_JOIN)) { + case join: ExternalLookupJoin if provablyEmpty(join) => emptyJoinResult(join) + case join: ExternalLookupJoin => prepare(join) + } + // Prepared lookup joins no longer contain their lookup relation, so a lookup-only table that + // is still read here cannot run. Pre-order visits a join before the relation it reads. + prepared.foreach { + // Classic Dataset.join(right, usingColumns) first analyzes Join(left, right, joinType, None), + // including CheckAnalysis, to disambiguate a self-join, and only then builds the USING join. + // That plan cannot be told apart from a join without a condition, so analysis keeps it, and + // a missing condition is reported here, on the final plan. + case Join(_, right, _, None, _) if readsLookupOnlyTable(right) => + throw QueryCompilationErrors.unsupportedExternalLookupJoinError("MISSING_JOIN_CONDITION") + case node => + ExternalLookupRelations.lookupOnlyTable(node).foreach { table => + throw QueryCompilationErrors.unsupportedBatchReadError(table) + } + } + prepared + } + + private def readsLookupOnlyTable(plan: LogicalPlan): Boolean = plan match { + case node @ (_: Project | _: Filter) => readsLookupOnlyTable(node.children.head) + case node => ExternalLookupRelations.lookupOnlyTable(node).isDefined + } + + // The optimizer folds `WHERE false` on the lookup side into an empty relation, and a condition + // such as `s.id = r.id AND false` into `false`. Such a join needs no connector call. + private def provablyEmpty(join: ExternalLookupJoin): Boolean = + join.right.maxRows.contains(0L) || (join.condition match { + case Some(Literal(value, BooleanType)) => value == null || value == false + case _ => false + }) + + private def emptyJoinResult(join: ExternalLookupJoin): LogicalPlan = join.joinType match { + case LeftOuter => + val nullLookupColumns = join.right.output.map { attr => + Alias(Literal(null, attr.dataType), attr.name)(attr.exprId) + } + Project(join.left.output ++ nullLookupColumns, join.left) + case _ => LocalRelation(join.output, Seq.empty, join.isStreaming) + } + + private def prepare(join: ExternalLookupJoin): PreparedExternalLookupJoin = { + val relation = extractLookupPlan(join.right).getOrElse { + throw QueryCompilationErrors.unsupportedExternalLookupJoinError("LOOKUP_PLAN") + } + val tableName = relation.table.name() + val supportsRead = relation.table match { + case table: SupportsRead + if table.capabilities().contains(TableCapability.EXTERNAL_LOOKUP) => table + case _ => + throw QueryCompilationErrors.unsupportedExternalLookupJoinError( + "TABLE", Map("table" -> toSQLId(tableName))) + } + val scanBuilder = supportsRead.newScanBuilder(relation.options) + val lookupBuilder = scanBuilder match { + case builder: SupportsExternalLookup => builder + case _ => + throw QueryCompilationErrors.unsupportedExternalLookupJoinError( + "SCAN_BUILDER", Map("table" -> toSQLId(tableName))) + } + + val conjuncts = join.condition.toSeq.flatMap(splitConjunctivePredicates) + val offeredCandidates = mutable.ArrayBuffer.empty[OfferedLookupPredicate] + val residual = mutable.ArrayBuffer.empty[Expression] + conjuncts.foreach { + case equality @ EqualTo(leftExpr, rightAttr: Attribute) + if leftExpr.references.subsetOf(join.left.outputSet) && + join.right.outputSet.contains(rightAttr) => + offeredCandidates += OfferedLookupPredicate(leftExpr, rightAttr, equality, "=", false) + case equality @ EqualTo(rightAttr: Attribute, leftExpr) + if leftExpr.references.subsetOf(join.left.outputSet) && + join.right.outputSet.contains(rightAttr) => + offeredCandidates += OfferedLookupPredicate(leftExpr, rightAttr, equality, "=", false) + case other => residual += other Review Comment: > [!NOTE] > **DISCLAIMER:** This comment was written by an LLM, Claude Opus 5.5 (`claude-opus-5-5`, 1M context), posting on behalf of @HeartSaVioR. Fixed in [01060c7f72e](https://github.com/apache/spark/pull/59910/commits/01060c7f72e4e7256c177939e0e2306e463dfae0), as suggested. A conjunct with a scalar Python UDF that references the lookup side never goes into the residual. In an inner join, it is evaluated in a `Filter` above the prepared join, where `ExtractPythonUDFs` extracts it. A left outer join fails during planning with `UNSUPPORTED_FEATURE.PYTHON_UDF_IN_ON_CLAUSE`, as `ExtractPythonUDFFromJoinCondition` does for other joins. A Python UDF over only the stream side stays in the lookup join, where `ExtractPythonUDFs` already extracts it. The new test uses the integrated Python UDF test utilities and covers both join types. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ExternalLookupJoinEvaluatorFactory.scala: ########## @@ -0,0 +1,410 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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.apache.spark.sql.execution.joins + +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicBoolean + +import scala.collection.mutable + +import org.apache.spark.{PartitionEvaluator, PartitionEvaluatorFactory, SparkIllegalStateException, + TaskContext} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, CollationKey, Expression, + GenericInternalRow, JoinedRow, Predicate, UnsafeProjection, UnsafeRow} +import org.apache.spark.sql.catalyst.optimizer.NormalizeFloatingNumbers +import org.apache.spark.sql.catalyst.plans.{Inner, JoinType, LeftOuter} +import org.apache.spark.sql.connector.metric.CustomTaskMetric +import org.apache.spark.sql.connector.read.{ExternalLookupReader, ExternalLookupReaderFactory, + ExternalLookupResultCursor} +import org.apache.spark.sql.execution.metric.{CustomMetrics, SQLMetric} +import org.apache.spark.util.Utils + +/** + * Creates the partition evaluators of an [[ExternalLookupJoinExec]]. The factory is sent to the + * tasks, so it holds only what they need; the plan node and its scan stay on the driver. + * + * An evaluator joins its partition one batch of up to `maxParameterRowsPerRequest` stream rows at + * a time. For each batch it sends one request with the distinct parameters that can match, and + * streams the connector's result rows into output rows. + */ +class ExternalLookupJoinEvaluatorFactory( + lookupParameters: Seq[Expression], + nullSafeParameters: Seq[Boolean], + streamOutput: Seq[Attribute], + lookupOutput: Seq[Attribute], + output: Seq[Attribute], + joinType: JoinType, + residualCondition: Option[Expression], + readerFactory: ExternalLookupReaderFactory, + maxParameterRowsPerRequest: Int, + metrics: ExternalLookupJoinMetrics, + readerMetrics: ExternalLookupReaderMetrics) + extends PartitionEvaluatorFactory[InternalRow, InternalRow] { + + require(lookupParameters.length == nullSafeParameters.length, + "Every external lookup parameter must define its null semantics") + + override def createEvaluator(): PartitionEvaluator[InternalRow, InternalRow] = + new ExternalLookupJoinEvaluator + + private class ExternalLookupJoinEvaluator extends PartitionEvaluator[InternalRow, InternalRow] { + override def eval( + partitionIndex: Int, + inputs: Iterator[InternalRow]*): Iterator[InternalRow] = { + assert(inputs.length == 1) + new PartitionJoinIterator(inputs.head, partitionIndex) + } + } + + /** Joins one partition batch by batch, and closes its reader after the last batch. */ + private class PartitionJoinIterator(input: Iterator[InternalRow], partitionIndex: Int) + extends Iterator[InternalRow] { + private val reader = new LookupReader + private val parameters = new LookupParameters(partitionIndex) + private val joinOutput = new JoinOutput(partitionIndex) + private var batch: Iterator[InternalRow] = Iterator.empty + private var inputExhausted = false + + override def hasNext: Boolean = { + tryAdvance() + batch.hasNext + } + + override def next(): InternalRow = { + tryAdvance() + if (!batch.hasNext) { + throw new SparkIllegalStateException( + errorClass = "INTERNAL_ERROR", + messageParameters = Map("message" -> "External lookup join exhausted")) + } + batch.next() + } + + /** Reads more only when nothing is buffered; safe to call repeatedly. */ + private def tryAdvance(): Unit = { + while (!batch.hasNext && !inputExhausted) { + val rows = readBatch() + if (rows.nonEmpty) batch = new BatchJoinIterator(rows, parameters, joinOutput, reader) + inputExhausted = !input.hasNext + } + if (inputExhausted && !batch.hasNext) reader.close() + } + + private def readBatch(): Seq[InternalRow] = { + val rows = mutable.ArrayBuffer.empty[InternalRow] + while (rows.length < maxParameterRowsPerRequest && input.hasNext) { + rows += input.next().copy() + } + rows.toSeq Review Comment: > [!NOTE] > **DISCLAIMER:** This comment was written by an LLM, Claude Opus 5.5 (`claude-opus-5-5`, 1M context), posting on behalf of @HeartSaVioR. Fixed in [01060c7f72e](https://github.com/apache/spark/pull/59910/commits/01060c7f72e4e7256c177939e0e2306e463dfae0). `readBatch` returns an `Array[InternalRow]`, and `BatchJoinIterator` takes the array, so reading a row by position is constant time. Results and request sizes are unchanged. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/ExternalLookupJoinEvaluatorFactory.scala: ########## @@ -0,0 +1,410 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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.apache.spark.sql.execution.joins + +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicBoolean + +import scala.collection.mutable + +import org.apache.spark.{PartitionEvaluator, PartitionEvaluatorFactory, SparkIllegalStateException, + TaskContext} +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, CollationKey, Expression, + GenericInternalRow, JoinedRow, Predicate, UnsafeProjection, UnsafeRow} +import org.apache.spark.sql.catalyst.optimizer.NormalizeFloatingNumbers +import org.apache.spark.sql.catalyst.plans.{Inner, JoinType, LeftOuter} +import org.apache.spark.sql.connector.metric.CustomTaskMetric +import org.apache.spark.sql.connector.read.{ExternalLookupReader, ExternalLookupReaderFactory, + ExternalLookupResultCursor} +import org.apache.spark.sql.execution.metric.{CustomMetrics, SQLMetric} +import org.apache.spark.util.Utils + +/** + * Creates the partition evaluators of an [[ExternalLookupJoinExec]]. The factory is sent to the + * tasks, so it holds only what they need; the plan node and its scan stay on the driver. + * + * An evaluator joins its partition one batch of up to `maxParameterRowsPerRequest` stream rows at + * a time. For each batch it sends one request with the distinct parameters that can match, and + * streams the connector's result rows into output rows. + */ +class ExternalLookupJoinEvaluatorFactory( + lookupParameters: Seq[Expression], + nullSafeParameters: Seq[Boolean], + streamOutput: Seq[Attribute], + lookupOutput: Seq[Attribute], + output: Seq[Attribute], + joinType: JoinType, + residualCondition: Option[Expression], + readerFactory: ExternalLookupReaderFactory, + maxParameterRowsPerRequest: Int, + metrics: ExternalLookupJoinMetrics, + readerMetrics: ExternalLookupReaderMetrics) + extends PartitionEvaluatorFactory[InternalRow, InternalRow] { + + require(lookupParameters.length == nullSafeParameters.length, + "Every external lookup parameter must define its null semantics") + + override def createEvaluator(): PartitionEvaluator[InternalRow, InternalRow] = + new ExternalLookupJoinEvaluator + + private class ExternalLookupJoinEvaluator extends PartitionEvaluator[InternalRow, InternalRow] { + override def eval( + partitionIndex: Int, + inputs: Iterator[InternalRow]*): Iterator[InternalRow] = { + assert(inputs.length == 1) + new PartitionJoinIterator(inputs.head, partitionIndex) + } + } + + /** Joins one partition batch by batch, and closes its reader after the last batch. */ + private class PartitionJoinIterator(input: Iterator[InternalRow], partitionIndex: Int) + extends Iterator[InternalRow] { + private val reader = new LookupReader + private val parameters = new LookupParameters(partitionIndex) + private val joinOutput = new JoinOutput(partitionIndex) + private var batch: Iterator[InternalRow] = Iterator.empty + private var inputExhausted = false + + override def hasNext: Boolean = { + tryAdvance() + batch.hasNext + } + + override def next(): InternalRow = { + tryAdvance() + if (!batch.hasNext) { + throw new SparkIllegalStateException( + errorClass = "INTERNAL_ERROR", + messageParameters = Map("message" -> "External lookup join exhausted")) + } + batch.next() + } + + /** Reads more only when nothing is buffered; safe to call repeatedly. */ + private def tryAdvance(): Unit = { + while (!batch.hasNext && !inputExhausted) { + val rows = readBatch() + if (rows.nonEmpty) batch = new BatchJoinIterator(rows, parameters, joinOutput, reader) + inputExhausted = !input.hasNext + } + if (inputExhausted && !batch.hasNext) reader.close() + } + + private def readBatch(): Seq[InternalRow] = { + val rows = mutable.ArrayBuffer.empty[InternalRow] + while (rows.length < maxParameterRowsPerRequest && input.hasNext) { + rows += input.next().copy() + } + rows.toSeq + } + } + + /** The stream rows that share a key, and the one request row sent for them. */ + private case class KeyEntry(requestRow: UnsafeRow, rowIndexes: mutable.ArrayBuffer[Int]) + + /** + * Joins one batch of stream rows. Rows whose parameters are equal under Spark equality share + * one request row, and every result row for a request row joins all the stream rows sharing it. + * A stream row is finished, which may emit its unmatched output row, once no more result rows + * can match it. + * + * Output rows are projected into a reused row on `next()`, so the consumer copies the rows it + * retains, as with other join iterators. + */ + private class BatchJoinIterator( + rows: Seq[InternalRow], + parameters: LookupParameters, + joinOutput: JoinOutput, + reader: LookupReader) extends Iterator[InternalRow] { + metrics.inputRows += rows.size + metrics.inputBatches += 1 + + private val matched = Array.fill(rows.length)(false) + private val finished = mutable.BitSet.empty + // Output rows not yet returned: a stream row index and the lookup row it is joined with. + private val pending = mutable.Queue.empty[(Int, InternalRow)] + + // One entry for each distinct key that can match, in the order of first occurrence. A row + // that can never match is finished right away, and its key is not requested. + private val keyEntries = mutable.LinkedHashMap.empty[UnsafeRow, KeyEntry] + rows.indices.foreach { index => + val key = parameters.key(rows(index)) + if (parameters.matchesNothing(key)) { + finish(index) + } else { + keyEntries.get(key) match { + case Some(entry) => entry.rowIndexes += index + case None => keyEntries.put( + key.copy(), + KeyEntry(parameters.requestRow(rows(index)), mutable.ArrayBuffer(index))) + } + } + } + private val requestedKeys = keyEntries.keys.toArray + + private val lookupStartNanos = System.nanoTime() + // No connector request, and no reader, when no key can match. + private var cursor: ExternalLookupResultCursor = + if (keyEntries.isEmpty) EmptyLookupCursor + else reader.lookup(keyEntries.values.map(_.requestRow).toArray) + + override def hasNext: Boolean = { + tryFillPending() + pending.nonEmpty + } + + override def next(): InternalRow = { + tryFillPending() + if (pending.isEmpty) throw new NoSuchElementException("external lookup result") + val (index, lookupRow) = pending.dequeue() + joinOutput.project(rows(index), lookupRow) + } + + /** Reads more only when nothing is buffered; safe to call repeatedly. */ + private def tryFillPending(): Unit = { + try { + while (pending.isEmpty && cursor.next()) { + reader.resultRowRead() + val key = requestedKey(cursor.inputOrdinal()) + // The connector may reuse its row, and several stream rows can join it. + joinResultRow(key, cursor.get().copy()) + } + if (pending.isEmpty) finishCursor() + } catch { + case error: Throwable => + Utils.tryWithSafeFinally(throw error) { + finishCursor() + } + } + } + + private def requestedKey(ordinal: Int): UnsafeRow = { + if (ordinal < 0 || ordinal >= requestedKeys.length) { + throw new SparkIllegalStateException( + errorClass = "INTERNAL_ERROR", + messageParameters = Map("message" -> + (s"External lookup returned invalid input ordinal $ordinal for " + + s"${requestedKeys.length} requested parameter rows"))) + } + requestedKeys(ordinal) + } + + private def joinResultRow(key: UnsafeRow, lookupRow: InternalRow): Unit = + keyEntries(key).rowIndexes.foreach(joinMatch(_, lookupRow)) + + private def joinMatch(index: Int, lookupRow: InternalRow): Unit = { + if (joinOutput.matches(rows(index), lookupRow)) { + matched(index) = true + joinOutput.matched(lookupRow).foreach(emit(index, _)) + } + } + + /** Closes the cursor once, and finishes every stream row that was waiting for its results. */ + private def finishCursor(): Unit = if (cursor ne EmptyLookupCursor) { + cursor = EmptyLookupCursor + try { + reader.closeCursor() + } finally { + metrics.lookupTime += System.nanoTime() - lookupStartNanos + reader.updateMetrics() + } + keyEntries.values.foreach(_.rowIndexes.foreach(finish)) + } + + private def finish(index: Int): Unit = if (!finished.contains(index)) { + finished += index + if (!matched(index)) joinOutput.unmatched.foreach(emit(index, _)) + } + + private def emit(index: Int, lookupRow: InternalRow): Unit = { + pending.enqueue(index -> lookupRow) + metrics.outputRows += 1 + } + } + + /** The cursor of a batch that sends no request, or whose cursor is finished. */ + private object EmptyLookupCursor extends ExternalLookupResultCursor { + override def next(): Boolean = false + override def inputOrdinal(): Int = throw new NoSuchElementException("external lookup result") + override def get(): InternalRow = throw new NoSuchElementException("external lookup result") + override def close(): Unit = {} + } + + /** + * Evaluates the lookup parameters of stream rows: the request row that is sent to the + * connector, and the key that identifies equal parameters under Spark equality, which + * normalizes floating-point values and collations. + */ + private class LookupParameters(partitionIndex: Int) { + private val requestProjection = UnsafeProjection.create(lookupParameters, streamOutput) + private val keyProjection = UnsafeProjection.create(lookupParameters.map { parameter => + NormalizeFloatingNumbers.normalize(CollationKey.injectCollationKey(parameter)) + }, streamOutput) + requestProjection.initialize(partitionIndex) + keyProjection.initialize(partitionIndex) + + def requestRow(streamRow: InternalRow): UnsafeRow = requestProjection(streamRow).copy() + + /** The key of a stream row, in a reused row: copy it to retain it. */ + def key(streamRow: InternalRow): UnsafeRow = keyProjection(streamRow) + + /** + * A null parameter that is not null-safe matches no lookup row, so it is not requested. + * The key has the request row's nulls. + */ + def matchesNothing(key: InternalRow): Boolean = + nullSafeParameters.indices.exists { index => + !nullSafeParameters(index) && key.isNullAt(index) + } + } + + /** Evaluates the residual condition on joined rows and builds output rows per join type. */ + private class JoinOutput(partitionIndex: Int) { + private val joinedLookupOutput = + if (joinType == LeftOuter) lookupOutput.map(_.withNullability(true)) else lookupOutput + private val joinedInput = streamOutput ++ joinedLookupOutput + private val joinedRow = new JoinedRow() + private val residual = residualCondition.map(Predicate.create(_, joinedInput)) + private val projection = UnsafeProjection.create(output, joinedInput) + private val nullLookupRow = new GenericInternalRow(lookupOutput.length) + residual.foreach(_.initialize(partitionIndex)) + projection.initialize(partitionIndex) + + def matches(streamRow: InternalRow, lookupRow: InternalRow): Boolean = + residual.forall(_.eval(joinedRow(streamRow, lookupRow))) + + /** The lookup row that the output for a match joins, if the join type emits one. */ + def matched(lookupRow: InternalRow): Option[InternalRow] = joinType match { + case Inner | LeftOuter => Some(lookupRow) + case _ => None + } + + /** The lookup row that the output for a stream row without a match joins, if any. */ + def unmatched: Option[InternalRow] = joinType match { + case LeftOuter => Some(nullLookupRow) + case _ => None + } + + /** The output row, projected into a reused row. */ + def project(streamRow: InternalRow, lookupRow: InternalRow): InternalRow = + projection(joinedRow(streamRow, lookupRow)) + } + + /** + * The connector reader of one partition and its active cursor. The reader is opened by the + * first request, and closed once, after the last batch or when the task completes. + */ + private class LookupReader { + private val taskReaders = readerMetrics.forCurrentTask() + private val closed = new AtomicBoolean(false) + private var reader: ExternalLookupReader = null + private var activeCursor: ExternalLookupResultCursor = null + private var resultRowsSinceMetricsUpdate = 0 + Option(TaskContext.get()).foreach(_.addTaskCompletionListener[Unit](_ => close())) + + /** Sends one request and returns the cursor over its result rows. */ + def lookup(requestRows: Array[UnsafeRow]): ExternalLookupResultCursor = { Review Comment: > [!NOTE] > **DISCLAIMER:** This comment was written by an LLM, Claude Opus 5.5 (`claude-opus-5-5`, 1M context), posting on behalf of @HeartSaVioR. Fixed in [01060c7f72e](https://github.com/apache/spark/pull/59910/commits/01060c7f72e4e7256c177939e0e2306e463dfae0), partly as suggested. `lookup()` now fails once the reader is closed, before it counts a request or opens a reader, so a closed reader is never reused and no reader is opened after the task completed. `ExternalLookupReaderMetrics` registers its cleanup listener outside the map update (`putIfAbsent`, then the listener), so a first call after task completion no longer removes the key inside `computeIfAbsent`. I did not add a lock around `lookup` and `close`. `close()` runs from the task-completion listener, so a synchronized close would block task completion behind a slow or hung connector request. `DataSourceRDD` does not lock its readers against the completion listener either. The new test computes the join in a task that completes before, or while, the join is consumed, and checks that the next request fails without creating a reader. ########## common/utils/src/main/resources/error/error-conditions.json: ########## @@ -2817,6 +2861,59 @@ ], "sqlState" : "42822" }, + "EXTERNAL_LOOKUP_JOIN_UNSUPPORTED" : { + "message" : [ + "The join cannot be executed as an external lookup join:" + ], + "subClass" : { + "CONFLICTING_JOIN_STRATEGY" : { + "message" : [ + "The join strategy hint <strategy> cannot be honored by an external lookup join over <relation>. Remove the <strategy> hint, or the `EXTERNAL_LOOKUP` hint if the table also supports batch reads." Review Comment: > [!NOTE] > **DISCLAIMER:** This comment was written by an LLM, Claude Opus 5.5 (`claude-opus-5-5`, 1M context), posting on behalf of @HeartSaVioR. Addressed in #59909 ([73b62b7c559](https://github.com/apache/spark/pull/59909/commits/73b62b7c559b90caf5ac9088417679591118c5d5)), which owns the error text. `CONFLICTING_JOIN_STRATEGY` now only advises removing the conflicting strategy hint, and `LOOKUP_PLAN` describes only the shapes supported now. -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
