HeartSaVioR commented on code in PR #59909: URL: https://github.com/apache/spark/pull/59909#discussion_r4419731062
########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PrepareExternalLookupJoin.scala: ########## @@ -0,0 +1,214 @@ +/* + * 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.{And, Attribute, EqualTo, Expression, + PredicateHelper} +import org.apache.spark.sql.catalyst.plans.logical.{ExternalLookupJoin, Filter, Join, 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.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 => 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 + } + + 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) 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 [73b62b7c559](https://github.com/apache/spark/pull/59909/commits/73b62b7c559b90caf5ac9088417679591118c5d5). The residual is now rebuilt by walking the conjuncts in written order and dropping only the selected candidates, so an unselected equality stays ahead of the conjuncts written after it. In the binding test, the condition has three candidates and one non-candidate. When the connector selects only the first candidate, the residual is the other three conjuncts in written order. When every candidate is selected, or none is offered, the residual is the same as before. ########## 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. Fixed in [73b62b7c559](https://github.com/apache/spark/pull/59909/commits/73b62b7c559b90caf5ac9088417679591118c5d5). `CONFLICTING_JOIN_STRATEGY` now only says to remove the conflicting hint. `LOOKUP_PLAN` now says that the lookup side must be a DSv2 table with no other operator between the table and the join. Both texts will be widened in the changes that add the hint and the other lookup-side shapes. ########## common/utils/src/main/resources/error/error-conditions.json: ########## @@ -2350,6 +2350,50 @@ ], "sqlState" : "KD010" }, + "DATA_SOURCE_EXTERNAL_LOOKUP_NOT_SUPPORTED" : { + "message" : [ + "<description>: External lookup is not supported." + ], + "sqlState" : "0A000" + }, + "DATA_SOURCE_INVALID_EXTERNAL_LOOKUP" : { + "message" : [ + "Data source table <table> violates the external lookup contract:" + ], + "subClass" : { + "DUPLICATE_SCAN_FIELD" : { + "message" : [ + "The lookup scan returned the field <field> more than once." + ] + }, + "DUPLICATE_SELECTED_PREDICATE" : { + "message" : [ + "`selectLookupPredicates()` returned the predicate <predicate> more than once." + ] + }, + "NON_POSITIVE_MAX_PARAMETER_ROWS" : { + "message" : [ + "`maxParameterRowsPerRequest()` returned <value>, but it must be positive." + ] + }, + "NULL_PREDICATE_SELECTION" : { + "message" : [ + "`selectLookupPredicates()` returned null or an array containing null." + ] + }, + "UNKNOWN_SCAN_FIELD" : { + "message" : [ + "The lookup scan returned the field <field>, which is not in the table output." 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 [73b62b7c559](https://github.com/apache/spark/pull/59909/commits/73b62b7c559b90caf5ac9088417679591118c5d5). The message now reads "The lookup scan returned the field <field>, which matches no table column by name and data type." The suite tests the type-mismatch case (`value INT` for a `STRING` column). ########## 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." + ] + }, + "JOIN_TYPE" : { + "message" : [ + "The join type <joinType> is not supported by an external lookup join over <relation>. Supported join types are <supportedJoinTypes>." + ] + }, + "LOOKUP_PLAN" : { + "message" : [ + "The lookup side must be a DSv2 table, optionally followed by projections and filters." + ] + }, + "LOOKUP_TABLE_ON_LEFT" : { + "message" : [ + "The lookup-only table <relation> must be on the right side of the join." + ] + }, + "MISSING_JOIN_CONDITION" : { + "message" : [ + "The join has no join condition." + ] + }, + "NO_LOOKUP_PREDICATE" : { + "message" : [ + "The data source of table <table> did not select any lookup predicate from the join condition." 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 [73b62b7c559](https://github.com/apache/spark/pull/59909/commits/73b62b7c559b90caf5ac9088417679591118c5d5), using the second option. When no conjunct has the form `lookupColumn = <stream-side expression>` with a bare lookup column, the rule fails with the new `EXTERNAL_LOOKUP_JOIN_UNSUPPORTED.NO_LOOKUP_CANDIDATE` before it calls the connector, including before `newScanBuilder`. The message names the accepted form. It also says that when the lookup column is cast to match the stream side, the stream-side expression should be cast to the lookup column's type instead. `NO_LOOKUP_PREDICATE` is now raised only when the connector selects none of a non-empty offer, and its text says the predicates were offered. The test joins a BIGINT stream key with an INT lookup column (`s.id = CAST(t.id AS BIGINT)`) through a stub that fails the test if its selection is called. ########## 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. Fixed in [73b62b7c559](https://github.com/apache/spark/pull/59909/commits/73b62b7c559b90caf5ac9088417679591118c5d5). The suite now extends `PlanTest` and doesn't start a session. -- 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]
