From c3ed090d44f57a5e552f4e99b76f1c3bf673c357 Mon Sep 17 00:00:00 2001 From: Bukhtawar Khan Date: Sat, 18 Apr 2026 00:54:13 +0530 Subject: [PATCH 1/2] Simplify the source sink interactions Signed-off-by: Bukhtawar Khan --- .../analytics/backend/ExchangeSink.java | 31 ++-- .../analytics/backend/ExchangeSource.java | 30 ++++ .../opensearch/analytics/exec/PlanWalker.java | 19 ++- .../analytics/exec/RowProducingSink.java | 43 ++---- .../analytics/exec/stage/DataConsumer.java | 27 ++++ .../analytics/exec/stage/DataProducer.java | 38 +++++ .../exec/stage/LocalStageExecution.java | 32 ++-- .../exec/stage/LocalStageScheduler.java | 2 +- .../exec/stage/PassThroughStageExecution.java | 34 ++-- .../analytics/exec/stage/ResponseCodec.java | 40 +++++ .../exec/stage/RowResponseCodec.java | 145 ++++++++++++++++++ .../stage/ShardFragmentStageExecution.java | 93 +++++------ .../stage/ShardFragmentStageScheduler.java | 18 ++- .../stage/SinkProvidingStageExecution.java | 25 ++- .../exec/stage/StageExecutionBuilder.java | 22 ++- 15 files changed, 434 insertions(+), 165 deletions(-) create mode 100644 sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ExchangeSource.java create mode 100644 sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/DataConsumer.java create mode 100644 sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/DataProducer.java create mode 100644 sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ResponseCodec.java create mode 100644 sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/RowResponseCodec.java diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ExchangeSink.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ExchangeSink.java index 0de25f9785e5d..2f1b8eb822a73 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ExchangeSink.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ExchangeSink.java @@ -11,8 +11,14 @@ import org.apache.arrow.vector.VectorSchemaRoot; /** - * Per-stage result accumulator that collects Arrow {@link VectorSchemaRoot} batches - * from shard executions and provides access to the accumulated result set. + * Write-only interface for feeding Arrow batches into a stage exchange. + * Producers (shard scan stages, local compute stages) call {@link #feed} + * to push data; they never read from the sink. + * + *

Implementations must be thread-safe — multiple shard response + * handlers may call {@link #feed} concurrently. + * + * @see ExchangeSource for the read-side counterpart */ public interface ExchangeSink { @@ -23,26 +29,7 @@ public interface ExchangeSink { void feed(VectorSchemaRoot batch); /** - * Signal that no more responses will be fed. + * Signal that no more batches will be fed. Releases resources. */ void close(); - - /** - * Return all accumulated rows in insertion order. - */ - Iterable readResult(); - - /** - * Return the total number of accumulated rows. - */ - long getRowCount(); - - /** - * Look up a cell value by column name and row index. - * - * @param column the column name - * @param rowIndex the zero-based row index - * @return the cell value, or {@code null} if the column is unknown or the row index is out of range - */ - Object getValueAt(String column, int rowIndex); } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ExchangeSource.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ExchangeSource.java new file mode 100644 index 0000000000000..fb3572f92dc40 --- /dev/null +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ExchangeSource.java @@ -0,0 +1,30 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.backend; + +/** + * Read-only interface for consuming accumulated results from a stage exchange. + * Consumers (parent stages, the walker's completion listener) read from this; + * they never write to it. + * + * @see ExchangeSink for the write-side counterpart + */ +public interface ExchangeSource { + + /** + * Return all accumulated rows in insertion order. + * Converts columnar Arrow batches to row-oriented {@code Object[]} arrays. + */ + Iterable readResult(); + + /** + * Return the total number of accumulated rows across all batches. + */ + long getRowCount(); +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/PlanWalker.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/PlanWalker.java index cc660ddf5aee5..dc289c0a6b473 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/PlanWalker.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/PlanWalker.java @@ -10,7 +10,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import org.opensearch.analytics.exec.stage.SinkProvidingStageExecution; +import org.opensearch.analytics.exec.stage.DataProducer; import org.opensearch.analytics.exec.stage.StageExecution; import org.opensearch.analytics.exec.stage.StageExecutionBuilder; import org.opensearch.analytics.planner.dag.Stage; @@ -36,8 +36,9 @@ * closures (no global dispatcher). * *

The walker is pure topology: it does not know about scheduler types, - * {@link SinkProvidingStageExecution}, or how a child's sink is resolved - * from its parent. That logic lives entirely in {@link StageExecutionBuilder}. + * {@link DataProducer}, {@link org.opensearch.analytics.exec.stage.DataConsumer}, + * or how a child's sink is resolved from its parent. That logic lives + * entirely in {@link StageExecutionBuilder}. * *

Lifecycle: constructed by {@link QueryScheduler#execute}, * tracked in the scheduler's pool by query id, removed on terminal. @@ -76,7 +77,7 @@ public void walk() { // Walk the DAG and build StageExecutions - this sets up stage control flow Stage rootStage = config.dag().rootStage(); final StageExecution rootExec = stageExecutionBuilder.buildRootExecution(rootStage, config); - wireCompletionListener((SinkProvidingStageExecution) rootExec); + wireCompletionListener(rootExec); executions.put(rootStage.getStageId(), rootExec); buildChildrenRecursively(rootExec, rootStage); @@ -166,10 +167,16 @@ private void buildChildrenRecursively(StageExecution parentExec, Stage parentSta } } - private void wireCompletionListener(SinkProvidingStageExecution rootExec) { + private void wireCompletionListener(StageExecution rootExec) { + if ((rootExec instanceof DataProducer) == false) { + throw new IllegalStateException( + "Root execution " + rootExec.getClass().getSimpleName() + " does not implement DataProducer" + ); + } + final DataProducer producer = (DataProducer) rootExec; rootExec.addStateListener((from, to) -> { switch (to) { - case SUCCEEDED -> fireTerminal(() -> completionListener.onResponse(rootExec.sink().readResult())); + case SUCCEEDED -> fireTerminal(() -> completionListener.onResponse(producer.outputSource().readResult())); case FAILED, CANCELLED -> { Exception failure = rootExec.getFailure(); if (config.parentTask() instanceof CancellableTask ct && ct.isCancelled()) { diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/RowProducingSink.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/RowProducingSink.java index b51605f23e483..f2ab30a634865 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/RowProducingSink.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/RowProducingSink.java @@ -11,32 +11,28 @@ import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; -import org.apache.arrow.vector.types.pojo.Field; import org.opensearch.analytics.backend.ExchangeSink; +import org.opensearch.analytics.backend.ExchangeSource; import java.util.ArrayList; import java.util.List; /** - * Default {@link ExchangeSink} implementation that collects Arrow - * {@link VectorSchemaRoot} batches. Converts back to {@code Object[]} rows - * on {@link #readResult()} for caller compatibility. + * Default exchange implementation that collects Arrow + * {@link VectorSchemaRoot} batches via {@link ExchangeSink#feed} and + * converts to {@code Object[]} rows on {@link ExchangeSource#readResult}. * - *

The sink takes ownership of fed batches and releases them on - * {@link #close()}. + *

Implements both {@link ExchangeSink} (write side for producers) and + * {@link ExchangeSource} (read side for consumers). The builder passes + * the {@link ExchangeSink} view to child stages and the walker reads + * results via the {@link ExchangeSource} view. */ -public class RowProducingSink implements ExchangeSink { +public class RowProducingSink implements ExchangeSink, ExchangeSource { private final List batches = new ArrayList<>(); - private final List fieldNames = new ArrayList<>(); @Override public void feed(VectorSchemaRoot batch) { - if (fieldNames.isEmpty() && batch.getSchema().getFields().isEmpty() == false) { - for (Field f : batch.getSchema().getFields()) { - fieldNames.add(f.getName()); - } - } batches.add(batch); } @@ -73,27 +69,6 @@ public long getRowCount() { return total; } - @Override - public Object getValueAt(String column, int rowIndex) { - int colIdx = fieldNames.indexOf(column); - if (colIdx < 0) return null; - - int offset = 0; - for (VectorSchemaRoot batch : batches) { - int batchRows = batch.getRowCount(); - if (rowIndex < offset + batchRows) { - return toJavaValue(batch.getVector(colIdx), rowIndex - offset); - } - offset += batchRows; - } - return null; - } - - /** - * Converts an Arrow vector value to a Java-native type. VarChar vectors - * return {@code org.apache.arrow.vector.util.Text} which callers don't - * expect — convert to {@code String}. Other types pass through as-is. - */ private static Object toJavaValue(FieldVector vector, int index) { if (vector.isNull(index)) return null; if (vector instanceof VarCharVector) { diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/DataConsumer.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/DataConsumer.java new file mode 100644 index 0000000000000..4d1cc23abdacc --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/DataConsumer.java @@ -0,0 +1,27 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.exec.stage; + +import org.opensearch.analytics.backend.ExchangeSink; + +/** + * Implemented by {@link StageExecution} types whose children write + * row batches into them. Used by the builder during construction to + * resolve the {@link ExchangeSink} a child stage should write into. + * + * @opensearch.internal + */ +public interface DataConsumer { + + /** + * Returns the {@link ExchangeSink} that the given child stage + * should write its output into. + */ + ExchangeSink inputSink(int childStageId); +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/DataProducer.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/DataProducer.java new file mode 100644 index 0000000000000..80cc9b26a4957 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/DataProducer.java @@ -0,0 +1,38 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.exec.stage; + +import org.opensearch.analytics.backend.ExchangeSink; +import org.opensearch.analytics.backend.ExchangeSource; + +/** + * Implemented by {@link StageExecution} types that write row batches + * into an output sink. The output sink is typically owned by the + * parent stage (resolved via {@link DataConsumer#inputSink}). + * + *

Exposes both the write-side ({@link ExchangeSink}) for producers + * feeding batches and the read-side ({@link ExchangeSource}) for the + * walker to consume final results. + * + * @opensearch.internal + */ +public interface DataProducer { + + /** + * Returns the write-side sink this stage feeds batches into. + */ + ExchangeSink outputSink(); + + /** + * Returns the read-side source for consuming accumulated results. + * For stages backed by a {@code RowProducingSink}, this returns the + * same object as {@link #outputSink()} (which implements both interfaces). + */ + ExchangeSource outputSource(); +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageExecution.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageExecution.java index 88debfdcbd551..8da2f2d1ef72d 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageExecution.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageExecution.java @@ -11,6 +11,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.analytics.backend.ExchangeSink; +import org.opensearch.analytics.backend.ExchangeSource; import org.opensearch.analytics.backend.LocalStageContext; import org.opensearch.analytics.planner.dag.Stage; import org.opensearch.core.action.ActionListener; @@ -20,6 +21,11 @@ * {@link LocalStageContext} lifecycle (start, finalize, fail, cancel) * and ensures the downstream listener is signaled exactly once. * + *

Implements {@link SinkProvidingStageExecution} as both a + * {@link DataConsumer} (children write into per-child input sinks via + * {@link #inputSink(int)}) and a {@link DataProducer} (output is + * produced by the backend into the parent's sink). + * *

Lifecycle: * {@code CREATED → RUNNING → (SUCCEEDED | FAILED | CANCELLED)} * @@ -33,10 +39,12 @@ final class LocalStageExecution extends AbstractStageExecution implements SinkPr private static final Logger logger = LogManager.getLogger(LocalStageExecution.class); private final LocalStageContext ctx; + private final ExchangeSink outputSinkRef; - public LocalStageExecution(Stage stage, LocalStageContext ctx) { + public LocalStageExecution(Stage stage, LocalStageContext ctx, ExchangeSink outputSink) { super(stage); this.ctx = ctx; + this.outputSinkRef = outputSink; logger.info("[LocalStage] CREATED stageId={} childCount={}", stage.getStageId(), stage.getChildStages().size()); } @@ -46,21 +54,21 @@ public LocalStageContext getCtx() { } @Override - public ExchangeSink sink(int childStageId) { + public ExchangeSink inputSink(int childStageId) { return ctx.sinkFor(childStageId); } - /** - * LocalStage routes per-child via {@link #sink(int)} / {@code ctx.sinkFor(childStageId)} - * — there is no single shared sink. Callers must use the childStageId-aware - * overload; this one is never reached because {@link #sink(int)} is overridden - * above, bypassing the default delegation in {@link SinkProvidingStageExecution}. - */ @Override - public ExchangeSink sink() { - throw new UnsupportedOperationException( - "LocalStageExecution has per-child input sinks — call sink(int childStageId) instead" - ); + public ExchangeSink outputSink() { + return outputSinkRef; + } + + @Override + public ExchangeSource outputSource() { + if (outputSinkRef instanceof ExchangeSource source) { + return source; + } + throw new UnsupportedOperationException("outputSink does not implement ExchangeSource"); } /** diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java index 01171f99aac81..f6b64cd0b2d82 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LocalStageScheduler.java @@ -111,7 +111,7 @@ private StageExecution buildComputeLocal(Stage stage, ExchangeSink sink, QueryCo } catch (Exception e) { throw new RuntimeException("Failed to create local stage context for stageId=" + stage.getStageId(), e); } - return new LocalStageExecution(stage, ctx); + return new LocalStageExecution(stage, ctx, sink); } /** diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/PassThroughStageExecution.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/PassThroughStageExecution.java index 6e0234d5049a3..6ab476fc6c781 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/PassThroughStageExecution.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/PassThroughStageExecution.java @@ -9,27 +9,31 @@ package org.opensearch.analytics.exec.stage; import org.opensearch.analytics.backend.ExchangeSink; +import org.opensearch.analytics.backend.ExchangeSource; +import org.opensearch.analytics.exec.RowProducingSink; import org.opensearch.analytics.planner.dag.Stage; /** * Sentinel {@link StageExecution} for LOCAL pass-through (root gather) - * stages. Owns an {@link ExchangeSink} and transitions synchronously through - * {@code CREATED → RUNNING → SUCCEEDED} on {@link #start()} — there is - * no real dispatch work. The driver calls {@code start()} only after all - * children have reached {@code SUCCEEDED}, which is the normal driver rule. + * stages. Owns a {@link RowProducingSink} and transitions synchronously + * through {@code CREATED → RUNNING → SUCCEEDED} on {@link #start()}. * - *

Keeps the registry fully populated so cancellation, EXPLAIN, and - * metrics traversal all work uniformly without pass-through special cases. + *

Implements both {@link DataConsumer} and {@link DataProducer} via + * {@link SinkProvidingStageExecution}: children write into the same sink + * that the root reads from. * * @opensearch.internal */ final class PassThroughStageExecution extends AbstractStageExecution implements SinkProvidingStageExecution { - private final ExchangeSink sink; + private final RowProducingSink ownedSink; public PassThroughStageExecution(Stage stage, ExchangeSink sink) { super(stage); - this.sink = sink; + if ((sink instanceof RowProducingSink) == false) { + throw new IllegalArgumentException("PassThroughStageExecution requires a RowProducingSink"); + } + this.ownedSink = (RowProducingSink) sink; } @Override @@ -44,7 +48,17 @@ public void cancel(String reason) { } @Override - public ExchangeSink sink() { - return sink; + public ExchangeSink inputSink(int childStageId) { + return ownedSink; + } + + @Override + public ExchangeSink outputSink() { + return ownedSink; + } + + @Override + public ExchangeSource outputSource() { + return ownedSink; } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ResponseCodec.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ResponseCodec.java new file mode 100644 index 0000000000000..528b3a93e2b1f --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ResponseCodec.java @@ -0,0 +1,40 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.exec.stage; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.opensearch.core.action.ActionResponse; + +/** + * Decodes a transport response into an Arrow {@link VectorSchemaRoot} for + * the coordinator-side sink. Implementations handle the specific wire + * format — {@code Object[]} rows (current), Arrow IPC (Flight), or any + * future format. + * + *

The codec is injected into {@link ShardFragmentStageExecution} at + * construction time by the scheduler. Swapping the codec swaps the + * serialization format without touching stage execution logic. + * + * @param the transport response type + * @opensearch.internal + */ +@FunctionalInterface +public interface ResponseCodec { + + /** + * Decodes a transport response into an Arrow {@link VectorSchemaRoot}. + * The returned VSR is owned by the caller (the sink). + * + * @param response the transport response + * @param allocator the buffer allocator for Arrow vectors + * @return a new VectorSchemaRoot; caller owns and must close it + */ + VectorSchemaRoot decode(R response, BufferAllocator allocator); +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/RowResponseCodec.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/RowResponseCodec.java new file mode 100644 index 0000000000000..d18ae5a372850 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/RowResponseCodec.java @@ -0,0 +1,145 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.exec.stage; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.SmallIntVector; +import org.apache.arrow.vector.TinyIntVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.opensearch.analytics.exec.action.FragmentExecutionResponse; + +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; + +/** + * {@link ResponseCodec} for the current row-oriented + * {@link FragmentExecutionResponse} wire format. Converts {@code Object[]} + * rows to Arrow {@link VectorSchemaRoot} via type inference. + * + *

This codec is the bridge that gets replaced when Arrow IPC transport + * lands. A future {@code ArrowIpcResponseCodec} would import IPC buffers + * directly — zero conversion. + * + * @opensearch.internal + */ +public final class RowResponseCodec implements ResponseCodec { + + /** Singleton instance — stateless, thread-safe. */ + public static final RowResponseCodec INSTANCE = new RowResponseCodec(); + + private RowResponseCodec() {} + + @Override + public VectorSchemaRoot decode(FragmentExecutionResponse response, BufferAllocator allocator) { + List fieldNames = response.getFieldNames(); + List rows = response.getRows(); + + if (allocator == null) { + allocator = new RootAllocator(); + } + + // Infer Arrow type per column from the first non-null value + List fields = new ArrayList<>(); + for (int col = 0; col < fieldNames.size(); col++) { + ArrowType arrowType = inferArrowType(rows, col); + fields.add(new Field(fieldNames.get(col), FieldType.nullable(arrowType), null)); + } + Schema schema = new Schema(fields); + + VectorSchemaRoot vsr = VectorSchemaRoot.create(schema, allocator); + try { + vsr.allocateNew(); + int rowCount = rows.size(); + for (int col = 0; col < fieldNames.size(); col++) { + FieldVector vector = vsr.getVector(col); + for (int r = 0; r < rowCount; r++) { + Object value = rows.get(r)[col]; + setVectorValue(vector, r, value); + } + vector.setValueCount(rowCount); + } + vsr.setRowCount(rowCount); + return vsr; + } catch (Exception e) { + vsr.close(); + throw e; + } + } + + /** + * Infers the Arrow type for a column by scanning rows for the first + * non-null value. Falls back to {@code Utf8} (VarChar) if all values + * are null or the Java type is unrecognized. + */ + static ArrowType inferArrowType(List rows, int col) { + for (Object[] row : rows) { + Object value = row[col]; + if (value == null) continue; + if (value instanceof Long) return new ArrowType.Int(64, true); + if (value instanceof Integer) return new ArrowType.Int(32, true); + if (value instanceof Short) return new ArrowType.Int(16, true); + if (value instanceof Byte) return new ArrowType.Int(8, true); + if (value instanceof Double) return new ArrowType.FloatingPoint(org.apache.arrow.vector.types.FloatingPointPrecision.DOUBLE); + if (value instanceof Float) return new ArrowType.FloatingPoint(org.apache.arrow.vector.types.FloatingPointPrecision.SINGLE); + if (value instanceof Boolean) return ArrowType.Bool.INSTANCE; + if (value instanceof CharSequence) return ArrowType.Utf8.INSTANCE; + if (value instanceof byte[]) return ArrowType.Binary.INSTANCE; + if (value instanceof Number) return new ArrowType.Int(64, true); + break; + } + return ArrowType.Utf8.INSTANCE; + } + + /** + * Sets a value on the appropriate Arrow vector type. Handles null by + * calling {@code setNull}. For typed vectors, casts the Java value to + * the expected type. + */ + static void setVectorValue(FieldVector vector, int index, Object value) { + if (value == null) { + vector.setNull(index); + return; + } + if (vector instanceof BigIntVector) { + ((BigIntVector) vector).setSafe(index, ((Number) value).longValue()); + } else if (vector instanceof IntVector) { + ((IntVector) vector).setSafe(index, ((Number) value).intValue()); + } else if (vector instanceof SmallIntVector) { + ((SmallIntVector) vector).setSafe(index, ((Number) value).shortValue()); + } else if (vector instanceof TinyIntVector) { + ((TinyIntVector) vector).setSafe(index, ((Number) value).byteValue()); + } else if (vector instanceof Float8Vector) { + ((Float8Vector) vector).setSafe(index, ((Number) value).doubleValue()); + } else if (vector instanceof Float4Vector) { + ((Float4Vector) vector).setSafe(index, ((Number) value).floatValue()); + } else if (vector instanceof BitVector) { + ((BitVector) vector).setSafe(index, ((Boolean) value) ? 1 : 0); + } else if (vector instanceof VarCharVector) { + ((VarCharVector) vector).setSafe(index, value.toString().getBytes(StandardCharsets.UTF_8)); + } else if (vector instanceof VarBinaryVector) { + ((VarBinaryVector) vector).setSafe(index, (byte[]) value); + } else { + throw new IllegalArgumentException("Unsupported Arrow vector type: " + vector.getClass().getSimpleName()); + } + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecution.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecution.java index 78221f493386e..87622c56e07ca 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecution.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecution.java @@ -8,36 +8,19 @@ package org.opensearch.analytics.exec.stage; -import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.RootAllocator; -import org.apache.arrow.vector.BigIntVector; -import org.apache.arrow.vector.BitVector; -import org.apache.arrow.vector.FieldVector; -import org.apache.arrow.vector.Float4Vector; -import org.apache.arrow.vector.Float8Vector; -import org.apache.arrow.vector.IntVector; -import org.apache.arrow.vector.SmallIntVector; -import org.apache.arrow.vector.TinyIntVector; -import org.apache.arrow.vector.VarBinaryVector; -import org.apache.arrow.vector.VarCharVector; import org.apache.arrow.vector.VectorSchemaRoot; -import org.apache.arrow.vector.types.pojo.ArrowType; -import org.apache.arrow.vector.types.pojo.Field; -import org.apache.arrow.vector.types.pojo.FieldType; -import org.apache.arrow.vector.types.pojo.Schema; import org.opensearch.analytics.backend.ExchangeSink; -import org.opensearch.analytics.exec.action.FragmentExecutionResponse; +import org.opensearch.analytics.backend.ExchangeSource; import org.opensearch.analytics.exec.AnalyticsSearchTransportService; import org.opensearch.analytics.exec.PendingExecutions; import org.opensearch.analytics.exec.QueryContext; -import org.opensearch.analytics.exec.action.ShardTarget; -import org.opensearch.analytics.exec.AnalyticsSearchTransportService; import org.opensearch.analytics.exec.StreamingResponseListener; import org.opensearch.analytics.exec.action.FragmentExecutionRequest; +import org.opensearch.analytics.exec.action.FragmentExecutionResponse; +import org.opensearch.analytics.exec.action.ShardTarget; import org.opensearch.analytics.planner.dag.Stage; +import org.opensearch.core.action.ActionResponse; -import java.nio.charset.StandardCharsets; -import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -47,49 +30,54 @@ /** * Per-stage execution for row-producing DATA_NODE stages (scans, filters, * partial aggregates). Dispatches shard requests via - * {@link AnalyticsSearchTransportService#dispatchFragment}, collects streaming - * {@link FragmentExecutionResponse} batches, and feeds them into the stage's - * {@link org.opensearch.analytics.backend.ExchangeSink}. + * {@link AnalyticsSearchTransportService#dispatchFragment}, decodes streaming + * responses through a {@link ResponseCodec}, and feeds the resulting Arrow + * batches into the stage's output {@link ExchangeSink}. + * + *

The codec abstracts the wire format: the current {@link RowResponseCodec} + * converts {@code Object[]} rows to Arrow; a future Arrow IPC codec would + * import IPC buffers directly with zero conversion. The stage execution logic + * is format-agnostic. * - *

Replaces the scan path that previously lived in the generic - * fan-out execution + sink-feeding handler. + *

Implements {@link DataProducer} because it writes batches into a sink + * owned by its parent stage. Does not implement {@link DataConsumer} because + * it is a leaf stage with no children. * *

Lifecycle: {@code CREATED → RUNNING → SUCCEEDED | FAILED | CANCELLED}. * Instances are one-shot: constructed, {@link #start()} called once, * listener signaled once, discarded. * - *

No {@code completedStages} tracking — that responsibility moves to - * the caller (PlanWalker / scheduler) in a later change. - * * @opensearch.internal */ -final class ShardFragmentStageExecution extends AbstractStageExecution implements SinkProvidingStageExecution { +final class ShardFragmentStageExecution extends AbstractStageExecution implements DataProducer { private final AtomicInteger inFlight = new AtomicInteger(0); - private final AtomicInteger completedTasks = new AtomicInteger(0); // Immutable config private final QueryContext config; - private final ExchangeSink sink; + private final ExchangeSink outputSink; private final List targets; private final Function requestBuilder; private final AnalyticsSearchTransportService dispatcher; + private final ResponseCodec responseCodec; private final Map pendingPerNode = new ConcurrentHashMap<>(); ShardFragmentStageExecution( Stage stage, QueryContext config, - ExchangeSink sink, + ExchangeSink outputSink, List targets, Function requestBuilder, - AnalyticsSearchTransportService dispatcher + AnalyticsSearchTransportService dispatcher, + ResponseCodec responseCodec ) { super(stage); this.config = config; - this.sink = sink; + this.outputSink = outputSink; this.targets = targets; this.requestBuilder = requestBuilder; this.dispatcher = dispatcher; + this.responseCodec = responseCodec; } @Override @@ -99,7 +87,6 @@ public void start() { transitionTo(StageExecution.State.SUCCEEDED); return; } - // TODO: Introduce Shard Filter & Termination Decider logic to this execution type if (transitionTo(StageExecution.State.RUNNING) == false) return; inFlight.set(targets.size()); for (ShardTarget target : targets) { @@ -116,14 +103,13 @@ public void onStreamResponse(FragmentExecutionResponse response, boolean isLast) config.searchExecutor().execute(() -> { if (isDone()) return; - // TODO: This serialization should not be required - VectorSchemaRoot vsr = scanResponseToArrow(response, config.bufferAllocator()); - sink.feed(vsr); + VectorSchemaRoot vsr = responseCodec.decode(response, config.bufferAllocator()); + outputSink.feed(vsr); metrics.addRowsProcessed(vsr.getRowCount()); if (isLast) { metrics.incrementTasksCompleted(); - onTaskCompletion(); + onShardTerminated(); } }); } @@ -132,37 +118,34 @@ public void onStreamResponse(FragmentExecutionResponse response, boolean isLast) public void onFailure(Exception e) { captureFailure(new RuntimeException("Stage " + stage.getStageId() + " failed", e)); metrics.incrementTasksFailed(); - onTaskCompletion(); + onShardTerminated(); } }, config.parentTask(), pending); } - private void onTaskCompletion() { - completedTasks.incrementAndGet(); + private void onShardTerminated() { if (inFlight.decrementAndGet() == 0) { - finishStageInternal(); + Exception captured = getFailure(); + transitionTo(captured != null ? StageExecution.State.FAILED : StageExecution.State.SUCCEEDED); } } - private void finishStageInternal() { - Exception captured = getFailure(); - StageExecution.State target = (captured != null) ? StageExecution.State.FAILED : StageExecution.State.SUCCEEDED; - transitionTo(target); - } - @Override public void cancel(String reason) { if (transitionTo(StageExecution.State.CANCELLED) == false) return; } @Override - public ExchangeSink sink() { - return sink; + public ExchangeSink outputSink() { + return outputSink; } - /** Returns the sink this execution writes batches into. */ - public ExchangeSink getSink() { - return sink; + @Override + public ExchangeSource outputSource() { + if (outputSink instanceof ExchangeSource source) { + return source; + } + throw new UnsupportedOperationException("outputSink does not implement ExchangeSource"); } private boolean isDone() { diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageScheduler.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageScheduler.java index e6b15028c7054..cf6857271b94c 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageScheduler.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/ShardFragmentStageScheduler.java @@ -13,6 +13,7 @@ import org.opensearch.analytics.exec.action.ShardTarget; import org.opensearch.analytics.exec.AnalyticsSearchTransportService; import org.opensearch.analytics.exec.action.FragmentExecutionRequest; +import org.opensearch.analytics.exec.action.FragmentExecutionResponse; import org.opensearch.analytics.planner.dag.Stage; import org.opensearch.analytics.planner.dag.StagePlan; import org.opensearch.cluster.service.ClusterService; @@ -27,16 +28,31 @@ * and doesn't care whether it is a root sink or a parent-provided child sink * — {@link StageExecutionBuilder} resolves that distinction before calling. * + *

Injects a {@link ResponseCodec} into the execution to decouple the wire + * format from stage logic. The default codec ({@link RowResponseCodec}) handles + * the current {@code Object[]} row format; a future Arrow IPC codec would be + * swapped in here. + * * @opensearch.internal */ final class ShardFragmentStageScheduler { private final ClusterService clusterService; private final AnalyticsSearchTransportService transport; + private final ResponseCodec responseCodec; ShardFragmentStageScheduler(ClusterService clusterService, AnalyticsSearchTransportService transport) { + this(clusterService, transport, RowResponseCodec.INSTANCE); + } + + ShardFragmentStageScheduler( + ClusterService clusterService, + AnalyticsSearchTransportService transport, + ResponseCodec responseCodec + ) { this.clusterService = clusterService; this.transport = transport; + this.responseCodec = responseCodec; } StageExecution createExecution(Stage stage, ExchangeSink sink, QueryContext config) { @@ -53,7 +69,7 @@ StageExecution createExecution(Stage stage, ExchangeSink sink, QueryContext conf planAlternatives ); - return new ShardFragmentStageExecution(stage, config, sink, targets, requestBuilder, transport); + return new ShardFragmentStageExecution(stage, config, sink, targets, requestBuilder, transport, responseCodec); } private static List buildPlanAlternatives(Stage stage) { diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/SinkProvidingStageExecution.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/SinkProvidingStageExecution.java index 7c3c3ca924747..ad82fbdf3bedf 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/SinkProvidingStageExecution.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/SinkProvidingStageExecution.java @@ -8,24 +8,17 @@ package org.opensearch.analytics.exec.stage; -import org.opensearch.analytics.backend.ExchangeSink; - /** - * Implemented by {@link StageExecution} types whose children may write - * row batches into them. Used by the walker during construction to thread - * a sink from parent to row-producing child without an instanceof cascade - * over concrete execution types. + * Combines {@link DataConsumer} and {@link DataProducer} for stages that + * both accept child input and produce output (root gather, local compute). + * + *

Stages that own a single shared sink (like {@link PassThroughStageExecution}) + * return the same sink from both {@link DataConsumer#inputSink(int)} and + * {@link DataProducer#outputSink()}. Stages with per-child routing (like + * {@link LocalStageExecution}) override {@link DataConsumer#inputSink(int)} + * to delegate to the backend. * * @opensearch.internal */ -public interface SinkProvidingStageExecution extends StageExecution { - /** - * Returns the {@link ExchangeSink} that the given child stage should write - * into. - */ - default ExchangeSink sink(int childStageId) { - return sink(); - } - - ExchangeSink sink(); +public interface SinkProvidingStageExecution extends StageExecution, DataConsumer, DataProducer { } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/StageExecutionBuilder.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/StageExecutionBuilder.java index ad54576f71e50..d1bbf099a758d 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/StageExecutionBuilder.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/StageExecutionBuilder.java @@ -35,7 +35,7 @@ * different {@code createExecution} signature. *

  • Output resolution. For row-producing stages, resolves the * {@link ExchangeSink} from the parent's - * {@link SinkProvidingStageExecution} contract. For manifest-producing + * {@link DataConsumer} contract. For manifest-producing * stages (future), it will resolve a manifest receiver from the parent's * corresponding contract.
    • * @@ -93,11 +93,11 @@ public StageExecutionBuilder( * Builds the root stage's execution. Creates a fresh {@link ExchangeSink} * internally — the walker doesn't see it. The walker reads the final * result by casting the returned execution to - * {@link SinkProvidingStageExecution} and calling - * {@code sink().readResult()}. + * {@link DataProducer} and calling + * {@code outputSink().readResult()}. */ public StageExecution buildRootExecution(Stage rootStage, QueryContext config) { - ExchangeSink rootSink = new org.opensearch.analytics.exec.RowProducingSink(); + org.opensearch.analytics.exec.RowProducingSink rootSink = new org.opensearch.analytics.exec.RowProducingSink(); return dispatchRowStage(rootStage, rootSink, config); } @@ -112,7 +112,7 @@ public StageExecution buildRootExecution(Stage rootStage, QueryContext config) { * * @throws IllegalStateException if a row-producing stage's * {@code parentExec} does not implement - * {@link SinkProvidingStageExecution} — this is a planner bug. + * {@link DataConsumer} — this is a planner bug. * @throws UnsupportedOperationException if {@code stage} produces a * manifest (shuffle-write or broadcast-write) — not yet implemented. */ @@ -145,16 +145,22 @@ private StageExecution dispatchRowStage(Stage stage, ExchangeSink sink, QueryCon return shardFanOutScheduler.createExecution(stage, sink, config); } + /** + * Resolves the output sink for a child stage from its parent. The parent + * must implement {@link DataConsumer} to accept child input. The returned + * sink is the same object the child will write into via + * {@link DataProducer#outputSink()}. + */ private static ExchangeSink resolveRowSink(Stage stage, StageExecution parentExec) { - if (parentExec instanceof SinkProvidingStageExecution exec) { - return exec.sink(stage.getStageId()); + if (parentExec instanceof DataConsumer consumer) { + return consumer.inputSink(stage.getStageId()); } throw new IllegalStateException( "row-producing stage " + stage.getStageId() + " has parent " + parentExec.getClass().getSimpleName() - + " which does not receive row input" + + " which does not accept child input (does not implement DataConsumer)" ); } } From 29c28e8ddfb65fddeda1cf10f275bd3b6851a468 Mon Sep 17 00:00:00 2001 From: Bukhtawar Khan Date: Sat, 18 Apr 2026 01:03:14 +0530 Subject: [PATCH 2/2] Simplify the source sink interactions Signed-off-by: Bukhtawar Khan --- .../analytics/exec/ExecutionGraph.java | 102 ++++++++++++++++ .../opensearch/analytics/exec/PlanWalker.java | 113 +++++++++++------- .../analytics/exec/QueryScheduler.java | 61 ++++++---- .../analytics/exec/RowProducingSink.java | 28 +++++ 4 files changed, 234 insertions(+), 70 deletions(-) create mode 100644 sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ExecutionGraph.java diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ExecutionGraph.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ExecutionGraph.java new file mode 100644 index 0000000000000..901784bfbff48 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ExecutionGraph.java @@ -0,0 +1,102 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.analytics.exec; + +import org.opensearch.analytics.exec.stage.StageExecution; +import org.opensearch.analytics.exec.stage.StageMetrics; + +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +/** + * Inspectable snapshot of a fully-wired execution graph. Built by + * {@link PlanWalker#build()}, consumed by {@link PlanWalker#start()} + * or by EXPLAIN without executing. + * + *

    The graph holds all {@link StageExecution} instances with their + * state listeners already wired. No stage has been started yet — all + * are in {@link StageExecution.State#CREATED}. + * + * @opensearch.internal + */ +public class ExecutionGraph { + + private final Map executions; + private final StageExecution rootExecution; + private final List leaves; + private final String queryId; + + ExecutionGraph( + String queryId, + Map executions, + StageExecution rootExecution, + List leaves + ) { + this.queryId = queryId; + this.executions = executions; + this.rootExecution = rootExecution; + this.leaves = leaves; + } + + /** The query this graph belongs to. */ + public String queryId() { + return queryId; + } + + /** The root stage execution. */ + public StageExecution rootExecution() { + return rootExecution; + } + + /** All leaf executions (stages with no children). */ + public List leaves() { + return Collections.unmodifiableList(leaves); + } + + /** Lookup a stage execution by stage id. */ + public StageExecution executionFor(int stageId) { + return executions.get(stageId); + } + + /** All stage executions in the graph. */ + public Collection allExecutions() { + return Collections.unmodifiableCollection(executions.values()); + } + + /** Number of stages in the graph. */ + public int stageCount() { + return executions.size(); + } + + /** + * Returns a human-readable summary of the execution graph for + * EXPLAIN output. Lists each stage with its type, state, and + * child dependencies. + */ + public String explain() { + StringBuilder sb = new StringBuilder(); + sb.append("ExecutionGraph[queryId=").append(queryId); + sb.append(", stages=").append(executions.size()); + sb.append(", leaves=").append(leaves.size()).append("]\n"); + for (StageExecution exec : executions.values()) { + sb.append(" Stage ").append(exec.getStageId()); + sb.append(" [").append(exec.getClass().getSimpleName()).append("]"); + sb.append(" state=").append(exec.getState()); + StageMetrics m = exec.getMetrics(); + if (m.getStartTimeMs() > 0) { + sb.append(" elapsed=").append(m.getEndTimeMs() - m.getStartTimeMs()).append("ms"); + sb.append(" rows=").append(m.getRowsProcessed()); + } + sb.append("\n"); + } + return sb.toString(); + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/PlanWalker.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/PlanWalker.java index dc289c0a6b473..d87914c53cb13 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/PlanWalker.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/PlanWalker.java @@ -28,20 +28,21 @@ import java.util.concurrent.atomic.AtomicInteger; /** - * Per-query walker that owns the execution graph. Walks the DAG once, - * constructs all {@link StageExecution} instances via - * {@link StageExecutionBuilder#buildExecution} (including the root, wired through a - * virtual sink-holder parent), wires per-parent listeners inline - * during construction, and drives state transitions via local listener - * closures (no global dispatcher). + * Per-query walker that owns the execution graph. Two-phase lifecycle: + *

      + *
    1. {@link #build()} — walks the DAG, constructs all {@link StageExecution} + * instances, wires listeners, and returns an inspectable + * {@link ExecutionGraph}. No stages are started.
      2. + *
    2. {@link #start(ExecutionGraph)} — starts leaf stages, triggering + * the event-driven cascade.
      3. + *
    * - *

    The walker is pure topology: it does not know about scheduler types, - * {@link DataProducer}, {@link org.opensearch.analytics.exec.stage.DataConsumer}, - * or how a child's sink is resolved from its parent. That logic lives - * entirely in {@link StageExecutionBuilder}. + *

    The split enables EXPLAIN: call {@link #build()} to get the graph, + * inspect it via {@link ExecutionGraph#explain()}, and optionally call + * {@link #start(ExecutionGraph)} to execute. * - *

    Lifecycle: constructed by {@link QueryScheduler#execute}, - * tracked in the scheduler's pool by query id, removed on terminal. + *

    The legacy {@link #walk()} method calls both phases for backward + * compatibility. * * @opensearch.internal */ @@ -51,9 +52,9 @@ public class PlanWalker { private final QueryContext config; private final StageExecutionBuilder stageExecutionBuilder; - private final Map executions = new ConcurrentHashMap<>(); private final AtomicBoolean terminalFired = new AtomicBoolean(false); private final ActionListener> completionListener; + private volatile ExecutionGraph graph; public PlanWalker(QueryContext config, StageExecutionBuilder stageExecutionBuilder, ActionListener> listener) { this.config = config; @@ -62,39 +63,56 @@ public PlanWalker(QueryContext config, StageExecutionBuilder stageExecutionBuild } /** - * Walks the DAG, builds all executions, wires per-parent listeners, - * wires the root terminal listener, and starts leaves. + * Phase 1: Build the execution graph without starting any stages. + * All stages are in {@link StageExecution.State#CREATED} state. + * Listeners are wired. The graph is inspectable for EXPLAIN. * - *

    Four phases: - *

      - *
    1. Build root execution with a locally-owned {@link RowProducingSink}.
      2. - *
    2. Walk children recursively, wiring per-parent listeners inline.
      3. - *
    3. Wire the root's terminal listener.
      4. - *
    4. Start leaves.
      5. - *
    + * @return the fully-wired execution graph */ - public void walk() { - // Walk the DAG and build StageExecutions - this sets up stage control flow + public ExecutionGraph build() { + Map executions = new ConcurrentHashMap<>(); + Stage rootStage = config.dag().rootStage(); final StageExecution rootExec = stageExecutionBuilder.buildRootExecution(rootStage, config); wireCompletionListener(rootExec); executions.put(rootStage.getStageId(), rootExec); - buildChildrenRecursively(rootExec, rootStage); + buildChildrenRecursively(executions, rootExec, rootStage); + + List leaves = findLeaves(executions, rootStage); + + this.graph = new ExecutionGraph(config.queryId(), executions, rootExec, leaves); + return this.graph; + } - for (StageExecution leaf : findLeaves()) { + /** + * Phase 2: Start execution by dispatching leaf stages. + * Must be called after {@link #build()}. + * + * @param executionGraph the graph returned by {@link #build()} + */ + public void start(ExecutionGraph executionGraph) { + for (StageExecution leaf : executionGraph.leaves()) { leaf.start(); } } + /** + * Legacy single-call entry point. Builds the graph and starts + * execution in one shot. Equivalent to {@code start(build())}. + */ + public void walk() { + start(build()); + } + /** * Top-down cancel: iterates all executions and cancels any in - * {@code RUNNING} or {@code CREATED} state. Used by external - * cancellation (task cancel, timeout) via - * {@link QueryScheduler}. + * {@code RUNNING} or {@code CREATED} state. */ public void cancelAll(String reason) { - for (StageExecution exec : executions.values()) { + ExecutionGraph g = this.graph; + if (g == null) return; + for (StageExecution exec : g.allExecutions()) { StageExecution.State state = exec.getState(); if (state == StageExecution.State.RUNNING || state == StageExecution.State.CREATED) { try { @@ -104,6 +122,11 @@ public void cancelAll(String reason) { } } + /** Returns the built execution graph, or null if {@link #build()} hasn't been called. */ + public ExecutionGraph getGraph() { + return graph; + } + public String getQueryId() { return config.queryId(); } @@ -117,20 +140,26 @@ public Task getParentTask() { } public StageExecution executionFor(int stageId) { - return executions.get(stageId); + ExecutionGraph g = this.graph; + return g != null ? g.executionFor(stageId) : null; } public Collection activeExecutions() { - return executions.values().stream() + ExecutionGraph g = this.graph; + if (g == null) return List.of(); + return g.allExecutions().stream() .filter(e -> e.getState() == StageExecution.State.RUNNING) .toList(); } public Collection allExecutions() { - return executions.values(); + ExecutionGraph g = this.graph; + return g != null ? g.allExecutions() : List.of(); } - private void buildChildrenRecursively(StageExecution parentExec, Stage parentStage) { + // ─── Internal graph construction ──────────────────────────────────── + + private void buildChildrenRecursively(Map executions, StageExecution parentExec, Stage parentStage) { List children = parentStage.getChildStages(); if (children.isEmpty()) { return; @@ -142,8 +171,6 @@ private void buildChildrenRecursively(StageExecution parentExec, Stage parentSta StageExecution childExec = stageExecutionBuilder.buildExecution(child, parentExec, config); executions.put(child.getStageId(), childExec); - // Per-parent listener: this child → this specific parent. - // No global dispatch, no parentsByChild lookup, no re-deriving readiness. childExec.addStateListener((from, to) -> { switch (to) { case SUCCEEDED -> { @@ -162,8 +189,7 @@ private void buildChildrenRecursively(StageExecution parentExec, Stage parentSta default -> { } } }); - // Recurse into grandchildren - buildChildrenRecursively(childExec, child); + buildChildrenRecursively(executions, childExec, child); } } @@ -182,9 +208,6 @@ private void wireCompletionListener(StageExecution rootExec) { if (config.parentTask() instanceof CancellableTask ct && ct.isCancelled()) { fireTerminal(() -> completionListener.onFailure(new TaskCancelledException("query cancelled"))); } else if (failure != null) { - // The failure is already wrapped as "Stage N failed" at the point of origin - // (see ShardFragmentStageExecution.dispatchShardTask.onFailure). Forward as-is - // so the originating stage id is preserved through propagation. fireTerminal(() -> completionListener.onFailure(failure)); } else { fireTerminal(() -> completionListener.onFailure( @@ -196,13 +219,13 @@ private void wireCompletionListener(StageExecution rootExec) { }); } - private List findLeaves() { + private static List findLeaves(Map executions, Stage rootStage) { final List leaves = new ArrayList<>(); - collectLeaves(config.dag().rootStage(), leaves); + collectLeaves(executions, rootStage, leaves); return leaves; } - private void collectLeaves(Stage stage, List leaves) { + private static void collectLeaves(Map executions, Stage stage, List leaves) { if (stage.getChildStages().isEmpty()) { StageExecution exec = executions.get(stage.getStageId()); if (exec != null) { @@ -210,7 +233,7 @@ private void collectLeaves(Stage stage, List leaves) { } } else { for (Stage child : stage.getChildStages()) { - collectLeaves(child, leaves); + collectLeaves(executions, child, leaves); } } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryScheduler.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryScheduler.java index 1f0538f700251..772eb1c8e6ce9 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryScheduler.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryScheduler.java @@ -20,22 +20,17 @@ import java.util.concurrent.ConcurrentHashMap; /** - * Default {@link Scheduler} implementation. Pool manager for per-query - * {@link PlanWalker} instances. Constructs a walker on each - * {@link #execute} call, tracks it by query id, and removes it on - * terminal (success or failure). + * Default {@link Scheduler} implementation. Two-phase execution: + *
      + *
    1. {@link #plan(QueryContext)} — builds the execution graph without + * starting any stages. Returns an {@link ExecutionGraph} that can + * be inspected for EXPLAIN.
      2. + *
    2. {@link #execute(QueryContext, ActionListener)} — builds and starts + * execution in one call (the normal query path).
      3. + *
    * - *

    Responsibilities: - *

      - *
    • Walker construction: creates a {@link PlanWalker} with the - * per-query {@link QueryContext} and the shared {@link StageExecutionBuilder}.
      • - *
    • Pool tracking: maintains {@link #walkerPool} for - * future observability and concurrency limiting.
      • - *
    • Cancellation wiring: installs a cancel callback on the query - * task that calls {@link PlanWalker#cancelAll(String)}.
      • - *
    • Per-query cleanup: removes the walker from the pool on the - * terminal path before firing the caller's listener.
      • - *
    + *

    Also manages a pool of active {@link PlanWalker} instances for + * observability and cancellation. * * @opensearch.internal */ @@ -51,23 +46,39 @@ public QueryScheduler(StageExecutionBuilder stageExecutionBuilder) { this.stageExecutionBuilder = stageExecutionBuilder; } + /** + * Builds the execution graph without starting any stages. + * Use for EXPLAIN — inspect the returned graph, then discard. + * + * @param config the per-query context + * @return the fully-wired but unstarted execution graph + */ + public ExecutionGraph plan(QueryContext config) { + PlanWalker walker = new PlanWalker(config, stageExecutionBuilder, ActionListener.wrap(r -> {}, e -> {})); + return walker.build(); + } + @Override public void execute(QueryContext config, ActionListener> listener) { final String queryId = config.queryId(); - PlanWalker walker = getPlanWalker(config, listener, queryId); + PlanWalker walker = createWalker(config, listener, queryId); walkerPool.put(queryId, walker); final AnalyticsQueryTask queryTask = config.parentTask(); queryTask.setOnCancelCallback(() -> { - String reason = "task cancelled: " - + (queryTask.getReasonCancelled() != null ? queryTask.getReasonCancelled() : "unknown"); - logger.info("[QueryScheduler] AnalyticsQueryTask.onCancelled fired, reason={}", reason); - walker.cancelAll(reason); - }); - walker.walk(); + String reason = "task cancelled: " + + (queryTask.getReasonCancelled() != null ? queryTask.getReasonCancelled() : "unknown"); + logger.info("[QueryScheduler] AnalyticsQueryTask.onCancelled fired, reason={}", reason); + walker.cancelAll(reason); + }); + + // Two-phase: build graph, then start execution + ExecutionGraph graph = walker.build(); + logger.info("[QueryScheduler] ExecutionGraph built:\n{}", graph.explain()); + walker.start(graph); } - private PlanWalker getPlanWalker(QueryContext config, ActionListener> listener, String queryId) { + private PlanWalker createWalker(QueryContext config, ActionListener> listener, String queryId) { ActionListener> wrapped = ActionListener.wrap( result -> { walkerPool.remove(queryId); @@ -81,12 +92,12 @@ private PlanWalker getPlanWalker(QueryContext config, ActionListener activeWalkers() { return walkerPool.values(); } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/RowProducingSink.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/RowProducingSink.java index f2ab30a634865..c82cee7959409 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/RowProducingSink.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/RowProducingSink.java @@ -30,9 +30,15 @@ public class RowProducingSink implements ExchangeSink, ExchangeSource { private final List batches = new ArrayList<>(); + private final List fieldNames = new ArrayList<>(); @Override public void feed(VectorSchemaRoot batch) { + if (fieldNames.isEmpty() && batch.getSchema().getFields().isEmpty() == false) { + for (org.apache.arrow.vector.types.pojo.Field f : batch.getSchema().getFields()) { + fieldNames.add(f.getName()); + } + } batches.add(batch); } @@ -69,6 +75,28 @@ public long getRowCount() { return total; } + /** + * Look up a cell value by column name and row index. + * + * @param column the column name + * @param rowIndex the zero-based row index + * @return the cell value, or {@code null} if the column is unknown or the row index is out of range + */ + public Object getValueAt(String column, int rowIndex) { + int colIdx = fieldNames.indexOf(column); + if (colIdx < 0) return null; + + int offset = 0; + for (VectorSchemaRoot batch : batches) { + int batchRows = batch.getRowCount(); + if (rowIndex < offset + batchRows) { + return toJavaValue(batch.getVector(colIdx), rowIndex - offset); + } + offset += batchRows; + } + return null; + } + private static Object toJavaValue(FieldVector vector, int index) { if (vector.isNull(index)) return null; if (vector instanceof VarCharVector) {