sunchao commented on code in PR #56542: URL: https://github.com/apache/spark/pull/56542#discussion_r3714120869
########## core/src/main/scala/org/apache/spark/shard/ShardManagerMasterEndpoint.scala: ########## @@ -0,0 +1,277 @@ +/* + * 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.shard + +import java.util.{HashMap => JHashMap} +import java.util.concurrent.atomic.AtomicLong + +import scala.collection.mutable +import scala.concurrent.{ExecutionContext, ExecutionContextExecutorService, Future} +import scala.jdk.CollectionConverters._ +import scala.util.Random + +import org.apache.spark.{SparkConf, SparkContext} +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.{EXECUTOR_ID, SHARD_ID, SHARD_MANAGER_ID, SHARD_SET_ID} +import org.apache.spark.rpc.{IsolatedThreadSafeRpcEndpoint, RpcCallContext, RpcEndpointRef, RpcEnv} +import org.apache.spark.shard.ShardManagerMessages._ +import org.apache.spark.util.ThreadUtils + +/** + * Driver-side RPC endpoint that tracks shard locations across executors + * and coordinates replica placement for distributed map join. + */ +private[spark] class ShardManagerMasterEndpoint( + val rpcEnv: RpcEnv, + val isLocal: Boolean, + conf: SparkConf, + shardManagerInfo: mutable.Map[ShardManagerId, ShardManagerInfo], + isDriver: Boolean) + extends IsolatedThreadSafeRpcEndpoint + with Logging { + + private val EXEC_LOAD_WEIGHT = 10 + private val HOST_LOAD_WEIGHT = 5 + private val SAME_HOST_PENALTY = 100000L + private val JITTER_BOUND = 256 + + private val nextShardSetId = new AtomicLong(0) + + private val shardSetInfo = new mutable.HashMap[Long, ShardSetInfo] + private val shardSetLocations = + new JHashMap[Long, JHashMap[Int, mutable.HashSet[ShardManagerId]]] + + private implicit val askEc: ExecutionContextExecutorService = + ExecutionContext.fromExecutorService( + ThreadUtils.newDaemonCachedThreadPool("shard-manager-ask-thread-pool", 8)) + + override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = { + case RegisterShardManager(id, endpoint) => + context.reply(register(id, endpoint)) + + case NewShardSet(numShards, replicaCount) => + context.reply(newShardSet(numShards, replicaCount)) + + case UpdateShardInfo(shardManagerId, setId, shardId) => + updateShardInfo(shardManagerId, setId, shardId) + context.reply(true) + + case InstallReplicaSet(setId, shardId) => + installReplicaToWorkers(setId, shardId).onComplete { + case scala.util.Success(_) => context.reply(true) + case scala.util.Failure(e) => + logWarning(log"Replica installation failed for" + + log" (${MDC(SHARD_SET_ID, setId)}, ${MDC(SHARD_ID, shardId)})", e) + context.reply(true) Review Comment: **[P1] Propagate replica-installation failures to the build exchange.** This failure branch logs an unsuccessful replica RPC but still replies `true`; consequently `ShardManagerMaster.installReplicaSet` returns normally and `ShardExchangeExec` publishes a successful materialized stage without the requested replicas. For `replica_count=2`, losing the remaining primary then permanently breaks the join. The previous asynchronous-installation concern is therefore not fixed: complete the RPC exceptionally and fail/retry the exchange instead of acknowledging failed replication as success. ########## core/src/main/scala/org/apache/spark/scheduler/dynalloc/ExecutorMonitor.scala: ########## @@ -462,6 +472,18 @@ private[spark] class ExecutorMonitor( override def broadcastCleaned(broadcastId: Long): Unit = { } + override def shardSetCleaned(setId: Long): Unit = { + executors.asScala.foreach { case (_, exec) => + val toRemove = exec.nonRddCachedBlocks.collect { Review Comment: **[P2] Keep shard cleanup on the executor-monitor event thread.** `ContextCleaner` invokes `shardSetCleaned` on its dedicated cleaner thread, but this callback directly iterates and mutates each executor's mutable `nonRddCachedBlocks`/`cachedBlocks` while `onBlockUpdated` and other listener callbacks mutate the same structures on the listener-bus management thread. This violates the documented event-thread confinement and can corrupt shard tracking or remove/retain the wrong executor. Follow the adjacent `shuffleCleaned` implementation and post a shard-cleanup event to the listener bus. ########## core/src/main/scala/org/apache/spark/SparkContext.scala: ########## @@ -646,6 +646,7 @@ class SparkContext(config: SparkConf) extends Logging { _env.blockManager.initialize(_applicationId) FallbackStorage.registerBlockManagerIfNeeded( _env.blockManager.master, _conf, _hadoopConfiguration) + _env.initializeShardManager() Review Comment: **[P1] Register the driver shard endpoint before executors can start.** `_taskScheduler.start()` has already started the cluster backend at line 625 and can launch/register executors before this call creates `ShardManagerMaster`. Executor construction immediately calls its own `initializeShardManager()`, synchronously looks up/registers with that driver endpoint, and exits when the endpoint is not present. This makes feature-enabled cluster startup racy, particularly with Kubernetes/YARN or quickly allocated executors. Register the driver endpoint before the backend permits executor registration, or defer/retry executor initialization. ########## sql/core/src/main/java/org/apache/spark/sql/execution/BufferedShardRowMap.java: ########## @@ -0,0 +1,484 @@ +/* + * 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; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Iterator; +import java.util.LinkedList; +import java.util.Map; +import java.util.NoSuchElementException; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.PooledByteBufAllocator; +import org.apache.spark.internal.SparkLogger; +import org.apache.spark.internal.SparkLoggerFactory; +import org.apache.spark.memory.MemoryConsumer; +import org.apache.spark.memory.TaskMemoryManager; +import org.apache.spark.network.buffer.ManagedBuffer; +import org.apache.spark.network.buffer.NettyManagedBuffer; +import org.apache.spark.network.util.NettyUtils; +import org.apache.spark.sql.catalyst.expressions.UnsafeRow; +import org.apache.spark.unsafe.Platform; +import org.apache.spark.unsafe.UnsafeAlignedOffset; +import org.apache.spark.unsafe.array.ByteArrayMethods; +import org.apache.spark.unsafe.map.HashMapGrowthStrategy; +import org.apache.spark.unsafe.memory.MemoryBlock; + +/** + * Probe-side key batching buffer for distributed map join. + * + * Accumulates streamed-side keys into per-shard batches backed by off-heap + * Tungsten pages, then serializes them into Netty buffers for RPC lookup + * against build-side executors. + */ +public final class BufferedShardRowMap extends MemoryConsumer { + + private static final SparkLogger logger = SparkLoggerFactory.getLogger(BufferedShardRowMap.class); + + private final TaskMemoryManager taskMemoryManager; + + private final long setId; + private final int numShards; + private final int batchCapacity; + private final int mask; + private final int uaoSize; + + private final int numFields; + private final UnsafeRow valueUr; + + private final KeyValueBatch[] tailingBatches; + private final LinkedList<KeyValueBatch> pendingBatches; + + private final LinkedList<KeyValuePage> pageList; + + private final PooledByteBufAllocator alloc; + + private KeyValuePage currentPage = null; + + public BufferedShardRowMap(TaskMemoryManager taskMemoryManager, long setId, int numShards, + int numFields, UnsafeRow valueUr, int batchSize) { + super(taskMemoryManager, taskMemoryManager.getTungstenMemoryMode()); + this.taskMemoryManager = taskMemoryManager; + this.setId = setId; + this.numShards = numShards; + this.tailingBatches = new KeyValueBatch[numShards]; + this.pendingBatches = new LinkedList<>(); + this.pageList = new LinkedList<>(); + this.numFields = numFields; + this.valueUr = valueUr; + this.alloc = NettyUtils.getSharedPooledByteBufAllocator(true, true); + this.batchCapacity = Math.max((int) ByteArrayMethods.nextPowerOf2(batchSize), 32); + this.mask = HashMapGrowthStrategy.DOUBLING.nextCapacity(batchCapacity) - 1; + this.uaoSize = UnsafeAlignedOffset.getUaoSize(); + } + + public void putRow(int shard, + Object kbase, long koff, int klen, int khash, + Object vbase, long voff, int vlen) { + final int idx = shard % numShards; + if (tailingBatches[idx] == null) { + tailingBatches[idx] = new KeyValueBatch(shard); + } + KeyValueBatch batch = tailingBatches[idx]; + if (batch.capReached()) { Review Comment: **[P1] Bound buffered probe rows globally across all shards.** A batch becomes eligible for lookup only after an individual shard accumulates roughly 1,024 complete probe rows, so uniformly interleaved input can retain `numShards * 1024` full rows before the first RPC. For 32 shards and 32-KiB probe rows that is about 1 GiB per task; the retained Tungsten pages cannot spill, and the in-flight RPC limit is not reached until after buffering. Large/wide probe workloads therefore fail with task/executor OOM. Enforce a global buffered-byte budget and flush partially filled shard batches under pressure. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala: ########## @@ -258,6 +289,28 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] { } } + def createDistributedMapJoin(onlyLookingAtHint: Boolean) = { + if (hashJoinSupport) { Review Comment: **[P1] Gate DISTMAPJOIN on initialized shard infrastructure.** `spark.shard.enabled` defaults to `false`, and `SparkEnv.initializeShardManager` leaves `shardManager` null in that configuration, but this planner branch checks only `hashJoinSupport` before selecting `DistributedMapJoinExec` for a hint. A hinted query in an otherwise default application therefore reaches `ShardExchangeExec` and dereferences the null manager instead of ignoring the unavailable strategy/falling back. Check the feature gate before constructing the physical join and add a default-disabled regression test. ########## common/network-common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java: ########## @@ -163,25 +163,58 @@ private void processStreamRequest(final StreamRequest req) { } private void processRpcRequest(final RpcRequest req) { - try { - rpcHandler.receive(reverseClient, req.body().nioByteBuffer(), new RpcResponseCallback() { - @Override - public void onSuccess(ByteBuffer response) { - respond(new RpcResponse(req.requestId, new NioManagedBuffer(response))); + if (rpcHandler instanceof RpcHandler.ManagedRpcHandler) { Review Comment: **[P1] Preserve managed-RPC dispatch through authentication wrappers.** With `spark.authenticate=true`, `AuthServerBootstrap` wraps `NettyShardRpcServer` in `AuthRpcHandler` or `SaslRpcHandler`, neither of which implements `ManagedRpcHandler`. This check therefore selects the legacy `ByteBuffer` path; after authentication, the wrapper delegates to `NettyShardRpcServer.receive(ByteBuffer)`, which unconditionally throws `UnsupportedOperationException`. Consequently every shard lookup fails on authenticated clusters. Preserve authenticated managed dispatch and add an authentication-enabled integration test. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelationAdapter.scala: ########## @@ -0,0 +1,122 @@ +/* + * 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 io.netty.buffer.Unpooled + +import org.apache.spark.SparkEnv +import org.apache.spark.network.buffer.{ManagedBuffer, NettyManagedBuffer} +import org.apache.spark.network.util.NettyUtils +import org.apache.spark.shard.{ShardLookupAdapter, ShardManager} +import org.apache.spark.sql.catalyst.expressions.{BasePredicate, Expression, Predicate, UnsafeRow} +import org.apache.spark.sql.catalyst.types.DataTypeUtils +import org.apache.spark.sql.types.StructType + +/** + * Build-side RPC handler for distributed map join. + * + * Receives batched key lookups from probe-side executors, performs hash + * lookups against the local [[HashedRelation]], and returns matching rows. + * + * When a build-only filter is stored in the shard meta, it is loaded once + * per setId and evaluated server-side to reduce network transfer. + * + * Wire format (request): + * {{{ + * (setId:long)(shardId:int)(numKeyFields:int) + * [(keyLen:int)(keyBytes)]... + * }}} + */ +private[spark] class HashedRelationAdapter extends ShardLookupAdapter { + + private val INITIAL_RESPONSE_BUFFER_BYTES = 1 << 20 + private val alloc = NettyUtils.getSharedPooledByteBufAllocator(true, true) + + private val filterCache = new ConcurrentHashMap[Long, Option[BasePredicate]]() + private val cleanupRegistered = new AtomicBoolean(false) + + override def lookup(manager: ShardManager, reqMsg: ManagedBuffer): ManagedBuffer = { + if (cleanupRegistered.compareAndSet(false, true)) { + manager.registerCleanupCallback(setId => filterCache.remove(setId)) + } + val keysBuf = Unpooled.wrappedBuffer(reqMsg.nioByteBuffer()) + val setId = keysBuf.readLong() + val shard = keysBuf.readInt() + val numKeyFields = keysBuf.readInt() + + val keyUr = new UnsafeRow(numKeyFields) + val rel = manager.getLocalValue[HashedRelation](setId, shard).asReadOnlyCopy() + val valuesBuf = alloc.buffer(INITIAL_RESPONSE_BUFFER_BYTES) + valuesBuf.writeLong(setId) + valuesBuf.writeInt(shard) + + val advanceReadKey = UnsafeRowBufCodec.makeAdvanceRead(keyUr, keysBuf) + val advanceWrite = UnsafeRowBufCodec.makeAdvanceWrite(valuesBuf) + + val predicate = filterCache.computeIfAbsent(setId, _ => { Review Comment: **[P1] Do not share one mutable generated predicate across RPC workers.** `computeIfAbsent` stores one `BasePredicate` per shard set, while `NettyShardRpcServer` evaluates lookups concurrently on up to 16 worker threads. Generated predicates embed mutable fields; for example, a build-only expression using `elt(size(sequence(d.start,d.stop,d.step)), d.keep, d.drop)` can remain in the join because `sequence(..., step)` is throwable, and concurrent `elt` evaluations race on its generated `inputVal` field. This includes/excludes the wrong build rows and silently corrupts join results. Create a per-request/per-thread predicate instead of sharing its evaluator. ########## core/src/main/scala/org/apache/spark/network/ShardLookupServiceFactory.scala: ########## @@ -0,0 +1,73 @@ +/* + * 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.network + +import org.apache.spark.{SecurityManager, SparkConf} +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.CLASS_NAME +import org.apache.spark.network.netty.NettyShardLookupService +import org.apache.spark.rpc.RpcEndpointRef +import org.apache.spark.util.Utils + +/** + * Factory for creating [[ShardLookupService]] instances. + * + * The implementation class is controlled by `spark.shard.service`. + * Custom implementations (e.g. native backends) can be plugged in by setting + * this config to the FQCN of a [[ShardLookupService]] subclass with a + * constructor matching `(SparkConf, String, String, Int)`. + */ +private[spark] object ShardLookupServiceFactory extends Logging { + + private val SHARD_LOOKUP_SERVICE_CLASS_KEY = "spark.shard.service" + + private val DEFAULT_CLASS = + classOf[NettyShardLookupService].getName + + def create( + conf: SparkConf, + securityManager: SecurityManager, + bindAddress: String, + advertiseAddress: String, + port: Int, + numCores: Int, + masterEndpoint: RpcEndpointRef): ShardLookupService = { + val className = conf.get(SHARD_LOOKUP_SERVICE_CLASS_KEY, DEFAULT_CLASS) + if (className == DEFAULT_CLASS) { + new NettyShardLookupService(conf, securityManager, bindAddress, advertiseAddress, port, + numCores, masterEndpoint) + } else { + try { + logInfo(log"Creating custom ShardLookupService: ${MDC(CLASS_NAME, className)}") + Utils.classForName(className) + .getDeclaredConstructor( Review Comment: **[P2] Honor the documented constructor for custom shard services.** The factory documents `spark.shard.service` implementations with a `(SparkConf, String, String, Int)` constructor, but this reflective lookup requires seven arguments including `SecurityManager`, core count, and `RpcEndpointRef`. A custom implementation following the documented contract therefore throws `NoSuchMethodException` and is silently replaced by Netty, ignoring the explicitly configured backend. Support the advertised constructor, or document/enforce the actual contract without silently discarding the user's selection. ########## core/src/main/scala/org/apache/spark/shard/ShardManagerMasterEndpoint.scala: ########## @@ -0,0 +1,277 @@ +/* + * 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.shard + +import java.util.{HashMap => JHashMap} +import java.util.concurrent.atomic.AtomicLong + +import scala.collection.mutable +import scala.concurrent.{ExecutionContext, ExecutionContextExecutorService, Future} +import scala.jdk.CollectionConverters._ +import scala.util.Random + +import org.apache.spark.{SparkConf, SparkContext} +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.{EXECUTOR_ID, SHARD_ID, SHARD_MANAGER_ID, SHARD_SET_ID} +import org.apache.spark.rpc.{IsolatedThreadSafeRpcEndpoint, RpcCallContext, RpcEndpointRef, RpcEnv} +import org.apache.spark.shard.ShardManagerMessages._ +import org.apache.spark.util.ThreadUtils + +/** + * Driver-side RPC endpoint that tracks shard locations across executors + * and coordinates replica placement for distributed map join. + */ +private[spark] class ShardManagerMasterEndpoint( + val rpcEnv: RpcEnv, + val isLocal: Boolean, + conf: SparkConf, + shardManagerInfo: mutable.Map[ShardManagerId, ShardManagerInfo], + isDriver: Boolean) + extends IsolatedThreadSafeRpcEndpoint + with Logging { + + private val EXEC_LOAD_WEIGHT = 10 + private val HOST_LOAD_WEIGHT = 5 + private val SAME_HOST_PENALTY = 100000L + private val JITTER_BOUND = 256 + + private val nextShardSetId = new AtomicLong(0) + + private val shardSetInfo = new mutable.HashMap[Long, ShardSetInfo] + private val shardSetLocations = + new JHashMap[Long, JHashMap[Int, mutable.HashSet[ShardManagerId]]] + + private implicit val askEc: ExecutionContextExecutorService = + ExecutionContext.fromExecutorService( + ThreadUtils.newDaemonCachedThreadPool("shard-manager-ask-thread-pool", 8)) + + override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = { + case RegisterShardManager(id, endpoint) => + context.reply(register(id, endpoint)) + + case NewShardSet(numShards, replicaCount) => + context.reply(newShardSet(numShards, replicaCount)) + + case UpdateShardInfo(shardManagerId, setId, shardId) => + updateShardInfo(shardManagerId, setId, shardId) + context.reply(true) + + case InstallReplicaSet(setId, shardId) => + installReplicaToWorkers(setId, shardId).onComplete { + case scala.util.Success(_) => context.reply(true) + case scala.util.Failure(e) => + logWarning(log"Replica installation failed for" + + log" (${MDC(SHARD_SET_ID, setId)}, ${MDC(SHARD_ID, shardId)})", e) + context.reply(true) + } + + case GetLocations(setId, shardId) => + context.reply(getLocations(setId, shardId)) + + case RemoveShardSet(setId) => + removeShardSet(setId) + context.reply(true) + + case RemoveExecutor(execId) => + removeExecutor(execId) + context.reply(true) + + case StopShardManagerMaster => + context.reply(true) + stop() + } + + private def register( + idWithoutTopologyInfo: ShardManagerId, + managerEndpoint: RpcEndpointRef): ShardManagerId = { + val id = ShardManagerId( + idWithoutTopologyInfo.executorId, + idWithoutTopologyInfo.host, + idWithoutTopologyInfo.port, + None) + shardManagerInfo(id) = new ShardManagerInfo(id, managerEndpoint) + id + } + + private def removeShardManager(shardManagerId: ShardManagerId): Unit = { + val info = shardManagerInfo(shardManagerId) + shardManagerInfo.remove(shardManagerId) + + val iterator = info.shards.entrySet().iterator() + while (iterator.hasNext) { + val entry = iterator.next() + val setId = entry.getKey + val shardLocations = shardSetLocations.get(setId) + if (shardLocations != null) { + val valueIterator = entry.getValue.iterator + while (valueIterator.hasNext) { + val shardId = valueIterator.next() + val locations = shardLocations.get(shardId) + if (locations != null) { + locations -= shardManagerId + if (locations.isEmpty) { Review Comment: **[P1] Rebuild a materialized shard set when its final replica is lost.** The default `replica_count=1` means ordinary executor loss reaches this branch after a completely successful exchange, yet it only removes the location and logs. Subsequent lookups fail with an ordinary `SparkException`, while the AQE shard stage remains materialized with the same now-unrecoverable set ID; every task retry therefore fails rather than rebuilding the build side. Executor loss or finite cached-executor idle timeouts must invalidate/recompute the shard stage, or the strategy needs a stronger durability invariant. ########## project/MimaExcludes.scala: ########## @@ -44,6 +44,8 @@ object MimaExcludes { // [SPARK-54879] Add exitCode field to ApplicationAttemptInfo ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.status.api.v1.ApplicationAttemptInfo.tupled"), ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.status.api.v1.ApplicationAttemptInfo.curried"), + // [SPARK-57487][SQL] Distributed map join adds shardManagerFactory parameter to SparkEnv + ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.SparkEnv.this"), Review Comment: **[P1] Preserve the existing SparkEnv constructor instead of suppressing MiMa.** Adding `shardManagerFactory` removes the existing public `@DeveloperApi` 12-argument `SparkEnv` constructor, and this exclusion only hides the resulting binary-compatibility failure. Previously compiled extensions that instantiate `SparkEnv` fail with `NoSuchMethodError`, even when `spark.shard.enabled=false`; source callers also stop compiling. Retain a compatible constructor/overload and remove this exclusion rather than masking the incompatible JVM descriptor. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveHints.scala: ########## @@ -98,6 +98,45 @@ object ResolveHints { private def matchedIdentifier(identInHint: Seq[String], identInQuery: Seq[String]): Boolean = ResolveHints.matchedIdentifier(identInHint, identInQuery, resolver) + private def createHintInfo(ident: Seq[String], hint: UnresolvedHint): HintInfo = { + + if (DistributedMapJoinStrategy(None, None).hintAliases + .map(_.toUpperCase(Locale.ROOT)) + .contains(hint.name.toUpperCase(Locale.ROOT))) { + + def keyName(e: Expression): Option[String] = e match { + case UnresolvedAttribute(parts) => Some(parts.last.toLowerCase(Locale.ROOT)) + case _ => None + } + def intValue(e: Expression): Option[Int] = e match { + case IntegerLiteral(i) => Some(i) + case Literal(i: Int, _) => Some(i) + case _ => None + } + + val keyOpts = hint.parameters + .collectFirst { + case uf: UnresolvedFunction if matchedIdentifier(uf.nameParts, ident) => + uf.arguments + .collect { case EqualTo(lhs, rhs) => + for { + k <- keyName(lhs) + v <- intValue(rhs) Review Comment: **[P2] Validate shard and replica hint counts before planning.** Integer values are accepted without a positive-range check, so a valid-looking `DISTMAPJOIN(d(shard_count=0))` proceeds into `ShardDistribution` and `HashPartitioning(..., 0)`, causing deterministic zero-partition/modulo-by-zero execution failures instead of rejecting the malformed hint. `replica_count=0` is likewise silently accepted even though a primary shard is still created. Validate both options as positive during hint resolution and cover zero/negative values. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/DistributedMapJoinExec.scala: ########## @@ -0,0 +1,604 @@ +/* + * 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 + +import scala.annotation.tailrec +import scala.concurrent.ExecutionContextExecutorService + +import io.netty.buffer.Unpooled + +import org.apache.spark.{SparkEnv, SparkException, TaskContext} +import org.apache.spark.network.buffer.ManagedBuffer +import org.apache.spark.rdd.RDD +import org.apache.spark.shard.ShardSetRef +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{And, AttributeSet, BindReferences, Expression, GenericInternalRow, JoinedRow, Predicate, PredicateHelper, SortOrder, UnsafeProjection, UnsafeRow} +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} +import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, InnerLike, JoinType, LeftAnti, LeftOuter, LeftSemi, RightOuter} +import org.apache.spark.sql.catalyst.plans.logical.{DistributedMapJoinStrategy, JoinHint} +import org.apache.spark.sql.catalyst.plans.physical._ +import org.apache.spark.sql.errors.QueryExecutionErrors +import org.apache.spark.sql.execution.{BufferedShardRowMap, SparkPlan} +import org.apache.spark.sql.execution.adaptive.ShardQueryStageExec +import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShardExchangeExec} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.util.sketch.BloomFilter + +/** + * Physical operator for Distributed MapJoin. + * + * This strategy avoids full shuffle by building a distributed hash table service for the build + * side (medium-sized table), and the probe side performs batched RPC lookups to complete the + * join. + * + * Currently only supports: + * - Equi-join + * - BuildRight (right table as build side) + * - Explicit SQL hint: /*+distmapjoin(t(shard_count=5,replica_count=2))*/ + */ +case class DistributedMapJoinExec( + leftKeys: Seq[Expression], + rightKeys: Seq[Expression], + joinType: JoinType, + buildSide: BuildSide, + condition: Option[Expression], + left: SparkPlan, + right: SparkPlan, + hint: JoinHint, + isNullAwareAntiJoin: Boolean = false) + extends HashJoin with PredicateHelper { + + override lazy val metrics = Map( + "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) + + private val (numShards, replicaCount): (Int, Int) = { + val strategy = + (if (buildSide == BuildLeft) hint.leftHint else hint.rightHint).flatMap(_.strategy) + strategy match { + case Some(DistributedMapJoinStrategy(ns, rc)) => (ns.getOrElse(5), rc.getOrElse(1)) + case _ => (5, 1) + } + } + + @transient private lazy val (buildOnlyFilter, remainingCondition): + (Option[Expression], Option[Expression]) = { + condition match { + case Some(cond) => + val buildAttrs = AttributeSet(buildOutput) + val conjuncts = splitConjunctivePredicates(cond) + val (bo, rem) = conjuncts.partition(_.references.subsetOf(buildAttrs)) + (bo.reduceOption(And), rem.reduceOption(And)) + case None => (None, None) + } + } + + @transient override protected[this] lazy val boundCondition: InternalRow => Boolean = { + remainingCondition match { + case Some(cond) => + Predicate.create(cond, streamedPlan.output ++ buildPlan.output).eval _ + case None => + (_: InternalRow) => true + } + } + + override def supportCodegen: Boolean = false + + override def supportsColumnar: Boolean = false + + override def needCopyResult: Boolean = false + + override def requiredChildDistribution: Seq[Distribution] = { + val filter = buildOnlyFilter.map(BindReferences.bindReference(_, buildOutput)) + val filterSchema = buildOnlyFilter.map(_ => buildPlan.schema) + val sd = ShardDistribution(buildBoundKeys, numShards, replicaCount, filter, filterSchema) + buildSide match { + case BuildLeft => Seq(sd, UnspecifiedDistribution) + case BuildRight => Seq(UnspecifiedDistribution, sd) + } + } + + override def outputPartitioning: Partitioning = streamedPlan.outputPartitioning + + override def outputOrdering: Seq[SortOrder] = Nil + + override def inputRDDs(): Seq[RDD[InternalRow]] = + throw QueryExecutionErrors.executeCodePathUnsupportedError("DistributedMapJoin") + + override protected def prepareRelation(ctx: CodegenContext): HashedRelationInfo = + throw QueryExecutionErrors.executeCodePathUnsupportedError("DistributedMapJoin") + + override protected def withNewChildrenInternal( + newLeft: SparkPlan, + newRight: SparkPlan): SparkPlan = copy(left = newLeft, right = newRight) + + override protected def doExecute(): RDD[InternalRow] = { + val numOutputRows = longMetric("numOutputRows") + val setRef = resolveShardSetRef(buildPlan) + streamedPlan.execute().mapPartitionsInternal { streamedIter => + join(streamedIter, setRef.setId, numOutputRows) + } + } + + @tailrec + private def resolveShardSetRef(plan: SparkPlan): ShardSetRef = plan match { + case s: ShardExchangeExec => s.buildShardSet() + case s: ShardQueryStageExec => s.shardSetRef + case r: ReusedExchangeExec => resolveShardSetRef(r.child) + case other => + throw new IllegalStateException(s"Unexpected build plan for DistributedMapJoin: $other") + } + + private def streamedBloomFilter(setId: Long): BloomFilter = { + SparkEnv.get.shardManager.fetchBloomFilter[BloomFilter](setId)(bfInput => + BloomFilter.readFrom(bfInput)) + } + + private def join( + streamedIter: Iterator[InternalRow], + setId: Long, + numOutputRows: SQLMetric): Iterator[InternalRow] = { + + val joinedRow = new JoinedRow + + val (scanBatch, onMissing): ( + BatchMatchReader => Iterator[InternalRow], + InternalRow => Option[InternalRow]) = joinType match { + case _: InnerLike => + (innerJoinScan(_, joinedRow), (_: InternalRow) => None) + case LeftOuter | RightOuter => + val nullRow = new GenericInternalRow(buildOutput.length) + (outerJoinScan(_, joinedRow, nullRow), + (sr: InternalRow) => Some(joinedRow.withLeft(sr).withRight(nullRow))) + case LeftSemi => + (semiJoinScan(_, joinedRow), (_: InternalRow) => None) + case LeftAnti => + (antiJoinScan(_, joinedRow), (sr: InternalRow) => Some(sr)) + case _: ExistenceJoin => + val existsRow = new GenericInternalRow(Array[Any](null)) + (existenceJoinScan(_, joinedRow, existsRow), + (sr: InternalRow) => { existsRow.setBoolean(0, false); Some(joinedRow(sr, existsRow)) }) + case x => + throw new IllegalArgumentException( + s"DistributedMapJoin should not take $x as the JoinType") + } + + val iter = new LookupJoinIterator(streamedIter, setId, scanBatch, onMissing) + val resultProj = createResultProjection() + iter.map { row => + numOutputRows.add(1) + resultProj(row) + } + } + + // --------------------------------------------------------------------------- + // Scan methods: one per join type, each self-contained + // --------------------------------------------------------------------------- + + private def innerJoinScan( + reader: BatchMatchReader, + joinedRow: JoinedRow): Iterator[InternalRow] = { + new Iterator[InternalRow] { + private var _has = false + override def hasNext: Boolean = { if (!_has) _has = advance(); _has } + override def next(): InternalRow = { _has = false; joinedRow } + + private def advance(): Boolean = { + while (true) { + val br = reader.nextBuildRow() + if (br != null) { + if (boundCondition(joinedRow.withLeft(reader.curStreamed).withRight(br))) return true + } else if (!reader.advanceStreamedRow()) { + return false + } + } + false // unreachable + } + } + } + + private def outerJoinScan( + reader: BatchMatchReader, + joinedRow: JoinedRow, + nullRow: InternalRow): Iterator[InternalRow] = { + new Iterator[InternalRow] { + private var found = false + private var _has = false + override def hasNext: Boolean = { if (!_has) _has = advance(); _has } + override def next(): InternalRow = { _has = false; joinedRow } + + private def advance(): Boolean = { + while (true) { + val br = reader.nextBuildRow() + if (br != null) { + if (boundCondition(joinedRow.withLeft(reader.curStreamed).withRight(br))) { + found = true + return true + } + } else { + if (!found && reader.curStreamed != null) { + joinedRow.withLeft(reader.curStreamed).withRight(nullRow) + found = true + return true + } + if (!reader.advanceStreamedRow()) return false + found = false + } + } + false // unreachable + } + } + } + + private def semiJoinScan( + reader: BatchMatchReader, + joinedRow: JoinedRow): Iterator[InternalRow] = { + new Iterator[InternalRow] { + private var _has = false + override def hasNext: Boolean = { if (!_has) _has = advance(); _has } + override def next(): InternalRow = { _has = false; reader.curStreamed } + + private def advance(): Boolean = { + while (reader.advanceStreamedRow()) { + var br = reader.nextBuildRow() + while (br != null) { + if (boundCondition(joinedRow.withLeft(reader.curStreamed).withRight(br))) { + reader.skipRemainingBuildRows() + return true + } + br = reader.nextBuildRow() + } + } + false + } + } + } + + private def antiJoinScan( + reader: BatchMatchReader, + joinedRow: JoinedRow): Iterator[InternalRow] = { + new Iterator[InternalRow] { + private var _has = false + override def hasNext: Boolean = { if (!_has) _has = advance(); _has } + override def next(): InternalRow = { _has = false; reader.curStreamed } + + private def advance(): Boolean = { + while (reader.advanceStreamedRow()) { + if (findMatchingBuildRow(reader, joinedRow)) { + reader.skipRemainingBuildRows() + } else { + return true + } + } + false + } + } + } + + private def existenceJoinScan( + reader: BatchMatchReader, + joinedRow: JoinedRow, + existsRow: GenericInternalRow): Iterator[InternalRow] = { + new Iterator[InternalRow] { + private var _has = false + override def hasNext: Boolean = { if (!_has) _has = advance(); _has } + override def next(): InternalRow = { _has = false; joinedRow } + + private def advance(): Boolean = { + if (!reader.advanceStreamedRow()) return false + val exists = findMatchingBuildRow(reader, joinedRow) + if (exists) reader.skipRemainingBuildRows() + existsRow.setBoolean(0, exists) + joinedRow.withLeft(reader.curStreamed).withRight(existsRow) + true + } + } + } + + @tailrec + private def findMatchingBuildRow( + reader: BatchMatchReader, + joinedRow: JoinedRow): Boolean = { + val br = reader.nextBuildRow() + if (br == null) false + else if (boundCondition(joinedRow.withLeft(reader.curStreamed).withRight(br))) true + else findMatchingBuildRow(reader, joinedRow) + } + + // --------------------------------------------------------------------------- + // BatchMatchReader: reads batch response buffer, zero-copy + // --------------------------------------------------------------------------- + + private type PBatch = BufferedShardRowMap#KeyValueBatch + private type BBuffer = ManagedBuffer + + private class BatchMatchReader(batch: PBatch, buffer: BBuffer) extends AutoCloseable { + private val keyIter = batch.multiValuesIterator() + private val buf = Unpooled.wrappedBuffer(buffer.nioByteBuffer()) + private val buildUr: UnsafeRow = new UnsafeRow(buildOutput.length) + private val advanceRead = UnsafeRowBufCodec.makeAdvanceRead(buildUr, buf) + private var streamedIter: java.util.Iterator[UnsafeRow] = _ + private var buildRowsStart: Int = -1 + private var atSentinel = true + var curStreamed: UnsafeRow = _ + + assert(batch.getSetId == buf.readLong(), "setId mismatch") + assert(batch.getShard == buf.readInt(), "shardId mismatch") + + def advanceStreamedRow(): Boolean = { + if (streamedIter != null && streamedIter.hasNext) { + curStreamed = streamedIter.next() + rewindBuildRows() + true + } else { + skipRemainingBuildRows() + if (!keyIter.hasNext) return false + streamedIter = keyIter.next() + buildRowsStart = buf.readerIndex() + atSentinel = false + if (!streamedIter.hasNext) advanceStreamedRow() + else { curStreamed = streamedIter.next(); true } + } + } + + def nextBuildRow(): UnsafeRow = { + if (atSentinel) return null + val blen = buf.readInt() + if (blen == 0) { atSentinel = true; null } + else advanceRead(blen) + } + + def skipRemainingBuildRows(): Unit = { + while (!atSentinel) { + val blen = buf.readInt() + if (blen == 0) atSentinel = true + else buf.readerIndex(buf.readerIndex() + blen) + } + } + + private def rewindBuildRows(): Unit = { + buf.readerIndex(buildRowsStart) + atSentinel = false + } + + override def close(): Unit = { + batch.release() + buffer.release() + } + } + + // --------------------------------------------------------------------------- + // LookupJoinIterator: async pipeline management + // --------------------------------------------------------------------------- + + private class LookupJoinIterator( + streamedIter: Iterator[InternalRow], + setId: Long, + scanBatch: BatchMatchReader => Iterator[InternalRow], + onMissing: InternalRow => Option[InternalRow]) + extends Iterator[InternalRow] { + + private val maxInFlightNum = conf.distributedMapJoinMaxInFlightNum + private val keyGenerator: UnsafeProjection = UnsafeProjection.create(streamedBoundKeys) + private val valueGenerator: UnsafeProjection = UnsafeProjection.create(streamedPlan.schema) + private val shardGenerator = + UnsafeProjection.create( + HashPartitioning(streamedBoundKeys, numShards).partitionIdExpression :: Nil) + + private val bloom = streamedBloomFilter(setId) + private val probeUr: UnsafeRow = new UnsafeRow(streamedPlan.schema.length) + @volatile private var cancelled = false + + private val bufferedMap = { + val mm = TaskContext.get().taskMemoryManager() + val maxBatchSize = conf.distributedMapJoinMaxBatchSize + val map = + new BufferedShardRowMap( + mm, + setId, + numShards, + streamedBoundKeys.length, + probeUr, + maxBatchSize) + TaskContext.get().addTaskCompletionListener[Unit] { _ => + cancelled = true + if (currentReader != null) { + currentReader.close() + currentReader = null + } + var item = lookupQueue.poll() + while (item != null) { + item match { + case LookupSuccess(batch, buffer) => + batch.release() + buffer.release() + case _ => + } + item = lookupQueue.poll() + } + map.free() Review Comment: **[P1] Quiesce outstanding lookup work before freeing Tungsten pages.** `ShardManager.fetchRemoteBatch` first schedules `Future { batch.wrapKeysBuffer() }`, which reads addresses from this task's pages, but an early `LIMIT`, cancellation, or failure can run this completion listener and call `map.free()` while that serialization future is still queued/running. `TaskMemoryManager.freePage` clears the page table first, so the later serializer dereferences freed pages; the callback's later `cancelled` check cannot prevent it. The check-then-enqueue path also races queue drainage. Cancel/await outstanding work before releasing task memory and make completion ownership atomic. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/HashedRelationAdapter.scala: ########## @@ -0,0 +1,122 @@ +/* + * 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 io.netty.buffer.Unpooled + +import org.apache.spark.SparkEnv +import org.apache.spark.network.buffer.{ManagedBuffer, NettyManagedBuffer} +import org.apache.spark.network.util.NettyUtils +import org.apache.spark.shard.{ShardLookupAdapter, ShardManager} +import org.apache.spark.sql.catalyst.expressions.{BasePredicate, Expression, Predicate, UnsafeRow} +import org.apache.spark.sql.catalyst.types.DataTypeUtils +import org.apache.spark.sql.types.StructType + +/** + * Build-side RPC handler for distributed map join. + * + * Receives batched key lookups from probe-side executors, performs hash + * lookups against the local [[HashedRelation]], and returns matching rows. + * + * When a build-only filter is stored in the shard meta, it is loaded once + * per setId and evaluated server-side to reduce network transfer. + * + * Wire format (request): + * {{{ + * (setId:long)(shardId:int)(numKeyFields:int) + * [(keyLen:int)(keyBytes)]... + * }}} + */ +private[spark] class HashedRelationAdapter extends ShardLookupAdapter { + + private val INITIAL_RESPONSE_BUFFER_BYTES = 1 << 20 + private val alloc = NettyUtils.getSharedPooledByteBufAllocator(true, true) + + private val filterCache = new ConcurrentHashMap[Long, Option[BasePredicate]]() + private val cleanupRegistered = new AtomicBoolean(false) + + override def lookup(manager: ShardManager, reqMsg: ManagedBuffer): ManagedBuffer = { + if (cleanupRegistered.compareAndSet(false, true)) { + manager.registerCleanupCallback(setId => filterCache.remove(setId)) + } + val keysBuf = Unpooled.wrappedBuffer(reqMsg.nioByteBuffer()) + val setId = keysBuf.readLong() + val shard = keysBuf.readInt() + val numKeyFields = keysBuf.readInt() + + val keyUr = new UnsafeRow(numKeyFields) + val rel = manager.getLocalValue[HashedRelation](setId, shard).asReadOnlyCopy() + val valuesBuf = alloc.buffer(INITIAL_RESPONSE_BUFFER_BYTES) Review Comment: **[P2] Release the pooled response buffer when lookup construction fails.** A 1-MiB pooled buffer is allocated here, but any later exception during filter deserialization/evaluation, malformed-key decoding, or row serialization escapes without releasing it; `NettyShardRpcServer` catches the failure and releases only the request. For example, a valid existing-shard header followed by one extra byte passes this allocation and then throws from `readInt()`, leaking at least 1 MiB per failed request/retry. Transfer ownership only on successful return and release the response buffer on every exceptional path. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/joins/DistributedMapJoinExec.scala: ########## @@ -0,0 +1,604 @@ +/* + * 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 + +import scala.annotation.tailrec +import scala.concurrent.ExecutionContextExecutorService + +import io.netty.buffer.Unpooled + +import org.apache.spark.{SparkEnv, SparkException, TaskContext} +import org.apache.spark.network.buffer.ManagedBuffer +import org.apache.spark.rdd.RDD +import org.apache.spark.shard.ShardSetRef +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{And, AttributeSet, BindReferences, Expression, GenericInternalRow, JoinedRow, Predicate, PredicateHelper, SortOrder, UnsafeProjection, UnsafeRow} +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext +import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} +import org.apache.spark.sql.catalyst.plans.{ExistenceJoin, InnerLike, JoinType, LeftAnti, LeftOuter, LeftSemi, RightOuter} +import org.apache.spark.sql.catalyst.plans.logical.{DistributedMapJoinStrategy, JoinHint} +import org.apache.spark.sql.catalyst.plans.physical._ +import org.apache.spark.sql.errors.QueryExecutionErrors +import org.apache.spark.sql.execution.{BufferedShardRowMap, SparkPlan} +import org.apache.spark.sql.execution.adaptive.ShardQueryStageExec +import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShardExchangeExec} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.util.sketch.BloomFilter + +/** + * Physical operator for Distributed MapJoin. + * + * This strategy avoids full shuffle by building a distributed hash table service for the build + * side (medium-sized table), and the probe side performs batched RPC lookups to complete the + * join. + * + * Currently only supports: + * - Equi-join + * - BuildRight (right table as build side) + * - Explicit SQL hint: /*+distmapjoin(t(shard_count=5,replica_count=2))*/ + */ +case class DistributedMapJoinExec( + leftKeys: Seq[Expression], + rightKeys: Seq[Expression], + joinType: JoinType, + buildSide: BuildSide, + condition: Option[Expression], + left: SparkPlan, + right: SparkPlan, + hint: JoinHint, + isNullAwareAntiJoin: Boolean = false) + extends HashJoin with PredicateHelper { + + override lazy val metrics = Map( + "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) + + private val (numShards, replicaCount): (Int, Int) = { + val strategy = + (if (buildSide == BuildLeft) hint.leftHint else hint.rightHint).flatMap(_.strategy) + strategy match { + case Some(DistributedMapJoinStrategy(ns, rc)) => (ns.getOrElse(5), rc.getOrElse(1)) + case _ => (5, 1) + } + } + + @transient private lazy val (buildOnlyFilter, remainingCondition): + (Option[Expression], Option[Expression]) = { + condition match { + case Some(cond) => + val buildAttrs = AttributeSet(buildOutput) + val conjuncts = splitConjunctivePredicates(cond) + val (bo, rem) = conjuncts.partition(_.references.subsetOf(buildAttrs)) + (bo.reduceOption(And), rem.reduceOption(And)) + case None => (None, None) + } + } + + @transient override protected[this] lazy val boundCondition: InternalRow => Boolean = { + remainingCondition match { + case Some(cond) => + Predicate.create(cond, streamedPlan.output ++ buildPlan.output).eval _ + case None => + (_: InternalRow) => true + } + } + + override def supportCodegen: Boolean = false + + override def supportsColumnar: Boolean = false + + override def needCopyResult: Boolean = false + + override def requiredChildDistribution: Seq[Distribution] = { + val filter = buildOnlyFilter.map(BindReferences.bindReference(_, buildOutput)) + val filterSchema = buildOnlyFilter.map(_ => buildPlan.schema) + val sd = ShardDistribution(buildBoundKeys, numShards, replicaCount, filter, filterSchema) + buildSide match { + case BuildLeft => Seq(sd, UnspecifiedDistribution) + case BuildRight => Seq(UnspecifiedDistribution, sd) + } + } + + override def outputPartitioning: Partitioning = streamedPlan.outputPartitioning + + override def outputOrdering: Seq[SortOrder] = Nil + + override def inputRDDs(): Seq[RDD[InternalRow]] = + throw QueryExecutionErrors.executeCodePathUnsupportedError("DistributedMapJoin") + + override protected def prepareRelation(ctx: CodegenContext): HashedRelationInfo = + throw QueryExecutionErrors.executeCodePathUnsupportedError("DistributedMapJoin") + + override protected def withNewChildrenInternal( + newLeft: SparkPlan, + newRight: SparkPlan): SparkPlan = copy(left = newLeft, right = newRight) + + override protected def doExecute(): RDD[InternalRow] = { + val numOutputRows = longMetric("numOutputRows") + val setRef = resolveShardSetRef(buildPlan) + streamedPlan.execute().mapPartitionsInternal { streamedIter => + join(streamedIter, setRef.setId, numOutputRows) + } + } + + @tailrec + private def resolveShardSetRef(plan: SparkPlan): ShardSetRef = plan match { + case s: ShardExchangeExec => s.buildShardSet() + case s: ShardQueryStageExec => s.shardSetRef + case r: ReusedExchangeExec => resolveShardSetRef(r.child) + case other => + throw new IllegalStateException(s"Unexpected build plan for DistributedMapJoin: $other") + } + + private def streamedBloomFilter(setId: Long): BloomFilter = { + SparkEnv.get.shardManager.fetchBloomFilter[BloomFilter](setId)(bfInput => + BloomFilter.readFrom(bfInput)) + } + + private def join( + streamedIter: Iterator[InternalRow], + setId: Long, + numOutputRows: SQLMetric): Iterator[InternalRow] = { + + val joinedRow = new JoinedRow + + val (scanBatch, onMissing): ( + BatchMatchReader => Iterator[InternalRow], + InternalRow => Option[InternalRow]) = joinType match { + case _: InnerLike => + (innerJoinScan(_, joinedRow), (_: InternalRow) => None) + case LeftOuter | RightOuter => + val nullRow = new GenericInternalRow(buildOutput.length) + (outerJoinScan(_, joinedRow, nullRow), + (sr: InternalRow) => Some(joinedRow.withLeft(sr).withRight(nullRow))) + case LeftSemi => + (semiJoinScan(_, joinedRow), (_: InternalRow) => None) + case LeftAnti => + (antiJoinScan(_, joinedRow), (sr: InternalRow) => Some(sr)) + case _: ExistenceJoin => + val existsRow = new GenericInternalRow(Array[Any](null)) + (existenceJoinScan(_, joinedRow, existsRow), + (sr: InternalRow) => { existsRow.setBoolean(0, false); Some(joinedRow(sr, existsRow)) }) + case x => + throw new IllegalArgumentException( + s"DistributedMapJoin should not take $x as the JoinType") + } + + val iter = new LookupJoinIterator(streamedIter, setId, scanBatch, onMissing) + val resultProj = createResultProjection() + iter.map { row => + numOutputRows.add(1) + resultProj(row) + } + } + + // --------------------------------------------------------------------------- + // Scan methods: one per join type, each self-contained + // --------------------------------------------------------------------------- + + private def innerJoinScan( + reader: BatchMatchReader, + joinedRow: JoinedRow): Iterator[InternalRow] = { + new Iterator[InternalRow] { + private var _has = false + override def hasNext: Boolean = { if (!_has) _has = advance(); _has } + override def next(): InternalRow = { _has = false; joinedRow } + + private def advance(): Boolean = { + while (true) { + val br = reader.nextBuildRow() + if (br != null) { + if (boundCondition(joinedRow.withLeft(reader.curStreamed).withRight(br))) return true + } else if (!reader.advanceStreamedRow()) { + return false + } + } + false // unreachable + } + } + } + + private def outerJoinScan( + reader: BatchMatchReader, + joinedRow: JoinedRow, + nullRow: InternalRow): Iterator[InternalRow] = { + new Iterator[InternalRow] { + private var found = false + private var _has = false + override def hasNext: Boolean = { if (!_has) _has = advance(); _has } + override def next(): InternalRow = { _has = false; joinedRow } + + private def advance(): Boolean = { + while (true) { + val br = reader.nextBuildRow() + if (br != null) { + if (boundCondition(joinedRow.withLeft(reader.curStreamed).withRight(br))) { + found = true + return true + } + } else { + if (!found && reader.curStreamed != null) { + joinedRow.withLeft(reader.curStreamed).withRight(nullRow) + found = true + return true + } + if (!reader.advanceStreamedRow()) return false + found = false + } + } + false // unreachable + } + } + } + + private def semiJoinScan( + reader: BatchMatchReader, + joinedRow: JoinedRow): Iterator[InternalRow] = { + new Iterator[InternalRow] { + private var _has = false + override def hasNext: Boolean = { if (!_has) _has = advance(); _has } + override def next(): InternalRow = { _has = false; reader.curStreamed } + + private def advance(): Boolean = { + while (reader.advanceStreamedRow()) { + var br = reader.nextBuildRow() + while (br != null) { + if (boundCondition(joinedRow.withLeft(reader.curStreamed).withRight(br))) { + reader.skipRemainingBuildRows() + return true + } + br = reader.nextBuildRow() + } + } + false + } + } + } + + private def antiJoinScan( + reader: BatchMatchReader, + joinedRow: JoinedRow): Iterator[InternalRow] = { + new Iterator[InternalRow] { + private var _has = false + override def hasNext: Boolean = { if (!_has) _has = advance(); _has } + override def next(): InternalRow = { _has = false; reader.curStreamed } + + private def advance(): Boolean = { + while (reader.advanceStreamedRow()) { + if (findMatchingBuildRow(reader, joinedRow)) { + reader.skipRemainingBuildRows() + } else { + return true + } + } + false + } + } + } + + private def existenceJoinScan( + reader: BatchMatchReader, + joinedRow: JoinedRow, + existsRow: GenericInternalRow): Iterator[InternalRow] = { + new Iterator[InternalRow] { + private var _has = false + override def hasNext: Boolean = { if (!_has) _has = advance(); _has } + override def next(): InternalRow = { _has = false; joinedRow } + + private def advance(): Boolean = { + if (!reader.advanceStreamedRow()) return false + val exists = findMatchingBuildRow(reader, joinedRow) + if (exists) reader.skipRemainingBuildRows() + existsRow.setBoolean(0, exists) + joinedRow.withLeft(reader.curStreamed).withRight(existsRow) + true + } + } + } + + @tailrec + private def findMatchingBuildRow( + reader: BatchMatchReader, + joinedRow: JoinedRow): Boolean = { + val br = reader.nextBuildRow() + if (br == null) false + else if (boundCondition(joinedRow.withLeft(reader.curStreamed).withRight(br))) true + else findMatchingBuildRow(reader, joinedRow) + } + + // --------------------------------------------------------------------------- + // BatchMatchReader: reads batch response buffer, zero-copy + // --------------------------------------------------------------------------- + + private type PBatch = BufferedShardRowMap#KeyValueBatch + private type BBuffer = ManagedBuffer + + private class BatchMatchReader(batch: PBatch, buffer: BBuffer) extends AutoCloseable { + private val keyIter = batch.multiValuesIterator() + private val buf = Unpooled.wrappedBuffer(buffer.nioByteBuffer()) + private val buildUr: UnsafeRow = new UnsafeRow(buildOutput.length) + private val advanceRead = UnsafeRowBufCodec.makeAdvanceRead(buildUr, buf) + private var streamedIter: java.util.Iterator[UnsafeRow] = _ + private var buildRowsStart: Int = -1 + private var atSentinel = true + var curStreamed: UnsafeRow = _ + + assert(batch.getSetId == buf.readLong(), "setId mismatch") + assert(batch.getShard == buf.readInt(), "shardId mismatch") + + def advanceStreamedRow(): Boolean = { + if (streamedIter != null && streamedIter.hasNext) { + curStreamed = streamedIter.next() + rewindBuildRows() + true + } else { + skipRemainingBuildRows() + if (!keyIter.hasNext) return false + streamedIter = keyIter.next() + buildRowsStart = buf.readerIndex() + atSentinel = false + if (!streamedIter.hasNext) advanceStreamedRow() + else { curStreamed = streamedIter.next(); true } + } + } + + def nextBuildRow(): UnsafeRow = { + if (atSentinel) return null + val blen = buf.readInt() + if (blen == 0) { atSentinel = true; null } + else advanceRead(blen) + } + + def skipRemainingBuildRows(): Unit = { + while (!atSentinel) { + val blen = buf.readInt() + if (blen == 0) atSentinel = true + else buf.readerIndex(buf.readerIndex() + blen) + } + } + + private def rewindBuildRows(): Unit = { + buf.readerIndex(buildRowsStart) + atSentinel = false + } + + override def close(): Unit = { + batch.release() + buffer.release() + } + } + + // --------------------------------------------------------------------------- + // LookupJoinIterator: async pipeline management + // --------------------------------------------------------------------------- + + private class LookupJoinIterator( + streamedIter: Iterator[InternalRow], + setId: Long, + scanBatch: BatchMatchReader => Iterator[InternalRow], + onMissing: InternalRow => Option[InternalRow]) + extends Iterator[InternalRow] { + + private val maxInFlightNum = conf.distributedMapJoinMaxInFlightNum + private val keyGenerator: UnsafeProjection = UnsafeProjection.create(streamedBoundKeys) + private val valueGenerator: UnsafeProjection = UnsafeProjection.create(streamedPlan.schema) + private val shardGenerator = + UnsafeProjection.create( + HashPartitioning(streamedBoundKeys, numShards).partitionIdExpression :: Nil) + + private val bloom = streamedBloomFilter(setId) + private val probeUr: UnsafeRow = new UnsafeRow(streamedPlan.schema.length) + @volatile private var cancelled = false + + private val bufferedMap = { + val mm = TaskContext.get().taskMemoryManager() + val maxBatchSize = conf.distributedMapJoinMaxBatchSize + val map = + new BufferedShardRowMap( + mm, + setId, + numShards, + streamedBoundKeys.length, + probeUr, + maxBatchSize) + TaskContext.get().addTaskCompletionListener[Unit] { _ => + cancelled = true + if (currentReader != null) { + currentReader.close() + currentReader = null + } + var item = lookupQueue.poll() + while (item != null) { + item match { + case LookupSuccess(batch, buffer) => + batch.release() + buffer.release() + case _ => + } + item = lookupQueue.poll() + } + map.free() + } + map + } + + private sealed trait Lookup + private case class LookupSuccess(batch: PBatch, buffer: BBuffer) extends Lookup + private case class LookupFailure(batch: PBatch, cause: Throwable) extends Lookup + + private val lookupQueue = new util.concurrent.LinkedBlockingQueue[Lookup] + + private var inputExhausted = false + private var prepared = false + private var nextRowVal: InternalRow = _ + private var numInFlight = 0 + private var currentReader: BatchMatchReader = _ + private var currentBatchIter: Iterator[InternalRow] = _ + + override def hasNext: Boolean = { + if (!prepared) { + processNext() + } + prepared + } + + override def next(): InternalRow = { + if (!prepared && !hasNext) { + throw QueryExecutionErrors.noSuchElementExceptionError() + } + prepared = false + nextRowVal + } + + private def prepareNextRow(row: InternalRow): Unit = { + nextRowVal = row + prepared = true + } + + private def processNext(): Unit = { + if (currentBatchIter != null) { + iterateLookup() + } + + var hasLookup = true + while (!prepared && hasLookup) { + val ele = lookupQueue.poll() + if (ele == null) { + hasLookup = false + } else { + pollLookup(ele) + if (currentBatchIter != null) { + iterateLookup() + } + } + } + + while (!prepared && !inputExhausted && streamedIter.hasNext) { + val streamedRow = streamedIter.next() + val keyUr = keyGenerator(streamedRow) + if (keyUr.anyNull || !bloom.mightContain(keyUr.getBytes)) { + onMissing(streamedRow).foreach(prepareNextRow) + } else { + val shardId = shardGenerator(streamedRow).getInt(0) + val valueUr = streamedRow match { + case ur: UnsafeRow => ur + case r => valueGenerator(r) + } + processLookup(shardId, keyUr, valueUr) + } + } + + if (!prepared) { + if (!inputExhausted) { + inputExhausted = true + flushLookup(bufferedMap.tailingIterator()) + } + while (!prepared && numInFlight > 0) { + pollLookup(lookupQueue.poll(200, util.concurrent.TimeUnit.MILLISECONDS)) + if (currentBatchIter != null) { + iterateLookup() + } + } + } + } + + private def processLookup(shardId: Int, keyUr: UnsafeRow, valueUr: UnsafeRow): Unit = { + bufferedMap.putRow( + shardId, + keyUr.getBaseObject, + keyUr.getBaseOffset, + keyUr.getSizeInBytes, + keyUr.hashCode(), + valueUr.getBaseObject, + valueUr.getBaseOffset, + valueUr.getSizeInBytes) + + if (bufferedMap.hasPending) { + flushLookup(bufferedMap.pendingIterator()) + } + + while (numInFlight >= maxInFlightNum) { + pollLookup(lookupQueue.poll(200, util.concurrent.TimeUnit.MILLISECONDS)) + if (currentBatchIter != null) { + iterateLookup() + } + } + } + + private def flushLookup[T <: PBatch](iter: util.Iterator[T]): Unit = { + val manager = SparkEnv.get.shardManager + implicit val ec: ExecutionContextExecutorService = manager.lookupEc + while (iter.hasNext && !cancelled) { + while (numInFlight >= maxInFlightNum && !cancelled) { Review Comment: **[P1] Stop flushing once throttling prepares an output row.** With `maxInFlightNum=1` and three populated shard tail batches A/B/C, this throttle consumes A and `iterateLookup()` sets `prepared=true`/`nextRowVal=A`, but `flushLookup` continues, submits B, then overwrites `currentReader`, `currentBatchIter`, and `nextRowVal` while consuming B. A's row and any remaining duplicate matches silently disappear, and its reader is abandoned; the default limit of eight reproduces with ten populated shards. Preserve the prepared row/current reader and resume flushing only after the caller consumes them. -- 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]
