HeartSaVioR commented on code in PR #59910: URL: https://github.com/apache/spark/pull/59910#discussion_r4419820705
########## 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/commit/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/commit/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. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/ExternalLookupJoinSuite.scala: ########## @@ -0,0 +1,1358 @@ +/* + * 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.connector + +import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} + +import scala.collection.mutable +import scala.concurrent.duration._ +import scala.jdk.CollectionConverters._ +import scala.reflect.{classTag, ClassTag} + +import org.apache.spark.{SparkException, SparkIllegalArgumentException, + SparkUnsupportedOperationException} +import org.apache.spark.sql.{AnalysisException, DataFrame, Encoders, Row} +import org.apache.spark.sql.QueryTest +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.analysis.NoSuchTableException +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, PreparedExternalLookupJoin} +import org.apache.spark.sql.classic.Dataset +import org.apache.spark.sql.connector.catalog.{Identifier, SupportsRead, Table, TableCapability, + TableCatalog, TableChange, TableProvider} +import org.apache.spark.sql.connector.catalog.TableCapability.{BATCH_READ, EXTERNAL_LOOKUP} +import org.apache.spark.sql.connector.expressions.{NamedReference, Transform} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.metric.{CustomMetric, CustomSumMetric, CustomTaskMetric} +import org.apache.spark.sql.connector.read.{Batch, ExternalLookup, ExternalLookupReader, + ExternalLookupReaderFactory, ExternalLookupResultCursor, InputPartition, LookupParameter, + PartitionReader, PartitionReaderFactory, Scan, ScanBuilder, SupportsExternalLookup, + SupportsPushDownRequiredColumns, SupportsPushDownV2Filters} +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, ExternalLookupRelation} +import org.apache.spark.sql.execution.joins.ExternalLookupJoinExec +import org.apache.spark.sql.execution.streaming.runtime.{IncrementalExecution, MemoryStream, + StreamingQueryWrapper} +import org.apache.spark.sql.execution.streaming.sources.LowLatencyMemoryStream +import org.apache.spark.sql.functions.{coalesce, col, concat, length, lit} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.streaming.{StreamingQueryException, StreamRealTimeModeE2ESuiteBase} +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{DataType, IntegerType, StringType, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.unsafe.types.UTF8String + +class ExternalLookupJoinSuite extends QueryTest with SharedSparkSession { + + import testImplicits._ + + private val ExternalLookupCatalogName = "external_lookup_cat" + + private val streamRuns = new AtomicInteger() + + /** The rows of `df` as a streaming DataFrame: one micro-batch, in `numPartitions` partitions. */ + private def streamOf(df: DataFrame, numPartitions: Int = 1): DataFrame = { + val input = MemoryStream[Row](numPartitions)(Encoders.row(df.schema), spark.sqlContext) + input.addData(df.collect().toSeq) + input.toDF() + } + + private case class StreamRun(rows: Seq[Row], execution: IncrementalExecution) { + def plan: LogicalPlan = execution.optimizedPlan + } + + /** Runs the streaming `df` on a memory sink until it has processed all of its input. */ + private def runStream(df: DataFrame): StreamRun = { + val queryName = s"external_lookup_stream_${streamRuns.incrementAndGet()}" + val query = df.writeStream.format("memory").queryName(queryName).start() + try { + query.processAllAvailable() + StreamRun( + spark.table(queryName).collect().toSeq, + query.asInstanceOf[StreamingQueryWrapper].streamingQuery.lastExecution) + } finally { + query.stop() + spark.catalog.dropTempView(queryName) + } + } + + private def checkStream(df: DataFrame, expected: Seq[Row]): Unit = + QueryTest.sameRows(expected, runStream(df).rows).foreach(fail(_)) + + /** The error that the streaming `df` fails with. */ + private def streamFailure[T <: Throwable : ClassTag](df: DataFrame): T = { + val cause = intercept[StreamingQueryException](runStream(df)).getCause + assert(classTag[T].runtimeClass.isInstance(cause), cause) + cause.asInstanceOf[T] + } + + test("external lookup capability produces a unary external lookup join") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1).toDF("id")) + val joined = left.join(lookup, left("id") === lookup("id"), "inner") + + val plan = runStream(joined).plan + assert(plan.exists(_.isInstanceOf[PreparedExternalLookupJoin])) + assert(!plan.toString.contains("SimpleExternalLookupScan")) + } + + test("an external lookup join is a stream-static join with the lookup-only table on the right") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val batch = Seq(1).toDF("id") + // The left side must be a streaming query. + checkError( + exception = intercept[AnalysisException] { + batch.join(lookup, batch("id") === lookup("id")) + }, + condition = "EXTERNAL_LOOKUP_JOIN_UNSUPPORTED.STATIC_LEFT_SIDE", + parameters = Map("relation" -> "`SimpleExternalLookupTable`")) + // The sides are never swapped. + val stream = streamOf(batch) + checkError( + exception = intercept[AnalysisException] { + lookup.join(stream, lookup("id") === stream("id")) + }, + condition = "EXTERNAL_LOOKUP_JOIN_UNSUPPORTED.LOOKUP_TABLE_ON_LEFT", + parameters = Map("relation" -> "`SimpleExternalLookupTable`")) + } + + test("external lookup validates connector planning and request contracts") { + val left = streamOf(Seq(1).toDF("id")) + Seq( + ("nullKeySelection", "true", "NULL_PREDICATE_SELECTION", Map.empty[String, String]), + ("selectionMode", "nullElement", "NULL_PREDICATE_SELECTION", Map.empty[String, String]), + ("selectionMode", "duplicate", "DUPLICATE_SELECTED_PREDICATE", Map("predicate" -> "id = ?")), + ("selectionMode", "copy", "UNOFFERED_SELECTED_PREDICATE", Map("predicate" -> "id = ?")), + ("maxBatchSize", "0", "NON_POSITIVE_MAX_PARAMETER_ROWS", Map("value" -> "0"))).foreach { + case (option, value, subClass, parameters) => + val lookup = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .option(option, value) + .load() + checkError( + exception = streamFailure[AnalysisException]( + left.join(lookup, left("id") === lookup("id"))), + condition = s"DATA_SOURCE_INVALID_EXTERNAL_LOOKUP.$subClass", + parameters = Map("table" -> "`SimpleExternalLookupTable`") ++ parameters) + } + + val lookup = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .option("nullCursor", "true") + .load() + val error = streamFailure[SparkException](left.join(lookup, left("id") === lookup("id"))) + assert(error.getMessage.contains("External lookup returned a null result cursor")) + + def misconfigured(kind: String): DataFrame = spark.read + .format(classOf[MisconfiguredExternalLookupDataSource].getName) + .option("kind", kind) + .load() + val table = Map("table" -> "`MisconfiguredExternalLookupTable`") + Seq( + ("unknownField", "UNKNOWN_SCAN_FIELD", Map("field" -> "`unknown`")), + ("duplicateField", "DUPLICATE_SCAN_FIELD", Map("field" -> "`id`"))).foreach { + case (kind, subClass, parameters) => + val lookup = misconfigured(kind) + checkError( + exception = streamFailure[AnalysisException]( + left.join(lookup, left("id") === lookup("id"))), + condition = s"DATA_SOURCE_INVALID_EXTERNAL_LOOKUP.$subClass", + parameters = table ++ parameters) + } + + val noExternalLookup = misconfigured("noExternalLookup") + checkError( + exception = streamFailure[SparkUnsupportedOperationException]( + left.join(noExternalLookup, left("id") === noExternalLookup("id"))), + condition = "DATA_SOURCE_EXTERNAL_LOOKUP_NOT_SUPPORTED", + parameters = Map("description" -> "NoExternalLookupScan")) + + val invalidOrdinal = misconfigured("invalidOrdinal") + val ordinalError = streamFailure[SparkException]( + left.join(invalidOrdinal, left("id") === invalidOrdinal("id"))) + assert(ordinalError.getMessage.contains( + "External lookup returned invalid input ordinal 5 for 1 requested parameter rows")) + + // Planning failures that need no hint. A catalog can serve a table without SupportsRead. + val noLookupBuilder = misconfigured("noLookupBuilder") + val nonReadable = Dataset.ofRows(spark, DataSourceV2Relation.create( + NonReadableExternalLookupTable, None, None, CaseInsensitiveStringMap.empty)) + val simple = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + Seq( + (noLookupBuilder, left("id") === noLookupBuilder("id"), "SCAN_BUILDER", + "MisconfiguredExternalLookupTable"), + (nonReadable, left("id") === nonReadable("id"), "TABLE", "NonReadableExternalLookupTable"), + (simple, simple("department") === "sales", "NO_LOOKUP_PREDICATE", + "SimpleExternalLookupTable")).foreach { case (lookup, condition, subClass, name) => + checkError( + exception = streamFailure[AnalysisException](left.join(lookup, condition)), + condition = s"EXTERNAL_LOOKUP_JOIN_UNSUPPORTED.$subClass", + parameters = Map("table" -> s"`$name`")) + } + } + + test("external lookup finalizes expressions on the stream side after analysis") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq[java.lang.Integer](null, 1).toDF("id")) + + checkStream( + left.join( + lookup, + coalesce(left("id"), lit(0)) === lookup("id"), + "inner"), + Seq(Row(1, 1, "engineering"), Row(1, 1, "research"))) + } + + test("external lookup maps a selected predicate to its own occurrence") { + SimpleExternalLookupDataSource.resetCounters() + val lookup = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .option("selectionMode", "last") + .load() + val left = streamOf(Seq((1, 2), (2, 2), (2, 1)).toDF("a", "b")) + + // Both candidates have the shape `id = ?`; only the second one, bound to `b`, is selected. + val joined = + left.join(lookup, lookup("id") === left("a") && lookup("id") === left("b"), "left_outer") + assert(runStream(joined).rows.toSet === + Set(Row(1, 2, null, null), Row(2, 2, 2, "sales"), Row(2, 1, null, null))) + assert(SimpleExternalLookupDataSource.lookupCandidates.get() === + Seq(("=", "id", IntegerType), ("=", "id", IntegerType))) + assert(SimpleExternalLookupDataSource.lookupCandidateDescriptions.get() === + Seq("id = ?", "id = ?")) + assert(SimpleExternalLookupDataSource.requestedKeys.asScala.toSeq === Seq(2, 1)) + + // Selected in reverse order, each parameter row is `(b, a)`. + SimpleExternalLookupDataSource.resetCounters() + val reversed = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .option("selectionMode", "reversed") + .load() + val condition = reversed("id") === left("a") && reversed("id") === left("b") + assert(runStream(left.join(reversed, condition, "left_outer")).rows.toSet === + Set(Row(1, 2, null, null), Row(2, 2, 2, "sales"), Row(2, 1, null, null))) + assert(SimpleExternalLookupDataSource.requestedKeys.asScala.toSeq === Seq(2, 2, 1)) + } + + test("a parent projection of some lookup columns still plans an external lookup join") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1, 3).toDF("id")) + // Generic column pruning must not project the lookup relation, which would hide it from + // lookup planning; the join's requiredLookupOutput keeps it from doing so. + val joined = left.join(lookup, left("id") === lookup("id"), "left_outer") + .select(left("id"), lookup("department")) + + assert(runStream(joined).plan.exists(_.isInstanceOf[PreparedExternalLookupJoin])) + checkStream(joined, Seq(Row(1, "engineering"), Row(1, "research"), Row(3, null))) + } + + test("external lookup join supports inner and left outer joins") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq((1, "one"), (2, "two"), (3, "three")).toDF("id", "label")) + + checkStream( + left.join(lookup, left("id") === lookup("id"), "inner"), + Seq(Row(1, "one", 1, "engineering"), Row(1, "one", 1, "research"), + Row(2, "two", 2, "sales"))) + checkStream( + left.join(lookup, left("id") === lookup("id"), "left_outer"), + Seq(Row(1, "one", 1, "engineering"), Row(1, "one", 1, "research"), + Row(2, "two", 2, "sales"), Row(3, "three", null, null))) + + // Only comparisons of a lookup field with a stream-side value are offered. The connector + // selects `id = ?`, and Spark evaluates the other conjuncts on the joined rows. + SimpleExternalLookupDataSource.resetCounters() + val flagged = streamOf(Seq((1, true), (2, true), (1, false)).toDF("id", "flag")) + checkStream( + flagged.join( + lookup, + flagged("id") === lookup("id") && flagged("flag") && + lookup("department") === "engineering" && length(lookup("department")) > lookup("id"), + "left_outer"), + Seq(Row(1, true, 1, "engineering"), Row(2, true, null, null), Row(1, false, null, null))) + assert(SimpleExternalLookupDataSource.lookupCandidates.get() === + Seq(("=", "id", IntegerType), ("=", "department", StringType))) + } + + test("external lookup-only source joins without an explicit hint") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1).toDF("id")) + + assert(lookup.queryExecution.analyzed.isInstanceOf[ExternalLookupRelation]) + val joined = left.join(lookup, "id") + assert(runStream(joined).plan.exists(_.isInstanceOf[PreparedExternalLookupJoin])) + checkStream(joined, Seq(Row(1, "engineering"), Row(1, "research"))) + } + + test("an external lookup-only source cannot be read outside an eligible lookup join") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1).toDF("id")) + checkError( + exception = intercept[AnalysisException] { lookup.queryExecution.optimizedPlan }, + condition = "UNSUPPORTED_FEATURE.TABLE_OPERATION", + parameters = Map( + "tableName" -> "`SimpleExternalLookupTable`", + "operation" -> "batch scan")) + // A join that reads it but cannot be a lookup join fails with the reason. + checkError( + exception = streamFailure[AnalysisException](left.crossJoin(lookup)), + condition = "EXTERNAL_LOOKUP_JOIN_UNSUPPORTED.MISSING_JOIN_CONDITION", + parameters = Map.empty) + val fullOuter = intercept[AnalysisException] { + left.join(lookup, left("id") === lookup("id"), "full_outer").queryExecution.analyzed + } + assert(fullOuter.getCondition === "EXTERNAL_LOOKUP_JOIN_UNSUPPORTED.JOIN_TYPE") + assert(fullOuter.getMessageParameters.asScala.toMap - "supportedJoinTypes" === + Map("joinType" -> "FULL OUTER", "relation" -> "`SimpleExternalLookupTable`")) + } + + test("a catalog-resolved external lookup-only table cannot be read outside a lookup join") { + withSQLConf( + s"spark.sql.catalog.$ExternalLookupCatalogName" -> + classOf[ExternalLookupTestCatalog].getName) { + val table = s"$ExternalLookupCatalogName.ns.lookup" + val left = Seq(1).toDF("id") + Seq[() => DataFrame]( + () => spark.table(table), + () => sql(s"SELECT * FROM $table WHERE id = 1")).foreach { + df => + checkError( + exception = intercept[AnalysisException] { df().queryExecution.optimizedPlan }, + condition = "UNSUPPORTED_FEATURE.TABLE_OPERATION", + parameters = Map( + "tableName" -> "`SimpleExternalLookupTable`", + "operation" -> "batch scan")) + } + } + } + + test("external lookup join handles null, missing, duplicate, and partial batches") { + SimpleExternalLookupDataSource.resetCounters() + val lookup = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .option("maxBatchSize", "3") + .load() + val left = streamOf(Seq[java.lang.Integer](1, 1, null, 3, 2).toDF("id")) + + val rows = runStream(left.join(lookup, left("id") === lookup("id"), "left_outer")).rows + assert(rows.toSet === Set( + Row(1, 1, "engineering"), + Row(1, 1, "research"), + Row(null, null, null), + Row(3, null, null), + Row(2, 2, "sales"))) + assert(rows.count(_.get(0) == 1) === 4) + assert(SimpleExternalLookupDataSource.requestSizes.asScala.toSeq === Seq(1, 2)) + assert(SimpleExternalLookupDataSource.requestedKeys.asScala.toSeq === Seq(1, 3, 2)) + assert(SimpleExternalLookupDataSource.nullKeyRequests.get() === 0) + } + + test("external lookup join supports complete composite equi-keys") { + val lookup = spark.read + .format(classOf[CompositeExternalLookupDataSource].getName) + .load() + val left = streamOf(Seq((1, 1), (1, 2), (2, 1)).toDF("tenant", "id")) + + checkStream( + left.join( + lookup, + left("tenant") === lookup("tenant") && left("id") === lookup("id"), + "left_outer"), + Seq(Row(1, 1, 1, 1, "one-one"), Row(1, 2, null, null, null), + Row(2, 1, 2, 1, "two-one"))) + } + + test("an upstream limit stops one-to-many cursor consumption") { + SimpleExternalLookupDataSource.resetCounters() + SimpleExternalLookupDataSource.values.set(Map(1 -> (1 to 2000).map(_.toString))) + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1).toDF("id")) + + assert(runStream(left.join( + lookup, left("id") === lookup("id"), "inner") + .limit(1)).rows.size === 1) + assert(SimpleExternalLookupDataSource.resultRowsRead.get() === 1) + } + + test("external lookup samples connector metrics while consuming a large response") { + SimpleExternalLookupDataSource.resetCounters() + SimpleExternalLookupDataSource.values.set(Map(1 -> (1 to 250).map(_.toString))) + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1).toDF("id")) + + assert(runStream(left.join( + lookup, left("id") === lookup("id"), "inner")).rows.size === 250) + // Two periodic samples, followed by request-completion and reader-close samples. + assert(SimpleExternalLookupDataSource.metricSamples.get() >= 4) + } + + test("external lookup deduplicates floating-point keys within a request") { + SemanticKeyExternalLookupDataSource.requestSizes.clear() + val nan1 = java.lang.Double.longBitsToDouble(0x7ff8000000000001L) + val nan2 = java.lang.Double.longBitsToDouble(0x7ff8000000000002L) + // Requests hold two rows, so each pair of equivalent keys shares one request row. + val left = streamOf(Seq(-0.0d, 0.0d, nan1, nan2).toDF("key")) + val lookup = semanticLookup("double") + + assert(runStream(left.join(lookup, left("key") === lookup("key"))).rows.size === 4) + assert(SemanticKeyExternalLookupDataSource.requestSizes.asScala.toSeq === Seq(1, 1)) + } + + test("external lookup deduplicates keys using Spark collation equality") { + SemanticKeyExternalLookupDataSource.requestSizes.clear() + val left = streamOf(Seq("lookup", "LOOKUP").toDF("key")) + .selectExpr("CAST(key AS STRING COLLATE UTF8_LCASE) AS key") + val lookup = semanticLookup("string collate utf8_lcase") + + assert(runStream(left.join(lookup, left("key") === lookup("key"))).rows.size === 2) + assert(SemanticKeyExternalLookupDataSource.requestSizes.asScala.toSeq === Seq(1)) + } + + private def semanticLookup(keyType: String): DataFrame = { + spark.read + .format(classOf[SemanticKeyExternalLookupDataSource].getName) + .option("keyType", keyType) + .load() + } + + test("a join strategy hint on the stream side of a lookup-only join is rejected") { + SimpleExternalLookupDataSource.resetCounters() + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1).toDF("id")) + checkError( + exception = intercept[AnalysisException] { + left.hint("merge").join(lookup, left("id") === lookup("id")).queryExecution.analyzed + }, + condition = "EXTERNAL_LOOKUP_JOIN_UNSUPPORTED.CONFLICTING_JOIN_STRATEGY", + parameters = Map("strategy" -> "MERGE", "relation" -> "`SimpleExternalLookupTable`")) + + // A strategy hint on a different join of the same query does not conflict. + val other = Seq((1, "a")).toDF("id", "tag") + val nested = left.join(other.hint("merge"), "id") + val joined = nested.join(lookup, nested("id") === lookup("id")) + assert(runStream(joined).plan.exists(_.isInstanceOf[PreparedExternalLookupJoin])) + checkStream(joined, Seq(Row(1, "a", 1, "engineering"), Row(1, "a", 1, "research"))) + } + + test("a lookup join whose condition folds to false does not call the connector") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1, 3).toDF("id")) + Seq( + "inner" -> Seq.empty[Row], + "left_outer" -> Seq(Row(1, null, null), Row(3, null, null))).foreach { + case (joinType, expected) => + SimpleExternalLookupDataSource.resetCounters() + val joined = left.join(lookup, left("id") === lookup("id") && lit(false), joinType) + assert(!runStream(joined).plan.exists(_.isInstanceOf[PreparedExternalLookupJoin])) + checkStream(joined, expected) + assert(SimpleExternalLookupDataSource.lookupCandidates.get() === Seq.empty) + assert(SimpleExternalLookupDataSource.readerCount.get() === 0) + assert(SimpleExternalLookupDataSource.requestCount.get() === 0) + } + } + + test("external lookup reads VARCHAR columns and nested fields whose nullability differs") { + def lookup(kind: String): DataFrame = spark.read + .format(classOf[ReadSchemaExternalLookupDataSource].getName) + .option("kind", kind) + .load() + val left = streamOf(Seq(1, 2).toDF("id")) + + val varchar = lookup("varchar") + checkStream( + left.join(varchar, left("id") === varchar("id")), + Seq(Row(1, 1, "ab"), Row(2, 2, "ab"))) + val nested = lookup("nested") + checkStream( + left.join(nested, left("id") === nested("id")), + Seq(Row(1, 1, Row(10)), Row(2, 2, Row(20)))) + // A field whose type really differs is still rejected. + val wrongType = lookup("wrongType") + checkError( + exception = streamFailure[AnalysisException]( + left.join(wrongType, left("id") === wrongType("id"))), + condition = "DATA_SOURCE_INVALID_EXTERNAL_LOOKUP.UNKNOWN_SCAN_FIELD", + parameters = Map("table" -> "`ReadSchemaExternalLookupTable`", "field" -> "`name`")) + } + + test("external lookup merges the connector metrics of every reader in a task") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val lookup = spark.read.format(classOf[MetricsExternalLookupDataSource].getName).load() + // Two one-row input partitions, computed by one task under the coalesce, open two readers + // that send one request each. + val left = streamOf(spark.range(0, 2).selectExpr("CAST(id + 1 AS INT) AS id"), 2) + val joined = left.join(lookup, left("id") === lookup("id")).coalesce(1) + + val run = runStream(joined) + assert(run.rows.length === 2) + val lookupJoin = run.execution.executedPlan.collectFirst { + case exec: ExternalLookupJoinExec => exec + }.get + assert(lookupJoin.metrics("source.connectorRequests").value === 2) + // A connector metric named like an operator metric is kept apart from it. + assert(lookupJoin.metrics("source.numOutputRows").value === 7) + assert(lookupJoin.metrics("numOutputRows").value === 2) + } + } + + test("external lookup closes task-local readers after success and failure") { + def lookup(options: (String, String)*): DataFrame = { + var reader = spark.read.format(classOf[SimpleExternalLookupDataSource].getName) + options.foreach { case (key, value) => reader = reader.option(key, value) } + reader.load() + } + val left = streamOf(Seq(1).toDF("id")) + + SimpleExternalLookupDataSource.resetCounters() + val successfulLookup = lookup() + runStream(left.join(successfulLookup, left("id") === successfulLookup("id"))) + assert(SimpleExternalLookupDataSource.readerCount.get() === 1) + assert(SimpleExternalLookupDataSource.closeCount.get() === 1) + + SimpleExternalLookupDataSource.resetCounters() + val failingLookup = lookup("failRequests" -> "true") + val failure = streamFailure[SparkException]( + left.join(failingLookup, left("id") === failingLookup("id"))) + assert(failure.getMessage.contains("injected external lookup failure")) + assert(SimpleExternalLookupDataSource.readerCount.get() === 1) + assert(SimpleExternalLookupDataSource.closeCount.get() === 1) + } + + test("external lookup caps each inner-join request at the engine batch size") { + SimpleExternalLookupDataSource.resetCounters() + val lookup = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .option("maxBatchSize", "5000") + .load() + val left = streamOf((1 to 1500).toDF("id")) + + assert(runStream(left.join(lookup, left("id") === lookup("id"))).rows.size === 3) + assert(SimpleExternalLookupDataSource.requestSizes.asScala.toSeq === Seq(1000, 500)) + } + + test("external lookup does not create readers without requestable keys") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val empty = streamOf(Seq.empty[java.lang.Integer].toDF("id")) + val nullOnly = streamOf(Seq[java.lang.Integer](null).toDF("id")) + + Seq(empty, nullOnly).foreach { left => + SimpleExternalLookupDataSource.resetCounters() + runStream(left.join( + lookup, left("id") === lookup("id"), "left_outer")) + assert(SimpleExternalLookupDataSource.readerCount.get() === 0) + assert(SimpleExternalLookupDataSource.closeCount.get() === 0) + } + } + + test("external lookup cleanup attempts every resource and preserves the primary failure") { + val attempted = mutable.ArrayBuffer.empty[String] + val cursorFailure = new IllegalStateException("cursor close") + val metricsFailure = new IllegalStateException("metrics") + val readerFailure = new IllegalStateException("reader close") + + val failure = intercept[IllegalStateException] { + ExternalLookupJoinExec.closeResources( + { attempted += "cursor"; throw cursorFailure }, + { attempted += "metrics"; throw metricsFailure }, + { attempted += "reader"; throw readerFailure }) + } + assert(failure eq cursorFailure) + assert(attempted === Seq("cursor", "metrics", "reader")) + assert(failure.getSuppressed.toSeq === Seq(metricsFailure)) + assert(metricsFailure.getSuppressed.toSeq === Seq(readerFailure)) + + val lookupFailure = new IllegalStateException("lookup") + val cleanupFailure = new IllegalStateException("cleanup") + val preserved = intercept[IllegalStateException] { + org.apache.spark.util.Utils.tryWithSafeFinally(throw lookupFailure) { + ExternalLookupJoinExec.closeResources(throw cleanupFailure, (), ()) + } + } + assert(preserved eq lookupFailure) + assert(preserved.getSuppressed.toSeq === Seq(cleanupFailure)) + } + + test("external lookup responds to task cancellation and closes the reader") { + SimpleExternalLookupDataSource.resetCounters() + val lookup = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .option("responseDelayMillis", "60000") + .load() + val left = streamOf(Seq(1).toDF("id")) + // Stopping the query cancels the task that waits for the response. + val query = left.join(lookup, left("id") === lookup("id")) + .writeStream.format("memory").queryName("external_lookup_cancellation").start() + + assert(SimpleExternalLookupDataSource.requestStarted.get().await(10, TimeUnit.SECONDS)) + query.stop() + spark.catalog.dropTempView("external_lookup_cancellation") + assert(SimpleExternalLookupDataSource.readerCount.get() === 1) + eventually(timeout(10.seconds)) { + assert(SimpleExternalLookupDataSource.closeCount.get() === 1) + } + } + + test("external lookup join runs unchanged in micro-batch streaming") { + val input = MemoryStream[Int] + val left = input.toDF().toDF("id") + val lookup = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .load() + val joined = left.join(lookup, left("id") === lookup("id"), "left_outer") + val queryName = "external_lookup_mbm" + val query = joined.writeStream.format("memory").queryName(queryName).start() + try { + input.addData(1, 2, 3) + query.processAllAvailable() + checkAnswer( + spark.table(queryName), + Seq(Row(1, 1, "engineering"), Row(1, 1, "research"), + Row(2, 2, "sales"), Row(3, null, null))) + } finally { + query.stop() + spark.catalog.dropTempView(queryName) + } + } +} + +/** A catalog serving `ns.lookup` (lookup-only) and `ns.dual` (dual-capability). */ +class ExternalLookupTestCatalog extends TableCatalog { + private var catalogName: String = _ + + override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = + catalogName = name + override def name(): String = catalogName + override def listTables(namespace: Array[String]): Array[Identifier] = + Array("lookup", "dual").map(Identifier.of(Array("ns"), _)) + override def loadTable(ident: Identifier): Table = (ident.namespace().toSeq, ident.name()) match { + case (Seq("ns"), "lookup") => SimpleExternalLookupTable + case (Seq("ns"), "dual") => DualCapabilityExternalLookupTable + case _ => throw new NoSuchTableException(ident.namespace().mkString("."), ident.name()) + } + override def alterTable(ident: Identifier, changes: TableChange*): Table = + throw new UnsupportedOperationException() + override def dropTable(ident: Identifier): Boolean = throw new UnsupportedOperationException() + override def renameTable(oldIdent: Identifier, newIdent: Identifier): Unit = + throw new UnsupportedOperationException() +} + +class ExternalLookupJoinRealTimeModeSuite extends StreamRealTimeModeE2ESuiteBase { + import testImplicits._ + + test("real-time mode rejects external lookup join") { + SimpleExternalLookupDataSource.resetCounters() + val input = LowLatencyMemoryStream[Int](2) + val left = input.toDF().toDF("id") + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val joined = left.join(lookup, left("id") === lookup("id")) + .select(concat(col("department"), lit("-"), left("id").cast("string")).as("output")) + + val query = runStreamingQuery("external_lookup_rtm", joined) + try { + eventually(timeout(60.seconds)) { + checkError( + exception = query.exception.get.getCause.asInstanceOf[SparkIllegalArgumentException], + condition = "STREAMING_REAL_TIME_MODE.OPERATOR_OR_SINK_NOT_IN_ALLOWLIST", + parameters = Map( + "errorType" -> "operator", + "message" -> "org.apache.spark.sql.execution.joins.ExternalLookupJoinExec is")) + } + } finally { + query.stop() + } + assert(SimpleExternalLookupDataSource.readerCount.get() === 0) + } +} + +class SimpleExternalLookupDataSource extends TableProvider { + override def inferSchema(options: CaseInsensitiveStringMap): StructType = + SimpleExternalLookupDataSource.outputSchema + + override def getTable( + schema: StructType, + partitioning: Array[Transform], + properties: java.util.Map[String, String]): Table = SimpleExternalLookupTable +} + +object DualCapabilityExternalLookupTable extends Table with SupportsRead { + override def name(): String = "DualCapabilityExternalLookupTable" + override def schema(): StructType = SimpleExternalLookupDataSource.outputSchema + override def capabilities(): java.util.Set[TableCapability] = + java.util.EnumSet.of(BATCH_READ, EXTERNAL_LOOKUP) + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = + SimpleExternalLookupTable.newScanBuilder(options) +} + +object SimpleExternalLookupDataSource { + val requestCount = new AtomicInteger() + val readerCount = new AtomicInteger() + val closeCount = new AtomicInteger() + val requestStarted = new AtomicReference(new CountDownLatch(1)) + val requiredColumns = new AtomicReference(Seq.empty[String]) + val pushedPredicateCount = new AtomicInteger() + val pushedPredicateFields = new AtomicReference(Seq.empty[String]) + val requestSizes = new ConcurrentLinkedQueue[Int]() + val requestedKeys = new ConcurrentLinkedQueue[Int]() + val nullKeyRequests = new AtomicInteger(0) + val resultRowsRead = new AtomicInteger(0) + val metricSamples = new AtomicInteger(0) + val lookupCandidates = + new AtomicReference(Seq.empty[(String, String, DataType)]) + val lookupCandidateDescriptions = new AtomicReference(Seq.empty[String]) + val outputSchema: StructType = new StructType() + .add("id", "int", nullable = true) + .add("department", "string", nullable = false) + private val InitialValues: Map[Int, Seq[String]] = Map( + 1 -> Seq("engineering", "research"), + 2 -> Seq("sales")) + val values = new AtomicReference(InitialValues) + + def resetCounters(): Unit = { + requestCount.set(0) + readerCount.set(0) + closeCount.set(0) + requestStarted.set(new CountDownLatch(1)) + requiredColumns.set(Seq.empty) + pushedPredicateCount.set(0) + pushedPredicateFields.set(Seq.empty) + requestSizes.clear() + requestedKeys.clear() + nullKeyRequests.set(0) + resultRowsRead.set(0) + metricSamples.set(0) + lookupCandidates.set(Seq.empty) + lookupCandidateDescriptions.set(Seq.empty) + values.set(InitialValues) + } +} + +object SimpleExternalLookupTable extends Table with SupportsRead { + override def name(): String = "SimpleExternalLookupTable" + override def schema(): StructType = SimpleExternalLookupDataSource.outputSchema + override def capabilities(): java.util.Set[TableCapability] = + java.util.EnumSet.of(EXTERNAL_LOOKUP) + + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + val configuredMaxBatchSize = options.getInt("maxBatchSize", 100) + val responseDelayMillis = options.getLong("responseDelayMillis", 0) + val failRequests = options.getBoolean("failRequests", false) + val mutateAfterRequest = options.getBoolean("mutateAfterRequest", false) + val nullKeySelection = options.getBoolean("nullKeySelection", false) + val selectionMode = options.getOrDefault("selectionMode", "all") + val nullCursor = options.getBoolean("nullCursor", false) + new SupportsExternalLookup + with SupportsPushDownRequiredColumns + with SupportsPushDownV2Filters { + private var requiredSchema = SimpleExternalLookupDataSource.outputSchema + private var excludeResearch = false + private var acceptedPredicates = Array.empty[Predicate] + private var selectedLookupPredicates = Seq.empty[String] + + override def selectLookupPredicates(candidates: Array[Predicate]): Array[Predicate] = { + if (nullKeySelection) return null + SimpleExternalLookupDataSource.lookupCandidates.set(candidates.map { candidate => + val field = candidate.children()(0).asInstanceOf[NamedReference] + val parameter = candidate.children()(1).asInstanceOf[LookupParameter] + (candidate.name(), field.fieldNames().mkString("."), parameter.dataType()) + }.toSeq) + SimpleExternalLookupDataSource.lookupCandidateDescriptions.set( + candidates.map(_.toString).toSeq) + val supportedPredicates = Set("=", "<=>", "<", "<=", ">", ">=") + val supported = candidates.filter { candidate => + supportedPredicates.contains(candidate.name()) && + candidate.children()(0).asInstanceOf[NamedReference] + .fieldNames().sameElements(Array("id")) + } + val selected = selectionMode match { + case "all" => supported + case "last" => supported.takeRight(1) + case "reversed" => supported.reverse + case "duplicate" => supported.take(1) ++ supported.take(1) + // An equal-shaped predicate that is not one of the offered instances. + case "copy" => supported.take(1).map(p => new Predicate(p.name(), p.children())) + case "nullElement" => Array[Predicate](null) + } + selectedLookupPredicates = selected.filter(_ != null).map(_.name()).toSeq + selected + } + + override def pruneColumns(schema: StructType): Unit = { + requiredSchema = schema + SimpleExternalLookupDataSource.requiredColumns.set(schema.fieldNames.toSeq) + } + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + def fields(predicate: Predicate): Seq[NamedReference] = + predicate.children().toSeq.flatMap { + case reference: NamedReference => Seq(reference) + case child: Predicate => fields(child) + case _ => Seq.empty + } + val (accepted, rejected) = predicates.partition { predicate => + fields(predicate).exists(_.fieldNames().sameElements(Array("department"))) + } + if (accepted.nonEmpty) { + acceptedPredicates = accepted + excludeResearch = true + SimpleExternalLookupDataSource.pushedPredicateCount.addAndGet(accepted.length) + SimpleExternalLookupDataSource.pushedPredicateFields.set( + accepted.flatMap(fields).map(_.fieldNames().mkString(".")).toSeq) + } + rejected + } + override def pushedPredicates(): Array[Predicate] = acceptedPredicates + override def build(): Scan = SimpleExternalLookupScan( + configuredMaxBatchSize, responseDelayMillis, failRequests, + mutateAfterRequest, nullCursor, requiredSchema, excludeResearch, + selectedLookupPredicates) + } + } +} + +case class SimpleExternalLookupScan( + configuredMaxBatchSize: Int, + responseDelayMillis: Long, + failRequests: Boolean, + mutateAfterRequest: Boolean, + nullCursor: Boolean, + requiredSchema: StructType, + excludeResearch: Boolean, + selectedLookupPredicates: Seq[String]) extends Scan { + + override def readSchema(): StructType = requiredSchema + override def toBatch(): Batch = + SimpleExternalLookupBatch(requiredSchema.fieldNames.toSeq, excludeResearch) + override def toExternalLookup(): ExternalLookup = + TestExternalLookup(configuredMaxBatchSize, SimpleExternalLookupReaderFactory( + responseDelayMillis, + failRequests, + mutateAfterRequest, + nullCursor, + requiredSchema.fieldNames.toSeq, + excludeResearch, + selectedLookupPredicates)) +} + +case class SimpleExternalLookupBatch( + requiredColumns: Seq[String], + excludeResearch: Boolean) extends Batch { + override def planInputPartitions(): Array[InputPartition] = + Array(SimpleExternalLookupInputPartition) + + override def createReaderFactory(): PartitionReaderFactory = + SimpleExternalLookupBatchReaderFactory(requiredColumns, excludeResearch) +} + +case object SimpleExternalLookupInputPartition extends InputPartition + +case class SimpleExternalLookupBatchReaderFactory( + requiredColumns: Seq[String], + excludeResearch: Boolean) extends PartitionReaderFactory { + override def createReader(partition: InputPartition): PartitionReader[InternalRow] = + new PartitionReader[InternalRow] { + private val rows = SimpleExternalLookupDataSource.values.get().toSeq + .sortBy(_._1) + .iterator + .flatMap { case (id, departments) => + departments.iterator + .filterNot(department => excludeResearch && department == "research") + .map { department => + InternalRow.fromSeq(requiredColumns.map { + case "id" => Int.box(id) + case "department" => UTF8String.fromString(department) + }) + } + } + private var current: InternalRow = _ + + override def next(): Boolean = { + if (rows.hasNext) { + current = rows.next() + true + } else { + false + } + } + + override def get(): InternalRow = current + + override def close(): Unit = {} + } +} + +case class SimpleExternalLookupReaderFactory( + responseDelayMillis: Long, + failRequests: Boolean, + mutateAfterRequest: Boolean, + nullCursor: Boolean, + requiredColumns: Seq[String], + excludeResearch: Boolean, + selectedLookupPredicates: Seq[String]) extends ExternalLookupReaderFactory { + override def createReader(): ExternalLookupReader = new ExternalLookupReader { + SimpleExternalLookupDataSource.readerCount.incrementAndGet() + + override def lookupBatch(keys: Array[InternalRow]): ExternalLookupResultCursor = { + SimpleExternalLookupDataSource.requestStarted.get().countDown() + if (responseDelayMillis > 0) Thread.sleep(responseDelayMillis) + if (failRequests) throw new IllegalStateException("injected external lookup failure") + if (nullCursor) return null + SimpleExternalLookupDataSource.requestCount.incrementAndGet() + SimpleExternalLookupDataSource.requestSizes.add(keys.length) + val results = keys.iterator.zipWithIndex.flatMap { case (key, ordinal) => + if (key.isNullAt(0)) { + SimpleExternalLookupDataSource.nullKeyRequests.incrementAndGet() + } else { + val value = key.getInt(0) + SimpleExternalLookupDataSource.requestedKeys.add(value) + } + val values = SimpleExternalLookupDataSource.values.get().toSeq.map { case (id, values) => + Some(id) -> values + } :+ (None -> Seq("unknown")) + values.iterator.filter { case (id, _) => + selectedLookupPredicates.zipWithIndex.forall { case (predicate, index) => + val parameter = if (key.isNullAt(index)) None else Some(key.getInt(index)) + predicate match { + case "=" => id.nonEmpty && parameter.nonEmpty && id == parameter + case "<=>" => id == parameter + case "<" => id.nonEmpty && parameter.nonEmpty && id.get < parameter.get + case "<=" => id.nonEmpty && parameter.nonEmpty && id.get <= parameter.get + case ">" => id.nonEmpty && parameter.nonEmpty && id.get > parameter.get + case ">=" => id.nonEmpty && parameter.nonEmpty && id.get >= parameter.get + } + } + }.flatMap { case (id, departments) => departments.iterator.map(id -> _) } + .filterNot { case (_, department) => excludeResearch && department == "research" } + .map { case (id, department) => + ordinal -> InternalRow.fromSeq(requiredColumns.map { + case "id" => id.map(Int.box).orNull + case "department" => UTF8String.fromString(department) + }) + } + }.toArray.iterator + if (mutateAfterRequest) { + SimpleExternalLookupDataSource.values.updateAndGet { current => + current.updated(2, Seq("sales-v2")) + } + } + new TestExternalLookupResultCursor( + results, () => SimpleExternalLookupDataSource.resultRowsRead.incrementAndGet()) + } + + override def currentMetricsValues(): + Array[org.apache.spark.sql.connector.metric.CustomTaskMetric] = { + SimpleExternalLookupDataSource.metricSamples.incrementAndGet() + Array.empty + } + + override def close(): Unit = SimpleExternalLookupDataSource.closeCount.incrementAndGet() + } +} + +class CompositeExternalLookupDataSource extends TableProvider { + override def inferSchema(options: CaseInsensitiveStringMap): StructType = + CompositeExternalLookupDataSource.outputSchema + + override def getTable( + schema: StructType, + partitioning: Array[Transform], + properties: java.util.Map[String, String]): Table = CompositeExternalLookupTable +} + +object CompositeExternalLookupDataSource { + val outputSchema: StructType = new StructType() + .add("tenant", "int", nullable = false) + .add("id", "int", nullable = false) + .add("value", "string", nullable = false) + val values: Map[(Int, Int), String] = Map( + (1, 1) -> "one-one", + (1, 3) -> "one-three", + (2, 1) -> "two-one") + val selectedLookupPredicates = + new AtomicReference(Seq.empty[(String, String)]) +} + +object CompositeExternalLookupTable extends Table with SupportsRead { + override def name(): String = "CompositeExternalLookupTable" + override def schema(): StructType = CompositeExternalLookupDataSource.outputSchema + override def capabilities(): java.util.Set[TableCapability] = + java.util.EnumSet.of(EXTERNAL_LOOKUP) + + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = + new SupportsExternalLookup with SupportsPushDownRequiredColumns { + private var requiredSchema = CompositeExternalLookupDataSource.outputSchema + private var selectedLookupPredicates = Seq.empty[(String, String)] + + override def selectLookupPredicates(candidates: Array[Predicate]): Array[Predicate] = { + val supported = candidates.filter { predicate => + val field = predicate.children()(0).asInstanceOf[NamedReference] + .fieldNames().mkString(".") + (field == "tenant" && predicate.name() == "=") || + (field == "id" && Set("=", "<", "<=", ">", ">=").contains(predicate.name())) + } + val selected = supported.map { predicate => + predicate.name() -> predicate.children()(0).asInstanceOf[NamedReference] + .fieldNames().mkString(".") + }.toSeq + if (selected.contains("=" -> "tenant") && selected.exists(_._2 == "id")) { + selectedLookupPredicates = selected + CompositeExternalLookupDataSource.selectedLookupPredicates.set(selected) + supported + } else { + Array.empty + } + } + + override def pruneColumns(schema: StructType): Unit = requiredSchema = schema + override def build(): Scan = + CompositeExternalLookupScan(requiredSchema, selectedLookupPredicates) + } +} + +case class CompositeExternalLookupScan( + requiredSchema: StructType, + selectedLookupPredicates: Seq[(String, String)]) extends Scan { + override def readSchema(): StructType = requiredSchema + override def toExternalLookup(): ExternalLookup = + TestExternalLookup( + 100, + CompositeExternalLookupReaderFactory(requiredSchema, selectedLookupPredicates)) +} + +case class CompositeExternalLookupReaderFactory( + requiredSchema: StructType, + selectedLookupPredicates: Seq[(String, String)]) extends ExternalLookupReaderFactory { + override def createReader(): ExternalLookupReader = new ExternalLookupReader { + override def lookupBatch(keys: Array[InternalRow]): ExternalLookupResultCursor = { + val results = keys.iterator.zipWithIndex.flatMap { case (key, ordinal) => + CompositeExternalLookupDataSource.values.iterator.filter { case ((tenant, id), _) => + selectedLookupPredicates.zipWithIndex.forall { + case ((predicate, field), index) => + val lookupValue = if (field == "tenant") tenant else id + val parameter = key.getInt(index) + predicate match { + case "=" => lookupValue == parameter + case "<" => lookupValue < parameter + case "<=" => lookupValue <= parameter + case ">" => lookupValue > parameter + case ">=" => lookupValue >= parameter + } + } + }.map { case ((tenant, id), lookupValue) => + ordinal -> InternalRow.fromSeq(requiredSchema.fieldNames.toSeq.map { + case "tenant" => Int.box(tenant) + case "id" => Int.box(id) + case "value" => UTF8String.fromString(lookupValue) + }) + } + } + new TestExternalLookupResultCursor(results) + } + + override def close(): Unit = {} + } +} + +class TestExternalLookupResultCursor( + results: Iterator[(Int, InternalRow)], + onNext: () => Unit = () => ()) + extends ExternalLookupResultCursor { + private var current: (Int, InternalRow) = _ + + override def next(): Boolean = if (results.hasNext) { + current = results.next() + onNext() + true + } else { + current = null + false + } + + override def get(): InternalRow = current._2 + override def inputOrdinal(): Int = current._1 + override def close(): Unit = {} 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/commit/01060c7f72e4e7256c177939e0e2306e463dfae0), by extending the existing tests rather than adding new fixtures. - Cursor lifecycle: the test cursor records its close, and the reader records its own. After success, after a failing cursor, and after an upstream `LIMIT` stops reading, the cursor is closed before its reader. A failing request has no cursor, and its reader is still closed. - Metrics: the partial-batch test asserts `numInputRows`, `numInputBatches`, `numLookupRequests`, `numLookupParameterRows` (counted after deduplication) and `numOutputRows`, and that `lookupTime` is positive. - Redaction: the metrics test's scan description contains a secret, which `ExternalLookupJoinExec.simpleString` and the executed plan redact. - Micro-batches: the micro-batch test runs two micro-batches, and the connector changes a lookup value after the first one. The second micro-batch returns the new value. ########## 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 = { + metrics.requests += 1 + metrics.parameterRows += requestRows.length + if (reader == null) { + reader = readerFactory.createReader() + taskReaders.opened(reader) + } + activeCursor = Option(reader.lookupBatch(requestRows.asInstanceOf[Array[InternalRow]])) + .getOrElse { + throw new SparkIllegalStateException( + errorClass = "INTERNAL_ERROR", + messageParameters = Map("message" -> "External lookup returned a null result cursor")) 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/commit/01060c7f72e4e7256c177939e0e2306e463dfae0). The table name is passed to the evaluator factory, through `PreparedExternalLookupJoin` and `ExternalLookupJoinExec`. Three new subclasses of `DATA_SOURCE_INVALID_EXTERNAL_LOOKUP` name the table: `NULL_RESULT_CURSOR` (`lookupBatch()` returned null), `INVALID_INPUT_ORDINAL` (with the ordinal and the request's parameter row count) and `NULL_LOOKUP_READER` (`createReader()` returned null, which is now checked before the reader is registered). The contract test checks all three. ########## 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 => + resolver(attr.name, field.name) && + DataType.equalsIgnoreNullability(attr.dataType, field.dataType) + }.getOrElse { + throw QueryCompilationErrors.invalidDataSourceExternalLookupError( + tableName, "UNKNOWN_SCAN_FIELD", Map("field" -> toSQLId(Seq(field.name)))) + } + }.toSeq + lookupInput.groupBy(_.name).collectFirst { + case (name, fields) if fields.size > 1 => name + }.foreach { name => + throw QueryCompilationErrors.invalidDataSourceExternalLookupError( + tableName, "DUPLICATE_SCAN_FIELD", Map("field" -> toSQLId(Seq(name)))) + } + val externalLookup = scan.toExternalLookup() + val maxParameterRowsPerRequest = externalLookup.maxParameterRowsPerRequest() + if (maxParameterRowsPerRequest <= 0) { + throw QueryCompilationErrors.invalidDataSourceExternalLookupError( + tableName, + "NON_POSITIVE_MAX_PARAMETER_ROWS", + Map("value" -> maxParameterRowsPerRequest.toString)) + } + val readerFactory = externalLookup.createReaderFactory() + // TODO: prune lookup output to the columns required above the join. + val lookupOutput = lookupInput + + PreparedExternalLookupJoin( + join.left, + scan, + readerFactory, + maxParameterRowsPerRequest, + lookupInput, + lookupOutput, + selectedPredicates.map(_.parameter).toSeq, + selectedPredicates.map(_.nullSafe).toSeq, + join.joinType, + residual.reduceOption(And)) + } + + /** + * Maps the connector's selection back to the offered candidates. Returns the selected candidates + * in the connector's order, which defines the fields of each request parameter row, and the + * candidates it did not select. Selections are matched by object identity rather than by + * equality, so two candidates of the same shape, for example the same lookup field compared with + * two different join-side values, remain distinguishable. + */ + private def selectedCandidates( + tableName: String, + offered: Array[V2Predicate], + candidates: Iterable[OfferedLookupPredicate], + selected: Array[V2Predicate]): (Seq[OfferedLookupPredicate], Seq[OfferedLookupPredicate]) = { + if (selected == null || selected.contains(null)) { + throw QueryCompilationErrors.invalidDataSourceExternalLookupError( + tableName, "NULL_PREDICATE_SELECTION") + } + if (selected.isEmpty) { + throw QueryCompilationErrors.unsupportedExternalLookupJoinError( + "NO_LOOKUP_PREDICATE", Map("table" -> toSQLId(tableName))) 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/commit/73b62b7c559b90caf5ac9088417679591118c5d5)). When Spark can offer no candidate, it now fails before calling the connector with `NO_LOOKUP_CANDIDATE`, which names the accepted shape and suggests casting the stream side. `NO_LOOKUP_PREDICATE` remains for a connector that rejects a non-empty offer. In [01060c7f72e](https://github.com/apache/spark/commit/01060c7f72e4e7256c177939e0e2306e463dfae0), this PR's contract test checks that a BIGINT stream key joined to an INT lookup key fails with `NO_LOOKUP_CANDIDATE`. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/PrepareExternalLookupJoinSuite.scala: ########## @@ -0,0 +1,120 @@ +/* + * 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.connector + +import org.apache.spark.SparkException +import org.apache.spark.sql.AnalysisException +import org.apache.spark.sql.QueryTest +import org.apache.spark.sql.catalyst.analysis.ResolveExternalLookupJoin +import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, EqualTo, Literal} +import org.apache.spark.sql.catalyst.plans.{FullOuter, Inner, LeftOuter} +import org.apache.spark.sql.catalyst.plans.logical.{Join, JoinHint, LocalRelation, + PreparedExternalLookupJoin, Project} +import org.apache.spark.sql.connector.catalog.{SupportsRead, Table, TableCapability} +import org.apache.spark.sql.connector.catalog.TableCapability.EXTERNAL_LOOKUP +import org.apache.spark.sql.connector.expressions.FieldReference +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{ExternalLookup, ExternalLookupReaderFactory, Scan, + ScanBuilder, SupportsExternalLookup} +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, + PrepareExternalLookupJoin} +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{IntegerType, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +class PrepareExternalLookupJoinSuite extends QueryTest with SharedSparkSession { 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/commit/73b62b7c559b90caf5ac9088417679591118c5d5)), which owns that suite. It now extends `PlanTest`. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/ExternalLookupJoinSuite.scala: ########## @@ -0,0 +1,1358 @@ +/* + * 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.connector + +import java.util.concurrent.{ConcurrentLinkedQueue, CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} + +import scala.collection.mutable +import scala.concurrent.duration._ +import scala.jdk.CollectionConverters._ +import scala.reflect.{classTag, ClassTag} + +import org.apache.spark.{SparkException, SparkIllegalArgumentException, + SparkUnsupportedOperationException} +import org.apache.spark.sql.{AnalysisException, DataFrame, Encoders, Row} +import org.apache.spark.sql.QueryTest +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.analysis.NoSuchTableException +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, PreparedExternalLookupJoin} +import org.apache.spark.sql.classic.Dataset +import org.apache.spark.sql.connector.catalog.{Identifier, SupportsRead, Table, TableCapability, + TableCatalog, TableChange, TableProvider} +import org.apache.spark.sql.connector.catalog.TableCapability.{BATCH_READ, EXTERNAL_LOOKUP} +import org.apache.spark.sql.connector.expressions.{NamedReference, Transform} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.metric.{CustomMetric, CustomSumMetric, CustomTaskMetric} +import org.apache.spark.sql.connector.read.{Batch, ExternalLookup, ExternalLookupReader, + ExternalLookupReaderFactory, ExternalLookupResultCursor, InputPartition, LookupParameter, + PartitionReader, PartitionReaderFactory, Scan, ScanBuilder, SupportsExternalLookup, + SupportsPushDownRequiredColumns, SupportsPushDownV2Filters} +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, ExternalLookupRelation} +import org.apache.spark.sql.execution.joins.ExternalLookupJoinExec +import org.apache.spark.sql.execution.streaming.runtime.{IncrementalExecution, MemoryStream, + StreamingQueryWrapper} +import org.apache.spark.sql.execution.streaming.sources.LowLatencyMemoryStream +import org.apache.spark.sql.functions.{coalesce, col, concat, length, lit} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.streaming.{StreamingQueryException, StreamRealTimeModeE2ESuiteBase} +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{DataType, IntegerType, StringType, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.unsafe.types.UTF8String + +class ExternalLookupJoinSuite extends QueryTest with SharedSparkSession { + + import testImplicits._ + + private val ExternalLookupCatalogName = "external_lookup_cat" + + private val streamRuns = new AtomicInteger() + + /** The rows of `df` as a streaming DataFrame: one micro-batch, in `numPartitions` partitions. */ + private def streamOf(df: DataFrame, numPartitions: Int = 1): DataFrame = { + val input = MemoryStream[Row](numPartitions)(Encoders.row(df.schema), spark.sqlContext) + input.addData(df.collect().toSeq) + input.toDF() + } + + private case class StreamRun(rows: Seq[Row], execution: IncrementalExecution) { + def plan: LogicalPlan = execution.optimizedPlan + } + + /** Runs the streaming `df` on a memory sink until it has processed all of its input. */ + private def runStream(df: DataFrame): StreamRun = { + val queryName = s"external_lookup_stream_${streamRuns.incrementAndGet()}" + val query = df.writeStream.format("memory").queryName(queryName).start() + try { + query.processAllAvailable() + StreamRun( + spark.table(queryName).collect().toSeq, + query.asInstanceOf[StreamingQueryWrapper].streamingQuery.lastExecution) + } finally { + query.stop() + spark.catalog.dropTempView(queryName) + } + } + + private def checkStream(df: DataFrame, expected: Seq[Row]): Unit = + QueryTest.sameRows(expected, runStream(df).rows).foreach(fail(_)) + + /** The error that the streaming `df` fails with. */ + private def streamFailure[T <: Throwable : ClassTag](df: DataFrame): T = { + val cause = intercept[StreamingQueryException](runStream(df)).getCause + assert(classTag[T].runtimeClass.isInstance(cause), cause) + cause.asInstanceOf[T] + } + + test("external lookup capability produces a unary external lookup join") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1).toDF("id")) + val joined = left.join(lookup, left("id") === lookup("id"), "inner") + + val plan = runStream(joined).plan + assert(plan.exists(_.isInstanceOf[PreparedExternalLookupJoin])) + assert(!plan.toString.contains("SimpleExternalLookupScan")) + } + + test("an external lookup join is a stream-static join with the lookup-only table on the right") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val batch = Seq(1).toDF("id") + // The left side must be a streaming query. + checkError( + exception = intercept[AnalysisException] { + batch.join(lookup, batch("id") === lookup("id")) + }, + condition = "EXTERNAL_LOOKUP_JOIN_UNSUPPORTED.STATIC_LEFT_SIDE", + parameters = Map("relation" -> "`SimpleExternalLookupTable`")) + // The sides are never swapped. + val stream = streamOf(batch) + checkError( + exception = intercept[AnalysisException] { + lookup.join(stream, lookup("id") === stream("id")) + }, + condition = "EXTERNAL_LOOKUP_JOIN_UNSUPPORTED.LOOKUP_TABLE_ON_LEFT", + parameters = Map("relation" -> "`SimpleExternalLookupTable`")) + } + + test("external lookup validates connector planning and request contracts") { + val left = streamOf(Seq(1).toDF("id")) + Seq( + ("nullKeySelection", "true", "NULL_PREDICATE_SELECTION", Map.empty[String, String]), + ("selectionMode", "nullElement", "NULL_PREDICATE_SELECTION", Map.empty[String, String]), + ("selectionMode", "duplicate", "DUPLICATE_SELECTED_PREDICATE", Map("predicate" -> "id = ?")), + ("selectionMode", "copy", "UNOFFERED_SELECTED_PREDICATE", Map("predicate" -> "id = ?")), + ("maxBatchSize", "0", "NON_POSITIVE_MAX_PARAMETER_ROWS", Map("value" -> "0"))).foreach { + case (option, value, subClass, parameters) => + val lookup = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .option(option, value) + .load() + checkError( + exception = streamFailure[AnalysisException]( + left.join(lookup, left("id") === lookup("id"))), + condition = s"DATA_SOURCE_INVALID_EXTERNAL_LOOKUP.$subClass", + parameters = Map("table" -> "`SimpleExternalLookupTable`") ++ parameters) + } + + val lookup = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .option("nullCursor", "true") + .load() + val error = streamFailure[SparkException](left.join(lookup, left("id") === lookup("id"))) + assert(error.getMessage.contains("External lookup returned a null result cursor")) + + def misconfigured(kind: String): DataFrame = spark.read + .format(classOf[MisconfiguredExternalLookupDataSource].getName) + .option("kind", kind) + .load() + val table = Map("table" -> "`MisconfiguredExternalLookupTable`") + Seq( + ("unknownField", "UNKNOWN_SCAN_FIELD", Map("field" -> "`unknown`")), + ("duplicateField", "DUPLICATE_SCAN_FIELD", Map("field" -> "`id`"))).foreach { + case (kind, subClass, parameters) => + val lookup = misconfigured(kind) + checkError( + exception = streamFailure[AnalysisException]( + left.join(lookup, left("id") === lookup("id"))), + condition = s"DATA_SOURCE_INVALID_EXTERNAL_LOOKUP.$subClass", + parameters = table ++ parameters) + } + + val noExternalLookup = misconfigured("noExternalLookup") + checkError( + exception = streamFailure[SparkUnsupportedOperationException]( + left.join(noExternalLookup, left("id") === noExternalLookup("id"))), + condition = "DATA_SOURCE_EXTERNAL_LOOKUP_NOT_SUPPORTED", + parameters = Map("description" -> "NoExternalLookupScan")) + + val invalidOrdinal = misconfigured("invalidOrdinal") + val ordinalError = streamFailure[SparkException]( + left.join(invalidOrdinal, left("id") === invalidOrdinal("id"))) + assert(ordinalError.getMessage.contains( + "External lookup returned invalid input ordinal 5 for 1 requested parameter rows")) + + // Planning failures that need no hint. A catalog can serve a table without SupportsRead. + val noLookupBuilder = misconfigured("noLookupBuilder") + val nonReadable = Dataset.ofRows(spark, DataSourceV2Relation.create( + NonReadableExternalLookupTable, None, None, CaseInsensitiveStringMap.empty)) + val simple = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + Seq( + (noLookupBuilder, left("id") === noLookupBuilder("id"), "SCAN_BUILDER", + "MisconfiguredExternalLookupTable"), + (nonReadable, left("id") === nonReadable("id"), "TABLE", "NonReadableExternalLookupTable"), + (simple, simple("department") === "sales", "NO_LOOKUP_PREDICATE", + "SimpleExternalLookupTable")).foreach { case (lookup, condition, subClass, name) => + checkError( + exception = streamFailure[AnalysisException](left.join(lookup, condition)), + condition = s"EXTERNAL_LOOKUP_JOIN_UNSUPPORTED.$subClass", + parameters = Map("table" -> s"`$name`")) + } + } + + test("external lookup finalizes expressions on the stream side after analysis") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq[java.lang.Integer](null, 1).toDF("id")) + + checkStream( + left.join( + lookup, + coalesce(left("id"), lit(0)) === lookup("id"), + "inner"), + Seq(Row(1, 1, "engineering"), Row(1, 1, "research"))) + } + + test("external lookup maps a selected predicate to its own occurrence") { + SimpleExternalLookupDataSource.resetCounters() + val lookup = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .option("selectionMode", "last") + .load() + val left = streamOf(Seq((1, 2), (2, 2), (2, 1)).toDF("a", "b")) + + // Both candidates have the shape `id = ?`; only the second one, bound to `b`, is selected. + val joined = + left.join(lookup, lookup("id") === left("a") && lookup("id") === left("b"), "left_outer") + assert(runStream(joined).rows.toSet === + Set(Row(1, 2, null, null), Row(2, 2, 2, "sales"), Row(2, 1, null, null))) + assert(SimpleExternalLookupDataSource.lookupCandidates.get() === + Seq(("=", "id", IntegerType), ("=", "id", IntegerType))) + assert(SimpleExternalLookupDataSource.lookupCandidateDescriptions.get() === + Seq("id = ?", "id = ?")) + assert(SimpleExternalLookupDataSource.requestedKeys.asScala.toSeq === Seq(2, 1)) + + // Selected in reverse order, each parameter row is `(b, a)`. + SimpleExternalLookupDataSource.resetCounters() + val reversed = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .option("selectionMode", "reversed") + .load() + val condition = reversed("id") === left("a") && reversed("id") === left("b") + assert(runStream(left.join(reversed, condition, "left_outer")).rows.toSet === + Set(Row(1, 2, null, null), Row(2, 2, 2, "sales"), Row(2, 1, null, null))) + assert(SimpleExternalLookupDataSource.requestedKeys.asScala.toSeq === Seq(2, 2, 1)) + } + + test("a parent projection of some lookup columns still plans an external lookup join") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1, 3).toDF("id")) + // Generic column pruning must not project the lookup relation, which would hide it from + // lookup planning; the join's requiredLookupOutput keeps it from doing so. + val joined = left.join(lookup, left("id") === lookup("id"), "left_outer") + .select(left("id"), lookup("department")) + + assert(runStream(joined).plan.exists(_.isInstanceOf[PreparedExternalLookupJoin])) + checkStream(joined, Seq(Row(1, "engineering"), Row(1, "research"), Row(3, null))) + } + + test("external lookup join supports inner and left outer joins") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq((1, "one"), (2, "two"), (3, "three")).toDF("id", "label")) + + checkStream( + left.join(lookup, left("id") === lookup("id"), "inner"), + Seq(Row(1, "one", 1, "engineering"), Row(1, "one", 1, "research"), + Row(2, "two", 2, "sales"))) + checkStream( + left.join(lookup, left("id") === lookup("id"), "left_outer"), + Seq(Row(1, "one", 1, "engineering"), Row(1, "one", 1, "research"), + Row(2, "two", 2, "sales"), Row(3, "three", null, null))) + + // Only comparisons of a lookup field with a stream-side value are offered. The connector + // selects `id = ?`, and Spark evaluates the other conjuncts on the joined rows. + SimpleExternalLookupDataSource.resetCounters() + val flagged = streamOf(Seq((1, true), (2, true), (1, false)).toDF("id", "flag")) + checkStream( + flagged.join( + lookup, + flagged("id") === lookup("id") && flagged("flag") && + lookup("department") === "engineering" && length(lookup("department")) > lookup("id"), + "left_outer"), + Seq(Row(1, true, 1, "engineering"), Row(2, true, null, null), Row(1, false, null, null))) + assert(SimpleExternalLookupDataSource.lookupCandidates.get() === + Seq(("=", "id", IntegerType), ("=", "department", StringType))) + } + + test("external lookup-only source joins without an explicit hint") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1).toDF("id")) + + assert(lookup.queryExecution.analyzed.isInstanceOf[ExternalLookupRelation]) + val joined = left.join(lookup, "id") + assert(runStream(joined).plan.exists(_.isInstanceOf[PreparedExternalLookupJoin])) + checkStream(joined, Seq(Row(1, "engineering"), Row(1, "research"))) + } + + test("an external lookup-only source cannot be read outside an eligible lookup join") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1).toDF("id")) + checkError( + exception = intercept[AnalysisException] { lookup.queryExecution.optimizedPlan }, + condition = "UNSUPPORTED_FEATURE.TABLE_OPERATION", + parameters = Map( + "tableName" -> "`SimpleExternalLookupTable`", + "operation" -> "batch scan")) + // A join that reads it but cannot be a lookup join fails with the reason. + checkError( + exception = streamFailure[AnalysisException](left.crossJoin(lookup)), + condition = "EXTERNAL_LOOKUP_JOIN_UNSUPPORTED.MISSING_JOIN_CONDITION", + parameters = Map.empty) + val fullOuter = intercept[AnalysisException] { + left.join(lookup, left("id") === lookup("id"), "full_outer").queryExecution.analyzed + } + assert(fullOuter.getCondition === "EXTERNAL_LOOKUP_JOIN_UNSUPPORTED.JOIN_TYPE") + assert(fullOuter.getMessageParameters.asScala.toMap - "supportedJoinTypes" === + Map("joinType" -> "FULL OUTER", "relation" -> "`SimpleExternalLookupTable`")) + } + + test("a catalog-resolved external lookup-only table cannot be read outside a lookup join") { + withSQLConf( + s"spark.sql.catalog.$ExternalLookupCatalogName" -> + classOf[ExternalLookupTestCatalog].getName) { + val table = s"$ExternalLookupCatalogName.ns.lookup" + val left = Seq(1).toDF("id") + Seq[() => DataFrame]( + () => spark.table(table), + () => sql(s"SELECT * FROM $table WHERE id = 1")).foreach { + df => + checkError( + exception = intercept[AnalysisException] { df().queryExecution.optimizedPlan }, + condition = "UNSUPPORTED_FEATURE.TABLE_OPERATION", + parameters = Map( + "tableName" -> "`SimpleExternalLookupTable`", + "operation" -> "batch scan")) + } + } + } + + test("external lookup join handles null, missing, duplicate, and partial batches") { + SimpleExternalLookupDataSource.resetCounters() + val lookup = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .option("maxBatchSize", "3") + .load() + val left = streamOf(Seq[java.lang.Integer](1, 1, null, 3, 2).toDF("id")) + + val rows = runStream(left.join(lookup, left("id") === lookup("id"), "left_outer")).rows + assert(rows.toSet === Set( + Row(1, 1, "engineering"), + Row(1, 1, "research"), + Row(null, null, null), + Row(3, null, null), + Row(2, 2, "sales"))) + assert(rows.count(_.get(0) == 1) === 4) + assert(SimpleExternalLookupDataSource.requestSizes.asScala.toSeq === Seq(1, 2)) + assert(SimpleExternalLookupDataSource.requestedKeys.asScala.toSeq === Seq(1, 3, 2)) + assert(SimpleExternalLookupDataSource.nullKeyRequests.get() === 0) + } + + test("external lookup join supports complete composite equi-keys") { + val lookup = spark.read + .format(classOf[CompositeExternalLookupDataSource].getName) + .load() + val left = streamOf(Seq((1, 1), (1, 2), (2, 1)).toDF("tenant", "id")) + + checkStream( + left.join( + lookup, + left("tenant") === lookup("tenant") && left("id") === lookup("id"), + "left_outer"), + Seq(Row(1, 1, 1, 1, "one-one"), Row(1, 2, null, null, null), + Row(2, 1, 2, 1, "two-one"))) + } + + test("an upstream limit stops one-to-many cursor consumption") { + SimpleExternalLookupDataSource.resetCounters() + SimpleExternalLookupDataSource.values.set(Map(1 -> (1 to 2000).map(_.toString))) + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1).toDF("id")) + + assert(runStream(left.join( + lookup, left("id") === lookup("id"), "inner") + .limit(1)).rows.size === 1) + assert(SimpleExternalLookupDataSource.resultRowsRead.get() === 1) + } + + test("external lookup samples connector metrics while consuming a large response") { + SimpleExternalLookupDataSource.resetCounters() + SimpleExternalLookupDataSource.values.set(Map(1 -> (1 to 250).map(_.toString))) + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1).toDF("id")) + + assert(runStream(left.join( + lookup, left("id") === lookup("id"), "inner")).rows.size === 250) + // Two periodic samples, followed by request-completion and reader-close samples. + assert(SimpleExternalLookupDataSource.metricSamples.get() >= 4) + } + + test("external lookup deduplicates floating-point keys within a request") { + SemanticKeyExternalLookupDataSource.requestSizes.clear() + val nan1 = java.lang.Double.longBitsToDouble(0x7ff8000000000001L) + val nan2 = java.lang.Double.longBitsToDouble(0x7ff8000000000002L) + // Requests hold two rows, so each pair of equivalent keys shares one request row. + val left = streamOf(Seq(-0.0d, 0.0d, nan1, nan2).toDF("key")) + val lookup = semanticLookup("double") + + assert(runStream(left.join(lookup, left("key") === lookup("key"))).rows.size === 4) + assert(SemanticKeyExternalLookupDataSource.requestSizes.asScala.toSeq === Seq(1, 1)) + } + + test("external lookup deduplicates keys using Spark collation equality") { + SemanticKeyExternalLookupDataSource.requestSizes.clear() + val left = streamOf(Seq("lookup", "LOOKUP").toDF("key")) + .selectExpr("CAST(key AS STRING COLLATE UTF8_LCASE) AS key") + val lookup = semanticLookup("string collate utf8_lcase") + + assert(runStream(left.join(lookup, left("key") === lookup("key"))).rows.size === 2) + assert(SemanticKeyExternalLookupDataSource.requestSizes.asScala.toSeq === Seq(1)) + } + + private def semanticLookup(keyType: String): DataFrame = { + spark.read + .format(classOf[SemanticKeyExternalLookupDataSource].getName) + .option("keyType", keyType) + .load() + } + + test("a join strategy hint on the stream side of a lookup-only join is rejected") { + SimpleExternalLookupDataSource.resetCounters() + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1).toDF("id")) + checkError( + exception = intercept[AnalysisException] { + left.hint("merge").join(lookup, left("id") === lookup("id")).queryExecution.analyzed + }, + condition = "EXTERNAL_LOOKUP_JOIN_UNSUPPORTED.CONFLICTING_JOIN_STRATEGY", + parameters = Map("strategy" -> "MERGE", "relation" -> "`SimpleExternalLookupTable`")) + + // A strategy hint on a different join of the same query does not conflict. + val other = Seq((1, "a")).toDF("id", "tag") + val nested = left.join(other.hint("merge"), "id") + val joined = nested.join(lookup, nested("id") === lookup("id")) + assert(runStream(joined).plan.exists(_.isInstanceOf[PreparedExternalLookupJoin])) + checkStream(joined, Seq(Row(1, "a", 1, "engineering"), Row(1, "a", 1, "research"))) + } + + test("a lookup join whose condition folds to false does not call the connector") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val left = streamOf(Seq(1, 3).toDF("id")) + Seq( + "inner" -> Seq.empty[Row], + "left_outer" -> Seq(Row(1, null, null), Row(3, null, null))).foreach { + case (joinType, expected) => + SimpleExternalLookupDataSource.resetCounters() + val joined = left.join(lookup, left("id") === lookup("id") && lit(false), joinType) + assert(!runStream(joined).plan.exists(_.isInstanceOf[PreparedExternalLookupJoin])) + checkStream(joined, expected) + assert(SimpleExternalLookupDataSource.lookupCandidates.get() === Seq.empty) + assert(SimpleExternalLookupDataSource.readerCount.get() === 0) + assert(SimpleExternalLookupDataSource.requestCount.get() === 0) + } + } + + test("external lookup reads VARCHAR columns and nested fields whose nullability differs") { + def lookup(kind: String): DataFrame = spark.read + .format(classOf[ReadSchemaExternalLookupDataSource].getName) + .option("kind", kind) + .load() + val left = streamOf(Seq(1, 2).toDF("id")) + + val varchar = lookup("varchar") + checkStream( + left.join(varchar, left("id") === varchar("id")), + Seq(Row(1, 1, "ab"), Row(2, 2, "ab"))) + val nested = lookup("nested") + checkStream( + left.join(nested, left("id") === nested("id")), + Seq(Row(1, 1, Row(10)), Row(2, 2, Row(20)))) + // A field whose type really differs is still rejected. + val wrongType = lookup("wrongType") + checkError( + exception = streamFailure[AnalysisException]( + left.join(wrongType, left("id") === wrongType("id"))), + condition = "DATA_SOURCE_INVALID_EXTERNAL_LOOKUP.UNKNOWN_SCAN_FIELD", + parameters = Map("table" -> "`ReadSchemaExternalLookupTable`", "field" -> "`name`")) + } + + test("external lookup merges the connector metrics of every reader in a task") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val lookup = spark.read.format(classOf[MetricsExternalLookupDataSource].getName).load() + // Two one-row input partitions, computed by one task under the coalesce, open two readers + // that send one request each. + val left = streamOf(spark.range(0, 2).selectExpr("CAST(id + 1 AS INT) AS id"), 2) + val joined = left.join(lookup, left("id") === lookup("id")).coalesce(1) + + val run = runStream(joined) + assert(run.rows.length === 2) + val lookupJoin = run.execution.executedPlan.collectFirst { + case exec: ExternalLookupJoinExec => exec + }.get + assert(lookupJoin.metrics("source.connectorRequests").value === 2) + // A connector metric named like an operator metric is kept apart from it. + assert(lookupJoin.metrics("source.numOutputRows").value === 7) + assert(lookupJoin.metrics("numOutputRows").value === 2) + } + } + + test("external lookup closes task-local readers after success and failure") { + def lookup(options: (String, String)*): DataFrame = { + var reader = spark.read.format(classOf[SimpleExternalLookupDataSource].getName) + options.foreach { case (key, value) => reader = reader.option(key, value) } + reader.load() + } + val left = streamOf(Seq(1).toDF("id")) + + SimpleExternalLookupDataSource.resetCounters() + val successfulLookup = lookup() + runStream(left.join(successfulLookup, left("id") === successfulLookup("id"))) + assert(SimpleExternalLookupDataSource.readerCount.get() === 1) + assert(SimpleExternalLookupDataSource.closeCount.get() === 1) + + SimpleExternalLookupDataSource.resetCounters() + val failingLookup = lookup("failRequests" -> "true") + val failure = streamFailure[SparkException]( + left.join(failingLookup, left("id") === failingLookup("id"))) + assert(failure.getMessage.contains("injected external lookup failure")) + assert(SimpleExternalLookupDataSource.readerCount.get() === 1) + assert(SimpleExternalLookupDataSource.closeCount.get() === 1) + } + + test("external lookup caps each inner-join request at the engine batch size") { + SimpleExternalLookupDataSource.resetCounters() + val lookup = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .option("maxBatchSize", "5000") + .load() + val left = streamOf((1 to 1500).toDF("id")) + + assert(runStream(left.join(lookup, left("id") === lookup("id"))).rows.size === 3) + assert(SimpleExternalLookupDataSource.requestSizes.asScala.toSeq === Seq(1000, 500)) + } + + test("external lookup does not create readers without requestable keys") { + val lookup = spark.read.format(classOf[SimpleExternalLookupDataSource].getName).load() + val empty = streamOf(Seq.empty[java.lang.Integer].toDF("id")) + val nullOnly = streamOf(Seq[java.lang.Integer](null).toDF("id")) + + Seq(empty, nullOnly).foreach { left => + SimpleExternalLookupDataSource.resetCounters() + runStream(left.join( + lookup, left("id") === lookup("id"), "left_outer")) + assert(SimpleExternalLookupDataSource.readerCount.get() === 0) + assert(SimpleExternalLookupDataSource.closeCount.get() === 0) + } + } + + test("external lookup cleanup attempts every resource and preserves the primary failure") { + val attempted = mutable.ArrayBuffer.empty[String] + val cursorFailure = new IllegalStateException("cursor close") + val metricsFailure = new IllegalStateException("metrics") + val readerFailure = new IllegalStateException("reader close") + + val failure = intercept[IllegalStateException] { + ExternalLookupJoinExec.closeResources( + { attempted += "cursor"; throw cursorFailure }, + { attempted += "metrics"; throw metricsFailure }, + { attempted += "reader"; throw readerFailure }) + } + assert(failure eq cursorFailure) + assert(attempted === Seq("cursor", "metrics", "reader")) + assert(failure.getSuppressed.toSeq === Seq(metricsFailure)) + assert(metricsFailure.getSuppressed.toSeq === Seq(readerFailure)) + + val lookupFailure = new IllegalStateException("lookup") + val cleanupFailure = new IllegalStateException("cleanup") + val preserved = intercept[IllegalStateException] { + org.apache.spark.util.Utils.tryWithSafeFinally(throw lookupFailure) { + ExternalLookupJoinExec.closeResources(throw cleanupFailure, (), ()) + } + } + assert(preserved eq lookupFailure) + assert(preserved.getSuppressed.toSeq === Seq(cleanupFailure)) + } + + test("external lookup responds to task cancellation and closes the reader") { + SimpleExternalLookupDataSource.resetCounters() + val lookup = spark.read + .format(classOf[SimpleExternalLookupDataSource].getName) + .option("responseDelayMillis", "60000") + .load() + val left = streamOf(Seq(1).toDF("id")) + // Stopping the query cancels the task that waits for the response. + val query = left.join(lookup, left("id") === lookup("id")) + .writeStream.format("memory").queryName("external_lookup_cancellation").start() + + assert(SimpleExternalLookupDataSource.requestStarted.get().await(10, TimeUnit.SECONDS)) 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/commit/01060c7f72e4e7256c177939e0e2306e463dfae0). `beforeEach` resets the shared fixture, including the lookup values, so the tests no longer depend on their order. The cancellation test stops the query and drops the view in a `finally` block. ########## 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) => 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/commit/73b62b7c559b90caf5ac9088417679591118c5d5)). With the analyzed plan kept a `Join`, `FormExternalLookupJoin` reports `MISSING_JOIN_CONDITION` at the start of optimization, before `PushPredicateThroughJoin` can move a later filter into the join. The message now says that the condition must be written in the join. In [01060c7f72e](https://github.com/apache/spark/commit/01060c7f72e4e7256c177939e0e2306e463dfae0), this PR's test checks that `join().where()` and `crossJoin().where()` both report `MISSING_JOIN_CONDITION`, along with the plain `crossJoin`. -- 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]
