HeartSaVioR commented on code in PR #59910: URL: https://github.com/apache/spark/pull/59910#discussion_r4419822339
########## 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/pull/59910/commits/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/pull/59909/commits/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/pull/59910/commits/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]
