diff --git a/pixels-common/src/main/java/io/pixelsdb/pixels/common/task/TaskQueue.java b/pixels-common/src/main/java/io/pixelsdb/pixels/common/task/TaskQueue.java index a36fe38179..1609e7cb10 100644 --- a/pixels-common/src/main/java/io/pixelsdb/pixels/common/task/TaskQueue.java +++ b/pixels-common/src/main/java/io/pixelsdb/pixels/common/task/TaskQueue.java @@ -60,6 +60,10 @@ public TaskQueue(Collection tasks) public boolean offerAllPending(Collection tasks) { checkPendingTasks(tasks); + synchronized (this.allTasks) + { + this.allTasks.addAll(tasks); + } return this.pendingQueue.addAll(tasks); } @@ -80,6 +84,11 @@ private void checkPendingTasks(Collection tasks) */ public boolean offerPending(E task) { + checkPendingTasks(Collections.singletonList(task)); + synchronized (this.allTasks) + { + this.allTasks.add(task); + } return this.pendingQueue.offer(task); } @@ -150,6 +159,53 @@ public E removeNextExpired() public List getAllTasks() { - return ImmutableList.copyOf(this.allTasks); + synchronized (this.allTasks) + { + return ImmutableList.copyOf(this.allTasks); + } + } + + public int getTotalTaskCount() + { + synchronized (this.allTasks) + { + return this.allTasks.size(); + } + } + + public int getPendingTaskCount() + { + return this.pendingQueue.size(); + } + + public int getRunningTaskCount() + { + return this.runningTasks.size(); + } + + public int getCompletedTaskCount() + { + return getTaskStatusCount(Task.Status.COMPLETE); + } + + public int getFailedTaskCount() + { + return getTaskStatusCount(Task.Status.FAILED); + } + + public int getTaskStatusCount(Task.Status status) + { + synchronized (this.allTasks) + { + int count = 0; + for (E task : this.allTasks) + { + if (task.getStatus() == status) + { + count++; + } + } + return count; + } } } diff --git a/pixels-common/src/main/java/io/pixelsdb/pixels/common/task/Worker.java b/pixels-common/src/main/java/io/pixelsdb/pixels/common/task/Worker.java index 094ec56f39..cf9b0033c4 100644 --- a/pixels-common/src/main/java/io/pixelsdb/pixels/common/task/Worker.java +++ b/pixels-common/src/main/java/io/pixelsdb/pixels/common/task/Worker.java @@ -88,6 +88,21 @@ public void terminate() } } + /** + * Whether this worker has explicitly finished or been terminated. + * + * This is different from {@link #isAlive()}: a lease can expire while a + * physical invocation is still running and must still be considered by + * coordinator admission or drain control. + */ + public boolean isTerminated() + { + synchronized (this.lease) + { + return this.terminated; + } + } + /** * Extent the lease of this worker if the worker is alive. This method should only be called on the coordinator. * @return the new start time (milliseconds since the Unix epoch) of the extended lease diff --git a/pixels-common/src/main/java/io/pixelsdb/pixels/common/turbo/WorkerType.java b/pixels-common/src/main/java/io/pixelsdb/pixels/common/turbo/WorkerType.java index 64eb8f0c59..f352a373f5 100644 --- a/pixels-common/src/main/java/io/pixelsdb/pixels/common/turbo/WorkerType.java +++ b/pixels-common/src/main/java/io/pixelsdb/pixels/common/turbo/WorkerType.java @@ -30,10 +30,13 @@ public enum WorkerType UNKNOWN, // The first enum value is the default value. SCAN, SCAN_STREAMING, PARTITION, PARTITION_STREAMING, + PARTITION_S3QS, BROADCAST_JOIN, BROADCAST_JOIN_STREAMING, BROADCAST_CHAIN_JOIN, BROADCAST_CHAIN_JOIN_STREAMING, PARTITIONED_JOIN, PARTITIONED_JOIN_STREAMING, + PARTITIONED_JOIN_S3QS, PARTITIONED_CHAIN_JOIN, PARTITIONED_CHAIN_JOIN_STREAMING, + PARTITIONED_CHAIN_JOIN_S3QS, SORT, SORTED_JOIN, AGGREGATION; diff --git a/pixels-common/src/main/resources/pixels.properties b/pixels-common/src/main/resources/pixels.properties index 3e86fb9680..cf914b3fc0 100644 --- a/pixels-common/src/main/resources/pixels.properties +++ b/pixels-common/src/main/resources/pixels.properties @@ -197,6 +197,8 @@ aggr.partition.size.rows=1280000 ### pixels-turbo - query execution ### executor.input.storage.scheme=s3 executor.intermediate.storage.scheme=s3 +# the storage scheme dedicated to the shuffle transport path, empty means disabled +executor.shuffle.storage.scheme= executor.intermediate.folder=/pixels-turbo/intermediate/ executor.output.storage.scheme=output-storage-scheme-dummy executor.output.folder=output-folder-dummy diff --git a/pixels-planner/pom.xml b/pixels-planner/pom.xml index 7d695d8a86..33ef52d5c0 100644 --- a/pixels-planner/pom.xml +++ b/pixels-planner/pom.xml @@ -31,6 +31,10 @@ io.pixelsdb pixels-executor + + io.pixelsdb + pixels-storage-s3qs + @@ -57,7 +61,6 @@ io.pixelsdb pixels-storage-s3 - test true @@ -96,4 +99,4 @@ - \ No newline at end of file + diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/PixelsPlanner.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/PixelsPlanner.java index f4d9b68481..c208496fa8 100644 --- a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/PixelsPlanner.java +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/PixelsPlanner.java @@ -36,6 +36,8 @@ import io.pixelsdb.pixels.executor.join.JoinAlgorithm; import io.pixelsdb.pixels.executor.join.JoinType; import io.pixelsdb.pixels.executor.predicate.TableScanFilter; +import io.pixelsdb.pixels.planner.coordinate.CoordinatedPlanExecution; +import io.pixelsdb.pixels.planner.coordinate.PlanCoordinatorFactory; import io.pixelsdb.pixels.planner.plan.PlanOptimizer; import io.pixelsdb.pixels.planner.plan.logical.*; import io.pixelsdb.pixels.planner.plan.logical.Table; @@ -63,6 +65,7 @@ public class PixelsPlanner private static final StorageInfo InputStorageInfo; private static final StorageInfo IntermediateStorageInfo; private static final StorageInfo StreamStorageInfo; + private static final StorageInfo ShuffleStorageInfo; private static final String IntermediateFolder; private static final int IntraWorkerParallelism; private static final ExchangeMethod EnabledExchangeMethod; @@ -92,9 +95,9 @@ public class PixelsPlanner ConfigFactory.Instance().getProperty("executor.input.storage.scheme")); InputStorageInfo = StorageInfoBuilder.BuildFromConfig(inputStorageScheme); - Storage.Scheme interStorageScheme = EnabledExchangeMethod == ExchangeMethod.batch ? - Storage.Scheme.from(ConfigFactory.Instance().getProperty("executor.intermediate.storage.scheme")) : - Storage.Scheme.valueOf("httpstream"); + Storage.Scheme interStorageScheme = EnabledExchangeMethod == ExchangeMethod.stream ? + Storage.Scheme.valueOf("httpstream") : + Storage.Scheme.from(ConfigFactory.Instance().getProperty("executor.intermediate.storage.scheme")); IntermediateStorageInfo = StorageInfoBuilder.BuildFromConfig(interStorageScheme); String interStorageFolder = ConfigFactory.Instance().getProperty("executor.intermediate.folder"); if (!interStorageFolder.endsWith("/")) @@ -106,6 +109,7 @@ public class PixelsPlanner .getProperty("executor.intra.worker.parallelism")); StreamStorageInfo = StorageInfoBuilder.BuildFromConfig(Storage.Scheme.httpstream); + ShuffleStorageInfo = getShuffleStorageInfo(); } /** @@ -161,6 +165,15 @@ else if (this.rootTable.getTableType() == Table.TableType.AGGREGATED) } } + /** + * Create a physical plan execution whose coordinator is registered before + * any runtime worker can be invoked. + */ + public CoordinatedPlanExecution createPlanExecution() throws IOException, MetadataException + { + return PlanCoordinatorFactory.Instance().createPlanExecution(this.transId, getRootOperator()); + } + public double getScanSize() { return scanSize; @@ -171,6 +184,84 @@ public static String getIntermediateFolderForTrans(long transId) return IntermediateFolder + transId + "/"; } + private static StorageInfo getShuffleStorageInfo() + { + String shuffleStorageScheme = ConfigFactory.Instance().getProperty("executor.shuffle.storage.scheme"); + if (shuffleStorageScheme == null || shuffleStorageScheme.trim().isEmpty()) + { + return null; + } + + Storage.Scheme scheme = Storage.Scheme.from(shuffleStorageScheme.trim()); + checkArgument(scheme == Storage.Scheme.s3qs, + "executor.shuffle.storage.scheme only supports s3qs currently: %s", shuffleStorageScheme); + return StorageInfoBuilder.BuildFromConfig(scheme); + } + + private static boolean isS3QSShuffleEnabled() + { + return EnabledExchangeMethod == ExchangeMethod.s3qs && + ShuffleStorageInfo != null && + ShuffleStorageInfo.getScheme() == Storage.Scheme.s3qs; + } + + private static boolean isStreamExchange() + { + return EnabledExchangeMethod == ExchangeMethod.stream; + } + + private PartitionedJoinOperator createPartitionedJoinOperator( + String name, List smallPartitionInputs, List largePartitionInputs, + List joinInputs, JoinAlgorithm joinAlgo) + { + switch (EnabledExchangeMethod) + { + case batch: + return new PartitionedJoinBatchOperator(name, smallPartitionInputs, largePartitionInputs, + joinInputs, joinAlgo); + case stream: + return new PartitionedJoinStreamOperator(name, smallPartitionInputs, largePartitionInputs, + joinInputs, joinAlgo); + case s3qs: + checkArgument(ShuffleStorageInfo != null && ShuffleStorageInfo.getScheme() == Storage.Scheme.s3qs, + "executor.exchange.method=s3qs requires executor.shuffle.storage.scheme=s3qs"); + return new PartitionedJoinS3QSOperator(name, smallPartitionInputs, largePartitionInputs, + joinInputs, joinAlgo); + default: + throw new UnsupportedOperationException("unsupported exchange method '" + EnabledExchangeMethod + "'"); + } + } + + private ShuffleInfo createShuffleInfo(Table inputTable, String outputBase, int numPartitions, int producerCount) + { + String objectPathPrefix = outputBase.endsWith("/") ? outputBase : outputBase + "/"; + String shuffleSuffix = Integer.toUnsignedString(objectPathPrefix.hashCode()); + String shuffleId = transId + "/" + inputTable.getSchemaName() + "/" + inputTable.getTableName() + + "/" + shuffleSuffix; + ImmutableList.Builder queues = ImmutableList.builderWithExpectedSize(numPartitions); + for (int i = 0; i < numPartitions; ++i) + { + queues.add(new ShuffleQueueInfo(i, buildShuffleQueueName(shuffleSuffix, i), null)); + } + return new ShuffleInfo(shuffleId, ShuffleStorageInfo, objectPathPrefix, numPartitions, + producerCount, numPartitions, 20, queues.build()); + } + + private String buildShuffleQueueName(String shuffleSuffix, int partitionId) + { + return "pixels-shuffle-" + transId + "-" + shuffleSuffix + "-" + partitionId; + } + + private static ShuffleInfo getShuffleInfo(List partitionInputs) + { + if (partitionInputs == null || partitionInputs.isEmpty()) + { + return null; + } + OutputInfo output = partitionInputs.get(0).getOutput(); + return output == null ? null : output.getShuffleInfo(); + } + private ScanOperator getScanOperator(BaseTable scanTable) throws IOException, MetadataException { final String intermediateBase = getIntermediateFolderForTrans(transId) + @@ -208,9 +299,9 @@ private ScanOperator getScanOperator(BaseTable scanTable) throws IOException, Me scanInput.setOutput(new OutputInfo(folderName, IntermediateStorageInfo, true)); scanInputsBuilder.add(scanInput); } - return EnabledExchangeMethod == ExchangeMethod.batch ? - new ScanBatchOperator(scanTable.getTableName(), scanInputsBuilder.build()) : - new ScanStreamOperator(scanTable.getTableName(), scanInputsBuilder.build()); + return isStreamExchange() ? + new ScanStreamOperator(scanTable.getTableName(), scanInputsBuilder.build()) : + new ScanBatchOperator(scanTable.getTableName(), scanInputsBuilder.build()); } private AggregationOperator getAggregationOperator(AggregatedTable aggregatedTable) @@ -344,10 +435,10 @@ else if (originTable.getTableType() == Table.TableType.JOINED) finalAggrInputsBuilder.add(finalAggrInput); } - AggregationOperator aggregationOperator = EnabledExchangeMethod == ExchangeMethod.batch ? - new AggregationBatchOperator(aggregatedTable.getTableName(), - finalAggrInputsBuilder.build(), scanInputsBuilder.build()) : + AggregationOperator aggregationOperator = isStreamExchange() ? new AggregationStreamOperator(aggregatedTable.getTableName(), + finalAggrInputsBuilder.build(), scanInputsBuilder.build()) : + new AggregationBatchOperator(aggregatedTable.getTableName(), finalAggrInputsBuilder.build(), scanInputsBuilder.build()); aggregationOperator.setChild(joinOperator); @@ -448,10 +539,10 @@ private JoinOperator getMultiPipelineJoinOperator(JoinedTable joinedTable, Optio joinInputs.add(complete); } - SingleStageJoinOperator joinOperator = EnabledExchangeMethod == ExchangeMethod.batch ? - new SingleStageJoinBatchOperator(joinedTable.getTableName(), true, - joinInputs.build(), JoinAlgorithm.BROADCAST_CHAIN) : + SingleStageJoinOperator joinOperator = isStreamExchange() ? new SingleStageJoinStreamOperator(joinedTable.getTableName(), true, + joinInputs.build(), JoinAlgorithm.BROADCAST_CHAIN) : + new SingleStageJoinBatchOperator(joinedTable.getTableName(), true, joinInputs.build(), JoinAlgorithm.BROADCAST_CHAIN); // The right operator must be set as the large child. joinOperator.setLargeChild(rightOperator); @@ -491,15 +582,11 @@ else if (rightOperator.getJoinAlgo() == JoinAlgorithm.PARTITIONED) joinInputs.add(chainJoinInput); } - PartitionedJoinOperator joinOperator = EnabledExchangeMethod == ExchangeMethod.batch ? - new PartitionedJoinBatchOperator(joinedTable.getTableName(), - rightJoinOperator.getSmallPartitionInputs(), - rightJoinOperator.getLargePartitionInputs(), - joinInputs.build(), JoinAlgorithm.PARTITIONED_CHAIN) : - new PartitionedJoinStreamOperator(joinedTable.getTableName(), - rightJoinOperator.getSmallPartitionInputs(), - rightJoinOperator.getLargePartitionInputs(), - joinInputs.build(), JoinAlgorithm.PARTITIONED_CHAIN); + PartitionedJoinOperator joinOperator = createPartitionedJoinOperator( + joinedTable.getTableName(), + rightJoinOperator.getSmallPartitionInputs(), + rightJoinOperator.getLargePartitionInputs(), + joinInputs.build(), JoinAlgorithm.PARTITIONED_CHAIN); // Set the children of the right operator as the children of the current join operator. joinOperator.setSmallChild(rightJoinOperator.getSmallChild()); joinOperator.setLargeChild(rightJoinOperator.getLargeChild()); @@ -543,11 +630,8 @@ else if (joinAlgo == JoinAlgorithm.PARTITIONED) List joinInputs = getPartitionedJoinInputs( joinedTable, parent, numPartition, leftTableInfo, rightTableInfo, null, null); - PartitionedJoinOperator joinOperator = EnabledExchangeMethod == ExchangeMethod.batch ? - new PartitionedJoinBatchOperator(joinedTable.getTableName(), - null, null, joinInputs, joinAlgo) : - new PartitionedJoinStreamOperator(joinedTable.getTableName(), - null, null, joinInputs, joinAlgo); + PartitionedJoinOperator joinOperator = createPartitionedJoinOperator( + joinedTable.getTableName(), null, null, joinInputs, joinAlgo); joinOperator.setSmallChild(leftOperator); joinOperator.setLargeChild(rightOperator); @@ -661,10 +745,10 @@ private JoinOperator getJoinOperator(JoinedTable joinedTable, Optional getPartitionInputs(Table inputTable, List partitionInputsBuilder = ImmutableList.builder(); int outputId = 0; + int producerCount = (inputSplits.size() + IntraWorkerParallelism - 1) / IntraWorkerParallelism; + ShuffleInfo shuffleInfo = isS3QSShuffleEnabled() ? + createShuffleInfo(inputTable, outputBase, numPartition, producerCount) : null; for (int i = 0; i < inputSplits.size();) { PartitionInput partitionInput = new PartitionInput(); @@ -1427,12 +1504,17 @@ private List getPartitionInputs(Table inputTable, List. + */ +package io.pixelsdb.pixels.planner.coordinate; + +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.planner.plan.physical.Operator; +import io.pixelsdb.pixels.planner.plan.physical.OperatorExecutor.OutputCollection; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutionException; +import java.io.IOException; +import java.io.UncheckedIOException; + +import static com.google.common.base.Preconditions.checkState; +import static java.util.Objects.requireNonNull; + +/** + * Owns coordinator registration for one physical query-plan execution. + * + * Callers should use try-with-resources so exceptional execution paths also + * remove the query coordinator. Normal output collection closes it + * automatically after the operator tree has completed. + */ +public class CoordinatedPlanExecution implements AutoCloseable +{ + private final long transId; + private final Operator rootOperator; + private final PlanCoordinator planCoordinator; + private final PlanCoordinatorFactory coordinatorFactory; + private final ShuffleResourceLifecycle shuffleResourceLifecycle; + private CompletableFuture[]> executionFuture; + private boolean executed; + private boolean resourcesPrepared; + private boolean closed; + + CoordinatedPlanExecution(long transId, Operator rootOperator, PlanCoordinator planCoordinator, + PlanCoordinatorFactory coordinatorFactory, + ShuffleResourceLifecycle shuffleResourceLifecycle) + { + this.transId = transId; + this.rootOperator = requireNonNull(rootOperator, "rootOperator is null"); + this.planCoordinator = requireNonNull(planCoordinator, "planCoordinator is null"); + this.coordinatorFactory = requireNonNull(coordinatorFactory, "coordinatorFactory is null"); + this.shuffleResourceLifecycle = + requireNonNull(shuffleResourceLifecycle, "shuffleResourceLifecycle is null"); + } + + public synchronized CompletableFuture[]> execute() + { + checkState(!closed, "plan execution is closed"); + checkState(!executed, "plan execution has already started"); + executed = true; + try + { + shuffleResourceLifecycle.prepare(planCoordinator.getShuffleInfos()); + resourcesPrepared = true; + executionFuture = rootOperator.execute(); + return executionFuture; + } + catch (IOException e) + { + CompletionException failure = + new CompletionException("failed to prepare query shuffle resources", e); + closeAfterFailure(failure); + throw failure; + } + catch (RuntimeException | Error e) + { + closeAfterFailure(e); + throw e; + } + } + + public synchronized OutputCollection collectOutputs() throws ExecutionException, InterruptedException + { + checkState(!closed, "plan execution is closed"); + checkState(executed, "plan execution has not started"); + OutputCollection outputs; + try + { + // Ensure the final stage has been started and its worker futures + // have been installed before traversing the operator output tree. + executionFuture.get(); + outputs = rootOperator.collectOutputs(); + } + catch (ExecutionException | InterruptedException | RuntimeException | Error e) + { + closeAfterFailure(e); + throw e; + } + close(); + return outputs; + } + + public PlanCoordinator getPlanCoordinator() + { + return planCoordinator; + } + + @Override + public synchronized void close() + { + if (closed) + { + return; + } + closed = true; + IOException cleanupFailure = null; + try + { + if (resourcesPrepared) + { + shuffleResourceLifecycle.cleanup(); + } + } + catch (IOException e) + { + cleanupFailure = e; + } + finally + { + coordinatorFactory.removePlanCoordinator(transId, planCoordinator); + } + if (cleanupFailure != null) + { + throw new UncheckedIOException("failed to clean query shuffle resources", cleanupFailure); + } + } + + private void closeAfterFailure(Throwable failure) + { + try + { + close(); + } + catch (RuntimeException cleanupFailure) + { + failure.addSuppressed(cleanupFailure); + } + } +} diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/CoordinatorEndpoint.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/CoordinatorEndpoint.java new file mode 100644 index 0000000000..0ae17d00e5 --- /dev/null +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/CoordinatorEndpoint.java @@ -0,0 +1,61 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of the License, + * or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public License + * along with Pixels. If not, see . + */ +package io.pixelsdb.pixels.planner.coordinate; + +import io.pixelsdb.pixels.common.utils.ConfigFactory; + +import static com.google.common.base.Preconditions.checkArgument; +import static java.util.Objects.requireNonNull; + +/** + * Network endpoint used by runtime workers to call back to the coordinator. + */ +public class CoordinatorEndpoint +{ + private final String host; + private final int port; + + public CoordinatorEndpoint(String host, int port) + { + this.host = requireNonNull(host, "host is null").trim(); + checkArgument(!this.host.isEmpty(), "host is empty"); + checkArgument(port > 0 && port <= 65535, "port is out of range"); + this.port = port; + } + + public static CoordinatorEndpoint fromConfig() + { + ConfigFactory config = ConfigFactory.Instance(); + String host = config.getProperty("worker.coordinate.server.host"); + String port = config.getProperty("worker.coordinate.server.port"); + checkArgument(port != null && !port.trim().isEmpty(), + "worker.coordinate.server.port is empty"); + return new CoordinatorEndpoint(host, Integer.parseInt(port)); + } + + public String getHost() + { + return host; + } + + public int getPort() + { + return port; + } +} diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/InvokerStageWorkerLauncher.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/InvokerStageWorkerLauncher.java new file mode 100644 index 0000000000..e5b36d2b5e --- /dev/null +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/InvokerStageWorkerLauncher.java @@ -0,0 +1,38 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of the License, + * or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public License + * along with Pixels. If not, see . + */ +package io.pixelsdb.pixels.planner.coordinate; + +import io.pixelsdb.pixels.common.turbo.InvokerFactory; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; + +import java.util.concurrent.CompletableFuture; + +/** + * Production stage worker launcher backed by the configured platform invoker. + */ +public class InvokerStageWorkerLauncher implements StageWorkerLauncher +{ + @Override + public CompletableFuture launch(WorkerType workerType, StageWorkerInput input) + { + return InvokerFactory.Instance().getInvoker(workerType).invoke(input); + } +} diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/PlanCoordinator.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/PlanCoordinator.java index dbb4d7cb9f..d0a1d2ac81 100644 --- a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/PlanCoordinator.java +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/PlanCoordinator.java @@ -19,8 +19,18 @@ */ package io.pixelsdb.pixels.planner.coordinate; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleQueueInfo; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; +import java.util.Objects; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicInteger; import static com.google.common.base.Preconditions.checkArgument; @@ -47,14 +57,43 @@ public class PlanCoordinator * It is only modified during coordinator initialization, thus there is no read-write conflict. */ private final Map stageDependencies = new HashMap<>(); + /** + * Runtime worker control is registered only for stages whose physical + * workers pull coordinator tasks, currently the S3QS shuffle stages. + */ + private final Map stageRuntimeControllers = new HashMap<>(); + /** + * Query-owned shuffle edges keyed by their query-unique shuffle id. + */ + private final Map shuffleInfos = new HashMap<>(); + private final CoordinatorEndpoint coordinatorEndpoint; + private final StageWorkerLauncher stageWorkerLauncher; /** * The assigner of stage id. */ private final AtomicInteger stageIdAssigner = new AtomicInteger(0); public PlanCoordinator(long transId) + { + this(transId, CoordinatorEndpoint.fromConfig(), new InvokerStageWorkerLauncher()); + } + + public PlanCoordinator(long transId, StageWorkerLauncher stageWorkerLauncher) + { + this(transId, CoordinatorEndpoint.fromConfig(), stageWorkerLauncher); + } + + public PlanCoordinator(long transId, CoordinatorEndpoint coordinatorEndpoint) + { + this(transId, coordinatorEndpoint, new InvokerStageWorkerLauncher()); + } + + public PlanCoordinator(long transId, CoordinatorEndpoint coordinatorEndpoint, + StageWorkerLauncher stageWorkerLauncher) { this.transId = transId; + this.coordinatorEndpoint = requireNonNull(coordinatorEndpoint, "coordinatorEndpoint is null"); + this.stageWorkerLauncher = requireNonNull(stageWorkerLauncher, "stageWorkerLauncher is null"); } /** @@ -85,6 +124,48 @@ public StageCoordinator getStageCoordinator(int stageId) return this.stageCoordinators.get(stageId); } + /** + * Register the physical execution description for a queued stage. Logical + * tasks must already exist in the corresponding StageCoordinator. + */ + public void addStageRuntimeController(StageExecutionDescriptor descriptor) + { + requireNonNull(descriptor, "descriptor is null"); + int stageId = descriptor.getStageId(); + StageCoordinator stageCoordinator = requireNonNull(this.stageCoordinators.get(stageId), + "stage coordinator is not found"); + checkArgument(!this.stageRuntimeControllers.containsKey(stageId), + "stage runtime controller already exists"); + this.stageRuntimeControllers.put(stageId, + new StageRuntimeController(stageCoordinator, descriptor, this.stageWorkerLauncher)); + } + + public StageRuntimeController getStageRuntimeController(int stageId) + { + return this.stageRuntimeControllers.get(stageId); + } + + /** + * Start the first worker set for a queued stage. The initial target is + * bounded by its current pending logical task count. + */ + public List> activateStage(int stageId) + { + StageCoordinator stageCoordinator = requireNonNull(this.stageCoordinators.get(stageId), + "stage coordinator is not found"); + return scaleStage(stageId, stageCoordinator.getPendingTaskCount()); + } + + /** + * Manually change the desired runtime worker count for a stage. + */ + public List> scaleStage(int stageId, int targetWorkerCount) + { + StageRuntimeController runtimeController = requireNonNull(this.stageRuntimeControllers.get(stageId), + "stage runtime controller is not found"); + return runtimeController.scaleTo(targetWorkerCount); + } + /** * Get this stage's dependency on the parent (downstream) stage. * @param stageId the id of this stage @@ -103,6 +184,36 @@ public long getTransId() return transId; } + public CoordinatorEndpoint getCoordinatorEndpoint() + { + return coordinatorEndpoint; + } + + /** + * Register one shuffle edge with this query. Producer and consumer inputs + * may reference separate but equivalent metadata objects, so duplicates + * are accepted only when all resource-defining fields agree. + */ + public void addShuffleInfo(ShuffleInfo shuffleInfo) + { + requireNonNull(shuffleInfo, "shuffleInfo is null"); + String shuffleId = requireNonNull(shuffleInfo.getShuffleId(), "shuffleId is null"); + checkArgument(!shuffleId.trim().isEmpty(), "shuffleId is empty"); + ShuffleInfo previous = this.shuffleInfos.get(shuffleId); + if (previous == null) + { + this.shuffleInfos.put(shuffleId, shuffleInfo); + return; + } + checkArgument(sameShuffleResource(previous, shuffleInfo), + "conflicting shuffle metadata for shuffleId " + shuffleId); + } + + public Collection getShuffleInfos() + { + return Collections.unmodifiableCollection(new ArrayList<>(this.shuffleInfos.values())); + } + /** * Assign an id for a stage. This should only be called in * {@link io.pixelsdb.pixels.planner.plan.physical.Operator#initPlanCoordinator(PlanCoordinator, int, boolean)} @@ -113,4 +224,37 @@ public int assignStageId() { return this.stageIdAssigner.getAndIncrement(); } + + private static boolean sameShuffleResource(ShuffleInfo left, ShuffleInfo right) + { + if (!Objects.equals(left.getObjectPathPrefix(), right.getObjectPathPrefix()) || + left.getNumPartitions() != right.getNumPartitions() || + left.getStorageInfo() == null || right.getStorageInfo() == null || + left.getStorageInfo().getScheme() != right.getStorageInfo().getScheme()) + { + return false; + } + List leftQueues = left.getQueues(); + List rightQueues = right.getQueues(); + if (leftQueues == null || rightQueues == null || leftQueues.size() != rightQueues.size()) + { + return false; + } + Map rightByPartition = new HashMap<>(); + for (ShuffleQueueInfo queue : rightQueues) + { + rightByPartition.put(queue.getPartitionId(), queue); + } + for (ShuffleQueueInfo leftQueue : leftQueues) + { + ShuffleQueueInfo rightQueue = rightByPartition.get(leftQueue.getPartitionId()); + if (rightQueue == null || + !Objects.equals(leftQueue.getQueueName(), rightQueue.getQueueName()) || + !Objects.equals(leftQueue.getQueueUrl(), rightQueue.getQueueUrl())) + { + return false; + } + } + return true; + } } diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/PlanCoordinatorFactory.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/PlanCoordinatorFactory.java index 71e9470305..775556be02 100644 --- a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/PlanCoordinatorFactory.java +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/PlanCoordinatorFactory.java @@ -24,6 +24,9 @@ import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import static com.google.common.base.Preconditions.checkState; +import static java.util.Objects.requireNonNull; + /** * The factory to create and manage the plan coordinators of queries. * @author hank @@ -56,12 +59,62 @@ private PlanCoordinatorFactory() */ public PlanCoordinator createPlanCoordinator(long transId, Operator planRootOperator) { - PlanCoordinator planCoordinator = new PlanCoordinator(transId); + return createPlanCoordinator(transId, planRootOperator, CoordinatorEndpoint.fromConfig()); + } + + public PlanCoordinator createPlanCoordinator(long transId, Operator planRootOperator, + CoordinatorEndpoint coordinatorEndpoint) + { + return createPlanCoordinator(transId, planRootOperator, coordinatorEndpoint, + new InvokerStageWorkerLauncher()); + } + + public PlanCoordinator createPlanCoordinator(long transId, Operator planRootOperator, + CoordinatorEndpoint coordinatorEndpoint, + StageWorkerLauncher stageWorkerLauncher) + { + requireNonNull(planRootOperator, "planRootOperator is null"); + checkState(!this.transIdToPlanCoordinator.containsKey(transId), + "plan coordinator already exists for transaction %s", transId); + PlanCoordinator planCoordinator = new PlanCoordinator(transId, coordinatorEndpoint, + requireNonNull(stageWorkerLauncher, "stageWorkerLauncher is null")); planRootOperator.initPlanCoordinator(planCoordinator, -1, false); - this.transIdToPlanCoordinator.put(transId, planCoordinator); + PlanCoordinator previous = this.transIdToPlanCoordinator.putIfAbsent(transId, planCoordinator); + checkState(previous == null, "plan coordinator already exists for transaction %s", transId); return planCoordinator; } + public CoordinatedPlanExecution createPlanExecution(long transId, Operator planRootOperator) + { + return createPlanExecution(transId, planRootOperator, CoordinatorEndpoint.fromConfig()); + } + + public CoordinatedPlanExecution createPlanExecution(long transId, Operator planRootOperator, + CoordinatorEndpoint coordinatorEndpoint) + { + return createPlanExecution(transId, planRootOperator, coordinatorEndpoint, + new S3QSShuffleResourceLifecycle()); + } + + public CoordinatedPlanExecution createPlanExecution(long transId, Operator planRootOperator, + CoordinatorEndpoint coordinatorEndpoint, + ShuffleResourceLifecycle shuffleResourceLifecycle) + { + return createPlanExecution(transId, planRootOperator, coordinatorEndpoint, + shuffleResourceLifecycle, new InvokerStageWorkerLauncher()); + } + + public CoordinatedPlanExecution createPlanExecution(long transId, Operator planRootOperator, + CoordinatorEndpoint coordinatorEndpoint, + ShuffleResourceLifecycle shuffleResourceLifecycle, + StageWorkerLauncher stageWorkerLauncher) + { + PlanCoordinator planCoordinator = + createPlanCoordinator(transId, planRootOperator, coordinatorEndpoint, stageWorkerLauncher); + return new CoordinatedPlanExecution(transId, planRootOperator, planCoordinator, this, + requireNonNull(shuffleResourceLifecycle, "shuffleResourceLifecycle is null")); + } + /** * Retrieve the plan coordinator of the query. * @param transId the transaction id of the query @@ -71,4 +124,14 @@ public PlanCoordinator getPlanCoordinator(long transId) { return this.transIdToPlanCoordinator.get(transId); } + + /** + * Remove only the coordinator owned by the given execution. This prevents + * a stale execution from deleting a newer coordinator with the same id. + */ + public boolean removePlanCoordinator(long transId, PlanCoordinator planCoordinator) + { + return this.transIdToPlanCoordinator.remove(transId, + requireNonNull(planCoordinator, "planCoordinator is null")); + } } diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/S3QSShuffleResourceLifecycle.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/S3QSShuffleResourceLifecycle.java new file mode 100644 index 0000000000..2a12ece202 --- /dev/null +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/S3QSShuffleResourceLifecycle.java @@ -0,0 +1,187 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public + * License along with Pixels. If not, see + * . + */ +package io.pixelsdb.pixels.planner.coordinate; + +import io.pixelsdb.pixels.common.physical.Storage; +import io.pixelsdb.pixels.common.physical.StorageFactory; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleQueueInfo; +import io.pixelsdb.pixels.storage.s3qs.S3QS; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collection; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; + +import static com.google.common.base.Preconditions.checkState; + +/** + * Provisions and cleans the AWS resources owned by one S3QS query execution. + * + * Workers still resolve the deterministic queue name locally. Provisioning it + * here first makes the query execution the owner and gives cleanup an exact + * list of queue URLs rather than relying on worker-JVM state. + */ +public class S3QSShuffleResourceLifecycle implements ShuffleResourceLifecycle +{ + private final List ownedQueues = new ArrayList<>(); + private final Set ownedObjectPrefixes = new LinkedHashSet<>(); + private S3QS s3qs; + private boolean prepared; + private boolean cleaned; + + @Override + public synchronized void prepare(Collection shuffleInfos) throws IOException + { + checkState(!prepared, "shuffle resources have already been prepared"); + prepared = true; + if (shuffleInfos == null || shuffleInfos.isEmpty()) + { + return; + } + + try + { + for (ShuffleInfo shuffleInfo : shuffleInfos) + { + if (shuffleInfo.getStorageInfo() == null || + shuffleInfo.getStorageInfo().getScheme() != Storage.Scheme.s3qs) + { + continue; + } + S3QS storage = getS3QS(); + ownedObjectPrefixes.add(shuffleInfo.getObjectPathPrefix()); + for (ShuffleQueueInfo queue : shuffleInfo.getQueues()) + { + String queueUrl = queue.getQueueUrl(); + if (queueUrl == null || queueUrl.trim().isEmpty()) + { + queueUrl = storage.createQueue(queue.getQueueName()); + queue.setQueueUrl(queueUrl); + ownedQueues.add(new OwnedQueue( + shuffleInfo.getShuffleId(), queue.getPartitionId(), queueUrl)); + } + } + } + } + catch (IOException | RuntimeException e) + { + try + { + cleanup(); + } + catch (IOException cleanupFailure) + { + e.addSuppressed(cleanupFailure); + } + throw e; + } + } + + @Override + public synchronized void cleanup() throws IOException + { + if (cleaned) + { + return; + } + cleaned = true; + IOException failure = null; + + if (s3qs != null) + { + for (OwnedQueue queue : ownedQueues) + { + try + { + s3qs.deleteQueue(queue.queueUrl); + } + catch (IOException e) + { + failure = addFailure(failure, e); + } + try + { + s3qs.unregisterQueue(queue.shuffleId, queue.partitionId); + } + catch (IOException e) + { + failure = addFailure(failure, e); + } + } + for (String objectPrefix : ownedObjectPrefixes) + { + try + { + s3qs.delete(objectPrefix, true); + } + catch (IOException e) + { + failure = addFailure(failure, e); + } + } + } + + if (failure != null) + { + throw failure; + } + } + + private S3QS getS3QS() throws IOException + { + if (s3qs == null) + { + Storage storage = StorageFactory.Instance().getStorage(Storage.Scheme.s3qs); + if (!(storage instanceof S3QS)) + { + throw new IOException("storage of scheme s3qs is not S3QS"); + } + s3qs = (S3QS) storage; + } + return s3qs; + } + + private static IOException addFailure(IOException current, IOException next) + { + if (current == null) + { + return next; + } + current.addSuppressed(next); + return current; + } + + private static class OwnedQueue + { + private final String shuffleId; + private final int partitionId; + private final String queueUrl; + + private OwnedQueue(String shuffleId, int partitionId, String queueUrl) + { + this.shuffleId = shuffleId; + this.partitionId = partitionId; + this.queueUrl = queueUrl; + } + } +} diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/ShuffleResourceLifecycle.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/ShuffleResourceLifecycle.java new file mode 100644 index 0000000000..2e6212971a --- /dev/null +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/ShuffleResourceLifecycle.java @@ -0,0 +1,39 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public + * License along with Pixels. If not, see + * . + */ +package io.pixelsdb.pixels.planner.coordinate; + +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleInfo; + +import java.io.IOException; +import java.util.Collection; + +/** + * Owns external resources used by the shuffle edges of one query execution. + * + * Implementations are created per query. This keeps AWS resource operations + * outside the general stage/task coordinator while preserving one lifecycle + * owner for prepare and cleanup. + */ +public interface ShuffleResourceLifecycle +{ + void prepare(Collection shuffleInfos) throws IOException; + + void cleanup() throws IOException; +} diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageCoordinator.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageCoordinator.java index f864d14342..15ac72934f 100644 --- a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageCoordinator.java +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageCoordinator.java @@ -32,6 +32,7 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import static com.google.common.base.Preconditions.checkArgument; @@ -96,6 +97,16 @@ public class StageCoordinator * The number of workers in the right child stage. */ private int rightChildWorkerNum; + /** + * Runtime capacity is unbounded until a StageRuntimeController takes + * ownership of this queued stage. + */ + private int desiredRuntimeWorkerCount = Integer.MAX_VALUE; + /** + * Draining workers may finish their claimed task batch but can not claim + * another one. + */ + private final Set drainingWorkerIds = ConcurrentHashMap.newKeySet(); private final Object lock = new Object(); /** @@ -203,6 +214,10 @@ public void addWorker(Worker worker) } } this.workers.add(worker); + if (this.isQueued && getAcceptingWorkerCountLocked() > this.desiredRuntimeWorkerCount) + { + this.drainingWorkerIds.add(worker.getWorkerId()); + } if (!this.isQueued && this.workers.size() == this.fixedWorkerNum) { this.lock.notifyAll(); @@ -225,6 +240,10 @@ public List getTasksToRun(long workerId) throws WorkerCoordinateException Worker worker = this.workerIdToWorkers.get(workerId); if (worker != null) { + if (this.drainingWorkerIds.contains(workerId)) + { + return Collections.emptyList(); + } List tasks = new ArrayList<>(WorkerTaskParallelism); for (int i = 0; i < WorkerTaskParallelism; ++i) { @@ -333,11 +352,150 @@ public int getStageId() return this.stageId; } + public boolean isQueued() + { + return this.isQueued; + } + public List> getWorkers() { return this.workers; } + public int getRegisteredWorkerCount() + { + return this.workers.size(); + } + + /** + * Set the target number of workers that may claim new task batches. + * + * Reducing the target marks the newest excess workers as draining. Raising + * the target does not reactivate draining workers because they may already + * have received an end-of-tasks response; the runtime controller launches + * replacements when needed. + */ + public void setDesiredRuntimeWorkerCount(int desiredRuntimeWorkerCount) + { + checkArgument(this.isQueued, "non-queued stage does not have runtime capacity"); + checkArgument(desiredRuntimeWorkerCount >= 0, "desiredRuntimeWorkerCount is negative"); + synchronized (this.lock) + { + this.desiredRuntimeWorkerCount = desiredRuntimeWorkerCount; + int workersToDrain = getAcceptingWorkerCountLocked() - desiredRuntimeWorkerCount; + for (int i = this.workers.size() - 1; i >= 0 && workersToDrain > 0; --i) + { + Worker worker = this.workers.get(i); + if (!worker.isTerminated() && !this.drainingWorkerIds.contains(worker.getWorkerId())) + { + this.drainingWorkerIds.add(worker.getWorkerId()); + workersToDrain--; + } + } + } + } + + public int getDesiredRuntimeWorkerCount() + { + synchronized (this.lock) + { + return this.desiredRuntimeWorkerCount; + } + } + + public int getActiveRegisteredWorkerCount() + { + synchronized (this.lock) + { + int count = 0; + for (Worker worker : this.workers) + { + if (!worker.isTerminated()) + { + count++; + } + } + return count; + } + } + + public int getAcceptingWorkerCount() + { + synchronized (this.lock) + { + return getAcceptingWorkerCountLocked(); + } + } + + public int getDrainingWorkerCount() + { + synchronized (this.lock) + { + int count = 0; + for (Worker worker : this.workers) + { + if (!worker.isTerminated() && this.drainingWorkerIds.contains(worker.getWorkerId())) + { + count++; + } + } + return count; + } + } + + public boolean isWorkerDraining(long workerId) + { + return this.drainingWorkerIds.contains(workerId); + } + + private int getAcceptingWorkerCountLocked() + { + int count = 0; + for (Worker worker : this.workers) + { + if (!worker.isTerminated() && !this.drainingWorkerIds.contains(worker.getWorkerId())) + { + count++; + } + } + return count; + } + + public int getTotalTaskCount() + { + checkArgument(this.isQueued && this.taskQueue != null, + "non-queued stage does not have task queue"); + return this.taskQueue.getTotalTaskCount(); + } + + public int getPendingTaskCount() + { + checkArgument(this.isQueued && this.taskQueue != null, + "non-queued stage does not have task queue"); + return this.taskQueue.getPendingTaskCount(); + } + + public int getRunningTaskCount() + { + checkArgument(this.isQueued && this.taskQueue != null, + "non-queued stage does not have task queue"); + return this.taskQueue.getRunningTaskCount(); + } + + public int getCompletedTaskCount() + { + checkArgument(this.isQueued && this.taskQueue != null, + "non-queued stage does not have task queue"); + return this.taskQueue.getCompletedTaskCount(); + } + + public int getFailedTaskCount() + { + checkArgument(this.isQueued && this.taskQueue != null, + "non-queued stage does not have task queue"); + return this.taskQueue.getFailedTaskCount(); + } + public void setDownStreamWorkerNum(int downStreamWorkerNum) { this.downStreamWorkerNum = downStreamWorkerNum; diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageExecutionDescriptor.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageExecutionDescriptor.java new file mode 100644 index 0000000000..743cc996c7 --- /dev/null +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageExecutionDescriptor.java @@ -0,0 +1,70 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of the License, + * or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public License + * along with Pixels. If not, see . + */ +package io.pixelsdb.pixels.planner.coordinate; + +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; + +import static java.util.Objects.requireNonNull; + +/** + * Immutable information required to start a worker for one queued stage. + * + * Logical task payloads are deliberately absent: workers use this bootstrap + * information to register and pull those payloads from StageCoordinator. + */ +public class StageExecutionDescriptor +{ + private final long transId; + private final long timestamp; + private final int stageId; + private final String operatorName; + private final WorkerType workerType; + private final CoordinatorEndpoint coordinatorEndpoint; + + public StageExecutionDescriptor(long transId, long timestamp, int stageId, + String operatorName, WorkerType workerType, + CoordinatorEndpoint coordinatorEndpoint) + { + this.transId = transId; + this.timestamp = timestamp; + this.stageId = stageId; + this.operatorName = requireNonNull(operatorName, "operatorName is null"); + this.workerType = requireNonNull(workerType, "workerType is null"); + this.coordinatorEndpoint = requireNonNull(coordinatorEndpoint, "coordinatorEndpoint is null"); + } + + public StageWorkerInput createWorkerInput() + { + StageWorkerInput input = new StageWorkerInput(transId, timestamp, stageId, operatorName, workerType); + input.setCoordinatorHost(coordinatorEndpoint.getHost()); + input.setCoordinatorPort(coordinatorEndpoint.getPort()); + return input; + } + + public int getStageId() + { + return stageId; + } + + public WorkerType getWorkerType() + { + return workerType; + } +} diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageRuntimeController.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageRuntimeController.java new file mode 100644 index 0000000000..f6ecc40e74 --- /dev/null +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageRuntimeController.java @@ -0,0 +1,129 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of the License, + * or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public License + * along with Pixels. If not, see . + */ +package io.pixelsdb.pixels.planner.coordinate; + +import io.pixelsdb.pixels.common.turbo.Output; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; + +import static com.google.common.base.Preconditions.checkArgument; +import static java.util.Objects.requireNonNull; + +/** + * Coordinator-owned control plane for physical worker attempts of one stage. + * + * StageCoordinator remains authoritative for logical tasks and registered + * workers. This class owns desired capacity and platform invocation futures. + */ +public class StageRuntimeController +{ + private final StageCoordinator stageCoordinator; + private final StageExecutionDescriptor descriptor; + private final StageWorkerLauncher workerLauncher; + private final List> activeAttempts = new ArrayList<>(); + private int desiredWorkerCount; + + public StageRuntimeController(StageCoordinator stageCoordinator, + StageExecutionDescriptor descriptor, + StageWorkerLauncher workerLauncher) + { + this.stageCoordinator = requireNonNull(stageCoordinator, "stageCoordinator is null"); + this.descriptor = requireNonNull(descriptor, "descriptor is null"); + this.workerLauncher = requireNonNull(workerLauncher, "workerLauncher is null"); + checkArgument(stageCoordinator.isQueued(), "runtime-controlled stage must be queued"); + checkArgument(stageCoordinator.getStageId() == descriptor.getStageId(), + "stage coordinator and execution descriptor have different stage ids"); + this.desiredWorkerCount = 0; + this.stageCoordinator.setDesiredRuntimeWorkerCount(0); + } + + /** + * Change the desired physical capacity and return all currently active + * platform invocation futures after applying the change. + */ + public synchronized List> scaleTo(int targetWorkerCount) + { + checkArgument(targetWorkerCount >= 0, "targetWorkerCount is negative"); + desiredWorkerCount = targetWorkerCount; + stageCoordinator.setDesiredRuntimeWorkerCount(targetWorkerCount); + + int unregisteredAttemptCount = Math.max(0, + activeAttempts.size() - stageCoordinator.getActiveRegisteredWorkerCount()); + int effectiveWorkerCount = stageCoordinator.getAcceptingWorkerCount() + unregisteredAttemptCount; + int workersToLaunch = Math.max(0, targetWorkerCount - effectiveWorkerCount); + + // A physical worker can only help with pending work. Running tasks are + // deliberately not migrated during scale-out. + workersToLaunch = Math.min(workersToLaunch, stageCoordinator.getPendingTaskCount()); + + List> visibleAttempts = new ArrayList<>(activeAttempts); + for (int i = 0; i < workersToLaunch; ++i) + { + CompletableFuture launchedAttempt = requireNonNull( + workerLauncher.launch(descriptor.getWorkerType(), descriptor.createWorkerInput()), + "worker launcher returned null future"); + CompletableFuture attempt = launchedAttempt.thenApply(output -> + { + if (output == null) + { + throw new CompletionException(new IllegalStateException( + "runtime worker returned null output for stage " + descriptor.getStageId())); + } + if (!output.isSuccessful()) + { + String errorMessage = output.getErrorMessage(); + throw new CompletionException(new IllegalStateException( + "runtime worker failed for stage " + descriptor.getStageId() + + (errorMessage == null || errorMessage.isEmpty() ? "" : ": " + errorMessage))); + } + return output; + }); + activeAttempts.add(attempt); + visibleAttempts.add(attempt); + attempt.whenComplete((output, error) -> removeCompletedAttempt(attempt)); + } + return Collections.unmodifiableList(visibleAttempts); + } + + public synchronized StageRuntimeStatus getStatus() + { + return new StageRuntimeStatus(desiredWorkerCount, activeAttempts.size(), + stageCoordinator.getActiveRegisteredWorkerCount(), + stageCoordinator.getAcceptingWorkerCount(), + stageCoordinator.getDrainingWorkerCount(), + stageCoordinator.getPendingTaskCount(), + stageCoordinator.getRunningTaskCount(), + stageCoordinator.getCompletedTaskCount(), + stageCoordinator.getFailedTaskCount()); + } + + public int getStageId() + { + return descriptor.getStageId(); + } + + private synchronized void removeCompletedAttempt(CompletableFuture attempt) + { + activeAttempts.remove(attempt); + } +} diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageRuntimeStatus.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageRuntimeStatus.java new file mode 100644 index 0000000000..24f49f290a --- /dev/null +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageRuntimeStatus.java @@ -0,0 +1,97 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of the License, + * or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public License + * along with Pixels. If not, see . + */ +package io.pixelsdb.pixels.planner.coordinate; + +/** + * Read-only snapshot of one coordinator-managed stage. + */ +public class StageRuntimeStatus +{ + private final int desiredWorkerCount; + private final int activeAttemptCount; + private final int activeRegisteredWorkerCount; + private final int acceptingWorkerCount; + private final int drainingWorkerCount; + private final int pendingTaskCount; + private final int runningTaskCount; + private final int completedTaskCount; + private final int failedTaskCount; + + public StageRuntimeStatus(int desiredWorkerCount, int activeAttemptCount, + int activeRegisteredWorkerCount, int acceptingWorkerCount, + int drainingWorkerCount, int pendingTaskCount, + int runningTaskCount, int completedTaskCount, + int failedTaskCount) + { + this.desiredWorkerCount = desiredWorkerCount; + this.activeAttemptCount = activeAttemptCount; + this.activeRegisteredWorkerCount = activeRegisteredWorkerCount; + this.acceptingWorkerCount = acceptingWorkerCount; + this.drainingWorkerCount = drainingWorkerCount; + this.pendingTaskCount = pendingTaskCount; + this.runningTaskCount = runningTaskCount; + this.completedTaskCount = completedTaskCount; + this.failedTaskCount = failedTaskCount; + } + + public int getDesiredWorkerCount() + { + return desiredWorkerCount; + } + + public int getActiveAttemptCount() + { + return activeAttemptCount; + } + + public int getActiveRegisteredWorkerCount() + { + return activeRegisteredWorkerCount; + } + + public int getAcceptingWorkerCount() + { + return acceptingWorkerCount; + } + + public int getDrainingWorkerCount() + { + return drainingWorkerCount; + } + + public int getPendingTaskCount() + { + return pendingTaskCount; + } + + public int getRunningTaskCount() + { + return runningTaskCount; + } + + public int getCompletedTaskCount() + { + return completedTaskCount; + } + + public int getFailedTaskCount() + { + return failedTaskCount; + } +} diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageWorkerLauncher.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageWorkerLauncher.java new file mode 100644 index 0000000000..3b830ae013 --- /dev/null +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/coordinate/StageWorkerLauncher.java @@ -0,0 +1,36 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of the License, + * or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public License + * along with Pixels. If not, see . + */ +package io.pixelsdb.pixels.planner.coordinate; + +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; + +import java.util.concurrent.CompletableFuture; + +/** + * Starts one physical worker attempt for a coordinator-managed stage. + * + * The interface keeps platform invocation outside the logical task queue and + * provides a narrow injection point for coordinator tests. + */ +public interface StageWorkerLauncher +{ + CompletableFuture launch(WorkerType workerType, StageWorkerInput input); +} diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/ExchangeMethod.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/ExchangeMethod.java index 88cdf76637..b5d4c4efce 100644 --- a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/ExchangeMethod.java +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/ExchangeMethod.java @@ -25,7 +25,7 @@ */ public enum ExchangeMethod { - batch, stream; + batch, stream, s3qs; /** * Case-insensitive parsing from String name to enum value. diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/PartitionedJoinOperator.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/PartitionedJoinOperator.java index 9333922155..d5ebe91aa3 100644 --- a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/PartitionedJoinOperator.java +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/PartitionedJoinOperator.java @@ -21,6 +21,7 @@ import com.alibaba.fastjson.JSON; import com.google.common.collect.ImmutableList; +import io.pixelsdb.pixels.common.physical.Storage; import io.pixelsdb.pixels.common.task.Task; import io.pixelsdb.pixels.common.turbo.Output; import io.pixelsdb.pixels.executor.join.JoinAlgorithm; @@ -28,6 +29,8 @@ import io.pixelsdb.pixels.planner.coordinate.StageCoordinator; import io.pixelsdb.pixels.planner.coordinate.StageDependency; import io.pixelsdb.pixels.planner.plan.physical.domain.InputSplit; +import io.pixelsdb.pixels.planner.plan.physical.domain.OutputInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleInfo; import io.pixelsdb.pixels.planner.plan.physical.input.JoinInput; import io.pixelsdb.pixels.planner.plan.physical.input.PartitionInput; @@ -52,6 +55,7 @@ public abstract class PartitionedJoinOperator extends SingleStageJoinOperator protected CompletableFuture[] largePartitionOutputs = null; protected int smallPartitionStageId; protected int largePartitionStageId; + protected PlanCoordinator planCoordinator; public PartitionedJoinOperator(String name, List smallPartitionInputs, List largePartitionInputs, @@ -136,6 +140,7 @@ public List getLargePartitionInputs() @Override public void initPlanCoordinator(PlanCoordinator planCoordinator, int parentStageId, boolean wideDependOnParent) { + this.planCoordinator = planCoordinator; this.parentStageId = parentStageId; this.joinStageId = planCoordinator.assignStageId(); for (JoinInput joinInput : this.joinInputs) @@ -185,17 +190,7 @@ public void initPlanCoordinator(PlanCoordinator planCoordinator, int parentStage partitionInput.setStageId(this.smallPartitionStageId); } StageDependency partitionStageDependency = new StageDependency(smallPartitionStageId, joinStageId, true); - List tasks = new ArrayList<>(); - int taskId = 0; - for (PartitionInput partitionInput : this.smallPartitionInputs) - { - List inputSplits = partitionInput.getTableInfo().getInputSplits(); - for (InputSplit inputSplit : inputSplits) - { - partitionInput.getTableInfo().setInputSplits(ImmutableList.of(inputSplit)); - tasks.add(new Task(taskId++, JSON.toJSONString(partitionInput))); - } - } + List tasks = createPartitionStageTasks(this.smallPartitionInputs); smallWorkerNum = this.smallPartitionInputs.size(); StageCoordinator partitionStageCoordinator = new StageCoordinator(smallPartitionStageId, tasks, 0); planCoordinator.addStageCoordinator(partitionStageCoordinator, partitionStageDependency); @@ -209,17 +204,7 @@ public void initPlanCoordinator(PlanCoordinator planCoordinator, int parentStage partitionInput.setStageId(this.largePartitionStageId); } StageDependency partitionStageDependency = new StageDependency(largePartitionStageId, joinStageId, true); - List tasks = new ArrayList<>(); - int taskId = 0; - for (PartitionInput partitionInput : this.largePartitionInputs) - { - List inputSplits = partitionInput.getTableInfo().getInputSplits(); - for (InputSplit inputSplit : inputSplits) - { - partitionInput.getTableInfo().setInputSplits(ImmutableList.of(inputSplit)); - tasks.add(new Task(taskId++, JSON.toJSONString(partitionInput))); - } - } + List tasks = createPartitionStageTasks(this.largePartitionInputs); StageCoordinator partitionStageCoordinator = new StageCoordinator(largePartitionStageId, tasks, smallWorkerNum); planCoordinator.addStageCoordinator(partitionStageCoordinator, partitionStageDependency); } @@ -230,6 +215,53 @@ public void initPlanCoordinator(PlanCoordinator planCoordinator, int parentStage } } + private List createPartitionStageTasks(List partitionInputs) + { + List tasks = new ArrayList<>(); + int nextTaskId = 0; + for (PartitionInput partitionInput : partitionInputs) + { + if (isS3QSProducerInput(partitionInput)) + { + PartitionInput taskPartitionInput = copyPartitionInput(partitionInput); + taskPartitionInput.setProducerTaskId(nextTaskId); + tasks.add(new Task(nextTaskId, JSON.toJSONString(taskPartitionInput))); + nextTaskId++; + continue; + } + + List inputSplits = partitionInput.getTableInfo().getInputSplits(); + for (InputSplit inputSplit : inputSplits) + { + PartitionInput taskPartitionInput = copyPartitionInput(partitionInput); + taskPartitionInput.getTableInfo().setInputSplits(ImmutableList.of(inputSplit)); + tasks.add(new Task(nextTaskId++, JSON.toJSONString(taskPartitionInput))); + } + } + return tasks; + } + + private PartitionInput copyPartitionInput(PartitionInput partitionInput) + { + return JSON.parseObject(JSON.toJSONString(partitionInput), PartitionInput.class); + } + + private boolean isS3QSProducerInput(PartitionInput partitionInput) + { + if (partitionInput == null) + { + return false; + } + OutputInfo outputInfo = partitionInput.getOutput(); + if (outputInfo == null) + { + return false; + } + ShuffleInfo shuffleInfo = outputInfo.getShuffleInfo(); + return shuffleInfo != null && shuffleInfo.getStorageInfo() != null && + shuffleInfo.getStorageInfo().getScheme() == Storage.Scheme.s3qs; + } + @Override public JoinOutputCollection collectOutputs() throws ExecutionException, InterruptedException { diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/PartitionedJoinS3QSOperator.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/PartitionedJoinS3QSOperator.java new file mode 100644 index 0000000000..becec6d077 --- /dev/null +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/PartitionedJoinS3QSOperator.java @@ -0,0 +1,410 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public + * License along with Pixels. If not, see + * . + */ +package io.pixelsdb.pixels.planner.plan.physical; + +import com.alibaba.fastjson.JSON; +import io.pixelsdb.pixels.common.task.Task; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.common.physical.Storage; +import io.pixelsdb.pixels.executor.join.JoinAlgorithm; +import io.pixelsdb.pixels.planner.coordinate.PlanCoordinator; +import io.pixelsdb.pixels.planner.coordinate.StageExecutionDescriptor; +import io.pixelsdb.pixels.planner.coordinate.StageCoordinator; +import io.pixelsdb.pixels.planner.plan.physical.domain.PartitionedTableInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleQueueInfo; +import io.pixelsdb.pixels.planner.plan.physical.input.JoinInput; +import io.pixelsdb.pixels.planner.plan.physical.input.PartitionInput; +import io.pixelsdb.pixels.planner.plan.physical.input.PartitionedChainJoinInput; +import io.pixelsdb.pixels.planner.plan.physical.input.PartitionedJoinInput; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ExecutionException; + +import static com.google.common.base.Preconditions.checkArgument; + +/** + * S3QS exchange scheduler for partitioned joins. + * + * In coordinator mode this operator starts platform workers with a lightweight + * stage bootstrap input. The actual PartitionInput/JoinInput payloads stay in + * StageCoordinator task queues and are pulled by the runtime workers. + */ +public class PartitionedJoinS3QSOperator extends PartitionedJoinOperator +{ + private static final Logger logger = LogManager.getLogger(PartitionedJoinS3QSOperator.class); + private static final CompletableFuture Completed = CompletableFuture.completedFuture(null); + + public PartitionedJoinS3QSOperator(String name, List smallPartitionInputs, + List largePartitionInputs, + List joinInputs, JoinAlgorithm joinAlgo) + { + super(name, smallPartitionInputs, largePartitionInputs, joinInputs, joinAlgo); + } + + @Override + public void initPlanCoordinator(PlanCoordinator planCoordinator, int parentStageId, boolean wideDependOnParent) + { + registerShuffleResources(planCoordinator); + super.initPlanCoordinator(planCoordinator, parentStageId, wideDependOnParent); + if (!smallPartitionInputs.isEmpty()) + { + registerStageRuntimeController(smallPartitionStageId, WorkerType.PARTITION_S3QS); + } + if (!largePartitionInputs.isEmpty()) + { + registerStageRuntimeController(largePartitionStageId, WorkerType.PARTITION_S3QS); + } + registerStageRuntimeController(joinStageId, getS3QSJoinWorkerType()); + } + + private void registerShuffleResources(PlanCoordinator planCoordinator) + { + for (PartitionInput input : smallPartitionInputs) + { + planCoordinator.addShuffleInfo(input.getOutput().getShuffleInfo()); + } + for (PartitionInput input : largePartitionInputs) + { + planCoordinator.addShuffleInfo(input.getOutput().getShuffleInfo()); + } + for (JoinInput input : joinInputs) + { + if (input instanceof PartitionedJoinInput) + { + PartitionedJoinInput joinInput = (PartitionedJoinInput) input; + planCoordinator.addShuffleInfo(joinInput.getSmallTable().getShuffleInfo()); + planCoordinator.addShuffleInfo(joinInput.getLargeTable().getShuffleInfo()); + } + else if (input instanceof PartitionedChainJoinInput) + { + PartitionedChainJoinInput joinInput = (PartitionedChainJoinInput) input; + planCoordinator.addShuffleInfo(joinInput.getSmallTable().getShuffleInfo()); + planCoordinator.addShuffleInfo(joinInput.getLargeTable().getShuffleInfo()); + } + } + } + + @Override + public CompletableFuture[]> execute() + { + return executePrev().handle((result, exception) -> + { + if (exception != null) + { + throw new CompletionException("failed to start the previous S3QS stages", exception); + } + validateJoinInputs(); + joinOutputs = invokeJoinRuntimeWorkers(); + + logger.debug("invoke S3QS join " + this.getName()); + return joinOutputs; + }); + } + + @Override + public CompletableFuture executePrev() + { + validatePartitionInputs(); + if (smallChild != null && largeChild != null) + { + checkArgument(smallPartitionInputs.isEmpty(), "smallPartitionInputs is not empty"); + checkArgument(largePartitionInputs.isEmpty(), "largePartitionInputs is not empty"); + smallChild.execute(); + largeChild.execute(); + } + else if (smallChild != null) + { + checkArgument(smallPartitionInputs.isEmpty(), "smallPartitionInputs is not empty"); + checkArgument(!largePartitionInputs.isEmpty(), "largePartitionInputs is empty"); + smallChild.execute(); + largePartitionOutputs = invokePartitionRuntimeWorkers(largePartitionStageId); + logger.debug("invoke large S3QS partition of " + this.getName()); + } + else if (largeChild != null) + { + checkArgument(!smallPartitionInputs.isEmpty(), "smallPartitionInputs is empty"); + checkArgument(largePartitionInputs.isEmpty(), "largePartitionInputs is not empty"); + smallPartitionOutputs = invokePartitionRuntimeWorkers(smallPartitionStageId); + logger.debug("invoke small S3QS partition of " + this.getName()); + largeChild.execute(); + } + else + { + checkArgument(!smallPartitionInputs.isEmpty(), "smallPartitionInputs is empty"); + checkArgument(!largePartitionInputs.isEmpty(), "largePartitionInputs is empty"); + smallPartitionOutputs = invokePartitionRuntimeWorkers(smallPartitionStageId); + logger.debug("invoke small S3QS partition of " + this.getName()); + largePartitionOutputs = invokePartitionRuntimeWorkers(largePartitionStageId); + logger.debug("invoke large S3QS partition of " + this.getName()); + } + return Completed; + } + + /** + * Wait for every worker attempt started by this S3QS operator before the + * query lifecycle is allowed to delete queues and object prefixes. The + * inherited collector then reads each future again and propagates the + * original worker failure. + */ + @Override + public JoinOutputCollection collectOutputs() throws ExecutionException, InterruptedException + { + waitForAllAttempts(joinOutputs, smallPartitionOutputs, largePartitionOutputs); + return super.collectOutputs(); + } + + @SafeVarargs + static void waitForAllAttempts(CompletableFuture[]... stageAttempts) + throws ExecutionException, InterruptedException + { + List> attempts = new ArrayList<>(); + for (CompletableFuture[] stage : stageAttempts) + { + if (stage != null) + { + for (CompletableFuture attempt : stage) + { + attempts.add(attempt); + } + } + } + if (!attempts.isEmpty()) + { + CompletableFuture.allOf(attempts.toArray(new CompletableFuture[0])) + .handle((ignored, failure) -> null).get(); + } + } + + private CompletableFuture[] invokePartitionRuntimeWorkers(int partitionStageId) + { + requireControlledStage(partitionStageId, "partition"); + List> outputFutures = planCoordinator.activateStage(partitionStageId); + return outputFutures.toArray(new CompletableFuture[0]); + } + + private CompletableFuture[] invokeJoinRuntimeWorkers() + { + requireControlledStage(joinStageId, "join"); + List> outputFutures = planCoordinator.activateStage(joinStageId); + return outputFutures.toArray(new CompletableFuture[0]); + } + + @Override + protected StageCoordinator createJoinStageCoordinator(StageCoordinator parentStageCoordinator, + int joinStageId, int workerNum) + { + List tasks = new ArrayList<>(joinInputs.size()); + for (int i = 0; i < joinInputs.size(); ++i) + { + tasks.add(new Task(i, JSON.toJSONString(joinInputs.get(i)))); + } + + int workerIndexBegin = 0; + if (parentStageCoordinator != null) + { + if (parentStageCoordinator.leftChildWorkerIsEmpty()) + { + parentStageCoordinator.setLeftChildWorkerNum(workerNum); + } + else + { + workerIndexBegin = parentStageCoordinator.getLeftChildWorkerNum(); + parentStageCoordinator.setRightChildWorkerNum(workerNum); + } + } + StageCoordinator joinStageCoordinator = new StageCoordinator(joinStageId, tasks, workerIndexBegin); + if (parentStageCoordinator != null) + { + joinStageCoordinator.setDownStreamWorkerNum(parentStageCoordinator.getFixedWorkerNum()); + } + return joinStageCoordinator; + } + + private StageCoordinator requireControlledStage(int stageId, String stageRole) + { + StageCoordinator stageCoordinator = getStageCoordinator(stageId); + if (stageCoordinator == null) + { + throw new IllegalStateException("S3QS " + stageRole + + " stage is not initialized in PlanCoordinator"); + } + if (!stageCoordinator.isQueued()) + { + throw new IllegalStateException("S3QS " + stageRole + " stage is not queued"); + } + if (planCoordinator.getStageRuntimeController(stageId) == null) + { + throw new IllegalStateException("S3QS " + stageRole + + " stage does not have a runtime controller"); + } + return stageCoordinator; + } + + private StageCoordinator getStageCoordinator(int stageId) + { + if (planCoordinator == null || stageId < 0) + { + return null; + } + return planCoordinator.getStageCoordinator(stageId); + } + + private WorkerType getS3QSJoinWorkerType() + { + if (joinAlgo == JoinAlgorithm.PARTITIONED) + { + return WorkerType.PARTITIONED_JOIN_S3QS; + } + if (joinAlgo == JoinAlgorithm.PARTITIONED_CHAIN) + { + return WorkerType.PARTITIONED_CHAIN_JOIN_S3QS; + } + throw new UnsupportedOperationException("join algorithm '" + joinAlgo + "' is unsupported"); + } + + private void registerStageRuntimeController(int stageId, WorkerType workerType) + { + planCoordinator.addStageRuntimeController(new StageExecutionDescriptor( + planCoordinator.getTransId(), getTimestamp(), stageId, getName(), workerType, + planCoordinator.getCoordinatorEndpoint())); + } + + private long getTimestamp() + { + if (!joinInputs.isEmpty()) + { + return joinInputs.get(0).getTimestamp(); + } + if (!smallPartitionInputs.isEmpty()) + { + return smallPartitionInputs.get(0).getTimestamp(); + } + if (!largePartitionInputs.isEmpty()) + { + return largePartitionInputs.get(0).getTimestamp(); + } + return 0L; + } + + private void validatePartitionInputs() + { + if (!smallPartitionInputs.isEmpty()) + { + checkS3QSPartitionInputs(smallPartitionInputs, "small"); + } + if (!largePartitionInputs.isEmpty()) + { + checkS3QSPartitionInputs(largePartitionInputs, "large"); + } + } + + private void checkS3QSPartitionInputs(List partitionInputs, String side) + { + for (PartitionInput partitionInput : partitionInputs) + { + checkArgument(partitionInput.getOutput() != null, + "%s partition input does not have output info", side); + checkS3QSShuffleInfo(partitionInput.getOutput().getShuffleInfo(), side + " partition input"); + checkArgument(partitionInput.getProducerTaskId() >= 0, + "%s partition input does not have producerTaskId", side); + } + } + + static void checkS3QSShuffleInfo(ShuffleInfo shuffleInfo, String source) + { + checkArgument(isS3QSShuffle(shuffleInfo), + "%s is not an explicit S3QS shuffle", source); + checkArgument(!isNullOrEmpty(shuffleInfo.getShuffleId()), + "%s does not have shuffleId", source); + checkArgument(!isNullOrEmpty(shuffleInfo.getObjectPathPrefix()), + "%s does not have objectPathPrefix", source); + checkArgument(shuffleInfo.getNumPartitions() > 0, + "%s does not have positive numPartitions", source); + checkArgument(shuffleInfo.getProducerCount() > 0, + "%s does not have positive producerCount", source); + checkArgument(shuffleInfo.getConsumerCount() > 0, + "%s does not have positive consumerCount", source); + checkArgument(shuffleInfo.getPollTimeoutSeconds() > 0, + "%s does not have positive pollTimeoutSeconds", source); + checkArgument(shuffleInfo.getQueues() != null, + "%s does not have partition queues", source); + checkArgument(shuffleInfo.getQueues().size() == shuffleInfo.getNumPartitions(), + "%s queue count does not match numPartitions", source); + + Set partitionIds = new HashSet<>(); + for (ShuffleQueueInfo queue : shuffleInfo.getQueues()) + { + checkArgument(queue != null, "%s has null queue info", source); + int partitionId = queue.getPartitionId(); + checkArgument(partitionId >= 0 && partitionId < shuffleInfo.getNumPartitions(), + "%s has out-of-range queue partitionId %s", source, partitionId); + checkArgument(partitionIds.add(partitionId), + "%s has duplicate queue partitionId %s", source, partitionId); + checkArgument(!isNullOrEmpty(queue.getQueueName()) || !isNullOrEmpty(queue.getQueueUrl()), + "%s partition %s does not have queueName or queueUrl", source, partitionId); + } + } + + private static boolean isNullOrEmpty(String value) + { + return value == null || value.trim().isEmpty(); + } + + private static boolean isS3QSShuffle(ShuffleInfo shuffleInfo) + { + return shuffleInfo != null && shuffleInfo.getStorageInfo() != null && + shuffleInfo.getStorageInfo().getScheme() == Storage.Scheme.s3qs; + } + + private void checkS3QSTable(PartitionedTableInfo tableInfo, String side) + { + checkArgument(tableInfo != null, "%s join table is null", side); + checkS3QSShuffleInfo(tableInfo.getShuffleInfo(), side + " join table"); + } + + private void validateJoinInputs() + { + for (JoinInput joinInput : joinInputs) + { + if (joinAlgo == JoinAlgorithm.PARTITIONED) + { + PartitionedJoinInput partitionedJoinInput = (PartitionedJoinInput) joinInput; + checkS3QSTable(partitionedJoinInput.getSmallTable(), "small"); + checkS3QSTable(partitionedJoinInput.getLargeTable(), "large"); + } + else if (joinAlgo == JoinAlgorithm.PARTITIONED_CHAIN) + { + PartitionedChainJoinInput chainJoinInput = (PartitionedChainJoinInput) joinInput; + checkS3QSTable(chainJoinInput.getSmallTable(), "small"); + checkS3QSTable(chainJoinInput.getLargeTable(), "large"); + } + } + } + +} diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/domain/OutputInfo.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/domain/OutputInfo.java index 946218eb30..68ae05e2b4 100644 --- a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/domain/OutputInfo.java +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/domain/OutputInfo.java @@ -37,6 +37,10 @@ public class OutputInfo * Whether the output file should be encoded. */ private boolean encoding; + /** + * The optional shuffle metadata for intermediate shuffle output. + */ + private ShuffleInfo shuffleInfo; /** * Default constructor for Jackson. @@ -87,4 +91,14 @@ public void setEncoding(boolean encoding) { this.encoding = encoding; } + + public ShuffleInfo getShuffleInfo() + { + return shuffleInfo; + } + + public void setShuffleInfo(ShuffleInfo shuffleInfo) + { + this.shuffleInfo = shuffleInfo; + } } diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/domain/PartitionedTableInfo.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/domain/PartitionedTableInfo.java index 2f0f59c31c..d529534b52 100644 --- a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/domain/PartitionedTableInfo.java +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/domain/PartitionedTableInfo.java @@ -41,6 +41,10 @@ public class PartitionedTableInfo extends TableInfo * The ids of the join-key columns of the table. */ private int[] keyColumnIds; + /** + * The optional shuffle metadata for reading partitioned intermediate data. + */ + private ShuffleInfo shuffleInfo; /** * Default constructor for Jackson. @@ -86,4 +90,14 @@ public void setKeyColumnIds(int[] keyColumnIds) { this.keyColumnIds = keyColumnIds; } + + public ShuffleInfo getShuffleInfo() + { + return shuffleInfo; + } + + public void setShuffleInfo(ShuffleInfo shuffleInfo) + { + this.shuffleInfo = shuffleInfo; + } } diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/domain/ShuffleInfo.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/domain/ShuffleInfo.java new file mode 100644 index 0000000000..46ab1f7b71 --- /dev/null +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/domain/ShuffleInfo.java @@ -0,0 +1,165 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public + * License along with Pixels. If not, see + * . + */ +package io.pixelsdb.pixels.planner.plan.physical.domain; + +import com.google.common.collect.ImmutableList; + +import java.util.List; + +/** + * The shared metadata of one shuffle edge between producer and consumer stages. + * + * @author Haoting Yan + * @create 2026-06-17 + */ +public class ShuffleInfo +{ + /** + * A query-unique id for this shuffle edge. + */ + private String shuffleId; + /** + * The storage used by this shuffle. For S3QS this should be Storage.Scheme.s3qs. + */ + private StorageInfo storageInfo; + /** + * The object prefix for shuffle data objects. + */ + private String objectPathPrefix; + /** + * The number of hash partitions produced by this shuffle. + */ + private int numPartitions; + /** + * The number of producer workers expected for this shuffle. + */ + private int producerCount; + /** + * The number of consumer workers expected for this shuffle. + */ + private int consumerCount; + /** + * The max long-poll time in seconds for queue consumers. + */ + private int pollTimeoutSeconds; + /** + * The queues assigned to hash partitions. + */ + private List queues; + + /** + * Default constructor for Jackson. + */ + public ShuffleInfo() { } + + public ShuffleInfo(String shuffleId, StorageInfo storageInfo, String objectPathPrefix, + int numPartitions, int producerCount, int consumerCount, + int pollTimeoutSeconds, List queues) + { + this.shuffleId = shuffleId; + this.storageInfo = storageInfo; + this.objectPathPrefix = objectPathPrefix; + this.numPartitions = numPartitions; + this.producerCount = producerCount; + this.consumerCount = consumerCount; + this.pollTimeoutSeconds = pollTimeoutSeconds; + this.queues = queues == null ? null : ImmutableList.copyOf(queues); + } + + public String getShuffleId() + { + return shuffleId; + } + + public void setShuffleId(String shuffleId) + { + this.shuffleId = shuffleId; + } + + public StorageInfo getStorageInfo() + { + return storageInfo; + } + + public void setStorageInfo(StorageInfo storageInfo) + { + this.storageInfo = storageInfo; + } + + public String getObjectPathPrefix() + { + return objectPathPrefix; + } + + public void setObjectPathPrefix(String objectPathPrefix) + { + this.objectPathPrefix = objectPathPrefix; + } + + public int getNumPartitions() + { + return numPartitions; + } + + public void setNumPartitions(int numPartitions) + { + this.numPartitions = numPartitions; + } + + public int getProducerCount() + { + return producerCount; + } + + public void setProducerCount(int producerCount) + { + this.producerCount = producerCount; + } + + public int getConsumerCount() + { + return consumerCount; + } + + public void setConsumerCount(int consumerCount) + { + this.consumerCount = consumerCount; + } + + public int getPollTimeoutSeconds() + { + return pollTimeoutSeconds; + } + + public void setPollTimeoutSeconds(int pollTimeoutSeconds) + { + this.pollTimeoutSeconds = pollTimeoutSeconds; + } + + public List getQueues() + { + return queues; + } + + public void setQueues(List queues) + { + this.queues = queues; + } +} diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/domain/ShuffleQueueInfo.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/domain/ShuffleQueueInfo.java new file mode 100644 index 0000000000..54264e737e --- /dev/null +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/domain/ShuffleQueueInfo.java @@ -0,0 +1,84 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public + * License along with Pixels. If not, see + * . + */ +package io.pixelsdb.pixels.planner.plan.physical.domain; + +/** + * The queue endpoint for one hash partition in a shuffle. + * + * @author Haoting Yan + * @create 2026-06-17 + */ +public class ShuffleQueueInfo +{ + /** + * The hash partition id consumed through this queue. + */ + private int partitionId; + /** + * The stable queue name for this shuffle partition. + */ + private String queueName; + /** + * The queue url used by shuffle producers and consumers after the queue is created or resolved. + */ + private String queueUrl; + + /** + * Default constructor for Jackson. + */ + public ShuffleQueueInfo() { } + + public ShuffleQueueInfo(int partitionId, String queueName, String queueUrl) + { + this.partitionId = partitionId; + this.queueName = queueName; + this.queueUrl = queueUrl; + } + + public int getPartitionId() + { + return partitionId; + } + + public void setPartitionId(int partitionId) + { + this.partitionId = partitionId; + } + + public String getQueueName() + { + return queueName; + } + + public void setQueueName(String queueName) + { + this.queueName = queueName; + } + + public String getQueueUrl() + { + return queueUrl; + } + + public void setQueueUrl(String queueUrl) + { + this.queueUrl = queueUrl; + } +} diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/input/PartitionInput.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/input/PartitionInput.java index 99999e59b8..7ecde3c841 100644 --- a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/input/PartitionInput.java +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/input/PartitionInput.java @@ -51,6 +51,10 @@ public class PartitionInput extends Input * The information about the hash partitioning. */ private PartitionInfo partitionInfo; + /** + * The stable logical producer task id for shuffle protocols. + */ + private int producerTaskId = -1; /** * Default constructor for Jackson. @@ -68,6 +72,7 @@ public PartitionInput(long transId, long timestamp, ScanTableInfo tableInfo, boo this.projection = projection; this.output = output; this.partitionInfo = partitionInfo; + this.producerTaskId = -1; } public ScanTableInfo getTableInfo() @@ -109,4 +114,14 @@ public void setPartitionInfo(PartitionInfo partitionInfo) { this.partitionInfo = partitionInfo; } + + public int getProducerTaskId() + { + return producerTaskId; + } + + public void setProducerTaskId(int producerTaskId) + { + this.producerTaskId = producerTaskId; + } } diff --git a/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/input/StageWorkerInput.java b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/input/StageWorkerInput.java new file mode 100644 index 0000000000..8b9f8cbdf7 --- /dev/null +++ b/pixels-planner/src/main/java/io/pixelsdb/pixels/planner/plan/physical/input/StageWorkerInput.java @@ -0,0 +1,79 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of the License, + * or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public License + * along with Pixels. If not, see . + */ +package io.pixelsdb.pixels.planner.plan.physical.input; + +import io.pixelsdb.pixels.common.turbo.Input; +import io.pixelsdb.pixels.common.turbo.WorkerType; + +/** + * Bootstrap input for a coordinator-dispatched stage worker. + * + * The real work item is not embedded in this input. The remote worker uses + * transId and stageId to register with the coordinator and pulls task payloads + * from the stage task queue. + */ +public class StageWorkerInput extends Input +{ + private String coordinatorHost; + private int coordinatorPort; + private WorkerType workerType; + + public StageWorkerInput() + { + super(0L, 0L); + } + + public StageWorkerInput(long transId, long timestamp, int stageId, String operatorName, WorkerType workerType) + { + super(transId, timestamp); + setStageId(stageId); + setOperatorName(operatorName); + this.workerType = workerType; + } + + public WorkerType getWorkerType() + { + return workerType; + } + + public void setWorkerType(WorkerType workerType) + { + this.workerType = workerType; + } + + public String getCoordinatorHost() + { + return coordinatorHost; + } + + public void setCoordinatorHost(String coordinatorHost) + { + this.coordinatorHost = coordinatorHost; + } + + public int getCoordinatorPort() + { + return coordinatorPort; + } + + public void setCoordinatorPort(int coordinatorPort) + { + this.coordinatorPort = coordinatorPort; + } +} diff --git a/pixels-planner/src/test/java/io/pixelsdb/pixels/planner/coordinate/TestCoordinatedPlanExecution.java b/pixels-planner/src/test/java/io/pixelsdb/pixels/planner/coordinate/TestCoordinatedPlanExecution.java new file mode 100644 index 0000000000..42c1f5dd0d --- /dev/null +++ b/pixels-planner/src/test/java/io/pixelsdb/pixels/planner/coordinate/TestCoordinatedPlanExecution.java @@ -0,0 +1,243 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of the License, + * or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public License + * along with Pixels. If not, see . + */ +package io.pixelsdb.pixels.planner.coordinate; + +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.common.physical.Storage; +import io.pixelsdb.pixels.planner.plan.physical.Operator; +import io.pixelsdb.pixels.planner.plan.physical.OperatorExecutor.OutputCollection; +import io.pixelsdb.pixels.planner.plan.physical.ScanOperator.ScanOutputCollection; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleQueueInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.StorageInfo; +import org.junit.Test; + +import java.io.IOException; +import java.util.Collection; +import java.util.Collections; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicLong; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; + +public class TestCoordinatedPlanExecution +{ + private static final AtomicLong NextTransId = new AtomicLong(12770000L); + + @Test + public void coordinatorRemainsRegisteredUntilOutputsAreCollected() throws Exception + { + long transId = NextTransId.getAndIncrement(); + PlanCoordinatorFactory factory = PlanCoordinatorFactory.Instance(); + RecordingOperator operator = new RecordingOperator(transId); + CoordinatedPlanExecution execution = factory.createPlanExecution(transId, operator, + new CoordinatorEndpoint("coordinator.internal", 19000)); + + assertSame(execution.getPlanCoordinator(), factory.getPlanCoordinator(transId)); + assertSame(execution.getPlanCoordinator(), operator.initializedCoordinator); + assertEquals("coordinator.internal", + execution.getPlanCoordinator().getCoordinatorEndpoint().getHost()); + + execution.execute().get(); + assertTrue(operator.coordinatorWasRegisteredDuringExecute); + assertSame(execution.getPlanCoordinator(), factory.getPlanCoordinator(transId)); + + OutputCollection outputs = execution.collectOutputs(); + assertSame(operator.outputs, outputs); + assertNull(factory.getPlanCoordinator(transId)); + } + + @Test + public void explicitCloseCleansUpExecutionThatDoesNotCollectOutputs() + { + long transId = NextTransId.getAndIncrement(); + PlanCoordinatorFactory factory = PlanCoordinatorFactory.Instance(); + CoordinatedPlanExecution execution = factory.createPlanExecution(transId, + new RecordingOperator(transId), new CoordinatorEndpoint("localhost", 18894)); + + execution.close(); + execution.close(); + + assertNull(factory.getPlanCoordinator(transId)); + } + + @Test + public void outputCollectionWaitsForFinalStageStartup() throws Exception + { + long transId = NextTransId.getAndIncrement(); + PlanCoordinatorFactory factory = PlanCoordinatorFactory.Instance(); + RecordingOperator operator = new RecordingOperator(transId); + operator.executionFuture = new CompletableFuture<>(); + CoordinatedPlanExecution execution = factory.createPlanExecution(transId, operator, + new CoordinatorEndpoint("localhost", 18894)); + execution.execute(); + + CompletableFuture collected = + CompletableFuture.supplyAsync(() -> + { + try + { + return execution.collectOutputs(); + } + catch (Exception e) + { + throw new RuntimeException(e); + } + }); + assertTrue(!collected.isDone()); + + operator.executionFuture.complete(new CompletableFuture[0]); + assertSame(operator.outputs, collected.get()); + assertNull(factory.getPlanCoordinator(transId)); + } + + @Test + public void shuffleResourcesArePreparedBeforeExecutionAndCleanedOnce() throws Exception + { + long transId = NextTransId.getAndIncrement(); + PlanCoordinatorFactory factory = PlanCoordinatorFactory.Instance(); + RecordingOperator operator = new RecordingOperator(transId); + operator.shuffleInfo = shuffleInfo("shuffle-lifecycle"); + RecordingShuffleResourceLifecycle lifecycle = new RecordingShuffleResourceLifecycle(); + operator.lifecycle = lifecycle; + CoordinatedPlanExecution execution = factory.createPlanExecution(transId, operator, + new CoordinatorEndpoint("localhost", 18894), lifecycle); + + execution.execute().get(); + assertTrue(operator.resourcesWerePreparedDuringExecute); + assertEquals(1, lifecycle.preparedShuffleCount); + + execution.collectOutputs(); + execution.close(); + assertEquals(1, lifecycle.cleanupCount); + assertNull(factory.getPlanCoordinator(transId)); + } + + @Test(expected = IllegalStateException.class) + public void duplicateTransactionCanNotReplaceActiveCoordinator() + { + long transId = NextTransId.getAndIncrement(); + PlanCoordinatorFactory factory = PlanCoordinatorFactory.Instance(); + CoordinatedPlanExecution execution = factory.createPlanExecution(transId, + new RecordingOperator(transId), new CoordinatorEndpoint("localhost", 18894)); + try + { + factory.createPlanExecution(transId, new RecordingOperator(transId), + new CoordinatorEndpoint("localhost", 18894)); + } + finally + { + execution.close(); + } + } + + @Test(expected = IllegalArgumentException.class) + public void endpointRejectsEmptyHost() + { + new CoordinatorEndpoint(" ", 18894); + } + + @Test(expected = IllegalArgumentException.class) + public void endpointRejectsInvalidPort() + { + new CoordinatorEndpoint("localhost", 0); + } + + private static ShuffleInfo shuffleInfo(String shuffleId) + { + return new ShuffleInfo(shuffleId, + new StorageInfo(Storage.Scheme.s3qs, null, null, null, null), + "bucket/" + shuffleId + "/", 1, 1, 1, 1, + Collections.singletonList(new ShuffleQueueInfo(0, shuffleId + "-p0", null))); + } + + private static class RecordingOperator extends Operator + { + private final long transId; + private final OutputCollection outputs = new ScanOutputCollection(); + private PlanCoordinator initializedCoordinator; + private boolean coordinatorWasRegisteredDuringExecute; + private boolean resourcesWerePreparedDuringExecute; + private ShuffleInfo shuffleInfo; + private RecordingShuffleResourceLifecycle lifecycle; + private CompletableFuture[]> executionFuture = + CompletableFuture.completedFuture(new CompletableFuture[0]); + + private RecordingOperator(long transId) + { + super("recording-plan"); + this.transId = transId; + } + + @Override + public void initPlanCoordinator(PlanCoordinator planCoordinator, int parentStageId, + boolean wideDependOnParent) + { + this.initializedCoordinator = planCoordinator; + if (shuffleInfo != null) + { + planCoordinator.addShuffleInfo(shuffleInfo); + } + } + + @Override + public CompletableFuture[]> execute() + { + coordinatorWasRegisteredDuringExecute = + PlanCoordinatorFactory.Instance().getPlanCoordinator(transId) == initializedCoordinator; + resourcesWerePreparedDuringExecute = lifecycle != null && lifecycle.prepared; + return executionFuture; + } + + @Override + public CompletableFuture executePrev() + { + return CompletableFuture.completedFuture(null); + } + + @Override + public OutputCollection collectOutputs() + { + return outputs; + } + } + + private static class RecordingShuffleResourceLifecycle implements ShuffleResourceLifecycle + { + private boolean prepared; + private int preparedShuffleCount; + private int cleanupCount; + + @Override + public void prepare(Collection shuffleInfos) + { + prepared = true; + preparedShuffleCount = shuffleInfos.size(); + } + + @Override + public void cleanup() throws IOException + { + ++cleanupCount; + } + } +} diff --git a/pixels-planner/src/test/java/io/pixelsdb/pixels/planner/coordinate/TestStageRuntimeController.java b/pixels-planner/src/test/java/io/pixelsdb/pixels/planner/coordinate/TestStageRuntimeController.java new file mode 100644 index 0000000000..10d814aa31 --- /dev/null +++ b/pixels-planner/src/test/java/io/pixelsdb/pixels/planner/coordinate/TestStageRuntimeController.java @@ -0,0 +1,229 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of the License, + * or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public License + * along with Pixels. If not, see . + */ +package io.pixelsdb.pixels.planner.coordinate; + +import io.pixelsdb.pixels.common.lease.Lease; +import io.pixelsdb.pixels.common.task.Task; +import io.pixelsdb.pixels.common.task.Worker; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.common.utils.Constants; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.PartitionOutput; +import org.junit.Test; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +public class TestStageRuntimeController +{ + @Test + public void planCoordinatorExposesStageActivationAndScaling() + { + StageCoordinator stage = stageCoordinator(2); + RecordingStageWorkerLauncher launcher = new RecordingStageWorkerLauncher(); + CoordinatorEndpoint endpoint = new CoordinatorEndpoint("coordinator.internal", 19000); + PlanCoordinator planCoordinator = new PlanCoordinator(1277L, endpoint, launcher); + planCoordinator.addStageCoordinator(stage, new StageDependency(17, -1, false)); + planCoordinator.addStageRuntimeController(new StageExecutionDescriptor( + 1277L, 1000L, 17, "s3qs-stage", WorkerType.PARTITION_S3QS, endpoint)); + + assertEquals(2, planCoordinator.activateStage(17).size()); + assertEquals(2, launcher.inputs.size()); + assertEquals("coordinator.internal", launcher.inputs.get(0).getCoordinatorHost()); + assertEquals(19000, launcher.inputs.get(0).getCoordinatorPort()); + + planCoordinator.scaleStage(17, 1); + assertEquals(1, planCoordinator.getStageRuntimeController(17) + .getStatus().getDesiredWorkerCount()); + } + + @Test + public void scaleOutIsIdempotentAndLaunchesOnlyMissingCapacity() + { + StageCoordinator stage = stageCoordinator(3); + RecordingStageWorkerLauncher launcher = new RecordingStageWorkerLauncher(); + StageRuntimeController controller = controller(stage, launcher); + + assertEquals(2, controller.scaleTo(2).size()); + assertEquals(2, launcher.inputs.size()); + assertEquals(2, controller.getStatus().getDesiredWorkerCount()); + assertEquals(2, controller.getStatus().getActiveAttemptCount()); + + assertEquals(2, controller.scaleTo(2).size()); + assertEquals(2, launcher.inputs.size()); + + assertEquals(3, controller.scaleTo(3).size()); + assertEquals(3, launcher.inputs.size()); + for (StageWorkerInput input : launcher.inputs) + { + assertEquals(17, input.getStageId()); + assertEquals(WorkerType.PARTITION_S3QS, input.getWorkerType()); + } + } + + @Test + public void scaleInDrainsRegisteredWorkerAndScaleOutLaunchesReplacement() throws Exception + { + StageCoordinator stage = stageCoordinator(3); + RecordingStageWorkerLauncher launcher = new RecordingStageWorkerLauncher(); + StageRuntimeController controller = controller(stage, launcher); + controller.scaleTo(2); + + Worker first = worker(1L); + Worker second = worker(2L); + stage.addWorker(first); + stage.addWorker(second); + assertEquals(2, stage.getAcceptingWorkerCount()); + + controller.scaleTo(1); + assertFalse(stage.isWorkerDraining(first.getWorkerId())); + assertTrue(stage.isWorkerDraining(second.getWorkerId())); + assertEquals(1, stage.getAcceptingWorkerCount()); + assertTrue(stage.getTasksToRun(second.getWorkerId()).isEmpty()); + + controller.scaleTo(2); + assertEquals(3, launcher.inputs.size()); + + Worker replacement = worker(3L); + stage.addWorker(replacement); + assertFalse(stage.isWorkerDraining(replacement.getWorkerId())); + assertEquals(2, stage.getAcceptingWorkerCount()); + assertEquals(1, stage.getDrainingWorkerCount()); + } + + @Test + public void workerRegisteringAfterScaleToZeroCanNotClaimTasks() throws Exception + { + StageCoordinator stage = stageCoordinator(1); + RecordingStageWorkerLauncher launcher = new RecordingStageWorkerLauncher(); + StageRuntimeController controller = controller(stage, launcher); + + controller.scaleTo(5); + assertEquals(1, launcher.inputs.size()); + controller.scaleTo(0); + + Worker lateWorker = new Worker<>(9L, + new Lease(System.currentTimeMillis() - 10000L, 1L), 0, + new CFWorkerInfo("localhost", -1, 1277L, 17, + Constants.PARTITION_OPERATOR_NAME, Collections.emptyList())); + stage.addWorker(lateWorker); + assertTrue(stage.isWorkerDraining(lateWorker.getWorkerId())); + assertEquals(0, stage.getAcceptingWorkerCount()); + assertTrue(stage.getTasksToRun(lateWorker.getWorkerId()).isEmpty()); + assertEquals(1, stage.getPendingTaskCount()); + } + + @Test + public void successfulWorkerOutputIsVisibleToOperator() throws Exception + { + StageCoordinator stage = stageCoordinator(1); + RecordingStageWorkerLauncher launcher = new RecordingStageWorkerLauncher(); + StageRuntimeController controller = controller(stage, launcher); + + CompletableFuture visibleAttempt = controller.scaleTo(1).get(0); + PartitionOutput output = new PartitionOutput(); + output.setSuccessful(true); + launcher.complete(0, output); + + assertSame(output, visibleAttempt.get()); + assertEquals(0, controller.getStatus().getActiveAttemptCount()); + } + + @Test + public void unsuccessfulWorkerOutputFailsTheVisibleAttempt() throws Exception + { + StageCoordinator stage = stageCoordinator(1); + RecordingStageWorkerLauncher launcher = new RecordingStageWorkerLauncher(); + StageRuntimeController controller = controller(stage, launcher); + + CompletableFuture visibleAttempt = controller.scaleTo(1).get(0); + PartitionOutput output = new PartitionOutput(); + output.setSuccessful(false); + output.setErrorMessage("partition task failed"); + launcher.complete(0, output); + + try + { + visibleAttempt.get(); + fail("unsuccessful worker output should fail the stage attempt"); + } + catch (ExecutionException e) + { + assertTrue(e.getCause() instanceof IllegalStateException); + assertTrue(e.getCause().getMessage().contains("partition task failed")); + } + assertEquals(0, controller.getStatus().getActiveAttemptCount()); + } + + private static StageRuntimeController controller(StageCoordinator stage, + RecordingStageWorkerLauncher launcher) + { + StageExecutionDescriptor descriptor = new StageExecutionDescriptor( + 1277L, 1000L, stage.getStageId(), "s3qs-stage", WorkerType.PARTITION_S3QS, + new CoordinatorEndpoint("localhost", 18894)); + return new StageRuntimeController(stage, descriptor, launcher); + } + + private static StageCoordinator stageCoordinator(int taskCount) + { + List tasks = new ArrayList<>(taskCount); + for (int i = 0; i < taskCount; ++i) + { + tasks.add(new Task(i, "task-" + i)); + } + return new StageCoordinator(17, tasks); + } + + private static Worker worker(long workerId) + { + return new Worker<>(workerId, new Lease(System.currentTimeMillis(), 60000L), 0, + new CFWorkerInfo("localhost", -1, 1277L, 17, + Constants.PARTITION_OPERATOR_NAME, Collections.emptyList())); + } + + private static class RecordingStageWorkerLauncher implements StageWorkerLauncher + { + private final List inputs = new ArrayList<>(); + private final List> futures = new ArrayList<>(); + + @Override + public CompletableFuture launch(WorkerType workerType, StageWorkerInput input) + { + inputs.add(input); + CompletableFuture future = new CompletableFuture<>(); + futures.add(future); + return future; + } + + private void complete(int attemptIndex, Output output) + { + futures.get(attemptIndex).complete(output); + } + } +} diff --git a/pixels-planner/src/test/java/io/pixelsdb/pixels/planner/plan/physical/TestPartitionedJoinS3QSOperator.java b/pixels-planner/src/test/java/io/pixelsdb/pixels/planner/plan/physical/TestPartitionedJoinS3QSOperator.java new file mode 100644 index 0000000000..7afbaf0b5b --- /dev/null +++ b/pixels-planner/src/test/java/io/pixelsdb/pixels/planner/plan/physical/TestPartitionedJoinS3QSOperator.java @@ -0,0 +1,420 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public + * License along with Pixels. If not, see + * . + */ +package io.pixelsdb.pixels.planner.plan.physical; + +import com.alibaba.fastjson.JSON; +import io.pixelsdb.pixels.common.lease.Lease; +import io.pixelsdb.pixels.common.physical.Storage; +import io.pixelsdb.pixels.common.task.Task; +import io.pixelsdb.pixels.common.task.Worker; +import io.pixelsdb.pixels.common.turbo.Input; +import io.pixelsdb.pixels.common.turbo.Invoker; +import io.pixelsdb.pixels.common.turbo.InvokerFactory; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.common.utils.Constants; +import io.pixelsdb.pixels.executor.join.JoinAlgorithm; +import io.pixelsdb.pixels.planner.coordinate.CFWorkerInfo; +import io.pixelsdb.pixels.planner.coordinate.CoordinatorEndpoint; +import io.pixelsdb.pixels.planner.coordinate.PlanCoordinator; +import io.pixelsdb.pixels.planner.coordinate.StageCoordinator; +import io.pixelsdb.pixels.planner.coordinate.StageRuntimeController; +import io.pixelsdb.pixels.planner.plan.physical.domain.InputInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.InputSplit; +import io.pixelsdb.pixels.planner.plan.physical.domain.OutputInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.PartitionedTableInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ScanTableInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleQueueInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.StorageInfo; +import io.pixelsdb.pixels.planner.plan.physical.input.JoinInput; +import io.pixelsdb.pixels.planner.plan.physical.input.PartitionInput; +import io.pixelsdb.pixels.planner.plan.physical.input.PartitionedJoinInput; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import org.junit.Test; + +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.Collections; +import java.util.EnumMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +/** + * @author Haoting Yan + * @create 2026-06-25 + */ +public class TestPartitionedJoinS3QSOperator +{ + @Test + public void acceptsCompleteShuffleInfo() + { + PartitionedJoinS3QSOperator.checkS3QSShuffleInfo(validShuffleInfo(), "test shuffle"); + } + + @Test(expected = IllegalArgumentException.class) + public void rejectsQueueWithoutNameOrUrl() + { + ShuffleInfo shuffleInfo = validShuffleInfo(); + shuffleInfo.setQueues(Collections.singletonList(new ShuffleQueueInfo(0, null, null))); + + PartitionedJoinS3QSOperator.checkS3QSShuffleInfo(shuffleInfo, "test shuffle"); + } + + @Test(expected = IllegalArgumentException.class) + public void rejectsPartitionInputWithoutProducerTaskId() + { + PartitionInput partitionInput = new PartitionInput(); + OutputInfo outputInfo = new OutputInfo("s3qs://bucket/shuffle/0/part", + new StorageInfo(Storage.Scheme.s3qs, null, null, null, null), true); + outputInfo.setShuffleInfo(validShuffleInfo()); + partitionInput.setOutput(outputInfo); + + PartitionedJoinS3QSOperator operator = new PartitionedJoinS3QSOperator("s3qs-join", + Collections.singletonList(partitionInput), Collections.emptyList(), Collections.emptyList(), + JoinAlgorithm.PARTITIONED); + + operator.executePrev(); + } + + @Test(expected = IllegalStateException.class) + public void executeRejectsS3QSPlanWithoutCoordinator() + { + PartitionedJoinS3QSOperator operator = new PartitionedJoinS3QSOperator("s3qs-join", + Collections.singletonList(partitionInput(0)), Collections.singletonList(partitionInput(1)), + Collections.singletonList(joinInput(validShuffleInfo())), JoinAlgorithm.PARTITIONED); + + operator.execute(); + } + + @Test + public void initPlanCoordinatorCreatesPartitionTasksWithS3QSMetadata() throws Exception + { + PartitionInput smallPartition = partitionInput(0); + smallPartition.setTableInfo(scanTableInfo("small")); + PartitionInput largePartition = partitionInput(1); + largePartition.setTableInfo(scanTableInfo("large")); + PartitionedJoinS3QSOperator operator = new PartitionedJoinS3QSOperator("s3qs-join", + Collections.singletonList(smallPartition), Collections.singletonList(largePartition), + Collections.singletonList(joinInput(validShuffleInfo())), JoinAlgorithm.PARTITIONED); + PlanCoordinator planCoordinator = new PlanCoordinator(100L); + + operator.initPlanCoordinator(planCoordinator, -1, false); + StageCoordinator smallStage = planCoordinator.getStageCoordinator(operator.smallPartitionStageId); + smallStage.setDesiredRuntimeWorkerCount(1); + Worker worker = new Worker<>(1L, new Lease(System.currentTimeMillis(), 60000L), + 0, new CFWorkerInfo("localhost", 8080, 100L, operator.smallPartitionStageId, + Constants.PARTITION_OPERATOR_NAME, Collections.emptyList())); + smallStage.addWorker(worker); + + List tasks = smallStage.getTasksToRun(worker.getWorkerId()); + assertEquals(1, tasks.size()); + PartitionInput taskInput = JSON.parseObject(tasks.get(0).getPayload(), PartitionInput.class); + assertEquals(0, taskInput.getProducerTaskId()); + assertEquals("shuffle-1", taskInput.getOutput().getShuffleInfo().getShuffleId()); + assertEquals(Storage.Scheme.s3qs, taskInput.getOutput().getShuffleInfo().getStorageInfo().getScheme()); + assertEquals(1, taskInput.getTableInfo().getInputSplits().size()); + } + + @Test + public void initPlanCoordinatorUsesOneCoordinatorTaskPerS3QSProducerInput() throws Exception + { + PartitionInput smallPartition = partitionInput(7); + smallPartition.setTableInfo(scanTableInfo("small", 3)); + PartitionInput largePartition = partitionInput(8); + largePartition.setTableInfo(scanTableInfo("large")); + PartitionedJoinS3QSOperator operator = new PartitionedJoinS3QSOperator("s3qs-join", + Collections.singletonList(smallPartition), Collections.singletonList(largePartition), + Collections.singletonList(joinInput(validShuffleInfo())), JoinAlgorithm.PARTITIONED); + PlanCoordinator planCoordinator = new PlanCoordinator(100L); + + operator.initPlanCoordinator(planCoordinator, -1, false); + StageCoordinator smallStage = planCoordinator.getStageCoordinator(operator.smallPartitionStageId); + smallStage.setDesiredRuntimeWorkerCount(1); + assertEquals(1, smallStage.getTotalTaskCount()); + assertEquals(1, smallStage.getPendingTaskCount()); + assertEquals(0, smallStage.getRunningTaskCount()); + + Worker worker = new Worker<>(1L, new Lease(System.currentTimeMillis(), 60000L), + 0, new CFWorkerInfo("localhost", 8080, 100L, operator.smallPartitionStageId, + Constants.PARTITION_OPERATOR_NAME, Collections.emptyList())); + smallStage.addWorker(worker); + + List tasks = smallStage.getTasksToRun(worker.getWorkerId()); + assertEquals(1, tasks.size()); + assertEquals(0, tasks.get(0).getTaskId()); + assertEquals(0, smallStage.getPendingTaskCount()); + assertEquals(1, smallStage.getRunningTaskCount()); + PartitionInput taskInput = JSON.parseObject(tasks.get(0).getPayload(), PartitionInput.class); + assertEquals(tasks.get(0).getTaskId(), taskInput.getProducerTaskId()); + assertEquals(3, taskInput.getTableInfo().getInputSplits().size()); + + smallStage.completeTask(tasks.get(0).getTaskId(), true); + assertEquals(0, smallStage.getRunningTaskCount()); + assertEquals(1, smallStage.getCompletedTaskCount()); + assertEquals(0, smallStage.getFailedTaskCount()); + } + + @Test + public void executeStartsS3QSRuntimeWorkersWhenPlanCoordinatorIsInitialized() throws Exception + { + RecordingInvoker partitionStageInvoker = new RecordingInvoker(); + RecordingInvoker joinStageInvoker = new RecordingInvoker(); + Map originalInvokers = replaceS3QSInvokers(partitionStageInvoker, joinStageInvoker); + try + { + ShuffleInfo smallShuffle = validShuffleInfo("small-shuffle"); + ShuffleInfo largeShuffle = validShuffleInfo("large-shuffle"); + PartitionInput smallPartition = partitionInput(7, smallShuffle); + smallPartition.setTableInfo(scanTableInfo("small", 2)); + PartitionInput largePartition = partitionInput(8, largeShuffle); + largePartition.setTableInfo(scanTableInfo("large", 2)); + PartitionedJoinInput joinInput = joinInput(smallShuffle, largeShuffle); + PartitionedJoinS3QSOperator operator = new PartitionedJoinS3QSOperator("s3qs-join", + Collections.singletonList(smallPartition), Collections.singletonList(largePartition), + Collections.singletonList(joinInput), JoinAlgorithm.PARTITIONED); + PlanCoordinator planCoordinator = + new PlanCoordinator(100L, new CoordinatorEndpoint("coordinator.internal", 19000)); + operator.initPlanCoordinator(planCoordinator, -1, false); + + CompletableFuture[]> future = operator.execute(); + CompletableFuture[] joinOutputs = future.get(); + + assertEquals(2, partitionStageInvoker.inputs.size()); + StageWorkerInput smallStageWorkerInput = (StageWorkerInput) partitionStageInvoker.inputs.get(0); + StageWorkerInput largeStageWorkerInput = (StageWorkerInput) partitionStageInvoker.inputs.get(1); + assertEquals(WorkerType.PARTITION_S3QS, smallStageWorkerInput.getWorkerType()); + assertEquals(WorkerType.PARTITION_S3QS, largeStageWorkerInput.getWorkerType()); + assertEquals(100L, smallStageWorkerInput.getTransId()); + assertEquals(100L, largeStageWorkerInput.getTransId()); + assertEquals(operator.smallPartitionStageId, smallStageWorkerInput.getStageId()); + assertEquals(operator.largePartitionStageId, largeStageWorkerInput.getStageId()); + assertEquals("coordinator.internal", smallStageWorkerInput.getCoordinatorHost()); + assertEquals(19000, smallStageWorkerInput.getCoordinatorPort()); + assertEquals(1, joinStageInvoker.inputs.size()); + StageWorkerInput joinStageWorkerInput = (StageWorkerInput) joinStageInvoker.inputs.get(0); + assertEquals(WorkerType.PARTITIONED_JOIN_S3QS, joinStageWorkerInput.getWorkerType()); + assertEquals(100L, joinStageWorkerInput.getTransId()); + assertEquals(operator.joinStageId, joinStageWorkerInput.getStageId()); + assertEquals("coordinator.internal", joinStageWorkerInput.getCoordinatorHost()); + assertEquals(19000, joinStageWorkerInput.getCoordinatorPort()); + assertEquals(1, joinOutputs.length); + + StageCoordinator smallStage = planCoordinator.getStageCoordinator(operator.smallPartitionStageId); + StageCoordinator largeStage = planCoordinator.getStageCoordinator(operator.largePartitionStageId); + StageCoordinator joinStage = planCoordinator.getStageCoordinator(operator.joinStageId); + StageRuntimeController smallRuntime = + planCoordinator.getStageRuntimeController(operator.smallPartitionStageId); + StageRuntimeController largeRuntime = + planCoordinator.getStageRuntimeController(operator.largePartitionStageId); + StageRuntimeController joinRuntime = + planCoordinator.getStageRuntimeController(operator.joinStageId); + assertNotNull(smallRuntime); + assertNotNull(largeRuntime); + assertNotNull(joinRuntime); + assertEquals(1, smallRuntime.getStatus().getDesiredWorkerCount()); + assertEquals(1, largeRuntime.getStatus().getDesiredWorkerCount()); + assertEquals(1, joinRuntime.getStatus().getDesiredWorkerCount()); + assertEquals(1, smallStage.getTotalTaskCount()); + assertEquals(1, largeStage.getTotalTaskCount()); + assertEquals(1, joinStage.getTotalTaskCount()); + assertEquals(1, smallStage.getPendingTaskCount()); + assertEquals(1, largeStage.getPendingTaskCount()); + assertEquals(1, joinStage.getPendingTaskCount()); + assertEquals(2, planCoordinator.getShuffleInfos().size()); + } + finally + { + restoreInvokers(originalInvokers); + } + } + + @Test + public void outputCollectionWaitsForAllKnownAttemptsAfterOneFails() throws Exception + { + CompletableFuture failedAttempt = new CompletableFuture<>(); + CompletableFuture runningAttempt = new CompletableFuture<>(); + + CompletableFuture collectionBarrier = CompletableFuture.runAsync(() -> + { + try + { + PartitionedJoinS3QSOperator.waitForAllAttempts( + new CompletableFuture[] {failedAttempt, runningAttempt}); + } + catch (ExecutionException | InterruptedException e) + { + throw new RuntimeException(e); + } + }); + + failedAttempt.completeExceptionally(new IllegalStateException("worker failed")); + assertFalse(collectionBarrier.isDone()); + + Output output = new Output() { }; + output.setSuccessful(true); + runningAttempt.complete(output); + collectionBarrier.get(); + assertTrue(collectionBarrier.isDone()); + } + + private static ShuffleInfo validShuffleInfo() + { + return validShuffleInfo("shuffle-1"); + } + + private static ShuffleInfo validShuffleInfo(String shuffleId) + { + return new ShuffleInfo(shuffleId, + new StorageInfo(Storage.Scheme.s3qs, null, null, null, null), + "s3qs://bucket/shuffle/" + shuffleId + "/", + 1, 1, 1, 1, + Collections.singletonList(new ShuffleQueueInfo(0, shuffleId + "-p0", null))); + } + + private static PartitionInput partitionInput(int producerTaskId) + { + return partitionInput(producerTaskId, validShuffleInfo()); + } + + private static PartitionInput partitionInput(int producerTaskId, ShuffleInfo shuffleInfo) + { + PartitionInput partitionInput = new PartitionInput(); + OutputInfo outputInfo = new OutputInfo("s3qs://bucket/shuffle/" + producerTaskId + "/part", + new StorageInfo(Storage.Scheme.s3qs, null, null, null, null), true); + outputInfo.setShuffleInfo(shuffleInfo); + partitionInput.setOutput(outputInfo); + partitionInput.setProducerTaskId(producerTaskId); + return partitionInput; + } + + private static ScanTableInfo scanTableInfo(String tableName) + { + return scanTableInfo(tableName, 1); + } + + private static ScanTableInfo scanTableInfo(String tableName, int splitCount) + { + List inputSplits = new ArrayList<>(splitCount); + for (int i = 0; i < splitCount; ++i) + { + inputSplits.add(new InputSplit(Collections.singletonList( + new InputInfo("s3://bucket/" + tableName + "-" + i + ".pxl", 0, 1)))); + } + return new ScanTableInfo(tableName, true, new String[] {"key"}, + new StorageInfo(Storage.Scheme.s3, null, null, null, null), + inputSplits, null); + } + + private static PartitionedJoinInput joinInput(ShuffleInfo shuffleInfo) + { + return joinInput(shuffleInfo, shuffleInfo); + } + + private static PartitionedJoinInput joinInput(ShuffleInfo smallShuffleInfo, ShuffleInfo largeShuffleInfo) + { + PartitionedJoinInput joinInput = new PartitionedJoinInput(); + joinInput.setSmallTable(tableInfo(smallShuffleInfo)); + joinInput.setLargeTable(tableInfo(largeShuffleInfo)); + return joinInput; + } + + private static PartitionedTableInfo tableInfo(ShuffleInfo shuffleInfo) + { + PartitionedTableInfo tableInfo = new PartitionedTableInfo(); + tableInfo.setStorageInfo(new StorageInfo(Storage.Scheme.s3qs, null, null, null, null)); + tableInfo.setShuffleInfo(shuffleInfo); + return tableInfo; + } + + private static Map replaceS3QSInvokers(Invoker partitionInvoker, Invoker joinInvoker) + throws Exception + { + Map invokerMap = invokerMap(); + Map original = new EnumMap<>(WorkerType.class); + original.put(WorkerType.PARTITION_S3QS, invokerMap.get(WorkerType.PARTITION_S3QS)); + original.put(WorkerType.PARTITIONED_JOIN_S3QS, invokerMap.get(WorkerType.PARTITIONED_JOIN_S3QS)); + invokerMap.put(WorkerType.PARTITION_S3QS, partitionInvoker); + invokerMap.put(WorkerType.PARTITIONED_JOIN_S3QS, joinInvoker); + return original; + } + + private static void restoreInvokers(Map original) throws Exception + { + Map invokerMap = invokerMap(); + for (Map.Entry entry : original.entrySet()) + { + if (entry.getValue() == null) + { + invokerMap.remove(entry.getKey()); + } + else + { + invokerMap.put(entry.getKey(), entry.getValue()); + } + } + } + + @SuppressWarnings("unchecked") + private static Map invokerMap() throws Exception + { + Field field = InvokerFactory.class.getDeclaredField("invokerMap"); + field.setAccessible(true); + return (Map) field.get(InvokerFactory.Instance()); + } + + private static class RecordingInvoker implements Invoker + { + private final List inputs = new ArrayList<>(); + + @Override + public Output parseOutput(String outputJson) + { + return null; + } + + @Override + public CompletableFuture invoke(Input input) + { + inputs.add(input); + return CompletableFuture.completedFuture(new Output() { }); + } + + @Override + public String getFunctionName() + { + return "recording"; + } + + @Override + public int getMemoryMB() + { + return 0; + } + } + +} diff --git a/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/PhysicalS3QSWriter.java b/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/PhysicalS3QSWriter.java index 9ba6da62f3..e2b8f161b3 100644 --- a/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/PhysicalS3QSWriter.java +++ b/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/PhysicalS3QSWriter.java @@ -38,6 +38,7 @@ public class PhysicalS3QSWriter implements PhysicalWriter private long position; private final DataOutputStream out; private S3Queue queue = null; + private S3QueueMessage message = null; public PhysicalS3QSWriter(Storage storage, String path) throws IOException { @@ -63,6 +64,12 @@ protected void setQueue(S3Queue queue) this.queue = queue; } + protected void setQueue(S3Queue queue, S3QueueMessage message) + { + this.queue = queue; + this.message = message; + } + /** * Tell the writer the offset of next write. * @@ -121,7 +128,14 @@ public void close() throws IOException this.out.close(); if (this.queue != null && !this.queue.isClosed()) { - this.queue.push(this.pathStr); + if (this.message == null) + { + this.queue.push(this.pathStr); + } + else + { + this.queue.push(this.message.setObjectPath(this.pathStr)); + } } } catch (IOException e) diff --git a/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3QS.java b/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3QS.java index 26330f1121..e0d9196cea 100644 --- a/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3QS.java +++ b/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3QS.java @@ -37,6 +37,7 @@ import java.lang.UnsupportedOperationException; import java.time.Duration; import java.util.*; +import java.util.concurrent.ConcurrentHashMap; /** * {@link S3QS} is to write and read the small intermediate files in data shuffling. It is compatible with S3, hence its @@ -55,11 +56,7 @@ public final class S3QS extends AbstractS3 { private static final String SchemePrefix = Scheme.s3qs.name() + "://"; - //TODO: need thread safe - public final HashSet producerSet; - // maybe can use array to improve - private final HashSet PartitionSet; - private final HashMap PartitionMap; + private final Map partitionQueues; private final int invisibleTime; private SqsClient sqs; @@ -71,9 +68,7 @@ public S3QS() public S3QS(int invisibleTime){ this.connect(); - this.producerSet = new HashSet<>(); - this.PartitionSet = new HashSet<>(); - this.PartitionMap = new HashMap<>(); + this.partitionQueues = new ConcurrentHashMap<>(); this.invisibleTime = invisibleTime; } @@ -111,60 +106,77 @@ public String ensureSchemePrefix(String path) throws IOException return SchemePrefix + path; } - public void addProducer(int workerId) + public PhysicalWriter offer(S3QueueMessage mesg) throws IOException { - if(!(this.producerSet).contains(workerId)) - { - this.producerSet.add(workerId); - } + S3Queue queue = getRegisteredQueue(queueKey(mesg)); + return queue.offer(mesg); } - //TODO: GC for files, objects and sqs. - //TODO: allow separated invisible timeout config - //producers in a shuffle offer their message to s3qs. - public PhysicalWriter offer(S3QueueMessage mesg) throws IOException + /** + * Publish a structured message that does not require writing an S3 object, + * such as a producer-end marker. + */ + public void publish(S3QueueMessage mesg) throws IOException { - S3Queue queue = null; + S3Queue queue = getRegisteredQueue(queueKey(mesg)); + queue.push(mesg); + } - //if current Partition is new one, create a new queue - if(!(this.PartitionSet).contains(mesg.getPartitionNum())) + public synchronized String registerQueue(String shuffleId, int partitionId, + String queueName, String queueUrl) throws IOException + { + ShuffleQueueKey key = queueKey(shuffleId, partitionId); + S3Queue queue = partitionQueues.get(key); + if (queue != null) { - String queueUrl = ""; - try + if (queueUrl != null && !queueUrl.trim().isEmpty() && + !queue.getQueueUrl().equals(queueUrl)) { - queueUrl = createQueue(sqs,invisibleTime, - mesg.getPartitionNum()+"-"+ - (String.valueOf(System.currentTimeMillis())) - ); + throw new IOException("queue " + key + " is already registered with a different URL"); } - catch (SqsException e) + return queue.getQueueUrl(); + } + + String resolvedQueueUrl = queueUrl; + if (resolvedQueueUrl == null || resolvedQueueUrl.trim().isEmpty()) + { + if (queueName == null || queueName.trim().isEmpty()) { - //TODO: if name is duplicated in aws try again later - throw new IOException(e); + throw new IOException("queue name is empty for " + key); } - if(!(queueUrl.isEmpty())) + try { - queue = openQueue(queueUrl); - PartitionSet.add(mesg.getPartitionNum()); - PartitionMap.put(mesg.getPartitionNum(), queue); + resolvedQueueUrl = createQueue(queueName); } - else + catch (RuntimeException e) { - throw new IOException("create new queue failed."); + throw new IOException("failed to create queue " + queueName + + " for " + key, e); } } - else + + queue = openQueue(resolvedQueueUrl); + partitionQueues.put(key, queue); + return resolvedQueueUrl; + } + + /** + * Remove one local queue wrapper without deleting the remote SQS queue. + * Remote resources are owned and deleted by the query lifecycle. + */ + public void unregisterQueue(String shuffleId, int partitionId) throws IOException + { + S3Queue queue = partitionQueues.remove(queueKey(shuffleId, partitionId)); + if (queue != null) { - queue = PartitionMap.get(mesg.getPartitionNum()); - if(queue.getStopInput()) throw new IOException("queue " + mesg.getPartitionNum() + " is closed."); + queue.close(); } - return queue.offer(mesg); } - private static String createQueue(SqsClient sqsClient,int invisibleTime, String queueName) { + public String createQueue(String queueName) throws IOException + { try { - CreateQueueRequest createQueueRequest = CreateQueueRequest.builder() .queueName(queueName) .attributes(Collections.singletonMap( @@ -172,81 +184,97 @@ private static String createQueue(SqsClient sqsClient,int invisibleTime, String )) .build(); - sqsClient.createQueue(createQueueRequest); + sqs.createQueue(createQueueRequest); - GetQueueUrlResponse getQueueUrlResponse = sqsClient + GetQueueUrlResponse getQueueUrlResponse = sqs .getQueueUrl(GetQueueUrlRequest.builder().queueName(queueName).build()); return getQueueUrlResponse.queueUrl(); + } + catch (SqsException e) + { + throw new IOException("fail to create sqs queue: " + queueName, e); + } + } + + public String getQueueUrl(String queueName) throws IOException + { + try + { + return sqs.getQueueUrl(GetQueueUrlRequest.builder().queueName(queueName).build()).queueUrl(); + } + catch (SqsException e) + { + throw new IOException("fail to get sqs queue url: " + queueName, e); + } + } + public void deleteQueue(String queueUrl) throws IOException + { + try + { + sqs.deleteQueue(DeleteQueueRequest.builder().queueUrl(queueUrl).build()); } catch (SqsException e) { - throw new RuntimeException("fail to create sqs queue: " + queueName, e); + throw new IOException("fail to delete sqs queue: " + queueUrl, e); } } public S3Queue openQueue(String queueUrl) { - return new S3Queue(this, queueUrl); + return new S3Queue(this, queueUrl, invisibleTime); } - /** - @return including writer and recipthanle(to sign a mesg in sqs). - 3 situation :closed(exception) / not ready or timeout or stopInput(null) / succeed(writer) - */ public Map.Entry poll(S3QueueMessage mesg, int timeoutSec) throws IOException { - S3Queue queue = PartitionMap.get(mesg.getPartitionNum()); - // queue close means consumer don't need to listen from it - if(queue.isClosed()) throw new IOException("queue " + mesg.getPartitionNum() + " is closed."); - if(queue == null) return null; - - //if there is no more input, just wait invisibletime. if timeout, no more message - //issue: dead message maybe only handle twice, I think that's reasonable - //TODO(OUT-OF-DATE): once a message dead, push it in dead message queue. when delete a message, delete - Map.Entry pair = null; + ShuffleQueueKey key = queueKey(mesg); + S3Queue queue = getRegisteredQueue(key); + if(queue.isClosed()) + { + throw new IOException("queue " + key + " is closed."); + } + try { - if (queue.getStopInput()) - { - pair = queue.poll(queue.getInvisibleTime()); - if (pair == null) - { - // no more message - queue.close(); //logical close, no effect to consumers - return null; //come back later and find queue is closed - } - } - else - { - pair = queue.poll(timeoutSec); - if (pair == null) return null; //upstream is working, come back later - } + return queue.poll(timeoutSec); } catch (TaskErrorException e) { - //clean up + throw new IOException("failed to poll queue " + key, e); } - - queue.addConsumer(mesg.getWorkerNum()); - return pair; } - public int finishWork(S3QueueMessage mesg) throws IOException + /** + * Poll a structured shuffle message from the partition queue. + * + * This is the preferred API for S3QS shuffle consumers. It lets callers + * distinguish DATA messages from producer-end markers before deciding + * whether to open an S3 object. + */ + public S3QueuePollResult pollMessage(String shuffleId, int partitionId, int timeoutSec) throws IOException { - //DONOT close queue here. we don't know whether downstream workers meet error - //if we check to close queue here, we maybe check many times. That is expensive - // Once queue closed, consumers in the queue will know in the next poll. - String receiptHandle = mesg.getReceiptHandle(); - S3Queue queue = PartitionMap.get(mesg.getPartitionNum()); - + ShuffleQueueKey key = queueKey(shuffleId, partitionId); + S3Queue queue = getRegisteredQueue(key); + if(queue.isClosed()) + { + throw new IOException("queue " + key + " is closed."); + } - if(queue == null) + try + { + return queue.pollMessage(timeoutSec); + } + catch (TaskErrorException e) { - //queue not exist: an error, or a timeout worker - throw new IOException("queue " + mesg.getPartitionNum() + " is closed."); + throw new IOException("failed to poll queue " + key, e); } + } + + public int finishWork(S3QueueMessage mesg) throws IOException + { + String receiptHandle = mesg.getReceiptHandle(); + S3Queue queue = getRegisteredQueue(queueKey(mesg)); try { queue.deleteMessage(receiptHandle); @@ -256,30 +284,38 @@ public int finishWork(S3QueueMessage mesg) throws IOException //TODO: log return 2; } + return 0; + } - // TODO(OUT-OF-DATE): when all the consumers exit with some messages staying in DLQ, the work occur an error - // though we can actually close in poll() function, we close here make sure sqs can close safely. - // thus, we can check which part of task failed in sqs. - if(queue.removeConsumer(mesg.getWorkerNum()) && queue.isClosed()) + private S3Queue getRegisteredQueue(ShuffleQueueKey key) throws IOException + { + S3Queue queue = partitionQueues.get(key); + if (queue == null) { - try - { - DeleteQueueRequest deleteQueueRequest = DeleteQueueRequest.builder() - .queueUrl(queue.getQueueUrl()) - .build(); - - sqs.deleteQueue(deleteQueueRequest); - - } - catch (SqsException e) - { - // TODO： log - return 1; - } + throw new IOException("queue is not registered for " + key); + } + return queue; + } + private static ShuffleQueueKey queueKey(S3QueueMessage message) throws IOException + { + if (message == null) + { + throw new IOException("queue message is null"); } + return queueKey(message.getShuffleId(), message.getPartitionId()); + } - return 0; + private static ShuffleQueueKey queueKey(String shuffleId, int partitionId) throws IOException + { + try + { + return new ShuffleQueueKey(shuffleId, partitionId); + } + catch (RuntimeException e) + { + throw new IOException("invalid shuffle queue identity", e); + } } @Override public DataInputStream open(String path) throws IOException @@ -335,13 +371,11 @@ public boolean directCopy(String src, String dest) @Override public void close() throws IOException { - for (S3Queue queue : PartitionMap.values()) + for (S3Queue queue : partitionQueues.values()) { queue.close(); } - this.producerSet.clear(); - this.PartitionSet.clear(); - this.PartitionMap.clear(); + this.partitionQueues.clear(); if (this.sqs != null) { this.sqs.close(); @@ -354,17 +388,15 @@ public void close() throws IOException public void refresh() throws IOException { - for (S3Queue queue : PartitionMap.values()) + for (S3Queue queue : partitionQueues.values()) { queue.close(); } - this.producerSet.clear(); - this.PartitionSet.clear(); - this.PartitionMap.clear(); + this.partitionQueues.clear(); } public SqsClient getSqsClient() { return sqs; } -} \ No newline at end of file +} diff --git a/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3Queue.java b/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3Queue.java index 73dc0bce55..d7b27f1f1e 100644 --- a/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3Queue.java +++ b/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3Queue.java @@ -34,6 +34,7 @@ import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; +import java.util.UUID; import static software.amazon.awssdk.services.sqs.model.MessageSystemAttributeName.APPROXIMATE_RECEIVE_COUNT; @@ -64,7 +65,7 @@ public class S3Queue implements Closeable } } - private final Queue> s3PathQueue = new ConcurrentLinkedQueue<>(); + private final Queue messageQueue = new ConcurrentLinkedQueue<>(); private final String queueUrl; @@ -72,15 +73,11 @@ public class S3Queue implements Closeable private final S3QS s3qs; - private final HashSetconsumerSet; - private final Lock lock = new ReentrantLock(); private boolean closed = false; - private boolean stopInput = false; - - private int invisibleTime; + private final int invisibleTime; public S3Queue(S3QS s3qs, String queueUrl, int invisibleTime) { @@ -88,7 +85,6 @@ public S3Queue(S3QS s3qs, String queueUrl, int invisibleTime) this.queueUrl = queueUrl; this.sqsClient = this.s3qs.getSqsClient(); this.invisibleTime = invisibleTime; - this.consumerSet = new HashSet<>(); } public S3Queue(S3QS s3qs, String queueUrl) { @@ -107,52 +103,21 @@ public int getInvisibleTime() } - //s3qs is the highest manager of a shuffle, so producerSet is in charge of s3qs. - private void removeProducer(int workerId) throws IOException - { - this.s3qs.producerSet.remove(workerId); - if(this.s3qs.producerSet.isEmpty()) - { - this.stopInput(); - //this.push(""); - } - } - - public void addConsumer(int workerId) - { - if(!(this.consumerSet).contains(workerId)) - { - this.consumerSet.add(workerId); - } - } - - //maybe unnecessary - public boolean removeConsumer(int workerId) - { - this.consumerSet.remove(workerId); - // TODO: close queue - return consumerSet.isEmpty(); - } - - //TODO(OUT-OF-DATE): Implement DLQ to handle bad message. - - /** - * Poll one object path from the SQS queue and create a physical reader for the object. - * Calling this method can receive a batch of object paths from SQS using long polling - * (the batch size is configured by s3qs.poll.batch.size in PIXELS_HOME/pixels.properties) - * and add the paths into a local in-memory queue. Thus reduces the receive-message requests sent to SQS. + * Poll one structured shuffle message from the SQS queue. + * + * This method does not open S3 objects. The caller should inspect the + * message type first and only create a reader for DATA messages. * * @param timeoutSec the max time in seconds to wait if the queue is currently empty, * a valid wait time should be between 1 and 20 seconds * @return null if the queue is still empty after timeout - * @throws IOException if fails to create the physical reader for the path + * @throws IOException if fails to parse the queue message */ - public Map.Entry poll(int timeoutSec) throws IOException, SqsException,TaskErrorException - + public S3QueuePollResult pollMessage(int timeoutSec) throws IOException, SqsException, TaskErrorException { - Map.Entry s3Path = this.s3PathQueue.poll(); - if (s3Path == null) + S3QueuePollResult result = this.messageQueue.poll(); + if (result == null) { if (timeoutSec < 1) { @@ -165,8 +130,7 @@ public Map.Entry poll(int timeoutSec) throws IOException this.lock.lock(); try { - // try poll from queue again to see if another thread has received the messages from sqs - while ((s3Path = this.s3PathQueue.poll()) == null) + while ((result = this.messageQueue.poll()) == null) { ReceiveMessageRequest request = ReceiveMessageRequest.builder() .queueUrl(queueUrl) @@ -178,27 +142,23 @@ public Map.Entry poll(int timeoutSec) throws IOException { for (Message message : response.messages()) { - String path = message.body(); - String receiptHandle = message.receiptHandle(); - this.s3PathQueue.add(new AbstractMap.SimpleEntry<>(path, receiptHandle)); - String countStr = message.attributes().get(APPROXIMATE_RECEIVE_COUNT); if (countStr == null) { - // If no value is returned, it may be because the property was not requested or the message does not contain that property. throw new TaskErrorException("ApproximateReceiveCount not returned"); } - else + int count = Integer.parseInt(countStr); + if (count > 2) { - int count = Integer.parseInt(countStr); - // because we can only promise two receipts can be handled - if (count > 2) throw new TaskErrorException("Dead message occurred"); + throw new TaskErrorException("Dead message occurred"); } + + S3QueueMessage queueMessage = S3QueueMessage.fromMessageBody(message.body()); + this.messageQueue.add(new S3QueuePollResult(message.receiptHandle(), queueMessage)); } } else { - // the sqs queue is also empty，timeout，invoker handle this situation, which means timeout. return null; } } @@ -208,9 +168,40 @@ public Map.Entry poll(int timeoutSec) throws IOException this.lock.unlock(); } } - String receiptHandle = s3Path.getValue(); - PhysicalReader reader = PhysicalReaderUtil.newPhysicalReader(this.s3qs, s3Path.getKey()); - return new AbstractMap.SimpleEntry(receiptHandle, reader); + return result; + } + + /** + * Poll one DATA message and create a physical reader for its S3 object. + * + * This method is kept as a compatibility adapter for existing tests and + * should not be used by the S3QS shuffle consumer path because control + * messages do not have S3 objects. + * + * Calling this method can receive a batch of object paths from SQS using long polling + * (the batch size is configured by s3qs.poll.batch.size in PIXELS_HOME/pixels.properties) + * and add the paths into a local in-memory queue. Thus reduces the receive-message requests sent to SQS. + * + * @param timeoutSec the max time in seconds to wait if the queue is currently empty, + * a valid wait time should be between 1 and 20 seconds + * @return null if the queue is still empty after timeout + * @throws IOException if fails to create the physical reader for the path + */ + public Map.Entry poll(int timeoutSec) throws IOException, SqsException,TaskErrorException + + { + S3QueuePollResult result = pollMessage(timeoutSec); + if (result == null) + { + return null; + } + S3QueueMessage message = result.getMessage(); + if (!message.isData()) + { + throw new IOException("polled a non-DATA message from compatibility poll: " + message.getMessageType()); + } + PhysicalReader reader = PhysicalReaderUtil.newPhysicalReader(this.s3qs, message.getObjectPath()); + return new AbstractMap.SimpleEntry(result.getReceiptHandle(), reader); } // do not need to delete local queue @@ -226,11 +217,19 @@ public void deleteMessage(String receiptHandle) throws SqsException protected void push(String objectPath) throws IOException + { + push(S3QueueMessage.data("", 0, 0, 0, 0L, objectPath)); + } + + /** + * Push one structured shuffle message into this SQS queue. + */ + protected void push(S3QueueMessage message) throws IOException { try { SendMessageRequest request = SendMessageRequest.builder() - .queueUrl(queueUrl).messageBody(objectPath).build(); + .queueUrl(queueUrl).messageBody(message.toMessageBody()).build(); sqsClient.sendMessage(request); } catch (Exception e) @@ -248,23 +247,16 @@ protected void push(String objectPath) throws IOException */ public PhysicalWriter offer(S3QueueMessage body) throws IOException { - //TODO: same name issue String objectPath = getObjectPath(body); PhysicalS3QSWriter writer = (PhysicalS3QSWriter) PhysicalWriterUtil .newPhysicalWriter(this.s3qs, objectPath, false); - writer.setQueue(this); - if(endWork(body)) removeProducer(body.getWorkerNum()); + writer.setQueue(this, body.setMessageType(S3QueueMessage.MessageType.DATA)); return writer; } private String getObjectPath(S3QueueMessage body) throws IOException { - return body.getObjectPath()+body.getPartitionNum() + "/"+ String.valueOf(System.currentTimeMillis()) ; - } - - private boolean endWork(S3QueueMessage body) throws IOException - { - return body.getEndWork(); + return body.getObjectPath() + body.getPartitionNum() + "/" + UUID.randomUUID(); } @@ -273,20 +265,10 @@ public boolean isClosed() return closed; } - public void stopInput() - { - this.stopInput = true; - } - - public boolean getStopInput() - { - return this.stopInput; - } - @Override public void close() throws IOException { - this.s3PathQueue.clear(); + this.messageQueue.clear(); this.closed = true; // do not close the s3qs storage and the sqs client as they are cached } diff --git a/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3QueueMessage.java b/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3QueueMessage.java index ad31fb6a93..9fca6cf093 100644 --- a/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3QueueMessage.java +++ b/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3QueueMessage.java @@ -19,86 +19,277 @@ */ package io.pixelsdb.pixels.storage.s3qs; +import java.io.IOException; +import java.io.StringReader; +import java.io.StringWriter; +import java.util.Properties; + /** - * Unified message format for communication between worker and s3qs. - * @author yanhaoting - * @create 2025-12-19 + * The structured message body sent through a S3QS partition queue. + * + * A producer task sends {@link MessageType#DATA} after an S3 object is durably + * written, and sends {@link MessageType#PRODUCER_END} to every partition when + * the producer task has no more data for this shuffle edge. */ public class S3QueueMessage { - private String objectPath; - private int workerNum = 0; - private int partitionNum = 0; - private boolean endWork = false;//specific to workers, indicating whether his work is finished - private String receiptHandle = "";//specific to consumers, indicating which mesg is consumed. - private long timestamp = System.currentTimeMillis(); - private String metadata = ""; - - public S3QueueMessage() + public enum MessageType { + DATA, + PRODUCER_END } + private static final int VERSION = 1; + private static final String EMPTY = ""; + + private int version = VERSION; + private MessageType messageType = MessageType.DATA; + private String shuffleId = EMPTY; + private int partitionId = 0; + private int producerId = 0; + private int producerAttemptId = 0; + private long sequenceId = 0L; + private long timestamp = System.currentTimeMillis(); + private String objectPath = EMPTY; + private long rowCount = -1L; + private long byteSize = -1L; + private String endReason = EMPTY; + private String receiptHandle = EMPTY; + private String metadata = EMPTY; + + public S3QueueMessage() { } + + /** + * Construct a DATA message using the given object path prefix. The final + * object path is filled by {@link S3Queue#offer(S3QueueMessage)}. + */ public S3QueueMessage(String objectPath) { this.objectPath = objectPath; this.timestamp = System.currentTimeMillis(); } - public String getObjectPath() + public static S3QueueMessage data(String shuffleId, int partitionId, int producerId, + int producerAttemptId, long sequenceId, String objectPath) + { + return new S3QueueMessage() + .setMessageType(MessageType.DATA) + .setShuffleId(shuffleId) + .setPartitionId(partitionId) + .setProducerId(producerId) + .setProducerAttemptId(producerAttemptId) + .setSequenceId(sequenceId) + .setObjectPath(objectPath); + } + + public static S3QueueMessage producerEnd(String shuffleId, int partitionId, int producerId, + int producerAttemptId, long sequenceId) + { + return new S3QueueMessage() + .setMessageType(MessageType.PRODUCER_END) + .setShuffleId(shuffleId) + .setPartitionId(partitionId) + .setProducerId(producerId) + .setProducerAttemptId(producerAttemptId) + .setSequenceId(sequenceId) + .setEndReason("NORMAL"); + } + + /** + * Serialize this message into a SQS message body. + */ + public String toMessageBody() throws IOException + { + Properties properties = new Properties(); + properties.setProperty("version", String.valueOf(version)); + properties.setProperty("messageType", messageType.name()); + properties.setProperty("shuffleId", nullToEmpty(shuffleId)); + properties.setProperty("partitionId", String.valueOf(partitionId)); + properties.setProperty("producerId", String.valueOf(producerId)); + properties.setProperty("producerAttemptId", String.valueOf(producerAttemptId)); + properties.setProperty("sequenceId", String.valueOf(sequenceId)); + properties.setProperty("timestamp", String.valueOf(timestamp)); + properties.setProperty("objectPath", nullToEmpty(objectPath)); + properties.setProperty("rowCount", String.valueOf(rowCount)); + properties.setProperty("byteSize", String.valueOf(byteSize)); + properties.setProperty("endReason", nullToEmpty(endReason)); + properties.setProperty("metadata", nullToEmpty(metadata)); + + StringWriter writer = new StringWriter(); + properties.store(writer, "s3qs-shuffle-message"); + return writer.toString(); + } + + /** + * Parse a structured S3QS shuffle message from a SQS message body. + */ + public static S3QueueMessage fromMessageBody(String body) throws IOException + { + Properties properties = new Properties(); + properties.load(new StringReader(body)); + S3QueueMessage message = new S3QueueMessage(); + message.setVersion(Integer.parseInt(properties.getProperty("version", String.valueOf(VERSION)))); + message.setMessageType(MessageType.valueOf(properties.getProperty("messageType", MessageType.DATA.name()))); + message.setShuffleId(properties.getProperty("shuffleId", EMPTY)); + message.setPartitionId(Integer.parseInt(properties.getProperty("partitionId", "0"))); + message.setProducerId(Integer.parseInt(properties.getProperty("producerId", "0"))); + message.setProducerAttemptId(Integer.parseInt(properties.getProperty("producerAttemptId", "0"))); + message.setSequenceId(Long.parseLong(properties.getProperty("sequenceId", "0"))); + message.setTimestamp(Long.parseLong(properties.getProperty("timestamp", "0"))); + message.setObjectPath(properties.getProperty("objectPath", EMPTY)); + message.setRowCount(Long.parseLong(properties.getProperty("rowCount", "-1"))); + message.setByteSize(Long.parseLong(properties.getProperty("byteSize", "-1"))); + message.setEndReason(properties.getProperty("endReason", EMPTY)); + message.setMetadata(properties.getProperty("metadata", EMPTY)); + return message; + } + + private static String nullToEmpty(String value) + { + return value == null ? EMPTY : value; + } + + public boolean isData() + { + return messageType == MessageType.DATA; + } + + public boolean isProducerEnd() + { + return messageType == MessageType.PRODUCER_END; + } + + public int getVersion() + { + return version; + } + + public S3QueueMessage setVersion(int version) + { + this.version = version; + return this; + } + + public MessageType getMessageType() + { + return messageType; + } + + public S3QueueMessage setMessageType(MessageType messageType) + { + this.messageType = messageType == null ? MessageType.DATA : messageType; + return this; + } + + public String getShuffleId() + { + return shuffleId; + } + + public S3QueueMessage setShuffleId(String shuffleId) + { + this.shuffleId = nullToEmpty(shuffleId); + return this; + } + + public String getObjectPath() { return objectPath; } public S3QueueMessage setObjectPath(String objectPath) { - this.objectPath = objectPath; + this.objectPath = nullToEmpty(objectPath); return this; } - public int getWorkerNum() { - return this.workerNum; + public int getPartitionId() + { + return partitionId; + } + + public S3QueueMessage setPartitionId(int partitionId) + { + this.partitionId = partitionId; + return this; + } + + public int getPartitionNum() + { + return partitionId; + } + + public S3QueueMessage setPartitionNum(int partitionNum) + { + return setPartitionId(partitionNum); + } + + public int getProducerId() + { + return producerId; + } + + public S3QueueMessage setProducerId(int producerId) + { + this.producerId = producerId; + return this; + } + + public int getWorkerNum() + { + return producerId; + } + + public S3QueueMessage setWorkerNum(int workerNum) + { + return setProducerId(workerNum); + } + + public int getProducerAttemptId() + { + return producerAttemptId; } - public S3QueueMessage setWorkerNum(int WorkerNum) + public S3QueueMessage setProducerAttemptId(int producerAttemptId) { - this.workerNum = WorkerNum; + this.producerAttemptId = producerAttemptId; return this; } - public int getPartitionNum() + public long getSequenceId() { - return this.partitionNum; + return sequenceId; } - public S3QueueMessage setPartitionNum(int PartitionNum) + public S3QueueMessage setSequenceId(long sequenceId) { - this.partitionNum = PartitionNum; + this.sequenceId = sequenceId; return this; } - public boolean getEndWork() + public boolean getEndWork() { - return this.endWork; + return isProducerEnd(); } public S3QueueMessage setEndwork(boolean endwork) { - this.endWork = endwork; + this.messageType = endwork ? MessageType.PRODUCER_END : MessageType.DATA; return this; } - public String getReceiptHandle() + public String getReceiptHandle() { - return this.receiptHandle; + return receiptHandle; } - public S3QueueMessage setReceiptHandle(String ReceiptHandle) + public S3QueueMessage setReceiptHandle(String receiptHandle) { - this.receiptHandle = ReceiptHandle; + this.receiptHandle = nullToEmpty(receiptHandle); return this; } - public long getTimestamp() + public long getTimestamp() { return timestamp; } @@ -109,14 +300,47 @@ public S3QueueMessage setTimestamp(long timestamp) return this; } - public String getMetadata() + public long getRowCount() + { + return rowCount; + } + + public S3QueueMessage setRowCount(long rowCount) + { + this.rowCount = rowCount; + return this; + } + + public long getByteSize() + { + return byteSize; + } + + public S3QueueMessage setByteSize(long byteSize) + { + this.byteSize = byteSize; + return this; + } + + public String getEndReason() + { + return endReason; + } + + public S3QueueMessage setEndReason(String endReason) + { + this.endReason = nullToEmpty(endReason); + return this; + } + + public String getMetadata() { return metadata; } public S3QueueMessage setMetadata(String metadata) { - this.metadata = metadata; + this.metadata = nullToEmpty(metadata); return this; } } diff --git a/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3QueuePollResult.java b/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3QueuePollResult.java new file mode 100644 index 0000000000..b6bb65cd22 --- /dev/null +++ b/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/S3QueuePollResult.java @@ -0,0 +1,49 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public + * License along with Pixels. If not, see + * . + */ +package io.pixelsdb.pixels.storage.s3qs; + +/** + * The result of polling a S3QS partition queue. + * + * It keeps the SQS receipt handle together with the structured shuffle message + * so the caller can decide whether to open a data object or record a control + * marker before acknowledging the message. + */ +public class S3QueuePollResult +{ + private final String receiptHandle; + private final S3QueueMessage message; + + public S3QueuePollResult(String receiptHandle, S3QueueMessage message) + { + this.receiptHandle = receiptHandle; + this.message = message; + } + + public String getReceiptHandle() + { + return receiptHandle; + } + + public S3QueueMessage getMessage() + { + return message; + } +} diff --git a/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/ShuffleQueueKey.java b/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/ShuffleQueueKey.java new file mode 100644 index 0000000000..4f36d1ea0a --- /dev/null +++ b/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/ShuffleQueueKey.java @@ -0,0 +1,82 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public + * License along with Pixels. If not, see + * . + */ +package io.pixelsdb.pixels.storage.s3qs; + +import java.util.Objects; + +import static com.google.common.base.Preconditions.checkArgument; +import static java.util.Objects.requireNonNull; + +/** + * Identifies one physical partition queue within one shuffle edge. + * + * Partition ids are only unique inside a shuffle. Including shuffleId keeps + * the two sides of a join and concurrent queries isolated in a warm worker JVM. + */ +public final class ShuffleQueueKey +{ + private final String shuffleId; + private final int partitionId; + + public ShuffleQueueKey(String shuffleId, int partitionId) + { + this.shuffleId = requireNonNull(shuffleId, "shuffleId is null").trim(); + checkArgument(!this.shuffleId.isEmpty(), "shuffleId is empty"); + checkArgument(partitionId >= 0, "partitionId is negative"); + this.partitionId = partitionId; + } + + public String getShuffleId() + { + return shuffleId; + } + + public int getPartitionId() + { + return partitionId; + } + + @Override + public boolean equals(Object other) + { + if (this == other) + { + return true; + } + if (!(other instanceof ShuffleQueueKey)) + { + return false; + } + ShuffleQueueKey that = (ShuffleQueueKey) other; + return partitionId == that.partitionId && shuffleId.equals(that.shuffleId); + } + + @Override + public int hashCode() + { + return Objects.hash(shuffleId, partitionId); + } + + @Override + public String toString() + { + return shuffleId + ":" + partitionId; + } +} diff --git a/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/exception/TaskErrorException.java b/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/exception/TaskErrorException.java index b7695a2227..54bd1080b4 100644 --- a/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/exception/TaskErrorException.java +++ b/pixels-storage/pixels-storage-s3qs/src/main/java/io/pixelsdb/pixels/storage/s3qs/exception/TaskErrorException.java @@ -20,7 +20,7 @@ package io.pixelsdb.pixels.storage.s3qs.exception; /** - * @author yanhaoting + * @author Haoting Yan * @create 2025-12-19 */ public class TaskErrorException extends Exception diff --git a/pixels-storage/pixels-storage-s3qs/src/test/java/io/pixelsdb/pixels/storage/s3qs/TestS3QS.java b/pixels-storage/pixels-storage-s3qs/src/test/java/io/pixelsdb/pixels/storage/s3qs/TestS3QS.java index 6fe62f0411..32fbef98c7 100644 --- a/pixels-storage/pixels-storage-s3qs/src/test/java/io/pixelsdb/pixels/storage/s3qs/TestS3QS.java +++ b/pixels-storage/pixels-storage-s3qs/src/test/java/io/pixelsdb/pixels/storage/s3qs/TestS3QS.java @@ -20,228 +20,429 @@ package io.pixelsdb.pixels.storage.s3qs; import io.pixelsdb.pixels.common.physical.PhysicalReader; +import io.pixelsdb.pixels.common.physical.PhysicalReaderUtil; import io.pixelsdb.pixels.common.physical.PhysicalWriter; -import io.pixelsdb.pixels.common.physical.Storage; -import io.pixelsdb.pixels.common.physical.StorageFactory; import org.junit.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.awssdk.services.sqs.SqsClient; +import software.amazon.awssdk.services.sqs.model.CreateQueueRequest; +import software.amazon.awssdk.services.sqs.model.DeleteMessageRequest; +import software.amazon.awssdk.services.sqs.model.GetQueueUrlRequest; +import software.amazon.awssdk.services.sqs.model.GetQueueUrlResponse; +import software.amazon.awssdk.services.sqs.model.Message; +import software.amazon.awssdk.services.sqs.model.MessageSystemAttributeName; +import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest; +import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse; +import software.amazon.awssdk.services.sqs.model.SendMessageRequest; import java.io.IOException; -import java.util.HashMap; -import java.util.concurrent.*; +import java.lang.reflect.Field; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assume.assumeTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static software.amazon.awssdk.services.sqs.model.MessageSystemAttributeName.APPROXIMATE_RECEIVE_COUNT; /** - * @author hank - * @create 2025-09-27 + * Unit tests for the current S3QS shuffle queue protocol. */ public class TestS3QS { - private final ExecutorService executor = Executors.newFixedThreadPool(2); + private static final String QUEUE_URL = "https://sqs.us-east-1.amazonaws.com/123456789012/pixels-shuffle-test"; @Test - public void TestWriterAndReader() throws IOException, InterruptedException, ExecutionException - { - S3QS s3qs = (S3QS) StorageFactory.Instance().getStorage(Storage.Scheme.s3qs); - Future future = this.executor.submit(() -> { - byte[] buffer = new byte[8 * 1024 * 1024]; - long startTime = System.currentTimeMillis(); - s3qs.addProducer(0); s3qs.addProducer(1); - for (int i = 0; i < 2; i++) - { - S3QueueMessage body = new S3QueueMessage() - .setObjectPath("pixels-turbo-intermediate/shuffle/") - .setWorkerNum(i) - .setPartitionNum(99) - .setEndwork(true); - try (PhysicalWriter writer = s3qs.offer(body)) - { - writer.append(buffer); - System.out.println("Wrote partition " + i); // 添加日志 - } - catch (IOException e) - { - System.err.println("Error in iteration " + i + ": " + e.getMessage()); - throw new RuntimeException(e); - } - } + public void testDataMessageRoundTrip() throws IOException + { + S3QueueMessage message = S3QueueMessage.data("shuffle-1", 3, 7, 2, 99L, + "bucket/shuffle/3/object") + .setTimestamp(12345L) + .setRowCount(1000L) + .setByteSize(4096L) + .setMetadata("source=unit"); - System.out.println("write finished in " + (System.currentTimeMillis() - startTime) + " ms"); + S3QueueMessage parsed = S3QueueMessage.fromMessageBody(message.toMessageBody()); - startTime = System.currentTimeMillis(); - for (int i = 0; i < 3; i++) - { - S3QueueMessage body = new S3QueueMessage() - .setObjectPath("pixels-turbo-intermediate/shuffle/") - .setPartitionNum(99) - .setWorkerNum(2+i); - PhysicalReader reader = null; - String receiptHandle = null; - try - { - HashMap.Entry pair = s3qs.poll(body,20); - if(pair != null) - { - reader = pair.getValue(); - receiptHandle = pair.getKey(); - body.setReceiptHandle(receiptHandle); - System.out.println("read object " + reader.getPath()); // 添加日志 - } - } - catch (IOException e) - { - System.err.println("Error in iteration " + i + ": " + e.getMessage()); - throw new RuntimeException(e); - } - finally - { - if(reader != null) - { - try - { - reader.close(); - } - catch (IOException e) - { - throw new RuntimeException(e); - } - try - { - s3qs.finishWork(body); - } - catch (IOException e) - { - throw new RuntimeException(e); - } - } - } - } + assertTrue(parsed.isData()); + assertFalse(parsed.isProducerEnd()); + assertEquals("shuffle-1", parsed.getShuffleId()); + assertEquals(3, parsed.getPartitionId()); + assertEquals(3, parsed.getPartitionNum()); + assertEquals(7, parsed.getProducerId()); + assertEquals(7, parsed.getWorkerNum()); + assertEquals(2, parsed.getProducerAttemptId()); + assertEquals(99L, parsed.getSequenceId()); + assertEquals(12345L, parsed.getTimestamp()); + assertEquals("bucket/shuffle/3/object", parsed.getObjectPath()); + assertEquals(1000L, parsed.getRowCount()); + assertEquals(4096L, parsed.getByteSize()); + assertEquals("source=unit", parsed.getMetadata()); + } - System.out.println("read finished in " + (System.currentTimeMillis() - startTime) + " ms"); - try - { - s3qs.refresh(); - } - catch (IOException e) - { - throw new RuntimeException(e); - } - }); + @Test + public void testProducerEndMessageRoundTrip() throws IOException + { + S3QueueMessage message = S3QueueMessage.producerEnd("shuffle-2", 5, 11, 4, 123L) + .setTimestamp(67890L) + .setMetadata("marker=true"); + + S3QueueMessage parsed = S3QueueMessage.fromMessageBody(message.toMessageBody()); + + assertTrue(parsed.isProducerEnd()); + assertFalse(parsed.isData()); + assertEquals(S3QueueMessage.MessageType.PRODUCER_END, parsed.getMessageType()); + assertEquals("shuffle-2", parsed.getShuffleId()); + assertEquals(5, parsed.getPartitionId()); + assertEquals(11, parsed.getProducerId()); + assertEquals(4, parsed.getProducerAttemptId()); + assertEquals(123L, parsed.getSequenceId()); + assertEquals(67890L, parsed.getTimestamp()); + assertEquals("NORMAL", parsed.getEndReason()); + assertEquals("marker=true", parsed.getMetadata()); + assertTrue(parsed.getEndWork()); + } + + @Test + public void testRegisterQueueWithResolvedUrlIsLocalAndIdempotent() throws Exception + { + SqsClient sqs = mock(SqsClient.class); + S3QS s3qs = newS3QS(sqs); + + String firstUrl = s3qs.registerQueue("shuffle-idempotent", 3, "ignored-name", QUEUE_URL); + String secondUrl = s3qs.registerQueue("shuffle-idempotent", 3, "another-name", QUEUE_URL); + + assertEquals(QUEUE_URL, firstUrl); + assertEquals(QUEUE_URL, secondUrl); + verify(sqs, never()).createQueue(any(CreateQueueRequest.class)); + } + + @Test(expected = IOException.class) + public void testRegisterQueueRejectsConflictingUrlForSameShufflePartition() throws Exception + { + S3QS s3qs = newS3QS(mock(SqsClient.class)); + s3qs.registerQueue("shuffle-conflict", 3, "ignored-name", QUEUE_URL); + s3qs.registerQueue("shuffle-conflict", 3, "another-name", "another-url"); + } + + @Test(expected = IOException.class) + public void testRegisterQueueRejectsEmptyShuffleId() throws Exception + { + S3QS s3qs = newS3QS(mock(SqsClient.class)); + s3qs.registerQueue(" ", 0, "ignored-name", QUEUE_URL); + } + @Test + public void testConcurrentRegistrationCreatesOneQueueForOneShufflePartition() throws Exception + { + SqsClient sqs = mock(SqsClient.class); + when(sqs.getQueueUrl(any(GetQueueUrlRequest.class))) + .thenReturn(GetQueueUrlResponse.builder().queueUrl(QUEUE_URL).build()); + S3QS s3qs = newS3QS(sqs); + ExecutorService executor = Executors.newFixedThreadPool(8); try { - future.get(); + Callable registration = + () -> s3qs.registerQueue("shuffle-concurrent", 0, "concurrent-queue", null); + List> results = executor.invokeAll(Collections.nCopies(16, registration)); + + for (Future result : results) + { + assertEquals(QUEUE_URL, result.get()); + } + verify(sqs, times(1)).createQueue(any(CreateQueueRequest.class)); } - catch (ExecutionException e) + finally { - //e.getCause().printStackTrace(); - throw e; + executor.shutdownNow(); } - - this.executor.shutdown(); - this.executor.awaitTermination(100, TimeUnit.HOURS); } @Test - public void testWriter() throws IOException, InterruptedException, ExecutionException + public void testSamePartitionInDifferentShufflesRoutesIndependently() throws Exception { - S3QS s3qs = (S3QS) StorageFactory.Instance().getStorage(Storage.Scheme.s3qs); - //S3Queue queue = s3qs.openQueue("https://sqs.us-east-2.amazonaws.com/970089764833/pixels-shuffle"); + String smallQueueUrl = QUEUE_URL + "-small"; + String largeQueueUrl = QUEUE_URL + "-large"; + SqsClient sqs = mock(SqsClient.class); + S3QS s3qs = newS3QS(sqs); + s3qs.registerQueue("small-shuffle", 0, "small", smallQueueUrl); + s3qs.registerQueue("large-shuffle", 0, "large", largeQueueUrl); - Future future = this.executor.submit(() -> - { - byte[] buffer = new byte[8 * 1024 * 1024]; - long startTime = System.currentTimeMillis(); - for (int i = 0; i < 2; i++) - { - S3QueueMessage body = new S3QueueMessage() - .setObjectPath("pixels-turbo-intermediate/shuffle/" + i + "/") - .setPartitionNum(2) - .setEndwork(false); - try (PhysicalWriter writer = s3qs.offer(body)) - { - writer.append(buffer); - System.out.println("Wrote partition " + i); // 添加日志 - } - catch (IOException e) - { - System.err.println("Error in iteration " + i + ": " + e.getMessage()); - throw new RuntimeException(e); - } - } - System.out.println("write finished in " + (System.currentTimeMillis() - startTime) + " ms"); + S3QueueMessage small = S3QueueMessage.producerEnd("small-shuffle", 0, 1, 0, 1L); + S3QueueMessage large = S3QueueMessage.producerEnd("large-shuffle", 0, 2, 0, 1L); + s3qs.publish(small); + s3qs.publish(large); - try + ArgumentCaptor sends = ArgumentCaptor.forClass(SendMessageRequest.class); + verify(sqs, times(2)).sendMessage(sends.capture()); + assertEquals(smallQueueUrl, sends.getAllValues().get(0).queueUrl()); + assertEquals(largeQueueUrl, sends.getAllValues().get(1).queueUrl()); + + when(sqs.receiveMessage(any(ReceiveMessageRequest.class))).thenAnswer(invocation -> + { + ReceiveMessageRequest request = invocation.getArgument(0); + if (smallQueueUrl.equals(request.queueUrl())) { - s3qs.refresh(); + return receive("small-receipt", small); } - catch (IOException e) + if (largeQueueUrl.equals(request.queueUrl())) { - throw new RuntimeException(e); + return receive("large-receipt", large); } + throw new AssertionError("unexpected queue URL: " + request.queueUrl()); }); + S3QueuePollResult smallResult = s3qs.pollMessage("small-shuffle", 0, 1); + S3QueuePollResult largeResult = s3qs.pollMessage("large-shuffle", 0, 1); + assertEquals("small-shuffle", smallResult.getMessage().getShuffleId()); + assertEquals("large-shuffle", largeResult.getMessage().getShuffleId()); + + smallResult.getMessage().setReceiptHandle(smallResult.getReceiptHandle()); + largeResult.getMessage().setReceiptHandle(largeResult.getReceiptHandle()); + assertEquals(0, s3qs.finishWork(smallResult.getMessage())); + assertEquals(0, s3qs.finishWork(largeResult.getMessage())); + verify(sqs).deleteMessage(DeleteMessageRequest.builder() + .queueUrl(smallQueueUrl).receiptHandle("small-receipt").build()); + verify(sqs).deleteMessage(DeleteMessageRequest.builder() + .queueUrl(largeQueueUrl).receiptHandle("large-receipt").build()); + } + + @Test + public void testPublishSendsStructuredProducerEndMessage() throws Exception + { + SqsClient sqs = mock(SqsClient.class); + S3QS s3qs = newS3QS(sqs); + s3qs.registerQueue("shuffle-3", 9, "ignored-name", QUEUE_URL); + + S3QueueMessage marker = S3QueueMessage.producerEnd("shuffle-3", 9, 17, 1, 88L); + s3qs.publish(marker); + + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(SendMessageRequest.class); + verify(sqs).sendMessage(requestCaptor.capture()); + SendMessageRequest request = requestCaptor.getValue(); + S3QueueMessage sent = S3QueueMessage.fromMessageBody(request.messageBody()); + + assertEquals(QUEUE_URL, request.queueUrl()); + assertTrue(sent.isProducerEnd()); + assertEquals("shuffle-3", sent.getShuffleId()); + assertEquals(9, sent.getPartitionId()); + assertEquals(17, sent.getProducerId()); + assertEquals(1, sent.getProducerAttemptId()); + assertEquals(88L, sent.getSequenceId()); + } + + @Test + public void testPollMessageParsesStructuredMessageAndReceiptHandle() throws Exception + { + SqsClient sqs = mock(SqsClient.class); + S3QS s3qs = newS3QS(sqs); + s3qs.registerQueue("shuffle-4", 4, "ignored-name", QUEUE_URL); + S3QueueMessage body = S3QueueMessage.data("shuffle-4", 4, 2, 0, 10L, + "bucket/shuffle/4/object"); + + Message sqsMessage = Message.builder() + .body(body.toMessageBody()) + .receiptHandle("receipt-1") + .attributes(receiveCount(1)) + .build(); + when(sqs.receiveMessage(any(ReceiveMessageRequest.class))) + .thenReturn(ReceiveMessageResponse.builder().messages(sqsMessage).build()); + + S3QueuePollResult result = s3qs.pollMessage("shuffle-4", 4, 99); + + assertNotNull(result); + assertEquals("receipt-1", result.getReceiptHandle()); + assertTrue(result.getMessage().isData()); + assertEquals("shuffle-4", result.getMessage().getShuffleId()); + assertEquals("bucket/shuffle/4/object", result.getMessage().getObjectPath()); + + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(ReceiveMessageRequest.class); + verify(sqs).receiveMessage(requestCaptor.capture()); + assertEquals(QUEUE_URL, requestCaptor.getValue().queueUrl()); + assertEquals(Integer.valueOf(20), requestCaptor.getValue().waitTimeSeconds()); + } + + @Test(expected = IOException.class) + public void testCompatibilityPollRejectsControlMessage() throws Exception + { + SqsClient sqs = mock(SqsClient.class); + S3QS s3qs = newS3QS(sqs); + s3qs.registerQueue("shuffle-5", 6, "ignored-name", QUEUE_URL); + S3QueueMessage body = S3QueueMessage.producerEnd("shuffle-5", 6, 2, 0, 10L); + + Message sqsMessage = Message.builder() + .body(body.toMessageBody()) + .receiptHandle("receipt-control") + .attributes(receiveCount(1)) + .build(); + when(sqs.receiveMessage(any(ReceiveMessageRequest.class))) + .thenReturn(ReceiveMessageResponse.builder().messages(sqsMessage).build()); + + s3qs.poll(new S3QueueMessage().setShuffleId("shuffle-5").setPartitionId(6), 1); + } + + @Test + public void testFinishWorkDeletesReceiptHandle() throws Exception + { + SqsClient sqs = mock(SqsClient.class); + S3QS s3qs = newS3QS(sqs); + s3qs.registerQueue("shuffle-ack", 8, "ignored-name", QUEUE_URL); + + int result = s3qs.finishWork(new S3QueueMessage() + .setShuffleId("shuffle-ack") + .setPartitionId(8) + .setReceiptHandle("receipt-finished")); + + assertEquals(0, result); + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(DeleteMessageRequest.class); + verify(sqs).deleteMessage(requestCaptor.capture()); + assertEquals(QUEUE_URL, requestCaptor.getValue().queueUrl()); + assertEquals("receipt-finished", requestCaptor.getValue().receiptHandle()); + } + + @Test(expected = IOException.class) + public void testUnregisteredPartitionCannotPublish() throws Exception + { + S3QS s3qs = newS3QS(mock(SqsClient.class)); + s3qs.publish(S3QueueMessage.producerEnd("shuffle-6", 12, 1, 0, 1L)); + } + + @Test + public void testAwsS3QSOfferPollReadAckAndCleanupIntegration() throws Exception + { + String bucket = System.getenv("PIXELS_S3QS_IT_BUCKET"); + String prefix = trimSlashes(System.getenv("PIXELS_S3QS_IT_PREFIX")); + String queuePrefix = System.getenv("PIXELS_S3QS_IT_QUEUE_PREFIX"); + assumeTrue("set PIXELS_S3QS_IT_BUCKET to run the AWS S3QS integration test", + bucket != null && !bucket.trim().isEmpty()); + assumeTrue("set PIXELS_S3QS_IT_QUEUE_PREFIX to run the AWS S3QS integration test", + queuePrefix != null && !queuePrefix.trim().isEmpty()); + + String suffix = UUID.randomUUID().toString().replace("-", ""); + String shuffleId = "it-" + suffix; + String queueUrl = null; + String objectRoot = bucket + "/" + (prefix.isEmpty() ? "" : prefix + "/") + shuffleId + "/"; + S3QS s3qs = new S3QS(30); + try { - future.get(); + queueUrl = s3qs.registerQueue(shuffleId, 0, queuePrefix + "-" + suffix, null); + byte[] payload = ("payload-" + suffix).getBytes(StandardCharsets.UTF_8); + S3QueueMessage dataMessage = S3QueueMessage.data(shuffleId, 0, 3, 0, 0L, objectRoot) + .setMetadata("schema=raw-bytes"); + + PhysicalWriter writer = s3qs.offer(dataMessage); + writer.append(payload); + writer.close(); + + S3QueuePollResult dataResult = pollUntilMessage(s3qs, shuffleId, 0, 5); + assertNotNull(dataResult); + assertTrue(dataResult.getMessage().isData()); + assertEquals(shuffleId, dataResult.getMessage().getShuffleId()); + assertEquals(0, dataResult.getMessage().getPartitionId()); + assertEquals(3, dataResult.getMessage().getProducerId()); + assertEquals("schema=raw-bytes", dataResult.getMessage().getMetadata()); + + PhysicalReader reader = PhysicalReaderUtil.newPhysicalReader(s3qs, + dataResult.getMessage().getObjectPath()); + ByteBuffer buffer = reader.readFully(payload.length); + byte[] actual = new byte[payload.length]; + buffer.get(actual); + reader.close(); + assertEquals(new String(payload, StandardCharsets.UTF_8), + new String(actual, StandardCharsets.UTF_8)); + + dataResult.getMessage().setReceiptHandle(dataResult.getReceiptHandle()); + assertEquals(0, s3qs.finishWork(dataResult.getMessage())); + + S3QueueMessage endMessage = S3QueueMessage.producerEnd(shuffleId, 0, 3, 0, 1L); + s3qs.publish(endMessage); + S3QueuePollResult endResult = pollUntilMessage(s3qs, shuffleId, 0, 5); + assertNotNull(endResult); + assertTrue(endResult.getMessage().isProducerEnd()); + endResult.getMessage().setReceiptHandle(endResult.getReceiptHandle()); + assertEquals(0, s3qs.finishWork(endResult.getMessage())); } - catch (ExecutionException e) + finally { - //e.getCause().printStackTrace(); - throw e; + if (queueUrl != null) + { + s3qs.deleteQueue(queueUrl); + } + s3qs.delete(objectRoot, true); } + } - this.executor.shutdown(); - this.executor.awaitTermination(100, TimeUnit.HOURS); + private static S3QS newS3QS(SqsClient sqs) throws Exception + { + S3QS s3qs = new S3QS(30); + Field sqsField = S3QS.class.getDeclaredField("sqs"); + sqsField.setAccessible(true); + sqsField.set(s3qs, sqs); + return s3qs; } - @Test - public void testMultiSQSWriter() throws IOException, InterruptedException, ExecutionException + private static Map receiveCount(int count) { - S3QS s3qs = (S3QS) StorageFactory.Instance().getStorage(Storage.Scheme.s3qs); + return Collections.singletonMap(APPROXIMATE_RECEIVE_COUNT, String.valueOf(count)); + } - Future future = this.executor.submit(() -> + private static S3QueuePollResult pollUntilMessage(S3QS s3qs, String shuffleId, + int partitionId, int attempts) throws IOException + { + for (int i = 0; i < attempts; ++i) { - byte[] buffer = new byte[8 * 1024 * 1024]; - long startTime = System.currentTimeMillis(); - for (int i = 0; i < 2; i++) + S3QueuePollResult result = s3qs.pollMessage(shuffleId, partitionId, 1); + if (result != null) { - S3QueueMessage body = new S3QueueMessage() - .setObjectPath("pixels-turbo-intermediate/shuffle/" + 9 + "/") - .setPartitionNum(9) - .setEndwork(false); - try (PhysicalWriter writer = s3qs.offer(body)) - { - writer.append(buffer); - System.out.println("Wrote partition " + i); // 添加日志 - } - catch (IOException e) - { - System.err.println("Error in iteration " + i + ": " + e.getMessage()); - throw new RuntimeException(e); - } + return result; } - System.out.println("write finished in " + (System.currentTimeMillis() - startTime) + " ms"); + } + return null; + } - }); + private static ReceiveMessageResponse receive(String receiptHandle, S3QueueMessage message) throws IOException + { + Message sqsMessage = Message.builder() + .body(message.toMessageBody()) + .receiptHandle(receiptHandle) + .attributes(receiveCount(1)) + .build(); + return ReceiveMessageResponse.builder().messages(sqsMessage).build(); + } - try - { - future.get(); - } - catch (ExecutionException e) + private static String trimSlashes(String value) + { + if (value == null) { - e.getCause().printStackTrace(); - throw e; + return ""; } - try + String result = value.trim(); + while (result.startsWith("/")) { - s3qs.close(); + result = result.substring(1); } - catch (IOException e) + while (result.endsWith("/")) { - throw new RuntimeException(e); + result = result.substring(0, result.length() - 1); } - - this.executor.shutdown(); - this.executor.awaitTermination(100, TimeUnit.HOURS); + return result; } } diff --git a/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionInvoker.java b/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionInvoker.java new file mode 100644 index 0000000000..72987229f7 --- /dev/null +++ b/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionInvoker.java @@ -0,0 +1,24 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.lambda; + +import com.alibaba.fastjson.JSON; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.planner.plan.physical.output.PartitionOutput; + +public class S3QSPartitionInvoker extends LambdaInvoker +{ + protected S3QSPartitionInvoker(String functionName) + { + super(functionName); + } + + @Override + public Output parseOutput(String outputJson) + { + return JSON.parseObject(outputJson, PartitionOutput.class); + } +} diff --git a/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionInvokerProvider.java b/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionInvokerProvider.java new file mode 100644 index 0000000000..5f5df890ac --- /dev/null +++ b/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionInvokerProvider.java @@ -0,0 +1,38 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.lambda; + +import io.pixelsdb.pixels.common.turbo.FunctionService; +import io.pixelsdb.pixels.common.turbo.Invoker; +import io.pixelsdb.pixels.common.turbo.InvokerProvider; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.common.utils.ConfigFactory; + +import static io.pixelsdb.pixels.common.turbo.FunctionService.lambda; +import static io.pixelsdb.pixels.common.turbo.WorkerType.PARTITION_S3QS; + +public class S3QSPartitionInvokerProvider implements InvokerProvider +{ + private static final ConfigFactory config = ConfigFactory.Instance(); + + @Override + public Invoker createInvoker() + { + return new S3QSPartitionInvoker(config.getProperty("s3qs.partition.worker.name")); + } + + @Override + public WorkerType workerType() + { + return PARTITION_S3QS; + } + + @Override + public boolean compatibleWith(FunctionService functionService) + { + return functionService.equals(lambda); + } +} diff --git a/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionedChainJoinInvoker.java b/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionedChainJoinInvoker.java new file mode 100644 index 0000000000..ba2f552409 --- /dev/null +++ b/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionedChainJoinInvoker.java @@ -0,0 +1,24 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.lambda; + +import com.alibaba.fastjson.JSON; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; + +public class S3QSPartitionedChainJoinInvoker extends LambdaInvoker +{ + protected S3QSPartitionedChainJoinInvoker(String functionName) + { + super(functionName); + } + + @Override + public Output parseOutput(String outputJson) + { + return JSON.parseObject(outputJson, JoinOutput.class); + } +} diff --git a/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionedChainJoinInvokerProvider.java b/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionedChainJoinInvokerProvider.java new file mode 100644 index 0000000000..c720429493 --- /dev/null +++ b/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionedChainJoinInvokerProvider.java @@ -0,0 +1,38 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.lambda; + +import io.pixelsdb.pixels.common.turbo.FunctionService; +import io.pixelsdb.pixels.common.turbo.Invoker; +import io.pixelsdb.pixels.common.turbo.InvokerProvider; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.common.utils.ConfigFactory; + +import static io.pixelsdb.pixels.common.turbo.FunctionService.lambda; +import static io.pixelsdb.pixels.common.turbo.WorkerType.PARTITIONED_CHAIN_JOIN_S3QS; + +public class S3QSPartitionedChainJoinInvokerProvider implements InvokerProvider +{ + private static final ConfigFactory config = ConfigFactory.Instance(); + + @Override + public Invoker createInvoker() + { + return new S3QSPartitionedChainJoinInvoker(config.getProperty("s3qs.partitioned.chain.join.worker.name")); + } + + @Override + public WorkerType workerType() + { + return PARTITIONED_CHAIN_JOIN_S3QS; + } + + @Override + public boolean compatibleWith(FunctionService functionService) + { + return functionService.equals(lambda); + } +} diff --git a/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionedJoinInvoker.java b/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionedJoinInvoker.java new file mode 100644 index 0000000000..90b5f98458 --- /dev/null +++ b/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionedJoinInvoker.java @@ -0,0 +1,24 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.lambda; + +import com.alibaba.fastjson.JSON; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; + +public class S3QSPartitionedJoinInvoker extends LambdaInvoker +{ + protected S3QSPartitionedJoinInvoker(String functionName) + { + super(functionName); + } + + @Override + public Output parseOutput(String outputJson) + { + return JSON.parseObject(outputJson, JoinOutput.class); + } +} diff --git a/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionedJoinInvokerProvider.java b/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionedJoinInvokerProvider.java new file mode 100644 index 0000000000..236f460efa --- /dev/null +++ b/pixels-turbo/pixels-invoker-lambda/src/main/java/io/pixelsdb/pixels/invoker/lambda/S3QSPartitionedJoinInvokerProvider.java @@ -0,0 +1,38 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.lambda; + +import io.pixelsdb.pixels.common.turbo.FunctionService; +import io.pixelsdb.pixels.common.turbo.Invoker; +import io.pixelsdb.pixels.common.turbo.InvokerProvider; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.common.utils.ConfigFactory; + +import static io.pixelsdb.pixels.common.turbo.FunctionService.lambda; +import static io.pixelsdb.pixels.common.turbo.WorkerType.PARTITIONED_JOIN_S3QS; + +public class S3QSPartitionedJoinInvokerProvider implements InvokerProvider +{ + private static final ConfigFactory config = ConfigFactory.Instance(); + + @Override + public Invoker createInvoker() + { + return new S3QSPartitionedJoinInvoker(config.getProperty("s3qs.partitioned.join.worker.name")); + } + + @Override + public WorkerType workerType() + { + return PARTITIONED_JOIN_S3QS; + } + + @Override + public boolean compatibleWith(FunctionService functionService) + { + return functionService.equals(lambda); + } +} diff --git a/pixels-turbo/pixels-invoker-lambda/src/main/resources/META-INF/services/io.pixelsdb.pixels.common.turbo.InvokerProvider b/pixels-turbo/pixels-invoker-lambda/src/main/resources/META-INF/services/io.pixelsdb.pixels.common.turbo.InvokerProvider index 8fccee7621..899eb705d7 100644 --- a/pixels-turbo/pixels-invoker-lambda/src/main/resources/META-INF/services/io.pixelsdb.pixels.common.turbo.InvokerProvider +++ b/pixels-turbo/pixels-invoker-lambda/src/main/resources/META-INF/services/io.pixelsdb.pixels.common.turbo.InvokerProvider @@ -1,9 +1,12 @@ io.pixelsdb.pixels.invoker.lambda.ScanInvokerProvider io.pixelsdb.pixels.invoker.lambda.PartitionInvokerProvider +io.pixelsdb.pixels.invoker.lambda.S3QSPartitionInvokerProvider io.pixelsdb.pixels.invoker.lambda.PartitionedJoinInvokerProvider +io.pixelsdb.pixels.invoker.lambda.S3QSPartitionedJoinInvokerProvider io.pixelsdb.pixels.invoker.lambda.PartitionedChainJoinInvokerProvider +io.pixelsdb.pixels.invoker.lambda.S3QSPartitionedChainJoinInvokerProvider io.pixelsdb.pixels.invoker.lambda.SortInvokerProvider io.pixelsdb.pixels.invoker.lambda.SortedJoinInvokerProvider io.pixelsdb.pixels.invoker.lambda.BroadcastJoinInvokerProvider io.pixelsdb.pixels.invoker.lambda.BroadcastChainJoinInvokerProvider -io.pixelsdb.pixels.invoker.lambda.AggregationInvokerProvider \ No newline at end of file +io.pixelsdb.pixels.invoker.lambda.AggregationInvokerProvider diff --git a/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionInvoker.java b/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionInvoker.java new file mode 100644 index 0000000000..b5f6e7d4d2 --- /dev/null +++ b/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionInvoker.java @@ -0,0 +1,25 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.spike; + +import com.alibaba.fastjson.JSON; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.planner.plan.physical.output.PartitionOutput; + +public class S3QSPartitionInvoker extends SpikeInvoker +{ + protected S3QSPartitionInvoker(String functionName) + { + super(functionName, WorkerType.PARTITION_S3QS); + } + + @Override + public Output parseOutput(String outputJson) + { + return JSON.parseObject(outputJson, PartitionOutput.class); + } +} diff --git a/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionInvokerProvider.java b/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionInvokerProvider.java new file mode 100644 index 0000000000..cc69cba33e --- /dev/null +++ b/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionInvokerProvider.java @@ -0,0 +1,37 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.spike; + +import io.pixelsdb.pixels.common.turbo.FunctionService; +import io.pixelsdb.pixels.common.turbo.Invoker; +import io.pixelsdb.pixels.common.turbo.InvokerProvider; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.common.utils.ConfigFactory; + +import static io.pixelsdb.pixels.common.turbo.FunctionService.spike; + +public class S3QSPartitionInvokerProvider implements InvokerProvider +{ + private static final ConfigFactory config = ConfigFactory.Instance(); + + @Override + public Invoker createInvoker() + { + return new S3QSPartitionInvoker(config.getProperty("s3qs.partition.worker.name")); + } + + @Override + public WorkerType workerType() + { + return WorkerType.PARTITION_S3QS; + } + + @Override + public boolean compatibleWith(FunctionService functionService) + { + return functionService.equals(spike); + } +} diff --git a/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionedChainJoinInvoker.java b/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionedChainJoinInvoker.java new file mode 100644 index 0000000000..41c94ee531 --- /dev/null +++ b/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionedChainJoinInvoker.java @@ -0,0 +1,25 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.spike; + +import com.alibaba.fastjson.JSON; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; + +public class S3QSPartitionedChainJoinInvoker extends SpikeInvoker +{ + protected S3QSPartitionedChainJoinInvoker(String functionName) + { + super(functionName, WorkerType.PARTITIONED_CHAIN_JOIN_S3QS); + } + + @Override + public Output parseOutput(String outputJson) + { + return JSON.parseObject(outputJson, JoinOutput.class); + } +} diff --git a/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionedChainJoinInvokerProvider.java b/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionedChainJoinInvokerProvider.java new file mode 100644 index 0000000000..6174697341 --- /dev/null +++ b/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionedChainJoinInvokerProvider.java @@ -0,0 +1,37 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.spike; + +import io.pixelsdb.pixels.common.turbo.FunctionService; +import io.pixelsdb.pixels.common.turbo.Invoker; +import io.pixelsdb.pixels.common.turbo.InvokerProvider; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.common.utils.ConfigFactory; + +import static io.pixelsdb.pixels.common.turbo.FunctionService.spike; + +public class S3QSPartitionedChainJoinInvokerProvider implements InvokerProvider +{ + private static final ConfigFactory config = ConfigFactory.Instance(); + + @Override + public Invoker createInvoker() + { + return new S3QSPartitionedChainJoinInvoker(config.getProperty("s3qs.partitioned.chain.join.worker.name")); + } + + @Override + public WorkerType workerType() + { + return WorkerType.PARTITIONED_CHAIN_JOIN_S3QS; + } + + @Override + public boolean compatibleWith(FunctionService functionService) + { + return functionService.equals(spike); + } +} diff --git a/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionedJoinInvoker.java b/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionedJoinInvoker.java new file mode 100644 index 0000000000..8017e450c7 --- /dev/null +++ b/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionedJoinInvoker.java @@ -0,0 +1,25 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.spike; + +import com.alibaba.fastjson.JSON; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; + +public class S3QSPartitionedJoinInvoker extends SpikeInvoker +{ + protected S3QSPartitionedJoinInvoker(String functionName) + { + super(functionName, WorkerType.PARTITIONED_JOIN_S3QS); + } + + @Override + public Output parseOutput(String outputJson) + { + return JSON.parseObject(outputJson, JoinOutput.class); + } +} diff --git a/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionedJoinInvokerProvider.java b/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionedJoinInvokerProvider.java new file mode 100644 index 0000000000..24f75ff951 --- /dev/null +++ b/pixels-turbo/pixels-invoker-spike/src/main/java/io/pixelsdb/pixels/invoker/spike/S3QSPartitionedJoinInvokerProvider.java @@ -0,0 +1,37 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.spike; + +import io.pixelsdb.pixels.common.turbo.FunctionService; +import io.pixelsdb.pixels.common.turbo.Invoker; +import io.pixelsdb.pixels.common.turbo.InvokerProvider; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.common.utils.ConfigFactory; + +import static io.pixelsdb.pixels.common.turbo.FunctionService.spike; + +public class S3QSPartitionedJoinInvokerProvider implements InvokerProvider +{ + private static final ConfigFactory config = ConfigFactory.Instance(); + + @Override + public Invoker createInvoker() + { + return new S3QSPartitionedJoinInvoker(config.getProperty("s3qs.partitioned.join.worker.name")); + } + + @Override + public WorkerType workerType() + { + return WorkerType.PARTITIONED_JOIN_S3QS; + } + + @Override + public boolean compatibleWith(FunctionService functionService) + { + return functionService.equals(spike); + } +} diff --git a/pixels-turbo/pixels-invoker-spike/src/main/resources/META-INF/services/io.pixelsdb.pixels.common.turbo.InvokerProvider b/pixels-turbo/pixels-invoker-spike/src/main/resources/META-INF/services/io.pixelsdb.pixels.common.turbo.InvokerProvider index d2fe5cef97..8f0f12e22d 100644 --- a/pixels-turbo/pixels-invoker-spike/src/main/resources/META-INF/services/io.pixelsdb.pixels.common.turbo.InvokerProvider +++ b/pixels-turbo/pixels-invoker-spike/src/main/resources/META-INF/services/io.pixelsdb.pixels.common.turbo.InvokerProvider @@ -2,9 +2,12 @@ io.pixelsdb.pixels.invoker.spike.AggregationInvokerProvider io.pixelsdb.pixels.invoker.spike.BroadcastChainJoinInvokerProvider io.pixelsdb.pixels.invoker.spike.BroadcastJoinInvokerProvider io.pixelsdb.pixels.invoker.spike.PartitionedChainJoinInvokerProvider +io.pixelsdb.pixels.invoker.spike.S3QSPartitionedChainJoinInvokerProvider io.pixelsdb.pixels.invoker.spike.PartitionedJoinInvokerProvider +io.pixelsdb.pixels.invoker.spike.S3QSPartitionedJoinInvokerProvider io.pixelsdb.pixels.invoker.spike.PartitionedJoinStreamInvokerProvider io.pixelsdb.pixels.invoker.spike.PartitionInvokerProvider +io.pixelsdb.pixels.invoker.spike.S3QSPartitionInvokerProvider io.pixelsdb.pixels.invoker.spike.PartitionStreamInvokerProvider io.pixelsdb.pixels.invoker.spike.ScanInvokerProvider -io.pixelsdb.pixels.invoker.spike.ScanStreamInvokerProvider \ No newline at end of file +io.pixelsdb.pixels.invoker.spike.ScanStreamInvokerProvider diff --git a/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionInvoker.java b/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionInvoker.java new file mode 100644 index 0000000000..c8c8ac00e4 --- /dev/null +++ b/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionInvoker.java @@ -0,0 +1,38 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.vhive; + +import com.alibaba.fastjson.JSON; +import com.google.common.util.concurrent.ListenableFuture; +import io.pixelsdb.pixels.common.turbo.Input; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.PartitionOutput; +import io.pixelsdb.pixels.turbo.TurboProto; + +import java.util.concurrent.CompletableFuture; + +public class S3QSPartitionInvoker extends VhiveInvoker +{ + protected S3QSPartitionInvoker(String functionName) + { + super(functionName); + } + + @Override + public Output parseOutput(String outputJson) + { + return JSON.parseObject(outputJson, PartitionOutput.class); + } + + @Override + public CompletableFuture invoke(Input input) + { + ListenableFuture future = + Vhive.Instance().getAsyncClient().partitionS3QS((StageWorkerInput) input); + return genCompletableFuture(future); + } +} diff --git a/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionInvokerProvider.java b/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionInvokerProvider.java new file mode 100644 index 0000000000..16944e3b5b --- /dev/null +++ b/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionInvokerProvider.java @@ -0,0 +1,35 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.vhive; + +import io.pixelsdb.pixels.common.turbo.FunctionService; +import io.pixelsdb.pixels.common.turbo.Invoker; +import io.pixelsdb.pixels.common.turbo.InvokerProvider; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.common.utils.ConfigFactory; + +public class S3QSPartitionInvokerProvider implements InvokerProvider +{ + private static final ConfigFactory config = ConfigFactory.Instance(); + + @Override + public Invoker createInvoker() + { + return new S3QSPartitionInvoker(config.getProperty("s3qs.partition.worker.name")); + } + + @Override + public WorkerType workerType() + { + return WorkerType.PARTITION_S3QS; + } + + @Override + public boolean compatibleWith(FunctionService functionService) + { + return functionService.equals(FunctionService.vhive); + } +} diff --git a/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionedChainJoinInvoker.java b/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionedChainJoinInvoker.java new file mode 100644 index 0000000000..2c83f7f40f --- /dev/null +++ b/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionedChainJoinInvoker.java @@ -0,0 +1,38 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.vhive; + +import com.alibaba.fastjson.JSON; +import com.google.common.util.concurrent.ListenableFuture; +import io.pixelsdb.pixels.common.turbo.Input; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; +import io.pixelsdb.pixels.turbo.TurboProto; + +import java.util.concurrent.CompletableFuture; + +public class S3QSPartitionedChainJoinInvoker extends VhiveInvoker +{ + protected S3QSPartitionedChainJoinInvoker(String functionName) + { + super(functionName); + } + + @Override + public Output parseOutput(String outputJson) + { + return JSON.parseObject(outputJson, JoinOutput.class); + } + + @Override + public CompletableFuture invoke(Input input) + { + ListenableFuture future = + Vhive.Instance().getAsyncClient().partitionChainJoinS3QS((StageWorkerInput) input); + return genCompletableFuture(future); + } +} diff --git a/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionedChainJoinInvokerProvider.java b/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionedChainJoinInvokerProvider.java new file mode 100644 index 0000000000..60315b897b --- /dev/null +++ b/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionedChainJoinInvokerProvider.java @@ -0,0 +1,35 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.vhive; + +import io.pixelsdb.pixels.common.turbo.FunctionService; +import io.pixelsdb.pixels.common.turbo.Invoker; +import io.pixelsdb.pixels.common.turbo.InvokerProvider; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.common.utils.ConfigFactory; + +public class S3QSPartitionedChainJoinInvokerProvider implements InvokerProvider +{ + private static final ConfigFactory config = ConfigFactory.Instance(); + + @Override + public Invoker createInvoker() + { + return new S3QSPartitionedChainJoinInvoker(config.getProperty("s3qs.partitioned.chain.join.worker.name")); + } + + @Override + public WorkerType workerType() + { + return WorkerType.PARTITIONED_CHAIN_JOIN_S3QS; + } + + @Override + public boolean compatibleWith(FunctionService functionService) + { + return functionService.equals(FunctionService.vhive); + } +} diff --git a/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionedJoinInvoker.java b/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionedJoinInvoker.java new file mode 100644 index 0000000000..385529561b --- /dev/null +++ b/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionedJoinInvoker.java @@ -0,0 +1,38 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.vhive; + +import com.alibaba.fastjson.JSON; +import com.google.common.util.concurrent.ListenableFuture; +import io.pixelsdb.pixels.common.turbo.Input; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; +import io.pixelsdb.pixels.turbo.TurboProto; + +import java.util.concurrent.CompletableFuture; + +public class S3QSPartitionedJoinInvoker extends VhiveInvoker +{ + protected S3QSPartitionedJoinInvoker(String functionName) + { + super(functionName); + } + + @Override + public Output parseOutput(String outputJson) + { + return JSON.parseObject(outputJson, JoinOutput.class); + } + + @Override + public CompletableFuture invoke(Input input) + { + ListenableFuture future = + Vhive.Instance().getAsyncClient().partitionJoinS3QS((StageWorkerInput) input); + return genCompletableFuture(future); + } +} diff --git a/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionedJoinInvokerProvider.java b/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionedJoinInvokerProvider.java new file mode 100644 index 0000000000..df9a7d6bae --- /dev/null +++ b/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/S3QSPartitionedJoinInvokerProvider.java @@ -0,0 +1,35 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.invoker.vhive; + +import io.pixelsdb.pixels.common.turbo.FunctionService; +import io.pixelsdb.pixels.common.turbo.Invoker; +import io.pixelsdb.pixels.common.turbo.InvokerProvider; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.common.utils.ConfigFactory; + +public class S3QSPartitionedJoinInvokerProvider implements InvokerProvider +{ + private static final ConfigFactory config = ConfigFactory.Instance(); + + @Override + public Invoker createInvoker() + { + return new S3QSPartitionedJoinInvoker(config.getProperty("s3qs.partitioned.join.worker.name")); + } + + @Override + public WorkerType workerType() + { + return WorkerType.PARTITIONED_JOIN_S3QS; + } + + @Override + public boolean compatibleWith(FunctionService functionService) + { + return functionService.equals(FunctionService.vhive); + } +} diff --git a/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/WorkerAsyncClient.java b/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/WorkerAsyncClient.java index 7da0bd827f..67d432598b 100644 --- a/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/WorkerAsyncClient.java +++ b/pixels-turbo/pixels-invoker-vhive/src/main/java/io/pixelsdb/pixels/invoker/vhive/WorkerAsyncClient.java @@ -129,6 +129,15 @@ public ListenableFuture partitionChainJoinStream return this.stub.process(request); } + public ListenableFuture partitionChainJoinS3QS(StageWorkerInput input) + { + TurboProto.vHiveWorkerRequest request = TurboProto.vHiveWorkerRequest.newBuilder() + .setWorkerType(String.valueOf(WorkerType.PARTITIONED_CHAIN_JOIN_S3QS)) + .setJson(JSON.toJSONString(input, SerializerFeature.DisableCircularReferenceDetect)) + .build(); + return this.stub.process(request); + } + public ListenableFuture partitionJoin(PartitionedJoinInput input) { TurboProto.vHiveWorkerRequest request = TurboProto.vHiveWorkerRequest.newBuilder() @@ -147,6 +156,15 @@ public ListenableFuture partitionJoinStreaming(P return this.stub.process(request); } + public ListenableFuture partitionJoinS3QS(StageWorkerInput input) + { + TurboProto.vHiveWorkerRequest request = TurboProto.vHiveWorkerRequest.newBuilder() + .setWorkerType(String.valueOf(WorkerType.PARTITIONED_JOIN_S3QS)) + .setJson(JSON.toJSONString(input, SerializerFeature.DisableCircularReferenceDetect)) + .build(); + return this.stub.process(request); + } + public ListenableFuture partition(PartitionInput input) { TurboProto.vHiveWorkerRequest request = TurboProto.vHiveWorkerRequest.newBuilder() @@ -165,6 +183,15 @@ public ListenableFuture partitionStreaming(Parti return this.stub.process(request); } + public ListenableFuture partitionS3QS(StageWorkerInput input) + { + TurboProto.vHiveWorkerRequest request = TurboProto.vHiveWorkerRequest.newBuilder() + .setWorkerType(String.valueOf(WorkerType.PARTITION_S3QS)) + .setJson(JSON.toJSONString(input, SerializerFeature.DisableCircularReferenceDetect)) + .build(); + return this.stub.process(request); + } + public ListenableFuture scan(ScanInput input) { TurboProto.vHiveWorkerRequest request = TurboProto.vHiveWorkerRequest.newBuilder() diff --git a/pixels-turbo/pixels-invoker-vhive/src/main/resources/META-INF/services/io.pixelsdb.pixels.common.turbo.InvokerProvider b/pixels-turbo/pixels-invoker-vhive/src/main/resources/META-INF/services/io.pixelsdb.pixels.common.turbo.InvokerProvider index df52eaf713..5b3afcf680 100644 --- a/pixels-turbo/pixels-invoker-vhive/src/main/resources/META-INF/services/io.pixelsdb.pixels.common.turbo.InvokerProvider +++ b/pixels-turbo/pixels-invoker-vhive/src/main/resources/META-INF/services/io.pixelsdb.pixels.common.turbo.InvokerProvider @@ -1,10 +1,13 @@ io.pixelsdb.pixels.invoker.vhive.ScanStreamInvokerProvider io.pixelsdb.pixels.invoker.vhive.ScanInvokerProvider io.pixelsdb.pixels.invoker.vhive.PartitionInvokerProvider +io.pixelsdb.pixels.invoker.vhive.S3QSPartitionInvokerProvider io.pixelsdb.pixels.invoker.vhive.PartitionStreamInvokerProvider io.pixelsdb.pixels.invoker.vhive.PartitionJoinInvokerProvider +io.pixelsdb.pixels.invoker.vhive.S3QSPartitionedJoinInvokerProvider io.pixelsdb.pixels.invoker.vhive.PartitionedJoinStreamInvokerProvider io.pixelsdb.pixels.invoker.vhive.PartitionChainJoinInvokerProvider +io.pixelsdb.pixels.invoker.vhive.S3QSPartitionedChainJoinInvokerProvider io.pixelsdb.pixels.invoker.vhive.PartitionChainJoinStreamInvokerProvider io.pixelsdb.pixels.invoker.vhive.BroadcastJoinStreamInvokerProvider io.pixelsdb.pixels.invoker.vhive.BroadcastJoinInvokerProvider diff --git a/pixels-turbo/pixels-worker-common/pom.xml b/pixels-turbo/pixels-worker-common/pom.xml index 2849f41f9e..f15943ceda 100644 --- a/pixels-turbo/pixels-worker-common/pom.xml +++ b/pixels-turbo/pixels-worker-common/pom.xml @@ -45,11 +45,23 @@ pixels-storage-s3 true + + io.pixelsdb + pixels-storage-s3qs + true + io.pixelsdb pixels-storage-redis true + + + + io.grpc + grpc-api + test + @@ -80,4 +92,4 @@ - \ No newline at end of file + diff --git a/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/BasePartitionWorker.java b/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/BasePartitionWorker.java index 77604d664e..f52f18b86b 100644 --- a/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/BasePartitionWorker.java +++ b/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/BasePartitionWorker.java @@ -20,6 +20,7 @@ package io.pixelsdb.pixels.worker.common; import com.alibaba.fastjson.JSON; +import io.pixelsdb.pixels.common.physical.PhysicalWriter; import io.pixelsdb.pixels.common.physical.Storage; import io.pixelsdb.pixels.core.PixelsReader; import io.pixelsdb.pixels.core.PixelsWriter; @@ -32,9 +33,12 @@ import io.pixelsdb.pixels.executor.scan.Scanner; import io.pixelsdb.pixels.planner.plan.physical.domain.InputInfo; import io.pixelsdb.pixels.planner.plan.physical.domain.InputSplit; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleInfo; import io.pixelsdb.pixels.planner.plan.physical.domain.StorageInfo; import io.pixelsdb.pixels.planner.plan.physical.input.PartitionInput; import io.pixelsdb.pixels.planner.plan.physical.output.PartitionOutput; +import io.pixelsdb.pixels.storage.s3qs.S3QS; +import io.pixelsdb.pixels.storage.s3qs.S3QueueMessage; import org.apache.logging.log4j.Logger; import java.util.*; @@ -102,6 +106,7 @@ public PartitionOutput process(PartitionInput event) WorkerCommon.initStorage(inputStorageInfo); WorkerCommon.initStorage(outputStorageInfo); + WorkerCommon.initOptionalShuffleStorage(event.getOutput().getShuffleInfo()); String[] columnsToRead = event.getTableInfo().getColumnsToRead(); TableScanFilter filter = JSON.parseObject(event.getTableInfo().getFilter(), TableScanFilter.class); @@ -150,30 +155,38 @@ public PartitionOutput process(PartitionInput event) TypeDescription resultSchema = WorkerCommon.getResultSchema(fileSchema, columnsToRead); writerSchema.set(resultSchema); } - PixelsWriter pixelsWriter = WorkerCommon.getWriter(writerSchema.get(), - WorkerCommon.getStorage(outputStorageInfo.getScheme()), outputPath, encoding, - true, Arrays.stream(keyColumnIds).boxed().collect(Collectors.toList())); Set hashValues = new HashSet<>(numPartition); - - for (int hash = 0; hash < numPartition; ++hash) + ShuffleInfo shuffleInfo = event.getOutput().getShuffleInfo(); + if (isS3QSShuffle(shuffleInfo)) + { + writeS3QSPartitionOutput(event, shuffleInfo, writerSchema.get(), partitioned, hashValues); + } + else { - ConcurrentLinkedQueue batches = partitioned.get(hash); - if (!batches.isEmpty()) + PixelsWriter pixelsWriter = WorkerCommon.getWriter(writerSchema.get(), + WorkerCommon.getStorage(outputStorageInfo.getScheme()), outputPath, encoding, + true, Arrays.stream(keyColumnIds).boxed().collect(Collectors.toList())); + + for (int hash = 0; hash < numPartition; ++hash) { - for (VectorizedRowBatch batch : batches) + ConcurrentLinkedQueue batches = partitioned.get(hash); + if (!batches.isEmpty()) { - pixelsWriter.addRowBatch(batch, hash); + for (VectorizedRowBatch batch : batches) + { + pixelsWriter.addRowBatch(batch, hash); + } + hashValues.add(hash); } - hashValues.add(hash); } + + pixelsWriter.close(); + workerMetrics.addWriteBytes(pixelsWriter.getCompletedBytes()); + workerMetrics.addNumWriteRequests(pixelsWriter.getNumWriteRequests()); } partitionOutput.addOutput(outputPath); partitionOutput.setHashValues(hashValues); - - pixelsWriter.close(); workerMetrics.addOutputCostNs(writeCostTimer.stop()); - workerMetrics.addWriteBytes(pixelsWriter.getCompletedBytes()); - workerMetrics.addNumWriteRequests(pixelsWriter.getNumWriteRequests()); partitionOutput.setDurationMs((int) (System.currentTimeMillis() - startTime)); WorkerCommon.setPerfMetrics(partitionOutput, workerMetrics); @@ -189,6 +202,92 @@ public PartitionOutput process(PartitionInput event) } } + static boolean isS3QSShuffle(ShuffleInfo shuffleInfo) + { + return shuffleInfo != null && + shuffleInfo.getStorageInfo() != null && + shuffleInfo.getStorageInfo().getScheme() == Storage.Scheme.s3qs; + } + + static S3QueueMessage createS3QSDataMessage(ShuffleInfo shuffleInfo, int partitionId, + int producerTaskId, int producerAttemptId, + long sequenceId, TypeDescription writerSchema) + { + return S3QueueMessage.data(shuffleInfo.getShuffleId(), partitionId, + producerTaskId, producerAttemptId, sequenceId, + shuffleInfo.getObjectPathPrefix()) + .setMetadata(writerSchema.toString()); + } + + static S3QueueMessage createS3QSProducerEndMessage(ShuffleInfo shuffleInfo, int partitionId, + int producerTaskId, int producerAttemptId, + long sequenceId, TypeDescription writerSchema) + { + return S3QueueMessage.producerEnd(shuffleInfo.getShuffleId(), partitionId, + producerTaskId, producerAttemptId, sequenceId) + .setMetadata(writerSchema.toString()); + } + + /** + * Write the partitioned output of one producer task into S3QS. + * + * Each non-empty partition is written as one S3 object and published as a + * DATA message. After data objects are written, this producer task sends a + * PRODUCER_END marker to every partition queue so consumers can determine + * partition completion without relying on live worker membership. + */ + private void writeS3QSPartitionOutput(PartitionInput event, ShuffleInfo shuffleInfo, TypeDescription writerSchema, + List> partitioned, + Set hashValues) throws Exception + { + if (event.getProducerTaskId() < 0) + { + throw new WorkerException("producerTaskId is not set for s3qs shuffle"); + } + Storage storage = WorkerCommon.getStorage(Storage.Scheme.s3qs); + if (!(storage instanceof S3QS)) + { + throw new WorkerException("storage of scheme s3qs is not S3QS"); + } + S3QS s3qs = (S3QS) storage; + int producerTaskId = event.getProducerTaskId(); + int producerAttemptId = 0; + long[] sequenceIds = new long[partitioned.size()]; + + for (int partitionId = 0; partitionId < partitioned.size(); ++partitionId) + { + ConcurrentLinkedQueue batches = partitioned.get(partitionId); + if (!batches.isEmpty()) + { + S3QueueMessage dataMessage = createS3QSDataMessage(shuffleInfo, partitionId, + producerTaskId, producerAttemptId, sequenceIds[partitionId]++, writerSchema); + PhysicalWriter physicalWriter = s3qs.offer(dataMessage); + PixelsWriter pixelsWriter = WorkerCommon.getWriter(writerSchema, storage, + physicalWriter.getPath(), physicalWriter, false, null); + try + { + for (VectorizedRowBatch batch : batches) + { + pixelsWriter.addRowBatch(batch); + } + } + finally + { + pixelsWriter.close(); + } + hashValues.add(partitionId); + workerMetrics.addWriteBytes(pixelsWriter.getCompletedBytes()); + workerMetrics.addNumWriteRequests(pixelsWriter.getNumWriteRequests()); + } + } + + for (int partitionId = 0; partitionId < partitioned.size(); ++partitionId) + { + s3qs.publish(createS3QSProducerEndMessage(shuffleInfo, partitionId, + producerTaskId, producerAttemptId, sequenceIds[partitionId]++, writerSchema)); + } + } + /** * Scan and partition the files in a query split. * diff --git a/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/BasePartitionedChainJoinWorker.java b/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/BasePartitionedChainJoinWorker.java index 12889e77a0..571c321f56 100644 --- a/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/BasePartitionedChainJoinWorker.java +++ b/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/BasePartitionedChainJoinWorker.java @@ -31,6 +31,7 @@ import io.pixelsdb.pixels.planner.plan.physical.domain.*; import io.pixelsdb.pixels.planner.plan.physical.input.PartitionedChainJoinInput; import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; +import io.pixelsdb.pixels.storage.s3qs.S3QueuePollResult; import org.apache.logging.log4j.Logger; import java.io.IOException; @@ -94,22 +95,30 @@ public JoinOutput process(PartitionedChainJoinInput event) checkArgument(chainTables.size() > 1, "there should be at least two chain tables"); requireNonNull(event.getSmallTable(), "leftTable is null"); + boolean leftS3QS = BasePartitionedJoinWorker.isS3QSShuffle(event.getSmallTable()); StorageInfo leftInputStorageInfo = requireNonNull(event.getSmallTable().getStorageInfo(), "leftInputStorageInfo is null"); - List leftPartitioned = requireNonNull( event.getSmallTable().getInputFiles(), - "leftPartitioned is null"); - checkArgument(leftPartitioned.size() > 0, "leftPartitioned is empty"); + List leftPartitioned = event.getSmallTable().getInputFiles(); + if (!leftS3QS) + { + requireNonNull(leftPartitioned, "leftPartitioned is null"); + checkArgument(leftPartitioned.size() > 0, "leftPartitioned is empty"); + } int leftParallelism = event.getSmallTable().getParallelism(); checkArgument(leftParallelism > 0, "leftParallelism is not positive"); String[] leftColumnsToRead = event.getSmallTable().getColumnsToRead(); int[] leftKeyColumnIds = event.getSmallTable().getKeyColumnIds(); requireNonNull(event.getLargeTable(), "rightTable is null"); + boolean rightS3QS = BasePartitionedJoinWorker.isS3QSShuffle(event.getLargeTable()); StorageInfo rightInputStorageInfo = requireNonNull(event.getLargeTable().getStorageInfo(), "rightInputStorageInfo is null"); - List rightPartitioned = requireNonNull(event.getLargeTable().getInputFiles(), - "rightPartitioned is null"); - checkArgument(rightPartitioned.size() > 0, "rightPartitioned is empty"); + List rightPartitioned = event.getLargeTable().getInputFiles(); + if (!rightS3QS) + { + requireNonNull(rightPartitioned, "rightPartitioned is null"); + checkArgument(rightPartitioned.size() > 0, "rightPartitioned is empty"); + } int rightParallelism = event.getLargeTable().getParallelism(); checkArgument(rightParallelism > 0, "rightParallelism is not positive"); String[] rightColumnsToRead = event.getLargeTable().getColumnsToRead(); @@ -154,15 +163,45 @@ public JoinOutput process(PartitionedChainJoinInput event) } WorkerCommon.initStorage(leftInputStorageInfo); WorkerCommon.initStorage(rightInputStorageInfo); + WorkerCommon.initOptionalShuffleStorage(event.getSmallTable().getShuffleInfo()); + WorkerCommon.initOptionalShuffleStorage(event.getLargeTable().getShuffleInfo()); WorkerCommon.initStorage(outputStorageInfo); // build the joiner. AtomicReference leftSchema = new AtomicReference<>(); AtomicReference rightSchema = new AtomicReference<>(); - WorkerCommon.getFileSchemaFromPaths(threadPool, - WorkerCommon.getStorage(leftInputStorageInfo.getScheme()), - WorkerCommon.getStorage(rightInputStorageInfo.getScheme()), - leftSchema, rightSchema, leftPartitioned, rightPartitioned); + Map> leftPendingMessages = new HashMap<>(); + Map> rightPendingMessages = new HashMap<>(); + if (!leftS3QS && !rightS3QS) + { + WorkerCommon.getFileSchemaFromPaths(threadPool, + WorkerCommon.getStorage(leftInputStorageInfo.getScheme()), + WorkerCommon.getStorage(rightInputStorageInfo.getScheme()), + leftSchema, rightSchema, leftPartitioned, rightPartitioned); + } + else + { + if (leftS3QS) + { + leftSchema.set(BasePartitionedJoinWorker.getS3QSFileSchema( + event.getSmallTable().getShuffleInfo(), hashValues, leftPendingMessages)); + } + else + { + leftSchema.set(WorkerCommon.getFileSchemaFromPaths( + WorkerCommon.getStorage(leftInputStorageInfo.getScheme()), leftPartitioned)); + } + if (rightS3QS) + { + rightSchema.set(BasePartitionedJoinWorker.getS3QSFileSchema( + event.getLargeTable().getShuffleInfo(), hashValues, rightPendingMessages)); + } + else + { + rightSchema.set(WorkerCommon.getFileSchemaFromPaths( + WorkerCommon.getStorage(rightInputStorageInfo.getScheme()), rightPartitioned)); + } + } /* * Issue #450: * For the left and the right partial partitioned files, the file schema is equal to the columns to read in normal cases. @@ -193,33 +232,42 @@ public JoinOutput process(PartitionedChainJoinInput event) // build the hash table for the left partitioned table. if (chainJoiner.getSmallTableSize() > 0) { - List leftFutures = new ArrayList<>(leftPartitioned.size()); - int leftSplitSize = leftPartitioned.size() / leftParallelism; - if (leftPartitioned.size() % leftParallelism > 0) + if (leftS3QS) { - leftSplitSize++; + BasePartitionedJoinWorker.buildHashTableS3QS(transId, timestamp, (HashJoiner) partitionJoiner, + event.getSmallTable().getShuffleInfo(), leftPendingMessages, leftColumnsToRead, + hashValues, workerMetrics); } - for (int i = 0; i < leftPartitioned.size(); i += leftSplitSize) + else { - List parts = new LinkedList<>(); - for (int j = i; j < i + leftSplitSize && j < leftPartitioned.size(); ++j) + List leftFutures = new ArrayList<>(leftPartitioned.size()); + int leftSplitSize = leftPartitioned.size() / leftParallelism; + if (leftPartitioned.size() % leftParallelism > 0) { - parts.add(leftPartitioned.get(j)); + leftSplitSize++; } - leftFutures.add(threadPool.submit(() -> { - try - { - BasePartitionedJoinWorker.buildHashTable(transId, timestamp, (HashJoiner) partitionJoiner, parts, leftColumnsToRead, - leftInputStorageInfo.getScheme(), hashValues, numPartition, workerMetrics); - } catch (Throwable e) + for (int i = 0; i < leftPartitioned.size(); i += leftSplitSize) + { + List parts = new LinkedList<>(); + for (int j = i; j < i + leftSplitSize && j < leftPartitioned.size(); ++j) { - throw new WorkerException("error during hash table construction", e); + parts.add(leftPartitioned.get(j)); } - })); - } - for (Future future : leftFutures) - { - future.get(); + leftFutures.add(threadPool.submit(() -> { + try + { + BasePartitionedJoinWorker.buildHashTable(transId, timestamp, (HashJoiner) partitionJoiner, parts, leftColumnsToRead, + leftInputStorageInfo.getScheme(), hashValues, numPartition, workerMetrics); + } catch (Throwable e) + { + throw new WorkerException("error during hash table construction", e); + } + })); + } + for (Future future : leftFutures) + { + future.get(); + } } logger.info("hash table size: " + partitionJoiner.getSmallTableSize() + ", duration (ns): " + (workerMetrics.getInputCostNs() + workerMetrics.getComputeCostNs())); @@ -227,46 +275,98 @@ public JoinOutput process(PartitionedChainJoinInput event) // scan the right table and do the join. if (partitionJoiner.getSmallTableSize() > 0) { - int rightSplitSize = rightPartitioned.size() / rightParallelism; - if (rightPartitioned.size() % rightParallelism > 0) + if (rightS3QS) { - rightSplitSize++; + if (partitionOutput) + { + joinWithRightTableAndPartitionS3QS(transId, timestamp, partitionJoiner, chainJoiner, + event.getLargeTable().getShuffleInfo(), rightPendingMessages, rightColumnsToRead, + hashValues, outputPartitionInfo, result, workerMetrics); + } + else + { + joinWithRightTableS3QS(transId, timestamp, partitionJoiner, chainJoiner, + event.getLargeTable().getShuffleInfo(), rightPendingMessages, rightColumnsToRead, + hashValues, result.get(0), workerMetrics); + } } - - for (int i = 0; i < rightPartitioned.size(); i += rightSplitSize) + else { - List parts = new LinkedList<>(); - for (int j = i; j < i + rightSplitSize && j < rightPartitioned.size(); ++j) + int rightSplitSize = rightPartitioned.size() / rightParallelism; + if (rightPartitioned.size() % rightParallelism > 0) { - parts.add(rightPartitioned.get(j)); + rightSplitSize++; } - threadPool.execute(() -> { - try - { - int numJoinedRows = partitionOutput ? - joinWithRightTableAndPartition( - transId, timestamp, partitionJoiner, chainJoiner, parts, rightColumnsToRead, - rightInputStorageInfo.getScheme(), hashValues, numPartition, - outputPartitionInfo, result, workerMetrics) : - joinWithRightTable(transId, timestamp, partitionJoiner, chainJoiner, parts, rightColumnsToRead, - rightInputStorageInfo.getScheme(), hashValues, numPartition, - result.get(0), workerMetrics); - } catch (Throwable e) + + for (int i = 0; i < rightPartitioned.size(); i += rightSplitSize) + { + List parts = new LinkedList<>(); + for (int j = i; j < i + rightSplitSize && j < rightPartitioned.size(); ++j) { - throw new WorkerException("error during hash join", e); + parts.add(rightPartitioned.get(j)); } - }); + threadPool.execute(() -> { + try + { + int numJoinedRows = partitionOutput ? + joinWithRightTableAndPartition( + transId, timestamp, partitionJoiner, chainJoiner, parts, rightColumnsToRead, + rightInputStorageInfo.getScheme(), hashValues, numPartition, + outputPartitionInfo, result, workerMetrics) : + joinWithRightTable(transId, timestamp, partitionJoiner, chainJoiner, parts, rightColumnsToRead, + rightInputStorageInfo.getScheme(), hashValues, numPartition, + result.get(0), workerMetrics); + } catch (Throwable e) + { + throw new WorkerException("error during hash join", e); + } + }); + } + threadPool.shutdown(); + try + { + while (!threadPool.awaitTermination(60, TimeUnit.SECONDS)) ; + } catch (InterruptedException e) + { + throw new WorkerException("interrupted while waiting for the termination of join", e); + } + } + } + else if (rightS3QS) + { + for (int hashValue : hashValues) + { + BasePartitionedJoinWorker.drainS3QSPartition( + event.getLargeTable().getShuffleInfo(), hashValue, rightPendingMessages, + message -> { }); } - threadPool.shutdown(); - try + } + } + else + { + if (leftS3QS) + { + for (int hashValue : hashValues) { - while (!threadPool.awaitTermination(60, TimeUnit.SECONDS)) ; - } catch (InterruptedException e) + BasePartitionedJoinWorker.drainS3QSPartition( + event.getSmallTable().getShuffleInfo(), hashValue, leftPendingMessages, + message -> { }); + } + } + if (rightS3QS) + { + for (int hashValue : hashValues) { - throw new WorkerException("interrupted while waiting for the termination of join", e); + BasePartitionedJoinWorker.drainS3QSPartition( + event.getLargeTable().getShuffleInfo(), hashValue, rightPendingMessages, + message -> { }); } } } + if (!threadPool.isShutdown()) + { + threadPool.shutdown(); + } String outputPath = outputFolder + outputInfo.getFileNames().get(0); try @@ -388,6 +488,95 @@ protected static Joiner buildChainJoiner( } } + protected static int joinWithRightTableS3QS( + long transId, long timestamp, Joiner partitionedJoiner, Joiner chainJoiner, ShuffleInfo shuffleInfo, + Map> pendingMessages, String[] rightCols, + List hashValues, ConcurrentLinkedQueue joinResult, + WorkerMetrics workerMetrics) throws Exception + { + final int[] joinedRows = {0}; + for (int partitionId : hashValues) + { + BasePartitionedJoinWorker.drainS3QSPartition(shuffleInfo, partitionId, pendingMessages, message -> + BasePartitionedJoinWorker.readS3QSDataObject( + transId, timestamp, message, rightCols, workerMetrics, rightRowBatch -> { + List partitionedJoinResults = + partitionedJoiner.join(rightRowBatch); + for (VectorizedRowBatch partitionedJoinResult : partitionedJoinResults) + { + if (!partitionedJoinResult.isEmpty()) + { + List chainJoinResults = + chainJoiner.join(partitionedJoinResult); + for (VectorizedRowBatch chainJoinResult : chainJoinResults) + { + if (!chainJoinResult.isEmpty()) + { + joinResult.add(chainJoinResult); + joinedRows[0] += chainJoinResult.size; + } + } + } + } + })); + } + return joinedRows[0]; + } + + protected static int joinWithRightTableAndPartitionS3QS( + long transId, long timestamp, Joiner partitionedJoiner, Joiner chainJoiner, ShuffleInfo shuffleInfo, + Map> pendingMessages, String[] rightCols, + List hashValues, PartitionInfo postPartitionInfo, + List> partitionResult, WorkerMetrics workerMetrics) + throws Exception + { + requireNonNull(postPartitionInfo, "outputPartitionInfo is null"); + Partitioner partitioner = new Partitioner(postPartitionInfo.getNumPartition(), + WorkerCommon.rowBatchSize, chainJoiner.getJoinedSchema(), postPartitionInfo.getKeyColumnIds()); + final int[] joinedRows = {0}; + for (int partitionId : hashValues) + { + BasePartitionedJoinWorker.drainS3QSPartition(shuffleInfo, partitionId, pendingMessages, message -> + BasePartitionedJoinWorker.readS3QSDataObject( + transId, timestamp, message, rightCols, workerMetrics, rightRowBatch -> { + List partitionedJoinResults = + partitionedJoiner.join(rightRowBatch); + for (VectorizedRowBatch partitionedJoinResult : partitionedJoinResults) + { + if (!partitionedJoinResult.isEmpty()) + { + List chainJoinResults = + chainJoiner.join(partitionedJoinResult); + for (VectorizedRowBatch chainJoinResult : chainJoinResults) + { + if (!chainJoinResult.isEmpty()) + { + Map parts = + partitioner.partition(chainJoinResult); + for (Map.Entry entry : parts.entrySet()) + { + partitionResult.get(entry.getKey()).add(entry.getValue()); + } + joinedRows[0] += chainJoinResult.size; + } + } + + } + } + })); + } + + VectorizedRowBatch[] tailBatches = partitioner.getRowBatches(); + for (int hash = 0; hash < tailBatches.length; ++hash) + { + if (!tailBatches[hash].isEmpty()) + { + partitionResult.get(hash).add(tailBatches[hash]); + } + } + return joinedRows[0]; + } + /** * Scan the partitioned file of the right table and do the join. * diff --git a/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/BasePartitionedJoinWorker.java b/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/BasePartitionedJoinWorker.java index 4db1a98c63..23a27b253a 100644 --- a/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/BasePartitionedJoinWorker.java +++ b/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/BasePartitionedJoinWorker.java @@ -30,9 +30,15 @@ import io.pixelsdb.pixels.executor.join.*; import io.pixelsdb.pixels.planner.plan.physical.domain.MultiOutputInfo; import io.pixelsdb.pixels.planner.plan.physical.domain.PartitionInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.PartitionedTableInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.InputInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleInfo; import io.pixelsdb.pixels.planner.plan.physical.domain.StorageInfo; import io.pixelsdb.pixels.planner.plan.physical.input.PartitionedJoinInput; import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; +import io.pixelsdb.pixels.storage.s3qs.S3QS; +import io.pixelsdb.pixels.storage.s3qs.S3QueueMessage; +import io.pixelsdb.pixels.storage.s3qs.S3QueuePollResult; import org.apache.logging.log4j.Logger; import java.io.IOException; @@ -83,20 +89,28 @@ public JoinOutput process(PartitionedJoinInput event) long transId = event.getTransId(); long timestamp = event.getTimestamp(); requireNonNull(event.getSmallTable(), "event.smallTable is null"); + boolean leftS3QS = isS3QSShuffle(event.getSmallTable()); StorageInfo leftInputStorageInfo = event.getSmallTable().getStorageInfo(); List leftPartitioned = event.getSmallTable().getInputFiles(); - requireNonNull(leftPartitioned, "leftPartitioned is null"); - checkArgument(leftPartitioned.size() > 0, "leftPartitioned is empty"); + if (!leftS3QS) + { + requireNonNull(leftPartitioned, "leftPartitioned is null"); + checkArgument(leftPartitioned.size() > 0, "leftPartitioned is empty"); + } int leftParallelism = event.getSmallTable().getParallelism(); checkArgument(leftParallelism > 0, "leftParallelism is not positive"); String[] leftColumnsToRead = event.getSmallTable().getColumnsToRead(); int[] leftKeyColumnIds = event.getSmallTable().getKeyColumnIds(); requireNonNull(event.getLargeTable(), "event.largeTable is null"); + boolean rightS3QS = isS3QSShuffle(event.getLargeTable()); StorageInfo rightInputStorageInfo = event.getLargeTable().getStorageInfo(); List rightPartitioned = event.getLargeTable().getInputFiles(); - requireNonNull(rightPartitioned, "rightPartitioned is null"); - checkArgument(rightPartitioned.size() > 0, "rightPartitioned is empty"); + if (!rightS3QS) + { + requireNonNull(rightPartitioned, "rightPartitioned is null"); + checkArgument(rightPartitioned.size() > 0, "rightPartitioned is empty"); + } int rightParallelism = event.getLargeTable().getParallelism(); checkArgument(rightParallelism > 0, "rightParallelism is not positive"); String[] rightColumnsToRead = event.getLargeTable().getColumnsToRead(); @@ -141,15 +155,45 @@ public JoinOutput process(PartitionedJoinInput event) WorkerCommon.initStorage(leftInputStorageInfo); WorkerCommon.initStorage(rightInputStorageInfo); + WorkerCommon.initOptionalShuffleStorage(event.getSmallTable().getShuffleInfo()); + WorkerCommon.initOptionalShuffleStorage(event.getLargeTable().getShuffleInfo()); WorkerCommon.initStorage(outputStorageInfo); // build the joiner. AtomicReference leftSchema = new AtomicReference<>(); AtomicReference rightSchema = new AtomicReference<>(); - WorkerCommon.getFileSchemaFromPaths(threadPool, - WorkerCommon.getStorage(leftInputStorageInfo.getScheme()), - WorkerCommon.getStorage(rightInputStorageInfo.getScheme()), - leftSchema, rightSchema, leftPartitioned, rightPartitioned); + Map> leftPendingMessages = new HashMap<>(); + Map> rightPendingMessages = new HashMap<>(); + if (!leftS3QS && !rightS3QS) + { + WorkerCommon.getFileSchemaFromPaths(threadPool, + WorkerCommon.getStorage(leftInputStorageInfo.getScheme()), + WorkerCommon.getStorage(rightInputStorageInfo.getScheme()), + leftSchema, rightSchema, leftPartitioned, rightPartitioned); + } + else + { + if (leftS3QS) + { + leftSchema.set(getS3QSFileSchema( + event.getSmallTable().getShuffleInfo(), hashValues, leftPendingMessages)); + } + else + { + leftSchema.set(WorkerCommon.getFileSchemaFromPaths( + WorkerCommon.getStorage(leftInputStorageInfo.getScheme()), leftPartitioned)); + } + if (rightS3QS) + { + rightSchema.set(getS3QSFileSchema( + event.getLargeTable().getShuffleInfo(), hashValues, rightPendingMessages)); + } + else + { + rightSchema.set(WorkerCommon.getFileSchemaFromPaths( + WorkerCommon.getStorage(rightInputStorageInfo.getScheme()), rightPartitioned)); + } + } /* * Issue #450: * For the left and the right partial partitioned files, the file schema is equal to the columns to read in normal cases. @@ -161,34 +205,42 @@ public JoinOutput process(PartitionedJoinInput event) WorkerCommon.getResultSchema(rightSchema.get(), rightColumnsToRead), rightColAlias, rightProjection, rightKeyColumnIds); // build the hash table for the left table. - List leftFutures = new ArrayList<>(leftPartitioned.size()); - int leftSplitSize = leftPartitioned.size() / leftParallelism; - if (leftPartitioned.size() % leftParallelism > 0) + if (leftS3QS) { - leftSplitSize++; + buildHashTableS3QS(transId, timestamp, (HashJoiner) joiner, event.getSmallTable().getShuffleInfo(), + leftPendingMessages, leftColumnsToRead, hashValues, workerMetrics); } - for (int i = 0; i < leftPartitioned.size(); i += leftSplitSize) + else { - List parts = new LinkedList<>(); - for (int j = i; j < i + leftSplitSize && j < leftPartitioned.size(); ++j) + List leftFutures = new ArrayList<>(leftPartitioned.size()); + int leftSplitSize = leftPartitioned.size() / leftParallelism; + if (leftPartitioned.size() % leftParallelism > 0) { - parts.add(leftPartitioned.get(j)); + leftSplitSize++; } - leftFutures.add(threadPool.submit(() -> { - try - { - buildHashTable(transId, timestamp, (HashJoiner) joiner, parts, leftColumnsToRead, leftInputStorageInfo.getScheme(), - hashValues, numPartition, workerMetrics); - } - catch (Throwable e) + for (int i = 0; i < leftPartitioned.size(); i += leftSplitSize) + { + List parts = new LinkedList<>(); + for (int j = i; j < i + leftSplitSize && j < leftPartitioned.size(); ++j) { - throw new WorkerException("error during hash table construction", e); + parts.add(leftPartitioned.get(j)); } - })); - } - for (Future future : leftFutures) - { - future.get(); + leftFutures.add(threadPool.submit(() -> { + try + { + buildHashTable(transId, timestamp, (HashJoiner) joiner, parts, leftColumnsToRead, leftInputStorageInfo.getScheme(), + hashValues, numPartition, workerMetrics); + } + catch (Throwable e) + { + throw new WorkerException("error during hash table construction", e); + } + })); + } + for (Future future : leftFutures) + { + future.get(); + } } logger.info("hash table size: " + joiner.getSmallTableSize() + ", duration (ns): " + (workerMetrics.getInputCostNs() + workerMetrics.getComputeCostNs())); @@ -209,50 +261,79 @@ public JoinOutput process(PartitionedJoinInput event) // scan the right table and do the join. if (joiner.getSmallTableSize() > 0) { - int rightSplitSize = rightPartitioned.size() / rightParallelism; - if (rightPartitioned.size() % rightParallelism > 0) + if (rightS3QS) { - rightSplitSize++; + if (partitionOutput) + { + joinWithRightTableAndPartitionS3QS(transId, timestamp, joiner, + event.getLargeTable().getShuffleInfo(), rightPendingMessages, rightColumnsToRead, + hashValues, outputPartitionInfo, result, workerMetrics); + } + else + { + joinWithRightTableS3QS(transId, timestamp, joiner, event.getLargeTable().getShuffleInfo(), + rightPendingMessages, rightColumnsToRead, hashValues, result.get(0), workerMetrics); + } } - - for (int i = 0; i < rightPartitioned.size(); i += rightSplitSize) + else { - List parts = new LinkedList<>(); - for (int j = i; j < i + rightSplitSize && j < rightPartitioned.size(); ++j) + int rightSplitSize = rightPartitioned.size() / rightParallelism; + if (rightPartitioned.size() % rightParallelism > 0) { - parts.add(rightPartitioned.get(j)); + rightSplitSize++; } - threadPool.execute(() -> { - try - { - int numJoinedRows = partitionOutput ? - joinWithRightTableAndPartition( - transId, timestamp, joiner, parts, rightColumnsToRead, - rightInputStorageInfo.getScheme(), hashValues, - numPartition, outputPartitionInfo, result, workerMetrics) : - joinWithRightTable(transId, timestamp, joiner, parts, rightColumnsToRead, - rightInputStorageInfo.getScheme(), hashValues, numPartition, - result.get(0), workerMetrics); - } catch (Throwable e) + + for (int i = 0; i < rightPartitioned.size(); i += rightSplitSize) + { + List parts = new LinkedList<>(); + for (int j = i; j < i + rightSplitSize && j < rightPartitioned.size(); ++j) { - throw new WorkerException("error during hash join", e); + parts.add(rightPartitioned.get(j)); } - }); - } - threadPool.shutdown(); - try - { - while (!threadPool.awaitTermination(60, TimeUnit.SECONDS)) ; - } catch (InterruptedException e) - { - throw new WorkerException("interrupted while waiting for the termination of join", e); - } + threadPool.execute(() -> { + try + { + int numJoinedRows = partitionOutput ? + joinWithRightTableAndPartition( + transId, timestamp, joiner, parts, rightColumnsToRead, + rightInputStorageInfo.getScheme(), hashValues, + numPartition, outputPartitionInfo, result, workerMetrics) : + joinWithRightTable(transId, timestamp, joiner, parts, rightColumnsToRead, + rightInputStorageInfo.getScheme(), hashValues, numPartition, + result.get(0), workerMetrics); + } catch (Throwable e) + { + throw new WorkerException("error during hash join", e); + } + }); + } + threadPool.shutdown(); + try + { + while (!threadPool.awaitTermination(60, TimeUnit.SECONDS)) ; + } catch (InterruptedException e) + { + throw new WorkerException("interrupted while waiting for the termination of join", e); + } - if (exceptionHandler.hasException()) + if (exceptionHandler.hasException()) + { + throw new WorkerException("error occurred threads, please check the stacktrace before this log record"); + } + } + } + else if (rightS3QS) + { + for (int hashValue : hashValues) { - throw new WorkerException("error occurred threads, please check the stacktrace before this log record"); + drainS3QSPartition(event.getLargeTable().getShuffleInfo(), hashValue, rightPendingMessages, + message -> { }); } } + if (!threadPool.isShutdown()) + { + threadPool.shutdown(); + } String outputPath = outputFolder + outputInfo.getFileNames().get(0); try @@ -357,6 +438,275 @@ public JoinOutput process(PartitionedJoinInput event) } } + protected interface S3QSDataHandler + { + void handle(S3QueueMessage message) throws Exception; + } + + protected interface RowBatchHandler + { + void handle(VectorizedRowBatch rowBatch) throws Exception; + } + + protected static boolean isS3QSShuffle(PartitionedTableInfo tableInfo) + { + ShuffleInfo shuffleInfo = tableInfo.getShuffleInfo(); + return shuffleInfo != null && + shuffleInfo.getStorageInfo() != null && + shuffleInfo.getStorageInfo().getScheme() == Storage.Scheme.s3qs; + } + + private static S3QS getS3QSStorage() + { + Storage storage = WorkerCommon.getStorage(Storage.Scheme.s3qs); + if (!(storage instanceof S3QS)) + { + throw new WorkerException("storage of scheme s3qs is not S3QS"); + } + return (S3QS) storage; + } + + protected static TypeDescription getS3QSFileSchema(ShuffleInfo shuffleInfo, List hashValues, + Map> pendingMessages) + throws IOException + { + checkArgument(!hashValues.isEmpty(), "hashValues is empty"); + S3QS s3qs = getS3QSStorage(); + int partitionId = hashValues.get(0); + while (true) + { + S3QueuePollResult result = pollS3QSMessage(s3qs, shuffleInfo, partitionId, pendingMessages); + if (result == null) + { + continue; + } + addPendingMessage(pendingMessages, partitionId, result); + S3QueueMessage message = result.getMessage(); + if (message.getMetadata() != null && !message.getMetadata().isEmpty()) + { + return TypeDescription.fromString(message.getMetadata()); + } + if (message.isData()) + { + try (PixelsReader pixelsReader = WorkerCommon.getReader( + message.getObjectPath(), WorkerCommon.getStorage(Storage.Scheme.s3qs))) + { + return pixelsReader.getFileSchema(); + } + } + } + } + + private static S3QueuePollResult pollS3QSMessage(S3QS s3qs, ShuffleInfo shuffleInfo, int partitionId, + Map> pendingMessages) + throws IOException + { + Queue pending = pendingMessages.get(partitionId); + if (pending != null) + { + S3QueuePollResult result = pending.poll(); + if (result != null) + { + return result; + } + } + S3QueuePollResult result = s3qs.pollMessage( + shuffleInfo.getShuffleId(), partitionId, shuffleInfo.getPollTimeoutSeconds()); + if (result != null) + { + validateS3QSMessage(shuffleInfo, partitionId, result.getMessage()); + } + return result; + } + + private static void addPendingMessage(Map> pendingMessages, int partitionId, + S3QueuePollResult result) + { + Queue pending = pendingMessages.get(partitionId); + if (pending == null) + { + pending = new LinkedList<>(); + pendingMessages.put(partitionId, pending); + } + pending.add(result); + } + + private static void validateS3QSMessage(ShuffleInfo shuffleInfo, int partitionId, S3QueueMessage message) + { + if (!shuffleInfo.getShuffleId().equals(message.getShuffleId())) + { + throw new WorkerException("unexpected s3qs shuffle id: " + message.getShuffleId()); + } + if (message.getPartitionId() != partitionId) + { + throw new WorkerException("unexpected s3qs partition id: " + message.getPartitionId() + + ", expected " + partitionId); + } + } + + /** + * Drain one S3QS partition queue until all producer end markers are seen and + * one final long poll returns empty. DATA messages are acknowledged only + * after the handler has processed the referenced S3 object. + */ + protected static void drainS3QSPartition(ShuffleInfo shuffleInfo, int partitionId, + Map> pendingMessages, + S3QSDataHandler dataHandler) throws Exception + { + S3QS s3qs = getS3QSStorage(); + Set endedProducers = new HashSet<>(shuffleInfo.getProducerCount()); + while (true) + { + S3QueuePollResult result = pollS3QSMessage(s3qs, shuffleInfo, partitionId, pendingMessages); + if (result == null) + { + if (endedProducers.size() >= shuffleInfo.getProducerCount()) + { + return; + } + continue; + } + + S3QueueMessage message = result.getMessage(); + message.setReceiptHandle(result.getReceiptHandle()); + if (message.isProducerEnd()) + { + endedProducers.add(message.getProducerId()); + s3qs.finishWork(message); + } + else if (message.isData()) + { + dataHandler.handle(message); + s3qs.finishWork(message); + } + else + { + throw new WorkerException("unsupported s3qs message type: " + message.getMessageType()); + } + } + } + + protected static void readS3QSDataObject(long transId, long timestamp, S3QueueMessage message, + String[] columnsToRead, WorkerMetrics workerMetrics, + RowBatchHandler rowBatchHandler) + throws Exception + { + WorkerMetrics.Timer readCostTimer = new WorkerMetrics.Timer(); + WorkerMetrics.Timer computeCostTimer = new WorkerMetrics.Timer(); + long readBytes = 0L; + int numReadRequests = 0; + readCostTimer.start(); + try (PixelsReader pixelsReader = WorkerCommon.getReader( + message.getObjectPath(), WorkerCommon.getStorage(Storage.Scheme.s3qs))) + { + readCostTimer.stop(); + PixelsReaderOption option = WorkerCommon.getReaderOption( + transId, timestamp, columnsToRead, new InputInfo(message.getObjectPath(), 0, -1)); + PixelsRecordReader recordReader = pixelsReader.read(option); + checkArgument(recordReader.isValid(), "failed to get record reader"); + + VectorizedRowBatch rowBatch; + computeCostTimer.start(); + do + { + rowBatch = recordReader.readBatch(WorkerCommon.rowBatchSize); + if (rowBatch.size > 0) + { + rowBatchHandler.handle(rowBatch); + } + } while (!rowBatch.endOfFile); + computeCostTimer.stop(); + computeCostTimer.minus(recordReader.getReadTimeNanos()); + readCostTimer.add(recordReader.getReadTimeNanos()); + readBytes += recordReader.getCompletedBytes(); + numReadRequests += recordReader.getNumReadRequests(); + } + workerMetrics.addReadBytes(readBytes); + workerMetrics.addNumReadRequests(numReadRequests); + workerMetrics.addInputCostNs(readCostTimer.getElapsedNs()); + workerMetrics.addComputeCostNs(computeCostTimer.getElapsedNs()); + } + + protected static void buildHashTableS3QS(long transId, long timestamp, HashJoiner joiner, ShuffleInfo shuffleInfo, + Map> pendingMessages, + String[] leftCols, List hashValues, + WorkerMetrics workerMetrics) throws Exception + { + for (int partitionId : hashValues) + { + drainS3QSPartition(shuffleInfo, partitionId, pendingMessages, message -> + readS3QSDataObject(transId, timestamp, message, leftCols, workerMetrics, + joiner::populateLeftTable)); + } + } + + protected static int joinWithRightTableS3QS( + long transId, long timestamp, Joiner joiner, ShuffleInfo shuffleInfo, + Map> pendingMessages, String[] rightCols, + List hashValues, ConcurrentLinkedQueue joinResult, + WorkerMetrics workerMetrics) throws Exception + { + final int[] joinedRows = {0}; + for (int partitionId : hashValues) + { + drainS3QSPartition(shuffleInfo, partitionId, pendingMessages, message -> + readS3QSDataObject(transId, timestamp, message, rightCols, workerMetrics, rowBatch -> { + List joinedBatches = joiner.join(rowBatch); + for (VectorizedRowBatch joined : joinedBatches) + { + if (!joined.isEmpty()) + { + joinResult.add(joined); + joinedRows[0] += joined.size; + } + } + })); + } + return joinedRows[0]; + } + + protected static int joinWithRightTableAndPartitionS3QS( + long transId, long timestamp, Joiner joiner, ShuffleInfo shuffleInfo, + Map> pendingMessages, String[] rightCols, + List hashValues, PartitionInfo postPartitionInfo, + List> partitionResult, WorkerMetrics workerMetrics) + throws Exception + { + requireNonNull(postPartitionInfo, "outputPartitionInfo is null"); + Partitioner partitioner = new Partitioner(postPartitionInfo.getNumPartition(), + WorkerCommon.rowBatchSize, joiner.getJoinedSchema(), postPartitionInfo.getKeyColumnIds()); + final int[] joinedRows = {0}; + for (int partitionId : hashValues) + { + drainS3QSPartition(shuffleInfo, partitionId, pendingMessages, message -> + readS3QSDataObject(transId, timestamp, message, rightCols, workerMetrics, rowBatch -> { + List joinedBatches = joiner.join(rowBatch); + for (VectorizedRowBatch joined : joinedBatches) + { + if (!joined.isEmpty()) + { + Map parts = partitioner.partition(joined); + for (Map.Entry entry : parts.entrySet()) + { + partitionResult.get(entry.getKey()).add(entry.getValue()); + } + joinedRows[0] += joined.size; + } + } + })); + } + + VectorizedRowBatch[] tailBatches = partitioner.getRowBatches(); + for (int hash = 0; hash < tailBatches.length; ++hash) + { + if (!tailBatches[hash].isEmpty()) + { + partitionResult.get(hash).add(tailBatches[hash]); + } + } + return joinedRows[0]; + } + /** * Scan the partitioned file of the left table and populate the hash table for the join. * diff --git a/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/S3QSStageWorkerRunner.java b/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/S3QSStageWorkerRunner.java new file mode 100644 index 0000000000..287cc6080f --- /dev/null +++ b/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/S3QSStageWorkerRunner.java @@ -0,0 +1,223 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of the License, + * or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public License + * along with Pixels. If not, see . + */ +package io.pixelsdb.pixels.worker.common; + +import com.alibaba.fastjson.JSON; +import io.pixelsdb.pixels.common.task.Worker; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.common.utils.Constants; +import io.pixelsdb.pixels.planner.coordinate.CFWorkerInfo; +import io.pixelsdb.pixels.planner.coordinate.TaskBatch; +import io.pixelsdb.pixels.planner.coordinate.TaskInfo; +import io.pixelsdb.pixels.planner.coordinate.WorkerCoordinateService; +import io.pixelsdb.pixels.planner.plan.physical.input.PartitionInput; +import io.pixelsdb.pixels.planner.plan.physical.input.PartitionedChainJoinInput; +import io.pixelsdb.pixels.planner.plan.physical.input.PartitionedJoinInput; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; +import io.pixelsdb.pixels.planner.plan.physical.output.PartitionOutput; + +import java.net.InetAddress; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +import static java.util.Objects.requireNonNull; + +/** + * Runs S3QS stage workers through the coordinator task protocol. + * + * This class deliberately lives above the concrete Base*Worker implementations: + * platform wrappers only adapt Lambda/vHive/Spike requests, this runner owns the + * register -> pull task -> execute -> complete loop, and Base*Worker still owns + * the physical S3QS read/write work for one task payload. + */ +public class S3QSStageWorkerRunner +{ + private final WorkerContext context; + private final WorkerCoordinateService workerCoordinateService; + + public S3QSStageWorkerRunner(WorkerContext context) + { + this.context = requireNonNull(context, "context is null"); + this.workerCoordinateService = null; + } + + S3QSStageWorkerRunner(WorkerContext context, WorkerCoordinateService workerCoordinateService) + { + this.context = requireNonNull(context, "context is null"); + this.workerCoordinateService = requireNonNull(workerCoordinateService, "workerCoordinateService is null"); + } + + public PartitionOutput runPartition(StageWorkerInput input) + { + return run(input, Constants.PARTITION_OPERATOR_NAME, WorkerType.PARTITION_S3QS, + PartitionInput.class, taskInput -> new BasePartitionWorker(context).process(taskInput), + new PartitionOutput()); + } + + public JoinOutput runPartitionedJoin(StageWorkerInput input) + { + return run(input, Constants.PARTITION_JOIN_OPERATOR_NAME, WorkerType.PARTITIONED_JOIN_S3QS, + PartitionedJoinInput.class, taskInput -> new BasePartitionedJoinWorker(context).process(taskInput), + new JoinOutput()); + } + + public JoinOutput runPartitionedChainJoin(StageWorkerInput input) + { + return run(input, Constants.PARTITION_JOIN_OPERATOR_NAME, WorkerType.PARTITIONED_CHAIN_JOIN_S3QS, + PartitionedChainJoinInput.class, + taskInput -> new BasePartitionedChainJoinWorker(context).process(taskInput), new JoinOutput()); + } + + private O run(StageWorkerInput input, String operatorName, WorkerType expectedWorkerType, + Class taskInputClass, TaskExecutor taskExecutor, O aggregateOutput) + { + requireNonNull(input, "input is null"); + if (input.getWorkerType() != expectedWorkerType) + { + throw new WorkerException("unexpected S3QS stage worker type: " + input.getWorkerType()); + } + + long startTimeMs = System.currentTimeMillis(); + aggregateOutput.setStartTimeMs(startTimeMs); + aggregateOutput.setRequestId(context.getRequestId()); + aggregateOutput.setSuccessful(true); + aggregateOutput.setErrorMessage(""); + + WorkerCoordinateService coordinateService = getWorkerCoordinateService(input); + Worker runtimeWorker = null; + try + { + CFWorkerInfo workerInfo = new CFWorkerInfo(InetAddress.getLocalHost().getHostAddress(), -1, + input.getTransId(), input.getStageId(), operatorName, new ArrayList<>()); + runtimeWorker = coordinateService.registerWorker(workerInfo); + + TaskBatch taskBatch = coordinateService.getTasksToExecute(runtimeWorker.getWorkerId()); + while (!taskBatch.isEndOfTasks()) + { + List taskInfos = taskBatch.getTasks(); + for (TaskInfo taskInfo : taskInfos) + { + executeTask(taskInfo, taskInputClass, taskExecutor, aggregateOutput); + } + coordinateService.completeTasks(runtimeWorker.getWorkerId(), taskInfos); + taskBatch = coordinateService.getTasksToExecute(runtimeWorker.getWorkerId()); + } + } + catch (Throwable e) + { + aggregateOutput.setSuccessful(false); + aggregateOutput.setErrorMessage(e.getMessage()); + throw new WorkerException("failed to run S3QS stage worker", e); + } + finally + { + if (runtimeWorker != null) + { + try + { + coordinateService.terminateWorker(runtimeWorker.getWorkerId()); + } + catch (Throwable ignored) + { + // The worker is already done from the caller's perspective. + } + } + aggregateOutput.setDurationMs((int) (System.currentTimeMillis() - startTimeMs)); + } + return aggregateOutput; + } + + private WorkerCoordinateService getWorkerCoordinateService(StageWorkerInput input) + { + if (workerCoordinateService != null) + { + return workerCoordinateService; + } + return new WorkerCoordinateService(input.getCoordinatorHost(), input.getCoordinatorPort()); + } + + private void executeTask(TaskInfo taskInfo, Class taskInputClass, + TaskExecutor taskExecutor, O aggregateOutput) + { + try + { + I taskInput = JSON.parseObject(taskInfo.getPayload(), taskInputClass); + O taskOutput = taskExecutor.execute(taskInput); + taskInfo.setSuccess(taskOutput.isSuccessful()); + mergeOutput(aggregateOutput, taskOutput); + } + catch (Throwable e) + { + taskInfo.setSuccess(false); + throw new WorkerException("failed to execute S3QS coordinator task " + taskInfo.getTaskId(), e); + } + } + + private void mergeOutput(Output aggregateOutput, Output taskOutput) + { + if (taskOutput == null) + { + return; + } + aggregateOutput.setSuccessful(aggregateOutput.isSuccessful() && taskOutput.isSuccessful()); + if (!taskOutput.isSuccessful()) + { + aggregateOutput.setErrorMessage(taskOutput.getErrorMessage()); + } + for (String output : taskOutput.getOutputs()) + { + aggregateOutput.addOutput(output); + } + aggregateOutput.setNumReadRequests(aggregateOutput.getNumReadRequests() + taskOutput.getNumReadRequests()); + aggregateOutput.setNumWriteRequests(aggregateOutput.getNumWriteRequests() + taskOutput.getNumWriteRequests()); + aggregateOutput.setTotalReadBytes(aggregateOutput.getTotalReadBytes() + taskOutput.getTotalReadBytes()); + aggregateOutput.setTotalWriteBytes(aggregateOutput.getTotalWriteBytes() + taskOutput.getTotalWriteBytes()); + aggregateOutput.setCumulativeInputCostMs( + aggregateOutput.getCumulativeInputCostMs() + taskOutput.getCumulativeInputCostMs()); + aggregateOutput.setCumulativeComputeCostMs( + aggregateOutput.getCumulativeComputeCostMs() + taskOutput.getCumulativeComputeCostMs()); + aggregateOutput.setCumulativeOutputCostMs( + aggregateOutput.getCumulativeOutputCostMs() + taskOutput.getCumulativeOutputCostMs()); + + if (aggregateOutput instanceof PartitionOutput && taskOutput instanceof PartitionOutput) + { + PartitionOutput aggregatePartitionOutput = (PartitionOutput) aggregateOutput; + PartitionOutput taskPartitionOutput = (PartitionOutput) taskOutput; + Set hashValues = aggregatePartitionOutput.getHashValues(); + if (hashValues == null) + { + hashValues = new HashSet<>(); + aggregatePartitionOutput.setHashValues(hashValues); + } + if (taskPartitionOutput.getHashValues() != null) + { + hashValues.addAll(taskPartitionOutput.getHashValues()); + } + } + } + + private interface TaskExecutor + { + O execute(I input); + } +} diff --git a/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/WorkerCommon.java b/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/WorkerCommon.java index aa42b6f0d9..2bb93d153c 100644 --- a/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/WorkerCommon.java +++ b/pixels-turbo/pixels-worker-common/src/main/java/io/pixelsdb/pixels/worker/common/WorkerCommon.java @@ -22,6 +22,7 @@ import com.google.common.collect.ImmutableList; import io.pixelsdb.pixels.common.physical.Storage; import io.pixelsdb.pixels.common.physical.StorageFactory; +import io.pixelsdb.pixels.common.physical.PhysicalWriter; import io.pixelsdb.pixels.common.turbo.Output; import io.pixelsdb.pixels.common.utils.ConfigFactory; import io.pixelsdb.pixels.core.*; @@ -29,7 +30,10 @@ import io.pixelsdb.pixels.core.reader.PixelsReaderOption; import io.pixelsdb.pixels.planner.plan.physical.domain.InputInfo; import io.pixelsdb.pixels.planner.plan.physical.domain.InputSplit; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleQueueInfo; import io.pixelsdb.pixels.planner.plan.physical.domain.StorageInfo; +import io.pixelsdb.pixels.storage.s3qs.S3QS; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -64,6 +68,7 @@ public class WorkerCommon protected static Storage minio; private static Storage redis; private static Storage stream; + private static Storage s3qs; public static final int rowBatchSize; protected static final int pixelStride; protected static final int rowGroupSize; @@ -106,6 +111,10 @@ else if (WorkerCommon.stream == null && storageInfo.getScheme() == Storage.Schem { WorkerCommon.stream = StorageFactory.Instance().getStorage(Storage.Scheme.httpstream); } + else if (WorkerCommon.s3qs == null && storageInfo.getScheme() == Storage.Scheme.s3qs) + { + WorkerCommon.s3qs = StorageFactory.Instance().getStorage(Storage.Scheme.s3qs); + } } catch (Throwable e) { throw new WorkerException("failed to initialize the storage of scheme " + storageInfo.getScheme(), e); @@ -124,10 +133,59 @@ public static Storage getStorage(Storage.Scheme scheme) return redis; case httpstream: return stream; + case s3qs: + return s3qs; } throw new UnsupportedOperationException("scheme " + scheme + " is not supported"); } + public static void initOptionalShuffleStorage(ShuffleInfo shuffleInfo) + { + if (shuffleInfo == null) + { + return; + } + StorageInfo storageInfo = requireNonNull(shuffleInfo.getStorageInfo(), "shuffle storageInfo is null"); + initStorage(storageInfo); + if (storageInfo.getScheme() == Storage.Scheme.s3qs) + { + registerS3QSShuffleQueues(shuffleInfo); + } + } + + private static void registerS3QSShuffleQueues(ShuffleInfo shuffleInfo) + { + List queues = requireNonNull(shuffleInfo.getQueues(), "shuffle queues is null"); + checkArgument(shuffleInfo.getNumPartitions() > 0, "shuffle numPartitions must be positive"); + checkArgument(queues.size() == shuffleInfo.getNumPartitions(), + "shuffle queues size does not match numPartitions"); + + Storage storage = getStorage(Storage.Scheme.s3qs); + if (!(storage instanceof S3QS)) + { + throw new WorkerException("storage of scheme s3qs is not S3QS"); + } + + S3QS s3qsStorage = (S3QS) storage; + try + { + for (ShuffleQueueInfo queue : queues) + { + requireNonNull(queue, "shuffle queue is null"); + checkArgument(queue.getPartitionId() >= 0 && queue.getPartitionId() < shuffleInfo.getNumPartitions(), + "shuffle queue partitionId is out of range"); + String queueUrl = s3qsStorage.registerQueue(shuffleInfo.getShuffleId(), queue.getPartitionId(), + queue.getQueueName(), queue.getQueueUrl()); + queue.setQueueUrl(queueUrl); + } + } + catch (IOException e) + { + throw new WorkerException("failed to register s3qs shuffle queues for " + + shuffleInfo.getShuffleId(), e); + } + } + /** * Read the schemas of the two joined tables, concurrently using the executor, thus * to reduce the latency of the schema reading. @@ -532,6 +590,40 @@ public static PixelsWriter getWriter(TypeDescription schema, Storage storage, return writer; } + /** + * Build a Pixels writer on top of an already-created physical writer. + * + * This is used by S3QS shuffle producers because the S3 object writer is + * created by S3QS.offer() together with the queue message that will be sent + * when the physical writer is closed. + */ + public static PixelsWriter getWriter(TypeDescription schema, Storage storage, String filePath, + PhysicalWriter physicalWriter, boolean isPartitioned, + List keyColumnIds) + { + requireNonNull(schema, "schema is null"); + requireNonNull(filePath, "fileName is null"); + requireNonNull(storage, "storage is null"); + requireNonNull(physicalWriter, "physicalWriter is null"); + checkArgument(!isPartitioned || keyColumnIds != null, + "keyColumnIds is null whereas isPartitioned is true"); + PixelsWriterImpl.Builder builder = PixelsWriterImpl.newBuilder() + .setSchema(schema) + .setPixelStride(pixelStride) + .setRowGroupSize(rowGroupSize) + .setStorage(storage) + .setPath(filePath) + .setPhysicalWriter(physicalWriter) + .setOverwrite(true) + .setEncodingLevel(EncodingLevel.EL2) + .setPartitioned(isPartitioned); + if (isPartitioned) + { + builder.setPartKeyColumnIds(keyColumnIds); + } + return builder.build(); + } + /** * Create the reader option for a record reader of the given input file. * diff --git a/pixels-turbo/pixels-worker-common/src/test/java/io/pixelsdb/pixels/worker/common/TestS3QSConsumerWiring.java b/pixels-turbo/pixels-worker-common/src/test/java/io/pixelsdb/pixels/worker/common/TestS3QSConsumerWiring.java new file mode 100644 index 0000000000..b452090103 --- /dev/null +++ b/pixels-turbo/pixels-worker-common/src/test/java/io/pixelsdb/pixels/worker/common/TestS3QSConsumerWiring.java @@ -0,0 +1,289 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public + * License along with Pixels. If not, see + * . + */ +package io.pixelsdb.pixels.worker.common; + +import io.pixelsdb.pixels.common.physical.Storage; +import io.pixelsdb.pixels.common.physical.PhysicalReader; +import io.pixelsdb.pixels.common.physical.PhysicalReaderUtil; +import io.pixelsdb.pixels.common.physical.PhysicalWriter; +import io.pixelsdb.pixels.core.TypeDescription; +import io.pixelsdb.pixels.executor.join.Joiner; +import io.pixelsdb.pixels.planner.plan.physical.domain.PartitionInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.PartitionedTableInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleQueueInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.StorageInfo; +import io.pixelsdb.pixels.storage.s3qs.S3QS; +import io.pixelsdb.pixels.storage.s3qs.S3QueueMessage; +import io.pixelsdb.pixels.storage.s3qs.S3QueuePollResult; +import org.junit.Test; +import software.amazon.awssdk.services.sqs.SqsClient; +import software.amazon.awssdk.services.sqs.model.DeleteMessageRequest; +import software.amazon.awssdk.services.sqs.model.Message; +import software.amazon.awssdk.services.sqs.model.MessageSystemAttributeName; +import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest; +import software.amazon.awssdk.services.sqs.model.ReceiveMessageResponse; + +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.lang.reflect.Field; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Queue; +import java.util.UUID; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assume.assumeTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static software.amazon.awssdk.services.sqs.model.MessageSystemAttributeName.APPROXIMATE_RECEIVE_COUNT; + +public class TestS3QSConsumerWiring +{ + @Test + public void s3qsConsumerPathRequiresExplicitShuffleInfo() + { + PartitionedTableInfo noShuffle = tableInfo(Storage.Scheme.s3qs, null); + assertFalse(BasePartitionedJoinWorker.isS3QSShuffle(noShuffle)); + + PartitionedTableInfo s3Shuffle = tableInfo(Storage.Scheme.s3qs, shuffleInfo(Storage.Scheme.s3)); + assertFalse(BasePartitionedJoinWorker.isS3QSShuffle(s3Shuffle)); + + PartitionedTableInfo s3qsShuffle = tableInfo(Storage.Scheme.s3, shuffleInfo(Storage.Scheme.s3qs)); + assertTrue(BasePartitionedJoinWorker.isS3QSShuffle(s3qsShuffle)); + } + + @Test + public void s3qsProbeHelpersAreNoOpWithoutAssignedPartitions() throws Exception + { + assertEquals(0, BasePartitionedJoinWorker.joinWithRightTableS3QS( + 1L, 1L, null, null, new HashMap>(), + new String[0], Collections.emptyList(), new ConcurrentLinkedQueue<>(), new WorkerMetrics())); + + assertEquals(0, BasePartitionedChainJoinWorker.joinWithRightTableS3QS( + 1L, 1L, null, null, null, new HashMap>(), + new String[0], Collections.emptyList(), new ConcurrentLinkedQueue<>(), new WorkerMetrics())); + } + + @Test + public void s3qsPostPartitionHelpersAreNoOpWithoutAssignedPartitions() throws Exception + { + PartitionInfo postPartition = new PartitionInfo(new int[] {0}, 1); + Joiner joiner = mockJoinerWithJoinedSchema(); + + assertEquals(0, BasePartitionedJoinWorker.joinWithRightTableAndPartitionS3QS( + 1L, 1L, joiner, null, new HashMap>(), + new String[0], Collections.emptyList(), postPartition, + Collections.singletonList(new ConcurrentLinkedQueue<>()), new WorkerMetrics())); + + assertEquals(0, BasePartitionedChainJoinWorker.joinWithRightTableAndPartitionS3QS( + 1L, 1L, null, joiner, null, new HashMap>(), + new String[0], Collections.emptyList(), postPartition, + Collections.singletonList(new ConcurrentLinkedQueue<>()), new WorkerMetrics())); + } + + @Test + public void s3qsProducerMessagesDrainThroughConsumerProtocol() throws Exception + { + SqsClient sqs = mock(SqsClient.class); + S3QS s3qs = newS3QS(sqs); + try + { + s3qs.registerQueue("shuffle-1", 0, "ignored-name", "queue-url-0"); + setWorkerCommonS3QS(s3qs); + + TypeDescription schema = TypeDescription.createStruct() + .addField("key", TypeDescription.createLong()); + ShuffleInfo shuffleInfo = new ShuffleInfo("shuffle-1", + new StorageInfo(Storage.Scheme.s3qs, null, null, null, null), + "s3://bucket/shuffle-1/", 1, 1, 1, 1, Collections.emptyList()); + S3QueueMessage data = BasePartitionWorker.createS3QSDataMessage( + shuffleInfo, 0, 7, 0, 0L, schema); + S3QueueMessage end = BasePartitionWorker.createS3QSProducerEndMessage( + shuffleInfo, 0, 7, 0, 1L, schema); + + when(sqs.receiveMessage(any(ReceiveMessageRequest.class))) + .thenReturn(receive("receipt-data", data)) + .thenReturn(receive("receipt-end", end)) + .thenReturn(ReceiveMessageResponse.builder().build()); + + AtomicInteger dataMessages = new AtomicInteger(); + BasePartitionedJoinWorker.drainS3QSPartition(shuffleInfo, 0, + new HashMap>(), + message -> { + assertTrue(message.isData()); + assertEquals(7, message.getProducerId()); + assertEquals(schema.toString(), message.getMetadata()); + dataMessages.incrementAndGet(); + }); + + assertEquals(1, dataMessages.get()); + verify(sqs).deleteMessage(DeleteMessageRequest.builder() + .queueUrl("queue-url-0").receiptHandle("receipt-data").build()); + verify(sqs).deleteMessage(DeleteMessageRequest.builder() + .queueUrl("queue-url-0").receiptHandle("receipt-end").build()); + } + finally + { + setWorkerCommonS3QS(null); + } + } + + @Test + public void awsWorkerCommonShuffleInfoWritesAndDrainsS3QS() throws Exception + { + String bucket = System.getenv("PIXELS_S3QS_IT_BUCKET"); + String prefix = trimSlashes(System.getenv("PIXELS_S3QS_IT_PREFIX")); + String queuePrefix = System.getenv("PIXELS_S3QS_IT_QUEUE_PREFIX"); + assumeTrue("set PIXELS_S3QS_IT_BUCKET to run the AWS worker-common S3QS integration test", + bucket != null && !bucket.trim().isEmpty()); + assumeTrue("set PIXELS_S3QS_IT_QUEUE_PREFIX to run the AWS worker-common S3QS integration test", + queuePrefix != null && !queuePrefix.trim().isEmpty()); + + String suffix = UUID.randomUUID().toString().replace("-", ""); + String shuffleId = "worker-it-" + suffix; + String objectRoot = bucket + "/" + (prefix.isEmpty() ? "" : prefix + "/") + shuffleId + "/"; + ShuffleQueueInfo queueInfo = new ShuffleQueueInfo(0, queuePrefix + "-" + suffix, null); + ShuffleInfo shuffleInfo = new ShuffleInfo(shuffleId, + new StorageInfo(Storage.Scheme.s3qs, null, null, null, null), + objectRoot, 1, 1, 1, 1, Collections.singletonList(queueInfo)); + + setWorkerCommonS3QS(null); + WorkerCommon.initOptionalShuffleStorage(shuffleInfo); + S3QS s3qs = (S3QS) WorkerCommon.getStorage(Storage.Scheme.s3qs); + byte[] payload = ("worker-payload-" + suffix).getBytes(StandardCharsets.UTF_8); + try + { + TypeDescription schema = TypeDescription.createStruct() + .addField("key", TypeDescription.createLong()); + S3QueueMessage data = BasePartitionWorker.createS3QSDataMessage( + shuffleInfo, 0, 0, 0, 0L, schema); + PhysicalWriter writer = s3qs.offer(data); + writer.append(payload); + writer.close(); + s3qs.publish(BasePartitionWorker.createS3QSProducerEndMessage( + shuffleInfo, 0, 0, 0, 1L, schema)); + + AtomicInteger dataMessages = new AtomicInteger(); + BasePartitionedJoinWorker.drainS3QSPartition(shuffleInfo, 0, + new HashMap>(), message -> { + PhysicalReader reader = PhysicalReaderUtil.newPhysicalReader(s3qs, message.getObjectPath()); + ByteBuffer buffer = reader.readFully(payload.length); + byte[] actual = new byte[payload.length]; + buffer.get(actual); + reader.close(); + assertArrayEquals(payload, actual); + assertEquals(schema.toString(), message.getMetadata()); + dataMessages.incrementAndGet(); + }); + + assertEquals(1, dataMessages.get()); + } + finally + { + if (queueInfo.getQueueUrl() != null) + { + s3qs.deleteQueue(queueInfo.getQueueUrl()); + } + s3qs.delete(objectRoot, true); + setWorkerCommonS3QS(null); + } + } + + private static Joiner mockJoinerWithJoinedSchema() + { + Joiner joiner = mock(Joiner.class); + TypeDescription joinedSchema = TypeDescription.createStruct() + .addField("key", TypeDescription.createLong()); + when(joiner.getJoinedSchema()).thenReturn(joinedSchema); + return joiner; + } + + private static PartitionedTableInfo tableInfo(Storage.Scheme tableScheme, ShuffleInfo shuffleInfo) + { + PartitionedTableInfo tableInfo = new PartitionedTableInfo(); + tableInfo.setStorageInfo(new StorageInfo(tableScheme, null, null, null, null)); + tableInfo.setShuffleInfo(shuffleInfo); + return tableInfo; + } + + private static ShuffleInfo shuffleInfo(Storage.Scheme shuffleScheme) + { + return new ShuffleInfo("shuffle-1", new StorageInfo(shuffleScheme, null, null, null, null), + "s3://bucket/shuffle-1/", 1, 1, 1, 1, Collections.emptyList()); + } + + private static S3QS newS3QS(SqsClient sqs) throws Exception + { + S3QS s3qs = new S3QS(30); + Field sqsField = S3QS.class.getDeclaredField("sqs"); + sqsField.setAccessible(true); + sqsField.set(s3qs, sqs); + return s3qs; + } + + private static void setWorkerCommonS3QS(S3QS s3qs) throws Exception + { + Field s3qsField = WorkerCommon.class.getDeclaredField("s3qs"); + s3qsField.setAccessible(true); + s3qsField.set(null, s3qs); + } + + private static ReceiveMessageResponse receive(String receiptHandle, S3QueueMessage body) throws Exception + { + Message message = Message.builder() + .body(body.toMessageBody()) + .receiptHandle(receiptHandle) + .attributes(receiveCount(1)) + .build(); + return ReceiveMessageResponse.builder().messages(message).build(); + } + + private static Map receiveCount(int count) + { + return Collections.singletonMap(APPROXIMATE_RECEIVE_COUNT, String.valueOf(count)); + } + + private static String trimSlashes(String value) + { + if (value == null) + { + return ""; + } + String result = value.trim(); + while (result.startsWith("/")) + { + result = result.substring(1); + } + while (result.endsWith("/")) + { + result = result.substring(0, result.length() - 1); + } + return result; + } +} diff --git a/pixels-turbo/pixels-worker-common/src/test/java/io/pixelsdb/pixels/worker/common/TestS3QSProducerWiring.java b/pixels-turbo/pixels-worker-common/src/test/java/io/pixelsdb/pixels/worker/common/TestS3QSProducerWiring.java new file mode 100644 index 0000000000..2ea17104a6 --- /dev/null +++ b/pixels-turbo/pixels-worker-common/src/test/java/io/pixelsdb/pixels/worker/common/TestS3QSProducerWiring.java @@ -0,0 +1,111 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public + * License along with Pixels. If not, see + * . + */ +package io.pixelsdb.pixels.worker.common; + +import io.pixelsdb.pixels.common.physical.Storage; +import io.pixelsdb.pixels.core.TypeDescription; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.StorageInfo; +import io.pixelsdb.pixels.storage.s3qs.S3QueueMessage; +import org.junit.Test; + +import java.util.Collections; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class TestS3QSProducerWiring +{ + @Test + public void s3qsProducerPathRequiresExplicitShuffleInfo() + { + assertFalse(BasePartitionWorker.isS3QSShuffle(null)); + assertFalse(BasePartitionWorker.isS3QSShuffle(shuffleInfo(Storage.Scheme.s3))); + assertTrue(BasePartitionWorker.isS3QSShuffle(shuffleInfo(Storage.Scheme.s3qs))); + } + + @Test + public void dataMessageCarriesProducerPartitionAndSchemaMetadata() throws Exception + { + TypeDescription schema = oneColumnSchema(); + S3QueueMessage message = BasePartitionWorker.createS3QSDataMessage( + shuffleInfo(Storage.Scheme.s3qs), 3, 7, 0, 11L, schema); + + assertTrue(message.isData()); + assertFalse(message.isProducerEnd()); + assertEquals("shuffle-1", message.getShuffleId()); + assertEquals(3, message.getPartitionId()); + assertEquals(7, message.getProducerId()); + assertEquals(0, message.getProducerAttemptId()); + assertEquals(11L, message.getSequenceId()); + assertEquals("s3://bucket/shuffle-1/", message.getObjectPath()); + assertEquals(schema.toString(), message.getMetadata()); + + S3QueueMessage parsed = S3QueueMessage.fromMessageBody(message.toMessageBody()); + assertTrue(parsed.isData()); + assertEquals(message.getShuffleId(), parsed.getShuffleId()); + assertEquals(message.getPartitionId(), parsed.getPartitionId()); + assertEquals(message.getProducerId(), parsed.getProducerId()); + assertEquals(message.getSequenceId(), parsed.getSequenceId()); + assertEquals(message.getObjectPath(), parsed.getObjectPath()); + assertEquals(message.getMetadata(), parsed.getMetadata()); + } + + @Test + public void producerEndMessageUsesSameProducerIdentityAndNormalEndReason() throws Exception + { + TypeDescription schema = oneColumnSchema(); + S3QueueMessage message = BasePartitionWorker.createS3QSProducerEndMessage( + shuffleInfo(Storage.Scheme.s3qs), 5, 9, 0, 2L, schema); + + assertFalse(message.isData()); + assertTrue(message.isProducerEnd()); + assertEquals("shuffle-1", message.getShuffleId()); + assertEquals(5, message.getPartitionId()); + assertEquals(9, message.getProducerId()); + assertEquals(0, message.getProducerAttemptId()); + assertEquals(2L, message.getSequenceId()); + assertEquals("NORMAL", message.getEndReason()); + assertEquals("", message.getObjectPath()); + assertEquals(schema.toString(), message.getMetadata()); + + S3QueueMessage parsed = S3QueueMessage.fromMessageBody(message.toMessageBody()); + assertTrue(parsed.isProducerEnd()); + assertEquals(message.getShuffleId(), parsed.getShuffleId()); + assertEquals(message.getPartitionId(), parsed.getPartitionId()); + assertEquals(message.getProducerId(), parsed.getProducerId()); + assertEquals(message.getSequenceId(), parsed.getSequenceId()); + assertEquals(message.getEndReason(), parsed.getEndReason()); + assertEquals(message.getMetadata(), parsed.getMetadata()); + } + + private static TypeDescription oneColumnSchema() + { + return TypeDescription.createStruct() + .addField("key", TypeDescription.createLong()); + } + + private static ShuffleInfo shuffleInfo(Storage.Scheme shuffleScheme) + { + return new ShuffleInfo("shuffle-1", new StorageInfo(shuffleScheme, null, null, null, null), + "s3://bucket/shuffle-1/", 8, 2, 4, 1, Collections.emptyList()); + } +} diff --git a/pixels-turbo/pixels-worker-common/src/test/java/io/pixelsdb/pixels/worker/common/TestS3QSStageWorkerRunner.java b/pixels-turbo/pixels-worker-common/src/test/java/io/pixelsdb/pixels/worker/common/TestS3QSStageWorkerRunner.java new file mode 100644 index 0000000000..be2fed8855 --- /dev/null +++ b/pixels-turbo/pixels-worker-common/src/test/java/io/pixelsdb/pixels/worker/common/TestS3QSStageWorkerRunner.java @@ -0,0 +1,158 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public + * License along with Pixels. If not, see + * . + */ +package io.pixelsdb.pixels.worker.common; + +import com.alibaba.fastjson.JSON; +import io.pixelsdb.pixels.common.lease.Lease; +import io.pixelsdb.pixels.common.task.Worker; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.common.utils.Constants; +import io.pixelsdb.pixels.planner.coordinate.CFWorkerInfo; +import io.pixelsdb.pixels.planner.coordinate.TaskBatch; +import io.pixelsdb.pixels.planner.coordinate.TaskInfo; +import io.pixelsdb.pixels.planner.coordinate.WorkerCoordinateService; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.PartitionOutput; +import org.apache.logging.log4j.LogManager; +import org.junit.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.InOrder; + +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.inOrder; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +public class TestS3QSStageWorkerRunner +{ + @Test + public void coordinatorLifecycleExecutesCompletesAndTerminatesTask() throws Exception + { + WorkerCoordinateService coordinateService = mock(WorkerCoordinateService.class); + WorkerContext context = new WorkerContext(LogManager.getLogger(TestS3QSStageWorkerRunner.class), + new WorkerMetrics(), "request-1277"); + S3QSStageWorkerRunner runner = new S3QSStageWorkerRunner(context, coordinateService); + + long workerId = 17L; + long transId = 1277L; + int stageId = 3; + Worker runtimeWorker = new Worker<>(workerId, + new Lease(System.currentTimeMillis(), 60000L), 0, + new CFWorkerInfo("localhost", -1, transId, stageId, + Constants.PARTITION_OPERATOR_NAME, Collections.emptyList())); + TaskInfo taskInfo = new TaskInfo(5, JSON.toJSONString("task-payload")); + List taskInfos = Collections.singletonList(taskInfo); + TaskBatch taskBatch = new TaskBatch(false, taskInfos); + TaskBatch endOfTasks = new TaskBatch(true, Collections.emptyList()); + + when(coordinateService.registerWorker(any(CFWorkerInfo.class))).thenReturn(runtimeWorker); + when(coordinateService.getTasksToExecute(workerId)).thenReturn(taskBatch, endOfTasks); + + PartitionOutput taskOutput = new PartitionOutput(); + taskOutput.setSuccessful(true); + taskOutput.addOutput("s3://bucket/shuffle/task-5"); + taskOutput.setNumReadRequests(2); + taskOutput.setNumWriteRequests(1); + taskOutput.setTotalReadBytes(128L); + taskOutput.setTotalWriteBytes(64L); + taskOutput.setHashValues(Collections.singleton(9)); + AtomicReference executedPayload = new AtomicReference<>(); + + Class taskExecutorClass = Class.forName( + "io.pixelsdb.pixels.worker.common.S3QSStageWorkerRunner$TaskExecutor"); + Object taskExecutor = Proxy.newProxyInstance(taskExecutorClass.getClassLoader(), + new Class[] {taskExecutorClass}, (proxy, method, args) -> + { + if (method.getName().equals("execute")) + { + executedPayload.set(args[0]); + return taskOutput; + } + throw new UnsupportedOperationException(method.getName()); + }); + Method run = S3QSStageWorkerRunner.class.getDeclaredMethod("run", + StageWorkerInput.class, String.class, WorkerType.class, Class.class, + taskExecutorClass, Output.class); + run.setAccessible(true); + + StageWorkerInput input = new StageWorkerInput(transId, 1000L, stageId, + "test-s3qs-stage", WorkerType.PARTITION_S3QS); + PartitionOutput aggregateOutput = (PartitionOutput) run.invoke(runner, input, + Constants.PARTITION_OPERATOR_NAME, WorkerType.PARTITION_S3QS, + String.class, taskExecutor, new PartitionOutput()); + + assertEquals("task-payload", executedPayload.get()); + assertTrue(taskInfo.isSuccess()); + assertTrue(aggregateOutput.isSuccessful()); + assertEquals("request-1277", aggregateOutput.getRequestId()); + assertEquals(Collections.singletonList("s3://bucket/shuffle/task-5"), aggregateOutput.getOutputs()); + assertEquals(2, aggregateOutput.getNumReadRequests()); + assertEquals(1, aggregateOutput.getNumWriteRequests()); + assertEquals(128L, aggregateOutput.getTotalReadBytes()); + assertEquals(64L, aggregateOutput.getTotalWriteBytes()); + assertEquals(Collections.singleton(9), aggregateOutput.getHashValues()); + + ArgumentCaptor workerInfo = ArgumentCaptor.forClass(CFWorkerInfo.class); + InOrder calls = inOrder(coordinateService); + calls.verify(coordinateService).registerWorker(workerInfo.capture()); + calls.verify(coordinateService).getTasksToExecute(workerId); + calls.verify(coordinateService).completeTasks(workerId, taskInfos); + calls.verify(coordinateService).getTasksToExecute(workerId); + calls.verify(coordinateService).terminateWorker(workerId); + + assertEquals(transId, workerInfo.getValue().getTransId()); + assertEquals(stageId, workerInfo.getValue().getStageId()); + assertEquals(Constants.PARTITION_OPERATOR_NAME, workerInfo.getValue().getOperatorName()); + } + + @Test + public void rejectsUnexpectedWorkerTypeBeforeRegistering() throws Exception + { + WorkerCoordinateService coordinateService = mock(WorkerCoordinateService.class); + WorkerContext context = new WorkerContext(LogManager.getLogger(TestS3QSStageWorkerRunner.class), + new WorkerMetrics(), "request-1277"); + S3QSStageWorkerRunner runner = new S3QSStageWorkerRunner(context, coordinateService); + StageWorkerInput input = new StageWorkerInput(1277L, 1000L, 3, + "test-s3qs-stage", WorkerType.PARTITIONED_JOIN_S3QS); + + try + { + runner.runPartition(input); + fail("expected WorkerException"); + } + catch (WorkerException e) + { + assertTrue(e.getMessage().contains("unexpected S3QS stage worker type")); + } + + verifyNoInteractions(coordinateService); + } +} diff --git a/pixels-turbo/pixels-worker-lambda/src/main/java/io/pixelsdb/pixels/worker/lambda/S3QSPartitionWorker.java b/pixels-turbo/pixels-worker-lambda/src/main/java/io/pixelsdb/pixels/worker/lambda/S3QSPartitionWorker.java new file mode 100644 index 0000000000..1fed3216e5 --- /dev/null +++ b/pixels-turbo/pixels-worker-lambda/src/main/java/io/pixelsdb/pixels/worker/lambda/S3QSPartitionWorker.java @@ -0,0 +1,34 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of the License, + * or (at your option) any later version. + */ +package io.pixelsdb.pixels.worker.lambda; + +import com.amazonaws.services.lambda.runtime.Context; +import com.amazonaws.services.lambda.runtime.RequestHandler; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.PartitionOutput; +import io.pixelsdb.pixels.worker.common.S3QSStageWorkerRunner; +import io.pixelsdb.pixels.worker.common.WorkerContext; +import io.pixelsdb.pixels.worker.common.WorkerMetrics; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +public class S3QSPartitionWorker implements RequestHandler +{ + private static final Logger logger = LogManager.getLogger(S3QSPartitionWorker.class); + private final WorkerMetrics workerMetrics = new WorkerMetrics(); + + @Override + public PartitionOutput handleRequest(StageWorkerInput event, Context context) + { + WorkerContext workerContext = new WorkerContext(logger, workerMetrics, context.getAwsRequestId()); + return new S3QSStageWorkerRunner(workerContext).runPartition(event); + } +} diff --git a/pixels-turbo/pixels-worker-lambda/src/main/java/io/pixelsdb/pixels/worker/lambda/S3QSPartitionedChainJoinWorker.java b/pixels-turbo/pixels-worker-lambda/src/main/java/io/pixelsdb/pixels/worker/lambda/S3QSPartitionedChainJoinWorker.java new file mode 100644 index 0000000000..45d44a21f1 --- /dev/null +++ b/pixels-turbo/pixels-worker-lambda/src/main/java/io/pixelsdb/pixels/worker/lambda/S3QSPartitionedChainJoinWorker.java @@ -0,0 +1,34 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of the License, + * or (at your option) any later version. + */ +package io.pixelsdb.pixels.worker.lambda; + +import com.amazonaws.services.lambda.runtime.Context; +import com.amazonaws.services.lambda.runtime.RequestHandler; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; +import io.pixelsdb.pixels.worker.common.S3QSStageWorkerRunner; +import io.pixelsdb.pixels.worker.common.WorkerContext; +import io.pixelsdb.pixels.worker.common.WorkerMetrics; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +public class S3QSPartitionedChainJoinWorker implements RequestHandler +{ + private static final Logger logger = LogManager.getLogger(S3QSPartitionedChainJoinWorker.class); + private final WorkerMetrics workerMetrics = new WorkerMetrics(); + + @Override + public JoinOutput handleRequest(StageWorkerInput event, Context context) + { + WorkerContext workerContext = new WorkerContext(logger, workerMetrics, context.getAwsRequestId()); + return new S3QSStageWorkerRunner(workerContext).runPartitionedChainJoin(event); + } +} diff --git a/pixels-turbo/pixels-worker-lambda/src/main/java/io/pixelsdb/pixels/worker/lambda/S3QSPartitionedJoinWorker.java b/pixels-turbo/pixels-worker-lambda/src/main/java/io/pixelsdb/pixels/worker/lambda/S3QSPartitionedJoinWorker.java new file mode 100644 index 0000000000..a1d7297dd5 --- /dev/null +++ b/pixels-turbo/pixels-worker-lambda/src/main/java/io/pixelsdb/pixels/worker/lambda/S3QSPartitionedJoinWorker.java @@ -0,0 +1,34 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of the License, + * or (at your option) any later version. + */ +package io.pixelsdb.pixels.worker.lambda; + +import com.amazonaws.services.lambda.runtime.Context; +import com.amazonaws.services.lambda.runtime.RequestHandler; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; +import io.pixelsdb.pixels.worker.common.S3QSStageWorkerRunner; +import io.pixelsdb.pixels.worker.common.WorkerContext; +import io.pixelsdb.pixels.worker.common.WorkerMetrics; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +public class S3QSPartitionedJoinWorker implements RequestHandler +{ + private static final Logger logger = LogManager.getLogger(S3QSPartitionedJoinWorker.class); + private final WorkerMetrics workerMetrics = new WorkerMetrics(); + + @Override + public JoinOutput handleRequest(StageWorkerInput event, Context context) + { + WorkerContext workerContext = new WorkerContext(logger, workerMetrics, context.getAwsRequestId()); + return new S3QSStageWorkerRunner(workerContext).runPartitionedJoin(event); + } +} diff --git a/pixels-turbo/pixels-worker-spike/pom.xml b/pixels-turbo/pixels-worker-spike/pom.xml index 444361ba2c..e992563efd 100644 --- a/pixels-turbo/pixels-worker-spike/pom.xml +++ b/pixels-turbo/pixels-worker-spike/pom.xml @@ -31,6 +31,15 @@ io.pixelsdb pixels-storage-localfs + + + com.google.protobuf + protobuf-java + io.pixelsdb spike-java-handler @@ -120,4 +129,4 @@ - \ No newline at end of file + diff --git a/pixels-turbo/pixels-worker-spike/src/main/java/io/pixelsdb/pixels/worker/spike/RequestHandlerImpl.java b/pixels-turbo/pixels-worker-spike/src/main/java/io/pixelsdb/pixels/worker/spike/RequestHandlerImpl.java index d183ae1fd4..44cbf6a18f 100644 --- a/pixels-turbo/pixels-worker-spike/src/main/java/io/pixelsdb/pixels/worker/spike/RequestHandlerImpl.java +++ b/pixels-turbo/pixels-worker-spike/src/main/java/io/pixelsdb/pixels/worker/spike/RequestHandlerImpl.java @@ -66,12 +66,24 @@ public SpikeWorker.CallWorkerFunctionResp execute(SpikeWorker.CallWorkerFunction new WorkerService<>(PartitionedChainJoinWorker.class, PartitionedChainJoinInput.class); return service.execute(workerRequest.getWorkerPayload(), request.getRequestId()); } + case PARTITIONED_CHAIN_JOIN_S3QS: + { + WorkerService service = + new WorkerService<>(S3QSPartitionedChainJoinWorker.class, StageWorkerInput.class); + return service.execute(workerRequest.getWorkerPayload(), request.getRequestId()); + } case PARTITIONED_JOIN: { WorkerService service = new WorkerService<>(PartitionedJoinWorker.class, PartitionedJoinInput.class); return service.execute(workerRequest.getWorkerPayload(), request.getRequestId()); } + case PARTITIONED_JOIN_S3QS: + { + WorkerService service = + new WorkerService<>(S3QSPartitionedJoinWorker.class, StageWorkerInput.class); + return service.execute(workerRequest.getWorkerPayload(), request.getRequestId()); + } case PARTITIONED_JOIN_STREAMING: { WorkerService service = @@ -84,6 +96,12 @@ public SpikeWorker.CallWorkerFunctionResp execute(SpikeWorker.CallWorkerFunction new WorkerService<>(PartitionWorker.class, PartitionInput.class); return service.execute(workerRequest.getWorkerPayload(), request.getRequestId()); } + case PARTITION_S3QS: + { + WorkerService service = + new WorkerService<>(S3QSPartitionWorker.class, StageWorkerInput.class); + return service.execute(workerRequest.getWorkerPayload(), request.getRequestId()); + } case PARTITION_STREAMING: { WorkerService service = diff --git a/pixels-turbo/pixels-worker-spike/src/main/java/io/pixelsdb/pixels/worker/spike/S3QSPartitionWorker.java b/pixels-turbo/pixels-worker-spike/src/main/java/io/pixelsdb/pixels/worker/spike/S3QSPartitionWorker.java new file mode 100644 index 0000000000..020cff05b8 --- /dev/null +++ b/pixels-turbo/pixels-worker-spike/src/main/java/io/pixelsdb/pixels/worker/spike/S3QSPartitionWorker.java @@ -0,0 +1,40 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.worker.spike; + +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.PartitionOutput; +import io.pixelsdb.pixels.worker.common.S3QSStageWorkerRunner; +import io.pixelsdb.pixels.worker.common.WorkerContext; + +public class S3QSPartitionWorker implements WorkerInterface +{ + private final WorkerContext context; + + public S3QSPartitionWorker(WorkerContext context) + { + this.context = context; + } + + @Override + public PartitionOutput handleRequest(StageWorkerInput input) + { + return new S3QSStageWorkerRunner(context).runPartition(input); + } + + @Override + public String getRequestId() + { + return this.context.getRequestId(); + } + + @Override + public WorkerType getWorkerType() + { + return WorkerType.PARTITION_S3QS; + } +} diff --git a/pixels-turbo/pixels-worker-spike/src/main/java/io/pixelsdb/pixels/worker/spike/S3QSPartitionedChainJoinWorker.java b/pixels-turbo/pixels-worker-spike/src/main/java/io/pixelsdb/pixels/worker/spike/S3QSPartitionedChainJoinWorker.java new file mode 100644 index 0000000000..673d2e3dbd --- /dev/null +++ b/pixels-turbo/pixels-worker-spike/src/main/java/io/pixelsdb/pixels/worker/spike/S3QSPartitionedChainJoinWorker.java @@ -0,0 +1,40 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.worker.spike; + +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; +import io.pixelsdb.pixels.worker.common.S3QSStageWorkerRunner; +import io.pixelsdb.pixels.worker.common.WorkerContext; + +public class S3QSPartitionedChainJoinWorker implements WorkerInterface +{ + private final WorkerContext context; + + public S3QSPartitionedChainJoinWorker(WorkerContext context) + { + this.context = context; + } + + @Override + public JoinOutput handleRequest(StageWorkerInput input) + { + return new S3QSStageWorkerRunner(context).runPartitionedChainJoin(input); + } + + @Override + public String getRequestId() + { + return this.context.getRequestId(); + } + + @Override + public WorkerType getWorkerType() + { + return WorkerType.PARTITIONED_CHAIN_JOIN_S3QS; + } +} diff --git a/pixels-turbo/pixels-worker-spike/src/main/java/io/pixelsdb/pixels/worker/spike/S3QSPartitionedJoinWorker.java b/pixels-turbo/pixels-worker-spike/src/main/java/io/pixelsdb/pixels/worker/spike/S3QSPartitionedJoinWorker.java new file mode 100644 index 0000000000..8ed1f4f62c --- /dev/null +++ b/pixels-turbo/pixels-worker-spike/src/main/java/io/pixelsdb/pixels/worker/spike/S3QSPartitionedJoinWorker.java @@ -0,0 +1,40 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.worker.spike; + +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; +import io.pixelsdb.pixels.worker.common.S3QSStageWorkerRunner; +import io.pixelsdb.pixels.worker.common.WorkerContext; + +public class S3QSPartitionedJoinWorker implements WorkerInterface +{ + private final WorkerContext context; + + public S3QSPartitionedJoinWorker(WorkerContext context) + { + this.context = context; + } + + @Override + public JoinOutput handleRequest(StageWorkerInput input) + { + return new S3QSStageWorkerRunner(context).runPartitionedJoin(input); + } + + @Override + public String getRequestId() + { + return this.context.getRequestId(); + } + + @Override + public WorkerType getWorkerType() + { + return WorkerType.PARTITIONED_JOIN_S3QS; + } +} diff --git a/pixels-turbo/pixels-worker-spike/src/test/java/io/pixelsdb/pixels/worker/spike/TestS3QSSpikeDebugFlow.java b/pixels-turbo/pixels-worker-spike/src/test/java/io/pixelsdb/pixels/worker/spike/TestS3QSSpikeDebugFlow.java new file mode 100644 index 0000000000..2b948d752d --- /dev/null +++ b/pixels-turbo/pixels-worker-spike/src/test/java/io/pixelsdb/pixels/worker/spike/TestS3QSSpikeDebugFlow.java @@ -0,0 +1,267 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public + * License along with Pixels. If not, see + * . + */ +package io.pixelsdb.pixels.worker.spike; + +import com.alibaba.fastjson.JSON; +import com.alibaba.fastjson.serializer.SerializerFeature; +import io.pixelsdb.pixels.common.physical.Storage; +import io.pixelsdb.pixels.common.turbo.SpikeWorkerRequest; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.executor.join.JoinAlgorithm; +import io.pixelsdb.pixels.planner.coordinate.PlanCoordinator; +import io.pixelsdb.pixels.planner.coordinate.PlanCoordinatorFactory; +import io.pixelsdb.pixels.planner.coordinate.WorkerCoordinateServer; +import io.pixelsdb.pixels.planner.plan.physical.PartitionedJoinS3QSOperator; +import io.pixelsdb.pixels.planner.plan.physical.domain.InputInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.InputSplit; +import io.pixelsdb.pixels.planner.plan.physical.domain.OutputInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.PartitionInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.PartitionedTableInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ScanTableInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleQueueInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.StorageInfo; +import io.pixelsdb.pixels.planner.plan.physical.input.JoinInput; +import io.pixelsdb.pixels.planner.plan.physical.input.PartitionInput; +import io.pixelsdb.pixels.planner.plan.physical.input.PartitionedJoinInput; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.spike.handler.SpikeWorker; +import org.junit.Ignore; +import org.junit.Test; + +import java.util.Collections; + +import static org.junit.Assert.assertEquals; + +/** + * Issue 1277 中用于排查 S3QS 与 Spike worker 分发链路的真实代码调试入口。 + * + * 该测试默认通过 @Ignore 关闭，因为它不会使用 mock 或桩对象替代执行路径， + * 而是会真正进入 worker 的生产执行链路： + * + * RequestHandlerImpl + * -> WorkerService + * -> S3QSPartitionWorker + * -> S3QSStageWorkerRunner + * -> WorkerCoordinateService / WorkerCoordinateServer + * -> StageCoordinator task queue + * -> BasePartitionWorker + * -> WorkerCommon / S3QS / S3 / SQS + * + * 因此它不是一个用于解释逻辑的轻量级单元测试，而是一个便于在 IDE 中下断点、 + * 观察真实调用栈和运行时对象状态的集成调试入口。启用该测试前，需要准备真实的 + * Pixels 输入文件、AWS S3/SQS 资源以及本地配置，否则 worker 会在读取输入、 + * 初始化 shuffle 存储或访问队列时失败。 + */ +public class TestS3QSSpikeDebugFlow +{ + @Ignore("仅用于调试的集成入口。需要真实 Pixels 输入数据、AWS S3/SQS 资源和运行配置。") + @Test + public void realSpikeRequestRunsS3QSPartitionStageWorkerThroughCoordinator() throws Exception + { + /* + * 必填环境变量： + * + * PIXELS_S3QS_DEBUG_INPUT_PATH + * 真实存在的 Pixels 文件路径，且必须能被当前配置的输入存储后端读取。 + * 示例：s3://bucket/path/source.pxl + * + * PIXELS_S3QS_DEBUG_OBJECT_PREFIX + * S3QS 写入 DATA 类型 S3QueueMessage 时使用的对象前缀。 + * 该前缀会作为 shuffle 中间数据的写入位置，通常应指向一个可清理的 + * 临时目录，避免与正式数据混用。 + * 示例：bucket/tmp/issue1277/debug-shuffle/ + * + * PIXELS_S3QS_DEBUG_QUEUE_NAME + * partition 0 对应的 SQS 队列名称。WorkerCommon 根据 ShuffleInfo + * 初始化 shuffle storage 时，S3QS 存储实现会通过 S3QS.registerQueue(...) + * 解析或创建该队列。 + * + * 可选环境变量： + * + * PIXELS_S3QS_DEBUG_INPUT_SCHEME + * 输入数据的存储协议，默认值为 s3。调试其它存储后端时，可以改成 + * Storage.Scheme 支持的协议名称。 + * + * PIXELS_S3QS_DEBUG_COLUMNS + * 需要读取的列名列表，使用逗号分隔，默认值为 c0。这里会同时影响 + * ScanTableInfo 的读取列和 PartitionInput 的投影数组长度。 + * + * 测试启动一个本地 WorkerCoordinateServer，然后把真实的 Spike worker + * 请求直接传入 RequestHandlerImpl。这样可以绕开外部 Spike 网络服务， + * 但仍然执行 worker 侧真实的请求反序列化、类型分发、WorkerService + * 调用以及后续 S3QS partition stage worker 逻辑。 + */ + long transId = 1277001L; + long timestamp = System.currentTimeMillis(); + int coordinatorPort = 18887; + + String inputPath = requireEnv("PIXELS_S3QS_DEBUG_INPUT_PATH"); + String objectPrefix = requireEnv("PIXELS_S3QS_DEBUG_OBJECT_PREFIX"); + String queueName = requireEnv("PIXELS_S3QS_DEBUG_QUEUE_NAME"); + Storage.Scheme inputScheme = Storage.Scheme.from( + getEnvOrDefault("PIXELS_S3QS_DEBUG_INPUT_SCHEME", "s3")); + String[] columnsToRead = getEnvOrDefault("PIXELS_S3QS_DEBUG_COLUMNS", "c0").split(","); + + WorkerCoordinateServer coordinateServer = new WorkerCoordinateServer(coordinatorPort); + Thread coordinateServerThread = new Thread(coordinateServer, "issue1277-worker-coordinate-server"); + coordinateServerThread.setDaemon(true); + coordinateServerThread.start(); + Thread.sleep(500L); + PlanCoordinator planCoordinator = null; + + try + { + PartitionInput smallPartitionInput = createProducerInput(transId, timestamp, + inputPath, inputScheme, columnsToRead, objectPrefix, queueName); + PartitionInput largePartitionInput = createProducerInput(transId, timestamp, + inputPath, inputScheme, columnsToRead, objectPrefix, queueName); + PartitionedJoinInput joinInput = createJoinInput(transId, timestamp, objectPrefix, queueName); + + PartitionedJoinS3QSOperator operator = new PartitionedJoinS3QSOperator("issue1277-s3qs-debug", + Collections.singletonList(smallPartitionInput), + Collections.singletonList(largePartitionInput), + Collections.singletonList(joinInput), + JoinAlgorithm.PARTITIONED); + planCoordinator = PlanCoordinatorFactory.Instance().createPlanCoordinator(transId, operator); + + /* + * 生产端 stage id 不是在构造 PartitionInput 时手工指定的，而是在 + * PlanCoordinatorFactory.createPlanCoordinator(...) 内部调用 + * initPlanCoordinator(...) 时分配的。该过程会把 stage id 回写到原始 + * PartitionInput 对象中，因此这里直接从 smallPartitionInput 读取， + * 确保发给 worker 的 StageWorkerInput 与协调器中的 stage 定义一致。 + */ + StageWorkerInput stageWorkerInput = new StageWorkerInput(transId, timestamp, + smallPartitionInput.getStageId(), operator.getName(), WorkerType.PARTITION_S3QS); + stageWorkerInput.setCoordinatorHost("localhost"); + stageWorkerInput.setCoordinatorPort(coordinatorPort); + + SpikeWorkerRequest workerRequest = new SpikeWorkerRequest(WorkerType.PARTITION_S3QS, + JSON.toJSONString(stageWorkerInput, SerializerFeature.DisableCircularReferenceDetect)); + SpikeWorker.CallWorkerFunctionReq request = SpikeWorker.CallWorkerFunctionReq.newBuilder() + .setRequestId(1277L) + .setPayload(JSON.toJSONString(workerRequest, SerializerFeature.DisableCircularReferenceDetect)) + .build(); + + /* + * 建议从下一行开始设置断点。执行 new RequestHandlerImpl().execute(request) + * 后，可以观察 Spike worker 请求进入真实 S3QS partition worker 的完整链路： + * + * RequestHandlerImpl.execute(request) + * -> switch PARTITION_S3QS，根据 WorkerType 选择 S3QS partition 分支 + * -> WorkerService + * -> S3QSPartitionWorker.handleRequest(...)，进入 S3QS 分区 worker + * -> S3QSStageWorkerRunner.runPartition(...)，启动 stage worker runner + * -> WorkerCoordinateService.registerWorker(...)，向本地协调器注册 worker + * -> WorkerCoordinateService.getTasksToExecute(...)，从协调器领取任务 + * -> BasePartitionWorker.process(PartitionInput)，执行 partition 输入处理 + * -> S3QS 生产端数据路径，将 shuffle 数据写入 S3QS/S3/SQS + */ + SpikeWorker.CallWorkerFunctionResp response = new RequestHandlerImpl().execute(request); + + assertEquals(1277L, response.getRequestId()); + System.out.println(response.getPayload()); + } + finally + { + if (planCoordinator != null) + { + PlanCoordinatorFactory.Instance().removePlanCoordinator(transId, planCoordinator); + } + coordinateServer.shutdown(); + } + } + + private static PartitionInput createProducerInput(long transId, long timestamp, String inputPath, + Storage.Scheme inputScheme, String[] columnsToRead, + String objectPrefix, String queueName) + { + ScanTableInfo tableInfo = new ScanTableInfo("issue1277_debug", true, columnsToRead, + new StorageInfo(inputScheme, null, null, null, null), + Collections.singletonList(new InputSplit(Collections.singletonList(new InputInfo(inputPath, 0, 1)))), + null); + ShuffleInfo shuffleInfo = createShuffleInfo(objectPrefix, queueName); + OutputInfo outputInfo = new OutputInfo(objectPrefix + "producer-output", + new StorageInfo(Storage.Scheme.s3qs, null, null, null, null), true); + outputInfo.setShuffleInfo(shuffleInfo); + PartitionInput partitionInput = new PartitionInput(transId, timestamp, tableInfo, + allProjected(columnsToRead.length), outputInfo, new PartitionInfo(new int[] {0}, 1)); + partitionInput.setProducerTaskId(0); + return partitionInput; + } + + private static PartitionedJoinInput createJoinInput(long transId, long timestamp, + String objectPrefix, String queueName) + { + PartitionedJoinInput joinInput = new PartitionedJoinInput(); + joinInput.setTransId(transId); + joinInput.setTimestamp(timestamp); + joinInput.setSmallTable(createPartitionedTableInfo(objectPrefix, queueName)); + joinInput.setLargeTable(createPartitionedTableInfo(objectPrefix, queueName)); + return joinInput; + } + + private static PartitionedTableInfo createPartitionedTableInfo(String objectPrefix, String queueName) + { + PartitionedTableInfo tableInfo = new PartitionedTableInfo(); + tableInfo.setStorageInfo(new StorageInfo(Storage.Scheme.s3qs, null, null, null, null)); + tableInfo.setShuffleInfo(createShuffleInfo(objectPrefix, queueName)); + return tableInfo; + } + + private static ShuffleInfo createShuffleInfo(String objectPrefix, String queueName) + { + return new ShuffleInfo("issue1277-debug-shuffle", + new StorageInfo(Storage.Scheme.s3qs, null, null, null, null), + objectPrefix, 1, 1, 1, 1, + Collections.singletonList(new ShuffleQueueInfo(0, queueName, null))); + } + + private static boolean[] allProjected(int columnCount) + { + boolean[] projection = new boolean[columnCount]; + for (int i = 0; i < projection.length; ++i) + { + projection[i] = true; + } + return projection; + } + + private static String requireEnv(String name) + { + String value = System.getenv(name); + if (value == null || value.trim().isEmpty()) + { + throw new IllegalStateException("missing environment variable: " + name); + } + return value; + } + + private static String getEnvOrDefault(String name, String defaultValue) + { + String value = System.getenv(name); + if (value == null || value.trim().isEmpty()) + { + return defaultValue; + } + return value; + } +} diff --git a/pixels-turbo/pixels-worker-spike/src/test/java/io/pixelsdb/pixels/worker/spike/TestS3QSSpikeEndToEnd.java b/pixels-turbo/pixels-worker-spike/src/test/java/io/pixelsdb/pixels/worker/spike/TestS3QSSpikeEndToEnd.java new file mode 100644 index 0000000000..75d57aded0 --- /dev/null +++ b/pixels-turbo/pixels-worker-spike/src/test/java/io/pixelsdb/pixels/worker/spike/TestS3QSSpikeEndToEnd.java @@ -0,0 +1,364 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + * + * Pixels is free software: you can redistribute it and/or modify + * it under the terms of the Affero GNU General Public License as + * published by the Free Software Foundation, either version 3 of + * the License, or (at your option) any later version. + * + * Pixels is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Affero GNU General Public License for more details. + * + * You should have received a copy of the Affero GNU General Public + * License along with Pixels. If not, see + * . + */ +package io.pixelsdb.pixels.worker.spike; + +import com.alibaba.fastjson.JSON; +import com.alibaba.fastjson.serializer.SerializerFeature; +import io.pixelsdb.pixels.common.physical.Storage; +import io.pixelsdb.pixels.common.physical.StorageFactory; +import io.pixelsdb.pixels.common.turbo.Output; +import io.pixelsdb.pixels.common.turbo.SpikeWorkerRequest; +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.core.PixelsReader; +import io.pixelsdb.pixels.core.PixelsWriter; +import io.pixelsdb.pixels.core.TypeDescription; +import io.pixelsdb.pixels.core.reader.PixelsReaderOption; +import io.pixelsdb.pixels.core.reader.PixelsRecordReader; +import io.pixelsdb.pixels.core.vector.LongColumnVector; +import io.pixelsdb.pixels.core.vector.VectorizedRowBatch; +import io.pixelsdb.pixels.executor.join.JoinAlgorithm; +import io.pixelsdb.pixels.executor.join.JoinType; +import io.pixelsdb.pixels.executor.predicate.TableScanFilter; +import io.pixelsdb.pixels.planner.coordinate.CoordinatedPlanExecution; +import io.pixelsdb.pixels.planner.coordinate.CoordinatorEndpoint; +import io.pixelsdb.pixels.planner.coordinate.PlanCoordinatorFactory; +import io.pixelsdb.pixels.planner.coordinate.S3QSShuffleResourceLifecycle; +import io.pixelsdb.pixels.planner.coordinate.StageWorkerLauncher; +import io.pixelsdb.pixels.planner.coordinate.WorkerCoordinateServer; +import io.pixelsdb.pixels.planner.plan.physical.PartitionedJoinS3QSOperator; +import io.pixelsdb.pixels.planner.plan.physical.domain.InputInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.InputSplit; +import io.pixelsdb.pixels.planner.plan.physical.domain.MultiOutputInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.OutputInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.PartitionInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.PartitionedJoinInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.PartitionedTableInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ScanTableInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.ShuffleQueueInfo; +import io.pixelsdb.pixels.planner.plan.physical.domain.StorageInfo; +import io.pixelsdb.pixels.planner.plan.physical.input.JoinInput; +import io.pixelsdb.pixels.planner.plan.physical.input.PartitionInput; +import io.pixelsdb.pixels.planner.plan.physical.input.PartitionedJoinInput; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; +import io.pixelsdb.pixels.planner.plan.physical.output.PartitionOutput; +import io.pixelsdb.pixels.storage.s3qs.S3QS; +import io.pixelsdb.pixels.worker.common.WorkerCommon; +import io.pixelsdb.spike.handler.SpikeWorker; +import org.junit.Test; + +import java.io.IOException; +import java.net.ServerSocket; +import java.util.Collections; +import java.util.HashSet; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.atomic.AtomicLong; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.fail; +import static org.junit.Assume.assumeTrue; + +/** + * Runs a complete coordinator-managed S3QS partitioned join against real AWS. + * + * Platform deployment and its network hop are intentionally replaced by an + * in-process StageWorkerLauncher. It still calls Spike's production request + * handler and therefore exercises the real worker wrappers, coordinator gRPC + * task protocol, partition/join workers, S3 object I/O, and SQS messaging. + */ +public class TestS3QSSpikeEndToEnd +{ + @Test + public void coordinatedSpikeWorkersShuffleAndJoinThroughAws() throws Exception + { + String bucket = System.getenv("PIXELS_S3QS_IT_BUCKET"); + String prefix = trimSlashes(System.getenv("PIXELS_S3QS_IT_PREFIX")); + String queuePrefix = System.getenv("PIXELS_S3QS_IT_QUEUE_PREFIX"); + assumeTrue("set PIXELS_S3QS_IT_BUCKET to run the AWS S3QS end-to-end test", + bucket != null && !bucket.trim().isEmpty()); + assumeTrue("set PIXELS_S3QS_IT_QUEUE_PREFIX to run the AWS S3QS end-to-end test", + queuePrefix != null && !queuePrefix.trim().isEmpty()); + + String suffix = UUID.randomUUID().toString().replace("-", "").substring(0, 16); + long transId = Math.abs(UUID.randomUUID().getMostSignificantBits()); + long timestamp = System.currentTimeMillis(); + int coordinatorPort = freePort(); + String testRoot = bucket + "/" + (prefix.isEmpty() ? "" : prefix + "/") + + "query-e2e-" + suffix + "/"; + String smallInputPath = testRoot + "input/small.pxl"; + String largeInputPath = testRoot + "input/large.pxl"; + String resultFolder = testRoot + "result/"; + String resultPath = resultFolder + "join.pxl"; + ShuffleInfo smallShuffle = shuffleInfo("small-" + suffix, testRoot + "shuffle/small/", + queuePrefix + "-small-" + suffix); + ShuffleInfo largeShuffle = shuffleInfo("large-" + suffix, testRoot + "shuffle/large/", + queuePrefix + "-large-" + suffix); + + Storage s3 = StorageFactory.Instance().getStorage(Storage.Scheme.s3); + writeInput(s3, smallInputPath, "small_value", new long[][] {{1, 10}, {2, 20}}); + writeInput(s3, largeInputPath, "large_value", new long[][] {{1, 100}, {3, 300}}); + + WorkerCoordinateServer coordinateServer = new WorkerCoordinateServer(coordinatorPort); + Thread serverThread = new Thread(coordinateServer, "issue1277-e2e-coordinate-server"); + serverThread.setDaemon(true); + serverThread.start(); + Thread.sleep(500L); + + CoordinatedPlanExecution execution = null; + try + { + PartitionInput smallProducer = producerInput(transId, timestamp, "small", smallInputPath, + "small_value", smallShuffle); + PartitionInput largeProducer = producerInput(transId, timestamp, "large", largeInputPath, + "large_value", largeShuffle); + PartitionedJoinInput consumer = joinInput(transId, timestamp, smallShuffle, largeShuffle, resultFolder); + PartitionedJoinS3QSOperator operator = new PartitionedJoinS3QSOperator( + "issue1277-s3qs-e2e", + Collections.singletonList(smallProducer), + Collections.singletonList(largeProducer), + Collections.singletonList(consumer), + JoinAlgorithm.PARTITIONED); + + execution = PlanCoordinatorFactory.Instance().createPlanExecution( + transId, operator, new CoordinatorEndpoint("localhost", coordinatorPort), + new S3QSShuffleResourceLifecycle(), new InProcessSpikeStageWorkerLauncher()); + execution.execute(); + execution.collectOutputs(); + + assertEquals(Collections.singleton("1:10:100"), readJoinedRows(s3, resultPath)); + assertShuffleObjectsDeleted(s3, smallShuffle.getObjectPathPrefix()); + assertShuffleObjectsDeleted(s3, largeShuffle.getObjectPathPrefix()); + S3QS s3qs = (S3QS) StorageFactory.Instance().getStorage(Storage.Scheme.s3qs); + assertQueueDeleted(s3qs, smallShuffle.getQueues().get(0).getQueueName()); + assertQueueDeleted(s3qs, largeShuffle.getQueues().get(0).getQueueName()); + } + finally + { + if (execution != null) + { + execution.close(); + } + coordinateServer.shutdown(); + s3.delete(testRoot, true); + } + } + + private static PartitionInput producerInput(long transId, long timestamp, String tableName, + String inputPath, String valueColumn, ShuffleInfo shuffleInfo) + { + String[] columns = {"key", valueColumn}; + ScanTableInfo tableInfo = new ScanTableInfo(tableName, true, columns, + new StorageInfo(Storage.Scheme.s3, null, null, null, null), + Collections.singletonList(new InputSplit( + Collections.singletonList(new InputInfo(inputPath, 0, 1)))), + JSON.toJSONString(TableScanFilter.empty("issue1277", tableName))); + OutputInfo output = new OutputInfo(shuffleInfo.getObjectPathPrefix(), + new StorageInfo(Storage.Scheme.s3qs, null, null, null, null), true); + output.setShuffleInfo(shuffleInfo); + PartitionInput input = new PartitionInput(transId, timestamp, tableInfo, + new boolean[] {true, true}, output, new PartitionInfo(new int[] {0}, 1)); + input.setProducerTaskId(0); + return input; + } + + private static PartitionedJoinInput joinInput(long transId, long timestamp, + ShuffleInfo smallShuffle, ShuffleInfo largeShuffle, + String resultFolder) + { + PartitionedTableInfo small = partitionedTable("small", "small_value", smallShuffle); + PartitionedTableInfo large = partitionedTable("large", "large_value", largeShuffle); + PartitionedJoinInfo joinInfo = new PartitionedJoinInfo( + JoinType.EQUI_INNER, + new String[] {"key", "small_value"}, + new String[] {"large_value"}, + new boolean[] {true, true}, + new boolean[] {false, true}, + false, null, 1, Collections.singletonList(0)); + MultiOutputInfo output = new MultiOutputInfo(resultFolder, + new StorageInfo(Storage.Scheme.s3, null, null, null, null), + true, Collections.singletonList("join.pxl")); + return new PartitionedJoinInput(transId, timestamp, small, large, joinInfo, + false, null, output); + } + + private static PartitionedTableInfo partitionedTable(String tableName, String valueColumn, + ShuffleInfo shuffleInfo) + { + PartitionedTableInfo tableInfo = new PartitionedTableInfo( + tableName, false, new String[] {"key", valueColumn}, + new StorageInfo(Storage.Scheme.s3qs, null, null, null, null), + Collections.emptyList(), 1, new int[] {0}); + tableInfo.setShuffleInfo(shuffleInfo); + return tableInfo; + } + + private static ShuffleInfo shuffleInfo(String shuffleId, String objectPrefix, String queueName) + { + return new ShuffleInfo(shuffleId, + new StorageInfo(Storage.Scheme.s3qs, null, null, null, null), + objectPrefix, 1, 1, 1, 1, + Collections.singletonList(new ShuffleQueueInfo(0, queueName, null))); + } + + private static void writeInput(Storage storage, String path, String valueColumn, long[][] rows) + throws IOException + { + TypeDescription schema = TypeDescription.createStruct() + .addField("key", TypeDescription.createLong()) + .addField(valueColumn, TypeDescription.createLong()); + VectorizedRowBatch batch = schema.createRowBatch(); + LongColumnVector key = (LongColumnVector) batch.cols[0]; + LongColumnVector value = (LongColumnVector) batch.cols[1]; + for (long[] row : rows) + { + int rowId = batch.size++; + key.vector[rowId] = row[0]; + value.vector[rowId] = row[1]; + } + PixelsWriter writer = WorkerCommon.getWriter(schema, storage, path, true, false, null); + writer.addRowBatch(batch); + writer.close(); + } + + private static Set readJoinedRows(Storage storage, String path) throws IOException + { + PixelsReaderOption option = new PixelsReaderOption() + .skipCorruptRecords(true) + .tolerantSchemaEvolution(true) + .includeCols(new String[] {"key", "small_value", "large_value"}); + Set rows = new HashSet<>(); + try (PixelsReader reader = WorkerCommon.getReader(path, storage); + PixelsRecordReader recordReader = reader.read(option)) + { + VectorizedRowBatch batch; + do + { + batch = recordReader.readBatch(32); + LongColumnVector key = (LongColumnVector) batch.cols[0]; + LongColumnVector smallValue = (LongColumnVector) batch.cols[1]; + LongColumnVector largeValue = (LongColumnVector) batch.cols[2]; + for (int i = 0; i < batch.size; ++i) + { + rows.add(key.vector[i] + ":" + smallValue.vector[i] + ":" + largeValue.vector[i]); + } + } + while (!batch.endOfFile); + } + return rows; + } + + private static void assertShuffleObjectsDeleted(Storage storage, String prefix) throws Exception + { + for (int i = 0; i < 10; ++i) + { + if (!storage.exists(prefix)) + { + return; + } + Thread.sleep(200L); + } + assertFalse("shuffle object prefix still exists: " + prefix, storage.exists(prefix)); + } + + private static void assertQueueDeleted(S3QS s3qs, String queueName) throws Exception + { + for (int i = 0; i < 20; ++i) + { + try + { + s3qs.getQueueUrl(queueName); + Thread.sleep(500L); + } + catch (IOException expected) + { + return; + } + } + fail("SQS queue still exists: " + queueName); + } + + private static int freePort() throws IOException + { + try (ServerSocket socket = new ServerSocket(0)) + { + return socket.getLocalPort(); + } + } + + private static String trimSlashes(String value) + { + if (value == null) + { + return ""; + } + String trimmed = value.trim(); + while (trimmed.startsWith("/")) + { + trimmed = trimmed.substring(1); + } + while (trimmed.endsWith("/")) + { + trimmed = trimmed.substring(0, trimmed.length() - 1); + } + return trimmed; + } + + private static class InProcessSpikeStageWorkerLauncher implements StageWorkerLauncher + { + private final AtomicLong requestIds = new AtomicLong(12770000L); + + @Override + public CompletableFuture launch(WorkerType workerType, StageWorkerInput input) + { + return CompletableFuture.supplyAsync(() -> + { + try + { + SpikeWorkerRequest workerRequest = new SpikeWorkerRequest(workerType, + JSON.toJSONString(input, SerializerFeature.DisableCircularReferenceDetect)); + SpikeWorker.CallWorkerFunctionReq request = SpikeWorker.CallWorkerFunctionReq.newBuilder() + .setRequestId(requestIds.getAndIncrement()) + .setPayload(JSON.toJSONString( + workerRequest, SerializerFeature.DisableCircularReferenceDetect)) + .build(); + String payload = new RequestHandlerImpl().execute(request).getPayload(); + if (workerType == WorkerType.PARTITION_S3QS) + { + return JSON.parseObject(payload, PartitionOutput.class); + } + if (workerType == WorkerType.PARTITIONED_JOIN_S3QS) + { + return JSON.parseObject(payload, JoinOutput.class); + } + throw new IllegalArgumentException("unsupported test worker type: " + workerType); + } + catch (Throwable e) + { + throw new CompletionException(e); + } + }); + } + } +} diff --git a/pixels-turbo/pixels-worker-vhive/src/main/java/io/pixelsdb/pixels/worker/vhive/S3QSPartitionWorker.java b/pixels-turbo/pixels-worker-vhive/src/main/java/io/pixelsdb/pixels/worker/vhive/S3QSPartitionWorker.java new file mode 100644 index 0000000000..79fa149571 --- /dev/null +++ b/pixels-turbo/pixels-worker-vhive/src/main/java/io/pixelsdb/pixels/worker/vhive/S3QSPartitionWorker.java @@ -0,0 +1,41 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.worker.vhive; + +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.PartitionOutput; +import io.pixelsdb.pixels.worker.common.S3QSStageWorkerRunner; +import io.pixelsdb.pixels.worker.common.WorkerContext; +import io.pixelsdb.pixels.worker.vhive.utils.RequestHandler; + +public class S3QSPartitionWorker implements RequestHandler +{ + private final WorkerContext context; + + public S3QSPartitionWorker(WorkerContext context) + { + this.context = context; + } + + @Override + public PartitionOutput handleRequest(StageWorkerInput input) + { + return new S3QSStageWorkerRunner(context).runPartition(input); + } + + @Override + public String getRequestId() + { + return context.getRequestId(); + } + + @Override + public WorkerType getWorkerType() + { + return WorkerType.PARTITION_S3QS; + } +} diff --git a/pixels-turbo/pixels-worker-vhive/src/main/java/io/pixelsdb/pixels/worker/vhive/S3QSPartitionedChainJoinWorker.java b/pixels-turbo/pixels-worker-vhive/src/main/java/io/pixelsdb/pixels/worker/vhive/S3QSPartitionedChainJoinWorker.java new file mode 100644 index 0000000000..e15788e300 --- /dev/null +++ b/pixels-turbo/pixels-worker-vhive/src/main/java/io/pixelsdb/pixels/worker/vhive/S3QSPartitionedChainJoinWorker.java @@ -0,0 +1,41 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.worker.vhive; + +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; +import io.pixelsdb.pixels.worker.common.S3QSStageWorkerRunner; +import io.pixelsdb.pixels.worker.common.WorkerContext; +import io.pixelsdb.pixels.worker.vhive.utils.RequestHandler; + +public class S3QSPartitionedChainJoinWorker implements RequestHandler +{ + private final WorkerContext context; + + public S3QSPartitionedChainJoinWorker(WorkerContext context) + { + this.context = context; + } + + @Override + public JoinOutput handleRequest(StageWorkerInput input) + { + return new S3QSStageWorkerRunner(context).runPartitionedChainJoin(input); + } + + @Override + public String getRequestId() + { + return context.getRequestId(); + } + + @Override + public WorkerType getWorkerType() + { + return WorkerType.PARTITIONED_CHAIN_JOIN_S3QS; + } +} diff --git a/pixels-turbo/pixels-worker-vhive/src/main/java/io/pixelsdb/pixels/worker/vhive/S3QSPartitionedJoinWorker.java b/pixels-turbo/pixels-worker-vhive/src/main/java/io/pixelsdb/pixels/worker/vhive/S3QSPartitionedJoinWorker.java new file mode 100644 index 0000000000..22dcd9d7eb --- /dev/null +++ b/pixels-turbo/pixels-worker-vhive/src/main/java/io/pixelsdb/pixels/worker/vhive/S3QSPartitionedJoinWorker.java @@ -0,0 +1,41 @@ +/* + * Copyright 2026 PixelsDB. + * + * This file is part of Pixels. + */ +package io.pixelsdb.pixels.worker.vhive; + +import io.pixelsdb.pixels.common.turbo.WorkerType; +import io.pixelsdb.pixels.planner.plan.physical.input.StageWorkerInput; +import io.pixelsdb.pixels.planner.plan.physical.output.JoinOutput; +import io.pixelsdb.pixels.worker.common.S3QSStageWorkerRunner; +import io.pixelsdb.pixels.worker.common.WorkerContext; +import io.pixelsdb.pixels.worker.vhive.utils.RequestHandler; + +public class S3QSPartitionedJoinWorker implements RequestHandler +{ + private final WorkerContext context; + + public S3QSPartitionedJoinWorker(WorkerContext context) + { + this.context = context; + } + + @Override + public JoinOutput handleRequest(StageWorkerInput input) + { + return new S3QSStageWorkerRunner(context).runPartitionedJoin(input); + } + + @Override + public String getRequestId() + { + return context.getRequestId(); + } + + @Override + public WorkerType getWorkerType() + { + return WorkerType.PARTITIONED_JOIN_S3QS; + } +} diff --git a/pixels-turbo/pixels-worker-vhive/src/main/java/io/pixelsdb/pixels/worker/vhive/WorkerServiceImpl.java b/pixels-turbo/pixels-worker-vhive/src/main/java/io/pixelsdb/pixels/worker/vhive/WorkerServiceImpl.java index 3a643eae9b..ee1f933723 100644 --- a/pixels-turbo/pixels-worker-vhive/src/main/java/io/pixelsdb/pixels/worker/vhive/WorkerServiceImpl.java +++ b/pixels-turbo/pixels-worker-vhive/src/main/java/io/pixelsdb/pixels/worker/vhive/WorkerServiceImpl.java @@ -76,6 +76,12 @@ public void process(TurboProto.vHiveWorkerRequest request, StreamObserver service = new ServiceImpl<>(S3QSPartitionedChainJoinWorker.class, StageWorkerInput.class); + service.execute(request, responseObserver); + break; + } case PARTITIONED_CHAIN_JOIN_STREAMING: { ServiceImpl service = new ServiceImpl<>(PartitionedChainJoinStreamWorker.class, PartitionedChainJoinInput.class); @@ -88,6 +94,12 @@ public void process(TurboProto.vHiveWorkerRequest request, StreamObserver service = new ServiceImpl<>(S3QSPartitionedJoinWorker.class, StageWorkerInput.class); + service.execute(request, responseObserver); + break; + } case PARTITIONED_JOIN_STREAMING: { ServiceImpl service = new ServiceImpl<>(PartitionedJoinStreamWorker.class, PartitionedJoinInput.class); @@ -100,6 +112,12 @@ public void process(TurboProto.vHiveWorkerRequest request, StreamObserver service = new ServiceImpl<>(S3QSPartitionWorker.class, StageWorkerInput.class); + service.execute(request, responseObserver); + break; + } case PARTITION_STREAMING: { ServiceImpl service = new ServiceImpl<>(PartitionStreamWorker.class, PartitionInput.class);