From 62311ad396ca5e40b238a3c094b584966c2a5574 Mon Sep 17 00:00:00 2001 From: Marc Handalian Date: Sun, 31 May 2026 19:37:49 -0700 Subject: [PATCH 1/3] analytics-engine: run query under framework AnalyticsQueryTask for HTTP-disconnect cancellation (#21917) Remove manual task register/unregister and rely on parent task in local transport call. Also update the explain endpoint to use our local transport call. Adds AnalyticsQueryTaskCleanupIT covering task cleanup on success and on cancel. Signed-off-by: Marc Handalian Co-authored-by: Sandesh Kumar --- .../analytics/exec/DefaultPlanExecutor.java | 49 ++-- .../exec/action/AnalyticsQueryRequest.java | 33 +++ .../exec/action/AnalyticsQueryResponse.java | 12 + .../AnalyticsQueryTaskCleanupIT.java | 259 ++++++++++++++++++ 4 files changed, 330 insertions(+), 23 deletions(-) create mode 100644 sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/analytics/cancellation/AnalyticsQueryTaskCleanupIT.java diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java index 1166e5dbd7245..c4f3ea99f949f 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/DefaultPlanExecutor.java @@ -29,7 +29,6 @@ import org.opensearch.analytics.exec.profile.QueryProfile; import org.opensearch.analytics.exec.profile.QueryProfileBuilder; import org.opensearch.analytics.exec.task.AnalyticsQueryTask; -import org.opensearch.analytics.exec.task.AnalyticsQueryTaskRequest; import org.opensearch.analytics.planner.CapabilityRegistry; import org.opensearch.analytics.planner.PlannerContext; import org.opensearch.analytics.planner.PlannerImpl; @@ -48,7 +47,6 @@ import org.opensearch.core.action.ActionListener; import org.opensearch.search.SearchService; import org.opensearch.tasks.Task; -import org.opensearch.tasks.TaskManager; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportService; import org.opensearch.transport.client.node.NodeClient; @@ -83,7 +81,6 @@ public class DefaultPlanExecutor extends HandledTransportAction listener) { - searchExecutor.execute(() -> { - try { - executeInternal(logicalFragment, queryCtx, true, listener); - } catch (Exception e) { - listener.onFailure(e); - } catch (AssertionError e) { - listener.onFailure(new IllegalStateException("Analytics-engine executor rejected the plan: " + e.getMessage(), e)); - } - }); + // Route through the framework action (profile=true) just like execute(), so a profiling + // query runs under the framework-provided, cancellable task rather than detached on + // searchExecutor. The SecurityFilter also evaluates index permissions on this path. + String[] indices = RelNodeUtils.extractIndices(logicalFragment); + AnalyticsQueryRequest request = new AnalyticsQueryRequest(logicalFragment, queryCtx, indices, true); + client.execute( + AnalyticsQueryAction.INSTANCE, + request, + ActionListener.wrap(resp -> listener.onResponse(resp.getProfiledResult()), listener::onFailure) + ); } /** @@ -165,6 +162,7 @@ public void executeWithProfile(RelNode logicalFragment, QueryRequestContext quer * the schema from. Otherwise a fresh {@code clusterService.state()} is read. */ private void executeInternal( + AnalyticsQueryTask queryTask, RelNode logicalFragment, QueryRequestContext queryCtx, boolean profile, @@ -191,11 +189,10 @@ private void executeInternal( final long planningTimeMs = profile ? java.util.concurrent.TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - planStartNanos) : 0; logger.debug("[DefaultPlanExecutor] QueryDAG:\n{}", dag); - final AnalyticsQueryTask queryTask = (AnalyticsQueryTask) taskManager.register( - "transport", - "analytics_query", - new AnalyticsQueryTaskRequest(dag.queryId(), null) - ); + // The task is the framework-provided task from doExecute (registered by + // HandledTransportAction before doExecute, unregistered when the listener completes). + // Using it — rather than self-registering a detached task — is what lets a client + // disconnect / explicit task cancel propagate into the running query. final BufferAllocator queryAllocator; final boolean ownsAllocator; if (perQueryBufferLimit <= 0) { @@ -233,9 +230,12 @@ private void executeInternal( : ActionListener.wrap(rows -> listener.onResponse(new ProfiledResult(rows, null, null)), listener::onFailure); final List outputColumnOrder = logicalFragment.getRowType().getFieldNames(); - ActionListener> batchesListener = ActionListener.runAfter( - ActionListener.wrap(batches -> rowsListener.onResponse(batchesToRows(batches, outputColumnOrder)), rowsListener::onFailure), - () -> taskManager.unregister(queryTask) + // No taskManager.unregister here: the framework (HandledTransportAction) unregisters the + // task it created for doExecute once this listener settles. Unregistering it ourselves + // would double-free a task we no longer own. + ActionListener> batchesListener = ActionListener.wrap( + batches -> rowsListener.onResponse(batchesToRows(batches, outputColumnOrder)), + rowsListener::onFailure ); TimeValue taskTimeout = queryTask.getCancelAfterTimeInterval(); @@ -288,11 +288,14 @@ protected void doExecute(Task task, AnalyticsQueryRequest request, ActionListene ContextAwareExecutor.wrap(searchExecutor, threadPool).execute(() -> { try { executeInternal( + (AnalyticsQueryTask) task, request.getPlan(), request.getQueryCtx(), - false, + request.isProfile(), ActionListener.wrap( - result -> convertingListener.onResponse(new AnalyticsQueryResponse(result.rows())), + result -> convertingListener.onResponse( + request.isProfile() ? new AnalyticsQueryResponse(result) : new AnalyticsQueryResponse(result.rows()) + ), convertingListener::onFailure ) ); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/AnalyticsQueryRequest.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/AnalyticsQueryRequest.java index c35643df1ae4a..903c1ddc51587 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/AnalyticsQueryRequest.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/AnalyticsQueryRequest.java @@ -14,10 +14,14 @@ import org.opensearch.action.IndicesRequest; import org.opensearch.action.support.IndicesOptions; import org.opensearch.analytics.QueryRequestContext; +import org.opensearch.analytics.exec.task.AnalyticsQueryTask; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.tasks.Task; import java.io.IOException; +import java.util.Map; /** * Request carrying a logical query plan for analytics engine execution. @@ -30,14 +34,29 @@ */ public class AnalyticsQueryRequest extends ActionRequest implements IndicesRequest.Replaceable { + /** + * Placeholder query id for the task created at {@link #createTask}. The real query id is + * derived from the planned {@code QueryDAG} later in {@code DefaultPlanExecutor.executeInternal}, + * which runs after the framework has already created this task — so none is available yet. + * It only affects the task's description in the tasks API; the executor identifies the query + * by its own {@code dag.queryId()}, not by this task's id. + */ + private static final String QUERY_ID_NOT_YET_ASSIGNED = "unassigned"; + private final transient RelNode plan; private final transient QueryRequestContext queryCtx; + private final boolean profile; private String[] indices; public AnalyticsQueryRequest(RelNode plan, QueryRequestContext queryCtx, String[] indices) { + this(plan, queryCtx, indices, false); + } + + public AnalyticsQueryRequest(RelNode plan, QueryRequestContext queryCtx, String[] indices, boolean profile) { this.plan = plan; this.queryCtx = queryCtx; this.indices = indices; + this.profile = profile; } public AnalyticsQueryRequest(StreamInput in) throws IOException { @@ -58,6 +77,10 @@ public QueryRequestContext getQueryCtx() { return queryCtx; } + public boolean isProfile() { + return profile; + } + @Override public String[] indices() { return indices; @@ -74,6 +97,16 @@ public IndicesOptions indicesOptions() { return IndicesOptions.strictExpandOpen(); } + @Override + public Task createTask(long id, String type, String action, TaskId parentTaskId, Map headers) { + // TODO: assign the real query id here instead of QUERY_ID_NOT_YET_ASSIGNED. The id is + // currently minted later in DAGBuilder.newQueryId() (random UUID), so the task and the + // DAG/context id don't coordinate. Mint the id at request construction and thread it + // through to DAGBuilder.build() so the task description and DAG id match. Requires + // restructuring queryId ownership — tracked as a follow-up. + return new AnalyticsQueryTask(id, type, action, QUERY_ID_NOT_YET_ASSIGNED, parentTaskId, headers); + } + @Override public ActionRequestValidationException validate() { return null; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/AnalyticsQueryResponse.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/AnalyticsQueryResponse.java index 3b2af61f29d73..e9bf486554b14 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/AnalyticsQueryResponse.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/AnalyticsQueryResponse.java @@ -8,6 +8,7 @@ package org.opensearch.analytics.exec.action; +import org.opensearch.analytics.exec.profile.ProfiledResult; import org.opensearch.core.action.ActionResponse; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; @@ -21,9 +22,16 @@ public class AnalyticsQueryResponse extends ActionResponse { private final transient Iterable rows; + private final transient ProfiledResult profiledResult; public AnalyticsQueryResponse(Iterable rows) { this.rows = rows; + this.profiledResult = null; + } + + public AnalyticsQueryResponse(ProfiledResult profiledResult) { + this.rows = profiledResult.rows(); + this.profiledResult = profiledResult; } public AnalyticsQueryResponse(StreamInput in) throws IOException { @@ -39,4 +47,8 @@ public void writeTo(StreamOutput out) throws IOException { public Iterable getRows() { return rows; } + + public ProfiledResult getProfiledResult() { + return profiledResult; + } } diff --git a/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/analytics/cancellation/AnalyticsQueryTaskCleanupIT.java b/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/analytics/cancellation/AnalyticsQueryTaskCleanupIT.java new file mode 100644 index 0000000000000..9edce1c646ff5 --- /dev/null +++ b/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/analytics/cancellation/AnalyticsQueryTaskCleanupIT.java @@ -0,0 +1,259 @@ +/* + * 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.cancellation; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.Version; +import org.opensearch.action.admin.cluster.node.tasks.cancel.CancelTasksResponse; +import org.opensearch.action.admin.cluster.node.tasks.list.ListTasksResponse; +import org.opensearch.action.admin.indices.create.CreateIndexResponse; +import org.opensearch.analytics.AnalyticsPlugin; +import org.opensearch.analytics.exec.action.AnalyticsQueryAction; +import org.opensearch.analytics.exec.action.FragmentExecutionAction; +import org.opensearch.arrow.allocator.ArrowBasePlugin; +import org.opensearch.arrow.flight.transport.FlightStreamPlugin; +import org.opensearch.be.datafusion.DataFusionPlugin; +import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.common.util.FeatureFlags; +import org.opensearch.composite.CompositeDataFormatPlugin; +import org.opensearch.index.engine.dataformat.stub.MockCommitterEnginePlugin; +import org.opensearch.parquet.ParquetDataFormatPlugin; +import org.opensearch.plugins.Plugin; +import org.opensearch.plugins.PluginInfo; +import org.opensearch.ppl.TestPPLPlugin; +import org.opensearch.ppl.action.PPLRequest; +import org.opensearch.ppl.action.PPLResponse; +import org.opensearch.ppl.action.UnifiedPPLExecuteAction; +import org.opensearch.test.OpenSearchIntegTestCase; +import org.opensearch.test.transport.MockTransportService; +import org.opensearch.transport.TransportService; + +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +/** + * Verifies the framework-task lifecycle for the analytics query action + * ({@link AnalyticsQueryAction}, {@code indices:data/read/analytics/query}). + * + *

PR 9 routes every analytics query through {@code client.execute(AnalyticsQueryAction, ...)} + * and runs it under the framework-provided, cancellable {@code AnalyticsQueryTask} — instead of a + * detached task self-registered inside {@code DefaultPlanExecutor.executeInternal}. Because the task + * is now framework-owned, {@code HandledTransportAction} registers it before {@code doExecute} and + * unregisters it when the listener settles, and a client disconnect / explicit cancel of that task + * propagates into the running query. + * + *

These tests assert the two cleanup guarantees that depend on dropping the old manual + * register/unregister: (1) a successful query leaves no residual analytics/query task (the framework + * unregistered it), and (2) cancelling the analytics/query task terminates the query and leaves no + * residual analytics/query or fragment tasks. + */ +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 2, numClientNodes = 0, supportsDedicatedMasters = false) +@com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope(com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope.Scope.TEST) +@com.carrotsearch.randomizedtesting.annotations.ThreadLeakLingering(linger = 5000) +@com.carrotsearch.randomizedtesting.annotations.ThreadLeakFilters(filters = org.opensearch.analytics.resilience.FlightTransportThreadLeakFilter.class) +public class AnalyticsQueryTaskCleanupIT extends OpenSearchIntegTestCase { + + private static final Logger logger = LogManager.getLogger(AnalyticsQueryTaskCleanupIT.class); + + private static final String INDEX = "analytics_task_cleanup_idx"; + private static final int NUM_SHARDS = 2; + private static final int DOCS_PER_SHARD = 50; + private static final int TOTAL_DOCS = NUM_SHARDS * DOCS_PER_SHARD; + private static final int VALUE = 7; + private static final long EXPECTED_SUM = (long) TOTAL_DOCS * VALUE; + private static final TimeValue QUERY_TIMEOUT = TimeValue.timeValueSeconds(30); + + @Override + protected Collection> nodePlugins() { + return List.of( + ArrowBasePlugin.class, + TestPPLPlugin.class, + CompositeDataFormatPlugin.class, + MockTransportService.TestPlugin.class, + MockCommitterEnginePlugin.class + ); + } + + @Override + protected Collection additionalNodePlugins() { + return List.of( + classpathPlugin(FlightStreamPlugin.class, List.of(ArrowBasePlugin.class.getName())), + classpathPlugin(AnalyticsPlugin.class, Collections.emptyList()), + classpathPlugin(ParquetDataFormatPlugin.class, Collections.emptyList()), + classpathPlugin(DataFusionPlugin.class, List.of(AnalyticsPlugin.class.getName())) + ); + } + + private static PluginInfo classpathPlugin(Class pluginClass, List extendedPlugins) { + return new PluginInfo( + pluginClass.getName(), + "classpath plugin", + "NA", + Version.CURRENT, + "1.8", + pluginClass.getName(), + null, + extendedPlugins, + false + ); + } + + @Override + protected Settings nodeSettings(int nodeOrdinal) { + return Settings.builder() + .put(super.nodeSettings(nodeOrdinal)) + .put(FeatureFlags.PLUGGABLE_DATAFORMAT_EXPERIMENTAL_FLAG, true) + .put(FeatureFlags.STREAM_TRANSPORT, true) + .build(); + } + + // ---------------------------------------------------------------- fixture + + private void createAndSeedIndex() { + Settings indexSettings = Settings.builder() + .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, NUM_SHARDS) + .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0) + .put("index.pluggable.dataformat.enabled", true) + .put("index.pluggable.dataformat", "composite") + .put("index.composite.primary_data_format", "parquet") + .putList("index.composite.secondary_data_formats") + .build(); + + CreateIndexResponse response = client().admin() + .indices() + .prepareCreate(INDEX) + .setSettings(indexSettings) + .setMapping("value", "type=integer") + .get(); + assertTrue("index creation must be acknowledged", response.isAcknowledged()); + ensureGreen(INDEX); + + for (int i = 0; i < TOTAL_DOCS; i++) { + client().prepareIndex(INDEX).setSource("value", VALUE).get(); + } + client().admin().indices().prepareRefresh(INDEX).get(); + client().admin().indices().prepareFlush(INDEX).get(); + + try { + assertBusy(() -> { + PPLResponse r = executePPL("source = " + INDEX + " | stats sum(value) as total"); + long actual = ((Number) r.getRows().get(0)[r.getColumns().indexOf("total")]).longValue(); + assertEquals("seed not yet visible", EXPECTED_SUM, actual); + }, 30, TimeUnit.SECONDS); + } catch (Exception e) { + throw new AssertionError("timed out waiting for seed visibility", e); + } + } + + private PPLResponse executePPL(String ppl) { + return client().execute(UnifiedPPLExecuteAction.INSTANCE, new PPLRequest(ppl)).actionGet(); + } + + private PPLResponse executePPL(String ppl, TimeValue timeout) { + return client().execute(UnifiedPPLExecuteAction.INSTANCE, new PPLRequest(ppl)).actionGet(timeout); + } + + private void assertNoResidualTasks(String action) throws Exception { + assertBusy(() -> { + ListTasksResponse tasks = client().admin().cluster().prepareListTasks().setActions(action).get(); + assertTrue("residual " + action + " tasks: " + tasks.getTasks(), tasks.getTasks().isEmpty()); + }, 10, TimeUnit.SECONDS); + } + + // ---------------------------------------------------------------- tests + + /** + * A successful query must leave NO residual {@code AnalyticsQueryAction} task. This is the + * regression guard for dropping the manual {@code taskManager.unregister}: the framework + * unregisters the task it created for {@code doExecute} once the listener settles, so neither a + * leak (we kept manual unregister AND framework unregister → double-free) nor a dangling task + * (we dropped unregister but the framework doesn't own it) is acceptable. + */ + public void testSuccessfulQueryLeavesNoResidualAnalyticsTask() throws Exception { + createAndSeedIndex(); + + PPLResponse response = executePPL("source = " + INDEX + " | stats sum(value) as total", QUERY_TIMEOUT); + long actual = ((Number) response.getRows().get(0)[response.getColumns().indexOf("total")]).longValue(); + assertEquals("query must return the correct sum", EXPECTED_SUM, actual); + + assertNoResidualTasks(AnalyticsQueryAction.NAME); + assertNoResidualTasks(FragmentExecutionAction.NAME); + } + + /** + * Cancelling the framework {@code AnalyticsQueryAction} task terminates the in-flight query + * (no hang) and leaves no residual analytics/query or fragment tasks. This only works because + * the query now runs under the framework-provided task; pre-PR-9 the query ran under a detached + * self-registered task, so cancelling the framework action had no effect on it. + */ + public void testCancelAnalyticsQueryTaskTerminatesQueryAndCleansUp() throws Exception { + createAndSeedIndex(); + + // Block one data node's shard handler so cancellation lands while the query is in-flight. + String victim = randomFrom(internalCluster().getDataNodeNames()); + MockTransportService mts = (MockTransportService) internalCluster().getInstance(TransportService.class, victim); + CountDownLatch released = new CountDownLatch(1); + mts.addRequestHandlingBehavior(FragmentExecutionAction.NAME, (handler, request, channel, task) -> { + try { + released.await(QUERY_TIMEOUT.seconds(), TimeUnit.SECONDS); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + } + handler.messageReceived(request, channel, task); + }); + + ExecutorService exec = Executors.newSingleThreadExecutor(); + try { + Future fut = exec.submit(() -> executePPL("source = " + INDEX + " | stats sum(value) as total")); + // Allow dispatch + shard handler entry so the analytics/query task is live. + assertBusy(() -> { + ListTasksResponse live = client().admin().cluster().prepareListTasks().setActions(AnalyticsQueryAction.NAME).get(); + assertFalse("analytics/query task should be running", live.getTasks().isEmpty()); + }, 10, TimeUnit.SECONDS); + + CancelTasksResponse cancel = client().admin() + .cluster() + .prepareCancelTasks() + .setActions(AnalyticsQueryAction.NAME) + .get(); + assertFalse( + "cancel must not report node failures", + cancel.getNodeFailures() != null && cancel.getNodeFailures().isEmpty() == false + ); + + released.countDown(); + + try { + fut.get(QUERY_TIMEOUT.seconds(), TimeUnit.SECONDS); + } catch (ExecutionException | TimeoutException e) { + // Expected: the query terminated due to cancellation rather than completing. + logger.info("query terminated as expected after analytics/query cancel: {}", e.getMessage()); + } + } finally { + released.countDown(); + mts.clearAllRules(); + exec.shutdownNow(); + exec.awaitTermination(5, TimeUnit.SECONDS); + } + + assertNoResidualTasks(AnalyticsQueryAction.NAME); + assertNoResidualTasks(FragmentExecutionAction.NAME); + } +} From 0e5bee48789fdd9fe4e165449a8eabecc436c6a4 Mon Sep 17 00:00:00 2001 From: Sandesh Kumar Date: Sun, 31 May 2026 19:42:55 -0700 Subject: [PATCH 2/3] [analytics-engine] Wire PPL percentile_approx end-to-end (#21920) Routes PPL percentile_approx(field, p) through the analytics-engine to DataFusion's builtin approx_percentile_cont as a single-stage gather-on- coordinator aggregate. PplAggregateCallRewriter strips the PPL type-flag arg and normalises SymbolFlag literals to VARCHAR (isthmus's LiteralConverter rejects unregistered enum classes). DataFusionFragmentConvertor's literal-arg normaliser rescales the percentile from PPL's [0, 100] to DataFusion's [0, 1] convention at substrait emission. OpenSearchAggregateRule skips backend-viability checks on null-FieldType metadata arg columns. Adds CoordinatorReduceIT.testPercentileApproxAcrossShards. Signed-off-by: Sandesh Kumar --- .../analytics/spi/AggregateFunctionTests.java | 18 +++ .../DataFusionAnalyticsBackendPlugin.java | 1 + .../DataFusionFragmentConvertor.java | 69 ++++++++++- .../datafusion/PplAggregateCallRewriter.java | 111 +++++++++++++++--- .../opensearch_aggregate_functions.yaml | 11 ++ .../rules/OpenSearchAggregateRule.java | 9 ++ .../analytics/qa/CoordinatorReduceIT.java | 25 ++++ 7 files changed, 225 insertions(+), 19 deletions(-) diff --git a/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/AggregateFunctionTests.java b/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/AggregateFunctionTests.java index ec1fa306372d9..e7c85f6339abe 100644 --- a/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/AggregateFunctionTests.java +++ b/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/AggregateFunctionTests.java @@ -164,6 +164,24 @@ public void testValuesReducerIsSelf() { assertEquals(integer, resolve(fields.get(0), integer)); } + // ── PERCENTILE_APPROX: state-bearing, no decomposition declared ── + // + // DataFusion's t-digest state is multi-field and doesn't fit the single-field + // IntermediateField shape. OpenSearchAggregateSplitRule skips the partial/final split + // for STATE_EXPANDING aggregates, so PERCENTILE_APPROX runs single-stage on the + // coordinator after a singleton gather. + public void testPercentileApproxHasNoDecomposition() { + assertFalse(AggregateFunction.PERCENTILE_APPROX.hasDecomposition()); + assertNull(AggregateFunction.PERCENTILE_APPROX.intermediateFields()); + assertEquals(AggregateFunction.Type.STATE_EXPANDING, AggregateFunction.PERCENTILE_APPROX.getType()); + assertEquals(SqlKind.OTHER, AggregateFunction.PERCENTILE_APPROX.getSqlKind()); + } + + public void testPercentileApproxResolvesByName() { + assertSame(AggregateFunction.PERCENTILE_APPROX, AggregateFunction.fromNameOrError("percentile_approx")); + assertSame(AggregateFunction.PERCENTILE_APPROX, AggregateFunction.fromNameOrError("PERCENTILE_APPROX")); + } + // ── fromSqlKind still works ── public void testFromSqlKindResolvesExistingEntries() { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java index 6d17dadae9c41..956594ff8bf98 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java @@ -409,6 +409,7 @@ public class DataFusionAnalyticsBackendPlugin implements AnalyticsSearchBackendP AggregateFunction.COUNT, AggregateFunction.AVG, AggregateFunction.APPROX_COUNT_DISTINCT, + AggregateFunction.PERCENTILE_APPROX, AggregateFunction.TAKE, AggregateFunction.FIRST, AggregateFunction.LAST, diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java index 01d8b6be6c6cf..133802cddccda 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java @@ -25,6 +25,7 @@ import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.calcite.rel.type.RelDataTypeField; +import org.apache.calcite.rex.RexBuilder; import org.apache.calcite.rex.RexLiteral; import org.apache.calcite.rex.RexNode; import org.apache.calcite.schema.ColumnStrategy; @@ -37,6 +38,7 @@ import org.apache.calcite.sql.fun.SqlStdOperatorTable; import org.apache.calcite.sql.type.OperandTypes; import org.apache.calcite.sql.type.ReturnTypes; +import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.sql.type.SqlTypeTransforms; import org.apache.calcite.util.ImmutableBitSet; import org.apache.calcite.util.Optionality; @@ -49,6 +51,8 @@ import org.opensearch.be.datafusion.planner.adapter.NumericConversionFunctionAdapter; import org.opensearch.be.datafusion.planner.adapter.TimeConversionFunctionAdapter; +import java.math.BigDecimal; +import java.math.MathContext; import java.util.ArrayList; import java.util.List; import java.util.Optional; @@ -304,6 +308,30 @@ public class DataFusionFragmentConvertor implements FragmentConvertor { ) { }; + /** + * PPL {@code percentile_approx(field, percentile)} → DataFusion's builtin + * {@code approx_percentile_cont(field, percentile)}. PPL's trailing field-type-flag + * arg is stripped by {@link PplAggregateCallRewriter} before binding; the percentile + * literal is rescaled from PPL's [0, 100] to DataFusion's [0, 1] convention via + * {@link LocalAggOp#normaliseLiteralArg} at substrait emission. + */ + static final LocalAggOp LOCAL_PERCENTILE_APPROX_OP = new LocalAggOp( + "approx_percentile_cont", + SqlKind.OTHER_FUNCTION, + ReturnTypes.ARG0.andThen(SqlTypeTransforms.FORCE_NULLABLE), + OperandTypes.ANY_ANY + ) { + @Override + public RexNode normaliseLiteralArg(int argIndex, RexLiteral lit, RexBuilder rexBuilder, RelDataTypeFactory typeFactory) { + if (argIndex == 1 && lit.getValue() instanceof BigDecimal bd) { + BigDecimal scaled = bd.divide(BigDecimal.valueOf(100), MathContext.DECIMAL64); + RelDataType doubleType = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.DOUBLE), true); + return rexBuilder.makeLiteral(scaled, doubleType); + } + return lit; + } + }; + /** BRAIN window stub for {@code patterns ... method=BRAIN mode=label}. */ static final SqlAggFunction LOCAL_INTERNAL_PATTERN_WINDOW_OP = new SqlAggFunction( "internal_pattern", @@ -342,6 +370,7 @@ public class DataFusionFragmentConvertor implements FragmentConvertor { FunctionMappings.s(LOCAL_ARRAY_AGG_OP, "array_agg"), FunctionMappings.s(LOCAL_LIST_MERGE_OP, "list_merge"), FunctionMappings.s(LOCAL_LIST_MERGE_DISTINCT_OP, "list_merge_distinct"), + FunctionMappings.s(LOCAL_PERCENTILE_APPROX_OP, "approx_percentile_cont"), FunctionMappings.s(LOCAL_INTERNAL_PATTERN_OP, "internal_pattern") ); @@ -604,6 +633,7 @@ public Optional convert( List projects = project.getProjects(); List args = fn.arguments(); List rewritten = null; + RexBuilder rexBuilder = project.getCluster().getRexBuilder(); for (int i = 0; i < args.size(); i++) { FunctionArg arg = args.get(i); if (!(arg instanceof io.substrait.expression.FieldReference fr)) continue; @@ -611,7 +641,10 @@ public Optional convert( if (offset == null || offset < 0 || offset >= projects.size()) continue; if (!(projects.get(offset) instanceof RexLiteral rexLit)) continue; if (rewritten == null) rewritten = new ArrayList<>(args); - rewritten.set(i, rexConverter.apply(rexLit)); + RexNode toConvert = call.getAggregation() instanceof LocalAggOp localOp + ? localOp.normaliseLiteralArg(i, rexLit, rexBuilder, typeFactory) + : rexLit; + rewritten.set(i, rexConverter.apply(toConvert)); } if (rewritten == null) return bound; return Optional.of(ImmutableAggregateFunctionInvocation.builder().from(fn).arguments(rewritten).build()); @@ -652,6 +685,40 @@ private static Integer simpleStructOffset(io.substrait.expression.FieldReference return sf.offset(); } + /** + * Local aggregate stub that may transform inlined literal args before substrait emission. + * Other local stubs without transformations stay as plain {@link SqlAggFunction}; the + * {@code convert()} override only invokes {@link #normaliseLiteralArg} when the call's + * operator is a {@code LocalAggOp}, so adding a new normalisation is purely a matter of + * subclassing here next to the op's declaration. + */ + abstract static class LocalAggOp extends SqlAggFunction { + LocalAggOp( + String name, + SqlKind kind, + org.apache.calcite.sql.type.SqlReturnTypeInference returnTypeInference, + org.apache.calcite.sql.type.SqlOperandTypeChecker operandTypeChecker + ) { + super( + name, + null, + kind, + returnTypeInference, + null, + operandTypeChecker, + SqlFunctionCategory.USER_DEFINED_FUNCTION, + false, + false, + Optionality.FORBIDDEN + ); + } + + /** Identity by default; override to transform the {@code argIndex}-th inlined literal arg. */ + public RexNode normaliseLiteralArg(int argIndex, RexLiteral lit, RexBuilder rexBuilder, RelDataTypeFactory typeFactory) { + return lit; + } + } + // ── Plan serde helpers ────────────────────────────────────────────────────── /** Decodes serialized Substrait bytes into a model-level {@link Plan}. */ diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java index c4a019c5b23e7..89de1e50a9ea8 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java @@ -12,8 +12,13 @@ import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.core.Aggregate; import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.core.Project; +import org.apache.calcite.rel.logical.LogicalProject; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexLiteral; +import org.apache.calcite.rex.RexNode; import org.apache.calcite.sql.SqlAggFunction; import org.apache.calcite.sql.type.SqlTypeName; @@ -21,7 +26,14 @@ import java.util.List; import java.util.Set; -/** Rewrites PPL state-expanding aggregates (TAKE/FIRST/LAST/LIST/VALUES/PATTERN) onto local stubs. */ +/** + * Rewrites PPL state-expanding aggregates (TAKE / FIRST / LAST / LIST / VALUES / PATTERN / + * PERCENTILE_APPROX) onto local stubs the substrait emitter binds via + * {@link DataFusionFragmentConvertor}'s ADDITIONAL_AGGREGATE_SIGS. Also normalises any + * RexLiteral{SqlTypeName.SYMBOL} in upstream Projects to VARCHAR — isthmus's + * LiteralConverter rejects unregistered Enum classes, and PPL's percentile_approx / + * median emit a SymbolFlag arg purely for type inference. + */ final class PplAggregateCallRewriter { private static final Set LOCAL_OPS = Set.of( @@ -31,6 +43,7 @@ final class PplAggregateCallRewriter { DataFusionFragmentConvertor.LOCAL_ARRAY_AGG_OP, DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP, DataFusionFragmentConvertor.LOCAL_LIST_MERGE_DISTINCT_OP, + DataFusionFragmentConvertor.LOCAL_PERCENTILE_APPROX_OP, DataFusionFragmentConvertor.LOCAL_INTERNAL_PATTERN_OP ); @@ -41,29 +54,65 @@ static RelNode rewrite(RelNode root) { @Override public RelNode visit(RelNode other) { RelNode visited = super.visit(other); - if (!(visited instanceof Aggregate agg)) { - return visited; - } - List oldCalls = agg.getAggCallList(); - List newCalls = new ArrayList<>(oldCalls.size()); - boolean changed = false; - for (AggregateCall call : oldCalls) { - AggregateCall rewritten = rewriteCall(agg, call); - if (rewritten == call) { - newCalls.add(call); - } else { - newCalls.add(rewritten); - changed = true; - } + if (visited instanceof Project p) { + return normaliseSymbolFlagLiterals(p); } - if (!changed) { - return visited; + if (visited instanceof Aggregate agg) { + return rewriteAggregate(agg); } - return agg.copy(agg.getTraitSet(), agg.getInput(), agg.getGroupSet(), agg.getGroupSets(), newCalls); + return visited; } }); } + private static RelNode rewriteAggregate(Aggregate agg) { + List oldCalls = agg.getAggCallList(); + List newCalls = new ArrayList<>(oldCalls.size()); + boolean changed = false; + for (AggregateCall call : oldCalls) { + AggregateCall rewritten = rewriteCall(agg, call); + if (rewritten == call) { + newCalls.add(call); + } else { + newCalls.add(rewritten); + changed = true; + } + } + if (!changed) { + return agg; + } + return agg.copy(agg.getTraitSet(), agg.getInput(), agg.getGroupSet(), agg.getGroupSets(), newCalls); + } + + /** Replace any RexLiteral{SymbolFlag} in {@code project}'s projection list with a VARCHAR literal of the symbol's name. */ + private static RelNode normaliseSymbolFlagLiterals(Project project) { + List oldProjects = project.getProjects(); + boolean hasSymbol = oldProjects.stream() + .anyMatch(p -> p instanceof RexLiteral lit && lit.getType().getSqlTypeName() == SqlTypeName.SYMBOL); + if (!hasSymbol) { + return project; + } + RelDataTypeFactory typeFactory = project.getCluster().getTypeFactory(); + RexBuilder rexBuilder = project.getCluster().getRexBuilder(); + RelDataType varcharType = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.VARCHAR), true); + List newProjects = new ArrayList<>(oldProjects.size()); + for (RexNode p : oldProjects) { + if (p instanceof RexLiteral lit && lit.getType().getSqlTypeName() == SqlTypeName.SYMBOL) { + String name = lit.getValue() == null ? "" : lit.getValue().toString(); + newProjects.add(rexBuilder.makeLiteral(name, varcharType)); + } else { + newProjects.add(p); + } + } + return LogicalProject.create( + project.getInput(), + project.getHints(), + newProjects, + project.getRowType().getFieldNames(), + project.getVariablesSet() + ); + } + private static AggregateCall rewriteCall(Aggregate agg, AggregateCall call) { SqlAggFunction aggregation = call.getAggregation(); if (LOCAL_OPS.contains(aggregation)) { @@ -101,6 +150,32 @@ private static AggregateCall rewriteCall(Aggregate agg, AggregateCall call) { targetOp = DataFusionFragmentConvertor.LOCAL_INTERNAL_PATTERN_OP; explicitReturnType = internalPatternReturnType(agg.getCluster().getTypeFactory()); } + case "PERCENTILE_APPROX" -> { + // Trim the PPL type-flag arg; the substrait emit-time literal-arg normaliser + // (DataFusionFragmentConvertor#normaliseLiteralArg) rescales the percentile. + if (call.getArgList().size() < 3) { + return call; + } + targetOp = DataFusionFragmentConvertor.LOCAL_PERCENTILE_APPROX_OP; + List trimmedArgList = new ArrayList<>(call.getArgList().subList(0, 2)); + RelDataType arg0Type = agg.getInput().getRowType().getFieldList().get(call.getArgList().get(0)).getType(); + RelDataType nullableArg0 = agg.getCluster().getTypeFactory().createTypeWithNullability(arg0Type, true); + return AggregateCall.create( + targetOp, + targetDistinct, + call.isApproximate(), + call.ignoreNulls(), + call.rexList, + trimmedArgList, + call.filterArg, + call.distinctKeys, + call.collation, + agg.getGroupCount(), + agg.getInput(), + nullableArg0, + call.getName() + ); + } default -> { return call; } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_aggregate_functions.yaml b/sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_aggregate_functions.yaml index 05130a163690d..f08c03f017c97 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_aggregate_functions.yaml +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/resources/opensearch_aggregate_functions.yaml @@ -92,6 +92,17 @@ aggregate_functions: - name: x value: any return: any + - name: approx_percentile_cont + description: >- + PPL `percentile_approx(field, percentile)` — t-digest based approximate + percentile. Maps to DataFusion's builtin `approx_percentile_cont`. + impls: + - args: + - name: x + value: any1 + - name: percentile + value: any2 + return: any1 - name: internal_pattern description: >- PPL `patterns ... method=BRAIN` (aggregation mode). Custom Rust UDAF in diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java index 352e13e98c5a5..e988eaf9364e3 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateRule.java @@ -138,6 +138,15 @@ private List resolveViableBackendsForCall(AggregateCall aggCall, List callViable = null; for (int fieldIndex : aggCall.getArgList()) { + // Skip metadata-only literal arg columns whose FieldType is null (e.g. SYMBOL — + // PPL's percentile_approx / median type-flag). Only data-field args need a + // backend viability check. + if (fieldIndex < childFieldStorageInfos.size()) { + FieldStorageInfo peek = childFieldStorageInfos.get(fieldIndex); + if (peek.isDerived() && peek.getFieldType() == null) { + continue; + } + } FieldStorageInfo storageInfo = FieldStorageInfo.resolve(childFieldStorageInfos, fieldIndex); FieldType fieldType = storageInfo.getFieldType(); diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CoordinatorReduceIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CoordinatorReduceIT.java index 30e9e00227cc7..4e3f00126706a 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CoordinatorReduceIT.java +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/CoordinatorReduceIT.java @@ -127,6 +127,31 @@ public void testDistinctCountAcrossShards() throws Exception { ); } + /** + * {@code stats percentile_approx(value, 50) as p} — t-digest approximate median. + * STATE_EXPANDING, so the split rule gathers to coordinator + single-stage. Maps to + * DataFusion's {@code approx_percentile_cont} via {@link + * org.opensearch.be.datafusion.PplAggregateCallRewriter}. + */ + public void testPercentileApproxAcrossShards() throws Exception { + String index = "coord_reduce_percentile_approx"; + createParquetBackedIndex(index); + indexVaryingValueDocs(index); + + Map result = executePpl("source = " + index + " | stats percentile_approx(value, 50) as p"); + List> rows = scalarRows(result, "p"); + + Object cell = rows.get(0).get(0); + assertNotNull("cell for 'p' must not be null — coordinator-reduce returned no value", cell); + double actual = ((Number) cell).doubleValue(); + int totalDocs = NUM_SHARDS * DOCS_PER_SHARD; + double expected = (totalDocs + 1) / 2.0; + assertTrue( + "percentile_approx(value, 50) should be approximately " + expected + " (±2.0), got " + actual, + Math.abs(actual - expected) <= 2.0 + ); + } + /** Single-shard {@code take(value, 3)} — bounded array of up to 3 values. */ public void testTakeSingleShard() throws Exception { String index = "coord_reduce_take_single"; From 7b7a65b36d873fba56ee2bdb34ec39f55467f056 Mon Sep 17 00:00:00 2001 From: Marc Handalian Date: Sun, 31 May 2026 19:49:42 -0700 Subject: [PATCH 3/3] Streaming coordinator-reduce: cancellation, memory-release, and correctness fixes Squashed benchmark-test work (PR #371), rebased on upstream/main and with review feedback applied. Excludes PR #21899 (flight backpressure) and PR #21887 (export-allocation cap). - Framework AnalyticsQueryTask for HTTP-disconnect cancellation - Reduce-input early cancel: tolerate receiver-drop, cancel stream when consumer satisfied - Fix streaming reduce deadlock: blocking_send instead of handle.block_on - RowProducingSink truncates result set at 10k rows to prevent OOM - Close session/senders/outStream on reduce teardown + child-FAILED to fix native leaks - Configurable datafusion.reduce.target_partitions setting --- .gitignore | 2 + .../analytics/spi/ExchangeSink.java | 10 ++ .../rust/src/api.rs | 38 ++++-- .../rust/src/cross_rt_stream.rs | 8 ++ .../rust/src/ffm.rs | 15 ++- .../rust/src/local_executor.rs | 10 +- .../rust/src/partition_stream.rs | 41 ++++--- .../AbstractDatafusionReduceSink.java | 21 +--- .../DataFusionFragmentConvertor.java | 11 -- .../be/datafusion/DataFusionPlugin.java | 16 +++ .../datafusion/DatafusionPartitionSender.java | 34 +++++- .../be/datafusion/DatafusionReduceSink.java | 113 +++++++++++------- .../be/datafusion/DatafusionSettings.java | 1 + .../be/datafusion/nativelib/NativeBridge.java | 28 ++++- .../datafusion/DatafusionReduceSinkTests.java | 19 +++ .../exec/AnalyticsSearchService.java | 8 +- .../exec/AnalyticsSearchTransportService.java | 33 +++-- .../exec/StreamingResponseListener.java | 5 +- .../LateMaterializationStageExecution.java | 5 +- .../analytics/exec/stage/StageExecution.java | 1 + .../PassThroughStageExecution.java | 12 +- .../coordinator/ReduceStageExecution.java | 24 ++-- .../shard/ShardFragmentStageExecution.java | 17 ++- .../ShardFragmentStageExecutionTests.java | 32 +++++ sandbox/qa/analytics-engine-rest/build.gradle | 1 - .../org/opensearch/analytics/qa/AliasIT.java | 36 ++++++ .../opensearch/analytics/qa/DataStreamIT.java | 20 ++++ .../analytics/qa/FilterDelegationIT.java | 16 +++ .../analytics/qa/IndexPatternUnionIT.java | 28 +++++ .../analytics/qa/MultiIndexQueryShapesIT.java | 70 +++++++++++ 30 files changed, 519 insertions(+), 156 deletions(-) diff --git a/.gitignore b/.gitignore index 308761cf3006f..0accb5de52da7 100644 --- a/.gitignore +++ b/.gitignore @@ -75,3 +75,5 @@ doc-tools/missing-doclet/bin/ **/Cargo.lock /sandbox/plugins/analytics-backend-datafusion/target/ /sandbox/libs/dataformat-native/rust/target +.docker-gradle-home/ +scripts/docker-crash/ diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSink.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSink.java index ea98e0e74d701..0b17aecccbfbe 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSink.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSink.java @@ -50,6 +50,16 @@ default void feed(VectorSchemaRoot batch, int sourceOrdinal) { feed(batch); } + /** + * Whether the downstream consumer has finished and will read no more batches (e.g. a reduce + * whose LimitExec satisfied its fetch). Producers may poll this after a {@link #feed} to stop + * early. Default {@code false}; best-effort — a {@code true} is monotonic, a {@code false} may + * race a concurrent completion. + */ + default boolean isConsumerDone() { + return false; + } + /** * Signal that no more batches will be fed. Releases resources. */ diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs index fa79a3edab101..1732dc1d11b83 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs @@ -323,7 +323,7 @@ pub fn create_global_runtime( let memory_pool = Arc::new(TrackConsumersPool::new( dynamic_pool, NonZeroUsize::new(5).unwrap(), - )); + )) as Arc; let (cache_manager_config, custom_cache_manager) = if cache_manager_ptr != 0 { let mgr = unsafe { *Box::from_raw(cache_manager_ptr as *mut CustomCacheManager) }; @@ -406,6 +406,18 @@ pub fn set_min_target_partitions(value: i64) { crate::query_budget::set_min_target_partitions(value.max(1) as usize); } +/// Initial target_partitions for coordinator-reduce sessions. Defaults to 4. +static REDUCE_TARGET_PARTITIONS: std::sync::atomic::AtomicUsize = + std::sync::atomic::AtomicUsize::new(4); + +pub fn set_reduce_target_partitions(value: i64) { + REDUCE_TARGET_PARTITIONS.store(value.max(1).min(32) as usize, std::sync::atomic::Ordering::Release); +} + +pub fn get_reduce_target_partitions() -> usize { + REDUCE_TARGET_PARTITIONS.load(std::sync::atomic::Ordering::Acquire) +} + /// Creates a native reader (ShardView) for the given path and files. /// /// Returns a heap-allocated pointer (as i64) to `ShardView`. @@ -908,6 +920,7 @@ pub async unsafe fn stream_next( } else { batch }; + let struct_array: StructArray = batch.into(); let array_data = struct_array.into_data(); let ffi_array = FFI_ArrowArray::new(&array_data); @@ -978,9 +991,10 @@ fn view_needs_gc(buffers: &[arrow::buffer::Buffer], bytes_used: usize) -> bool { /// `stream_ptr` must be 0 or a valid pointer returned by `execute_query`. pub unsafe fn stream_close(stream_ptr: i64) { if stream_ptr != 0 { - // Dropping the handle drops both the stream and the query context. - // The context's Drop impl marks the query completed in the registry. - let _ = Box::from_raw(stream_ptr as *mut QueryStreamHandle); + native_bridge_common::log_debug!("[stream-close] dropping QueryStreamHandle ptr={:#x}", stream_ptr); + let handle = Box::from_raw(stream_ptr as *mut QueryStreamHandle); + drop(handle); + native_bridge_common::log_debug!("[stream-close] QueryStreamHandle dropped ptr={:#x}", stream_ptr); } } @@ -1317,11 +1331,10 @@ pub(crate) fn base_schema_for_table(plan: &substrait::proto::Plan, table_name: & /// `create_global_runtime`. pub unsafe fn create_local_session(runtime_ptr: i64) -> Result { let runtime = &*(runtime_ptr as *const DataFusionRuntime); - // No phantom reservation at creation time — the schema isn't known yet. - // register_partition_stream acquires a schema-accurate phantom once the - // output schema is derived from the producer plan. let session = LocalSession::new(&runtime.runtime_env); - Ok(Box::into_raw(Box::new(session)) as i64) + let ptr = Box::into_raw(Box::new(session)) as i64; + native_bridge_common::log_debug!("[local-session] OPEN ptr={:#x}", ptr); + Ok(ptr) } /// Closes a `LocalSession`. Safe to call with 0 (no-op). @@ -1330,7 +1343,10 @@ pub unsafe fn create_local_session(runtime_ptr: i64) -> Result Result<(), DataFusionError> { +) -> Result { let sender = &*(sender_ptr as *const PartitionStreamSender); // Take ownership of the Java-allocated FFI structs. `from_raw` reads @@ -1525,7 +1541,7 @@ pub unsafe fn sender_send( let struct_array = StructArray::from(array_data); let batch = RecordBatch::from(struct_array); - sender.send_blocking(Ok(batch), io_handle) + Ok(sender.send_blocking(Ok(batch), io_handle)) } /// Closes a partition stream sender. Dropping the sender closes the mpsc, diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/cross_rt_stream.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/cross_rt_stream.rs index b44f53eb99b18..23225229b2994 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/cross_rt_stream.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/cross_rt_stream.rs @@ -36,6 +36,12 @@ pub struct CrossRtStream { phantom_corrector: Option>, } +impl Drop for CrossRtStream { + fn drop(&mut self) { + native_bridge_common::log_debug!("[cross-rt-stream] DROP"); + } +} + impl CrossRtStream { fn new_with_tx(f: F, schema: SchemaRef) -> Self where @@ -77,9 +83,11 @@ impl CrossRtStream { tokio::pin!(stream); while let Some(res) = stream.next().await { if tx_captured.send(res).await.is_err() { + native_bridge_common::log_debug!("[cross-rt-stream] CPU task: receiver gone, exiting"); return; } } + native_bridge_common::log_debug!("[cross-rt-stream] CPU task: stream exhausted normally"); }; let (abort_handle, join_fut) = exec.spawn_with_abort_handle(fut); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs index 7f83d64021fbc..31429b92fcb45 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs @@ -157,6 +157,11 @@ pub extern "C" fn df_set_min_target_partitions(value: i64) { api::set_min_target_partitions(value); } +#[no_mangle] +pub extern "C" fn df_set_reduce_target_partitions(value: i64) { + api::set_reduce_target_partitions(value); +} + /// Sets memory guard thresholds. Values are thresholds multiplied by 1000 /// (e.g., 700 = 0.70, 850 = 0.85, 950 = 0.95). #[no_mangle] @@ -542,12 +547,20 @@ pub unsafe extern "C" fn df_execute_local_plan( .map_err(|e| e.to_string()) } +/// Positive FFI sentinel from [`df_sender_send`] when the send was skipped because the consumer +/// finished (receiver dropped). Rides the success half of the return contract (`>= 0` = success), +/// so `checkResult` passes it through. MUST match `NativeBridge.SENDER_SEND_RECEIVER_DROPPED`. +pub const SENDER_SEND_RECEIVER_DROPPED: i64 = 1; + #[ffm_safe] #[no_mangle] pub unsafe extern "C" fn df_sender_send(sender_ptr: i64, array_ptr: i64, schema_ptr: i64) -> i64 { let mgr = get_rt_manager()?; api::sender_send(sender_ptr, array_ptr, schema_ptr, mgr.io_runtime.handle()) - .map(|_| 0) + .map(|outcome| match outcome { + crate::partition_stream::SendOutcome::Sent => 0, + crate::partition_stream::SendOutcome::ReceiverDropped => SENDER_SEND_RECEIVER_DROPPED, + }) .map_err(|e| e.to_string()) } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs index bf910273854eb..d0b58398c5602 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs @@ -71,7 +71,7 @@ impl LocalSession { pub fn new(runtime_env: &RuntimeEnv) -> Self { let runtime_env = Arc::new(runtime_env.clone()); let mut config = SessionConfig::new(); - config.options_mut().execution.target_partitions = 4; + config.options_mut().execution.target_partitions = crate::api::get_reduce_target_partitions(); let state = SessionStateBuilder::new() .with_config(config) .with_runtime_env(runtime_env) @@ -217,16 +217,15 @@ impl LocalSession { )) })?; let logical_plan = from_substrait_plan(&self.ctx.state(), &plan).await?; - log_debug!("DataFusion logical plan (reduce):\n{}", logical_plan.display_indent()); let dataframe = self.ctx.execute_logical_plan(logical_plan).await?; let physical_plan = dataframe.create_physical_plan().await?; let target_schema = crate::schema_coerce::coerce_inferred_schema(physical_plan.schema()); let physical_plan = crate::relabel_exec::wrap_if_relabel_needed(physical_plan, target_schema)?; - log_debug!("DataFusion physical plan (reduce):\n{}", displayable(physical_plan.as_ref()).indent(true)); let stripped = crate::agg_mode::apply_aggregate_mode( physical_plan, crate::agg_mode::Mode::Final, )?; + log_debug!("DataFusion physical plan (reduce):\n{}", displayable(stripped.as_ref()).indent(true)); self.prepared_plan = Some(stripped); Ok(()) } @@ -330,9 +329,8 @@ mod tests { let handle = Handle::current(); let producer = std::thread::spawn(move || { for chunk in &[vec![1i64, 2, 3], vec![4, 5, 6], vec![7, 8, 9]] { - sender - .send_blocking(Ok(i64_batch(&producer_schema, chunk)), &handle) - .expect("send"); + let outcome = sender.send_blocking(Ok(i64_batch(&producer_schema, chunk)), &handle); + assert!(matches!(outcome, crate::partition_stream::SendOutcome::Sent)); } drop(sender); // EOF }); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/partition_stream.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/partition_stream.rs index fc2ac5cc04c0b..ada38fa122b7b 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/partition_stream.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/partition_stream.rs @@ -60,6 +60,16 @@ pub struct PartitionStreamSender { schema: SchemaRef, } +/// Outcome of a blocking send. `ReceiverDropped` is the benign terminal case — the +/// DataFusion consumer finished (e.g. a `LimitExec` satisfied its fetch) and dropped the +/// receiver. Surfaced as a distinct variant so the FFM layer can signal it without the Java +/// side substring-matching an error message. +#[must_use] +pub enum SendOutcome { + Sent, + ReceiverDropped, +} + impl PartitionStreamSender { /// Returns the schema this sender was created with. pub fn schema(&self) -> &SchemaRef { @@ -73,16 +83,18 @@ impl PartitionStreamSender { /// bridge push without being async itself and without requiring the calling /// thread to be a Tokio worker. /// - /// Blocks while the channel is full (natural backpressure). Returns an - /// error only if the receiver has been dropped. + /// Blocks while the channel is full (natural backpressure). Returns + /// [`SendOutcome::ReceiverDropped`] only if the receiver has been dropped — + /// the sole failure mode of `mpsc::Sender::send`. pub fn send_blocking( &self, batch: Result, handle: &Handle, - ) -> Result<(), DataFusionError> { - handle.block_on(self.tx.send(batch)).map_err(|_| { - DataFusionError::Execution("partition stream receiver dropped before send".to_string()) - }) + ) -> SendOutcome { + match handle.block_on(self.tx.send(batch)) { + Ok(()) => SendOutcome::Sent, + Err(_) => SendOutcome::ReceiverDropped, + } } } @@ -259,9 +271,8 @@ mod tests { let sender_schema = Arc::clone(&schema); let producer = std::thread::spawn(move || { - sender - .send_blocking(Ok(test_batch(&sender_schema, &[7, 8, 9])), &handle) - .unwrap(); + let outcome = sender.send_blocking(Ok(test_batch(&sender_schema, &[7, 8, 9])), &handle); + assert!(matches!(outcome, SendOutcome::Sent)); drop(sender); }); @@ -280,14 +291,10 @@ mod tests { let (sender, receiver) = channel(Arc::clone(&schema)); drop(receiver); - let err = std::thread::spawn(move || { - sender - .send_blocking(Ok(test_batch(&schema, &[1])), &handle) - .unwrap_err() - }) - .join() - .unwrap(); - assert!(err.to_string().contains("receiver dropped")); + let outcome = std::thread::spawn(move || sender.send_blocking(Ok(test_batch(&schema, &[1])), &handle)) + .join() + .unwrap(); + assert!(matches!(outcome, SendOutcome::ReceiverDropped)); } #[tokio::test] diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java index 4a67308d1a239..177a95fdc3170 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java @@ -12,10 +12,8 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; -import org.opensearch.analytics.spi.CancellableExchangeSink; import org.opensearch.analytics.spi.ExchangeSinkContext; import org.opensearch.analytics.spi.ReducingExchangeSink; -import org.opensearch.be.datafusion.nativelib.NativeBridge; import org.opensearch.be.datafusion.nativelib.StreamHandle; import java.util.LinkedHashMap; @@ -43,7 +41,7 @@ * * @opensearch.internal */ -abstract class AbstractDatafusionReduceSink implements ReducingExchangeSink, CancellableExchangeSink { +abstract class AbstractDatafusionReduceSink implements ReducingExchangeSink { /** Single-input shortcut for the per-child table id; multi-input uses {@link #inputIdFor(int)}. */ static final String INPUT_ID = "input-0"; @@ -96,20 +94,6 @@ protected static String inputIdFor(int childStageId) { return "input-" + childStageId; } - /** - * Fires the Rust CancellationToken for this query so that {@code stream_next} - * returns the sentinel {@code 0} immediately and the drain unblocks without - * waiting for DataFusion to finish naturally. No-op when {@code taskId} is 0 - * (no context registered) or when already closed. - */ - @Override - public final void cancel() { - long taskId = ctx.taskId(); - if (taskId != 0L) { - NativeBridge.cancelQuery(taskId); - } - } - @Override public void feed(VectorSchemaRoot batch) { synchronized (feedLock) { @@ -163,7 +147,8 @@ protected final void drainOutputIntoDownstream(StreamHandle outStream) { try (CDataDictionaryProvider dictProvider = new CDataDictionaryProvider()) { DatafusionResultStream.BatchIterator it = new DatafusionResultStream.BatchIterator(outStream, alloc, dictProvider); while (it.hasNext()) { - ctx.downstream().feed(it.next().getArrowRoot()); + VectorSchemaRoot root = it.next().getArrowRoot(); + ctx.downstream().feed(root); } } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java index 133802cddccda..0f0890ddab630 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java @@ -110,19 +110,8 @@ public class DataFusionFragmentConvertor implements FragmentConvertor { SqlFunctionCategory.USER_DEFINED_FUNCTION ); - /** TopK reduce expression: evaluates opaque aggregate state to a sortable scalar. */ - static final SqlOperator REDUCE_EVAL_OP = new SqlFunction( - "reduce_eval", - SqlKind.OTHER_FUNCTION, - ReturnTypes.BIGINT_NULLABLE, - null, - OperandTypes.ANY_ANY, - SqlFunctionCategory.USER_DEFINED_FUNCTION - ); - private static final List ADDITIONAL_SCALAR_SIGS = List.of( FunctionMappings.s(DelegatedPredicateFunction.FUNCTION, DelegatedPredicateFunction.NAME), - FunctionMappings.s(REDUCE_EVAL_OP, "reduce_eval"), FunctionMappings.s(DelegationPossibleFunction.FUNCTION, DelegationPossibleFunction.NAME), FunctionMappings.s(SqlStdOperatorTable.ASCII, "ascii"), FunctionMappings.s(SqlStdOperatorTable.CHAR_LENGTH, "length"), diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java index 03f83c5da149b..251a49e52ea88 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionPlugin.java @@ -281,6 +281,20 @@ static String deriveSpillLimitDefault() { Setting.Property.Dynamic ); + /** + * Number of partitions used by the coordinator-reduce DataFusion plan. + * More partitions = more parallelism = more memory (each partition holds its own hash table). + * Lower values reduce peak memory at the cost of slower single-query latency. + */ + public static final Setting DATAFUSION_REDUCE_TARGET_PARTITIONS = Setting.intSetting( + "datafusion.reduce.target_partitions", + 4, + 1, + 32, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + private static final String SUPPORTED_FORMAT = "parquet"; /** @@ -345,6 +359,7 @@ public Collection createComponents( clusterService.getClusterSettings().addSettingsUpdateConsumer(DATAFUSION_MEMORY_POOL_LIMIT, this::updateMemoryPoolLimit); clusterService.getClusterSettings().addSettingsUpdateConsumer(DATAFUSION_SPILL_MEMORY_LIMIT, this::updateSpillMemoryLimit); clusterService.getClusterSettings().addSettingsUpdateConsumer(DATAFUSION_MIN_TARGET_PARTITIONS, this::updateMinTargetPartitions); + clusterService.getClusterSettings().addSettingsUpdateConsumer(DATAFUSION_REDUCE_TARGET_PARTITIONS, NativeBridge::setReduceTargetPartitions); clusterService.getClusterSettings() .addSettingsUpdateConsumer(DATAFUSION_MEMORY_GUARD_ADMISSION_THROTTLE_THRESHOLD, v -> updateMemoryGuardThresholds()); clusterService.getClusterSettings() @@ -356,6 +371,7 @@ public Collection createComponents( // Apply initial values NativeBridge.setMinTargetPartitions(DATAFUSION_MIN_TARGET_PARTITIONS.get(settings)); + NativeBridge.setReduceTargetPartitions(DATAFUSION_REDUCE_TARGET_PARTITIONS.get(settings)); NativeBridge.setMemoryGuardThresholds( DATAFUSION_MEMORY_GUARD_ADMISSION_THROTTLE_THRESHOLD.get(settings), DATAFUSION_MEMORY_GUARD_ADMISSION_REJECT_THRESHOLD.get(settings), diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionPartitionSender.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionPartitionSender.java index dd1a0d26989ce..ead1dd9414977 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionPartitionSender.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionPartitionSender.java @@ -8,6 +8,8 @@ package org.opensearch.be.datafusion; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; import org.opensearch.analytics.backend.jni.NativeHandle; import org.opensearch.be.datafusion.nativelib.NativeBridge; @@ -24,27 +26,53 @@ */ public final class DatafusionPartitionSender extends NativeHandle { + private static final Logger logger = LogManager.getLogger(DatafusionPartitionSender.class); private final ReentrantReadWriteLock lifecycle = new ReentrantReadWriteLock(); + /** + * Latched once a send reports {@link NativeBridge#SENDER_SEND_RECEIVER_DROPPED} — the + * consumer (e.g. a LimitExec above the ExchangeReducer) satisfied its fetch and tore down + * this channel's receiver. Monotonic; once set, no further batch on this channel will be + * consumed. Per-sender (not per-sink) so a multi-input reduce only stops the input whose + * receiver is actually gone. + */ + private volatile boolean receiverDropped; + public DatafusionPartitionSender(long senderPtr) { super(senderPtr); } - public void send(long arrayAddr, long schemaAddr) { + /** + * Sends one exported batch. Returns {@code 0} on a normal send or + * {@link NativeBridge#SENDER_SEND_RECEIVER_DROPPED} if the consumer already dropped the + * receiver (benign — the caller should discard the batch and stop feeding). + */ + public long send(long arrayAddr, long schemaAddr) { lifecycle.readLock().lock(); try { - NativeBridge.senderSend(getPointer(), arrayAddr, schemaAddr); + long rc = NativeBridge.senderSend(getPointer(), arrayAddr, schemaAddr); + if (rc == NativeBridge.SENDER_SEND_RECEIVER_DROPPED) { + receiverDropped = true; + } + return rc; } finally { lifecycle.readLock().unlock(); } } + /** True once the consumer dropped this channel's receiver (see {@link #receiverDropped}). */ + public boolean isReceiverDropped() { + return receiverDropped; + } + @Override public void close() { lifecycle.writeLock().lock(); try { - assert lifecycle.isWriteLockedByCurrentThread() : "close must hold the write lock across super.close()"; super.close(); + if (logger.isTraceEnabled()) { + logger.trace("[sender] closed, native EOF signalled, ptr={}", ptr); + } } finally { lifecycle.writeLock().unlock(); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java index 1a47b903329f2..68ea2290bc474 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java @@ -87,15 +87,14 @@ public DatafusionReduceSink(ExchangeSinkContext ctx, NativeRuntimeHandle runtime public DatafusionReduceSink(ExchangeSinkContext ctx, NativeRuntimeHandle runtimeHandle, DataFusionReduceState preparedState) { super(ctx, runtimeHandle, preparedState); + logger.debug("[reduce-sink] OPEN taskId={} hasPreparedState={} sessionPtr={}", + ctx.taskId(), preparedState != null, session != null ? session.getPointer() : 0); Map senders = new LinkedHashMap<>(childInputs.size()); long streamPtr = 0; StreamHandle outStreamLocal = null; boolean success = false; try { if (preparedState != null) { - // Plan was already prepared by FinalAggregateInstructionHandler. The handler - // registered senders + captured per-input schemas in ctx.childInputs() - // iteration order; re-index them by childStageId here for lookup during feed(). int i = 0; for (Map.Entry child : childInputs.entrySet()) { senders.put(child.getKey(), preparedState.senders().get(i)); @@ -103,12 +102,8 @@ public DatafusionReduceSink(ExchangeSinkContext ctx, NativeRuntimeHandle runtime i++; } streamPtr = NativeBridge.executeLocalPreparedPlan(session.getPointer(), ctx.taskId()); + logger.debug("[reduce-sink] ALLOC preparedPlan stream taskId={} streamPtr={}", ctx.taskId(), streamPtr); } else { - // Legacy path (non-aggregate reduce): register partitions and execute the - // fragment bytes directly. Used when no prior instruction prepared a plan. - // - // ctx.fragmentBytes() references each partition by its "input-" name - // (DataFusionFragmentConvertor names them this way during plan conversion). for (Map.Entry child : childInputs.entrySet()) { int childStageId = child.getKey(); byte[] producerPlanBytes = child.getValue(); @@ -118,9 +113,12 @@ public DatafusionReduceSink(ExchangeSinkContext ctx, NativeRuntimeHandle runtime producerPlanBytes ); senders.put(childStageId, new DatafusionPartitionSender(registered.pointer())); + logger.debug("[reduce-sink] ALLOC sender taskId={} childStageId={} senderPtr={}", + ctx.taskId(), childStageId, registered.pointer()); childSchemas.put(childStageId, ArrowSchemaIpc.fromBytes(registered.schemaIpc())); } streamPtr = NativeBridge.executeLocalPlan(session.getPointer(), ctx.fragmentBytes(), ctx.taskId()); + logger.debug("[reduce-sink] ALLOC localPlan stream taskId={} streamPtr={}", ctx.taskId(), streamPtr); } outStreamLocal = new StreamHandle(streamPtr, runtimeHandle); success = true; @@ -162,6 +160,19 @@ public void feed(VectorSchemaRoot batch) { feedToSender(sendersByChildStageId.values().iterator().next(), batch, childSchemas.values().iterator().next()); } + /** + * Single-input path only: true once the sole input's consumer dropped its receiver (e.g. a + * LimitExec satisfied its fetch). Multi-input shapes (join/union) feed via {@link #sinkForChild} + * and each producer observes early-termination on its own per-child wrapper + * ({@link ChildSink#isConsumerDone()}) — a single top-level answer can't be correct there (one + * dropped join side ≠ whole reduce done), so this conservatively returns false unless there is + * exactly one registered sender. + */ + @Override + public boolean isConsumerDone() { + return sendersByChildStageId.size() == 1 && sendersByChildStageId.values().iterator().next().isReceiverDropped(); + } + @Override public ExchangeSink sinkForChild(int childStageId) { DatafusionPartitionSender sender = sendersByChildStageId.get(childStageId); @@ -177,13 +188,17 @@ public ExchangeSink sinkForChild(int childStageId) { * Lock-free per-sender feed. Exports the batch via Arrow C Data outside any lock * (the allocator is thread-safe; multiple shard handlers can export concurrently), * then sends it through the supplied sender. The Rust mpsc::Sender is thread-safe, - * so multiple producers feeding the same sender is safe. If close() raced and - * already ran senderClose, the native side returns an error ("receiver dropped") - * which we catch and discard. + * so multiple producers feeding the same sender is safe. + * + *

Two teardown signals are handled distinctly, and neither fails the query: a benign + * receiver-drop (the consumer finished early) returns the {@link NativeBridge#SENDER_SEND_RECEIVER_DROPPED} + * code, while a concurrent {@link #close()} surfaces as an IllegalStateException from + * {@code getPointer()} before the native call. Both discard the batch. */ private void feedToSender(DatafusionPartitionSender sender, VectorSchemaRoot batch, Schema declaredSchema) { - // Best-effort fast path — skip export work if already closed. - if (closed) { + // Best-effort fast path — skip the export if the sink is closed or this input's consumer + // already dropped its receiver (nothing downstream will read another batch on it). + if (closed || sender.isReceiverDropped()) { batch.close(); return; } @@ -211,18 +226,28 @@ private void feedToSender(DatafusionPartitionSender sender, VectorSchemaRoot bat // close — see DatafusionPartitionSender. Throws IllegalStateException via // NativeHandle.getPointer() if the sender was closed (the close-race path). try { - sender.send(array.memoryAddress(), arrowSchema.memoryAddress()); + long rc = sender.send(array.memoryAddress(), arrowSchema.memoryAddress()); + if (rc == NativeBridge.SENDER_SEND_RECEIVER_DROPPED) { + // Consumer finished first (e.g. a LimitExec satisfied its fetch) and dropped the + // receiver while shards were still feeding. api::sender_send already consumed the + // FFI structs via from_raw, so the buffers are Rust's to drop — do NOT release() + // here (double-free). The sender latched the drop (see DatafusionPartitionSender), + // so subsequent feeds for this input short-circuit and the producer stream is + // cancelled by the shard listener via isConsumerDone(). + logger.trace("[ReduceSink] receiver dropped before send (consumer finished), discarding batch"); + return; + } feedCount.incrementAndGet(); } catch (IllegalStateException e) { - // Sender close raced our send — Rust didn't take ownership, so the FFI - // structs' release callbacks are still set. Invoke them explicitly to free - // the exported buffers back to the Java allocator. (ArrowArray.close / - // ArrowSchema.close in the finally below frees the wrapper but does NOT - // invoke the C release callback.) + // Sender close raced our send — getPointer() threw BEFORE the native call, + // so Rust never took ownership and the FFI structs' release callbacks are + // still set. Invoke them explicitly to free the exported buffers back to the + // Java allocator. (ArrowArray.close / ArrowSchema.close in the finally below + // frees the wrapper but does NOT invoke the C release callback.) array.release(); arrowSchema.release(); if (closed) { - logger.debug("[ReduceSink] send-after-close race caught, discarding batch"); + logger.trace("[ReduceSink] send-after-close race caught, discarding batch"); return; } throw e; @@ -281,6 +306,11 @@ public void feed(VectorSchemaRoot batch) { feedToSender(sender, batch, declaredSchema); } + @Override + public boolean isConsumerDone() { + return sender.isReceiverDropped(); + } + @Override public void close() { if (childClosed) { @@ -319,55 +349,41 @@ protected void feedBatchUnderLock(VectorSchemaRoot batch) { */ @Override protected Exception closeImpl() { - SinkState before = state.compareAndExchange(SinkState.READY, SinkState.DONE); - if (before == SinkState.REDUCING) { - // Drain parked — dropping senders/outStream now would panic in drop_in_place. - fireCancelQuery(); - return null; // reduce()'s finally calls closeImpl directly to tear down. - } - // before == READY (we just won) or DONE (reduce's finally calling us, or duplicate close). if (torndown.compareAndSet(false, true) == false) { return null; } + assert torndown.get() == true; + logger.debug("[reduce-sink] teardown taskId={} feedCount={}", ctx.taskId(), feedCount.get()); Exception failure = null; - // 1. Signal EOF on every sender (ChildSink may have closed some already; idempotent). - for (DatafusionPartitionSender sender : sendersByChildStageId.values()) { - try { - sender.close(); - } catch (Exception t) { - failure = accumulate(failure, t); - } - } try { outStream.close(); } catch (Exception t) { failure = accumulate(failure, t); } - if (preparedState == null) { - try { + try { + if (preparedState != null) { + preparedState.close(); + } else { session.close(); - } catch (Exception t) { - failure = accumulate(failure, t); } + } catch (Exception t) { + failure = accumulate(failure, t); } + logger.debug("[reduce-sink] teardown complete taskId={}", ctx.taskId()); return failure; } /** Test seam: overridden to count invocations without static mocking. */ void fireCancelQuery() { + logger.debug("[reduce-sink] fireCancelQuery: taskId={}", ctx.taskId()); NativeBridge.cancelQuery(ctx.taskId()); } - /** - * Drains inline on the caller (the reduce stage's task, on a virtual thread). - * Drain terminates on input EOF (all per-child wrappers closed via - * {@link #sinkForChild}) or on cancel-Err (an external {@link #close()} fired - * {@code cancel_query}). The {@code finally} runs {@code super.close()} so cleanup - * happens on the same thread as the drain — no race with concurrent close. - */ @Override public void reduce(ActionListener listener) { SinkState before = state.compareAndExchange(SinkState.READY, SinkState.REDUCING); + logger.debug("[reduce-sink] reduce() entered: taskId={} stateBefore={} feedCount={} senderCount={}", + ctx.taskId(), before, feedCount.get(), sendersByChildStageId.size()); if (before == SinkState.DONE) { listener.onFailure(new IllegalStateException("sink closed before reduce")); return; @@ -376,10 +392,13 @@ public void reduce(ActionListener listener) { Exception failure = null; try { drainOutputIntoDownstream(outStream); + logger.debug("[reduce-sink] drain returned normally: taskId={} totalFeedCount={}", ctx.taskId(), feedCount.get()); } catch (Exception e) { + logger.debug("[reduce-sink] drain threw: taskId={} feedCount={} error={}", ctx.taskId(), feedCount.get(), e.getMessage()); failure = e; } finally { state.set(SinkState.DONE); + logger.debug("[reduce-sink] closing in reduce finally: taskId={}", ctx.taskId()); try { Exception closeFailure = closeImpl(); if (closeFailure != null) { @@ -390,8 +409,10 @@ public void reduce(ActionListener listener) { } } if (failure == null) { + logger.debug("[reduce-sink] reduce() success: taskId={}", ctx.taskId()); listener.onResponse(null); } else { + logger.debug("[reduce-sink] reduce() failed: taskId={} error={}", ctx.taskId(), failure.getMessage()); listener.onFailure(failure); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java index 2e58008934f7c..cf3e9bbd817d9 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java @@ -231,6 +231,7 @@ public final class DatafusionSettings { DataFusionPlugin.DATAFUSION_MEMORY_POOL_LIMIT, DataFusionPlugin.DATAFUSION_SPILL_MEMORY_LIMIT, DataFusionPlugin.DATAFUSION_REDUCE_INPUT_MODE, + DataFusionPlugin.DATAFUSION_REDUCE_TARGET_PARTITIONS, DataFusionPlugin.DATAFUSION_MIN_TARGET_PARTITIONS, DataFusionPlugin.DATAFUSION_MEMORY_GUARD_ADMISSION_THROTTLE_THRESHOLD, DataFusionPlugin.DATAFUSION_MEMORY_GUARD_ADMISSION_REJECT_THRESHOLD, diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java index 2a457956c5e51..9a9d1c6b7dde8 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java @@ -77,6 +77,7 @@ public final class NativeBridge { */ private static final MethodHandle SET_SPILL_LIMIT; private static final MethodHandle SET_MIN_TARGET_PARTITIONS; + private static final MethodHandle SET_REDUCE_TARGET_PARTITIONS; private static final MethodHandle SET_MEMORY_GUARD_THRESHOLDS; private static final MethodHandle CREATE_READER; private static final MethodHandle CLOSE_READER; @@ -185,6 +186,11 @@ public final class NativeBridge { FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG) ); + SET_REDUCE_TARGET_PARTITIONS = linker.downcallHandle( + lib.find("df_set_reduce_target_partitions").orElseThrow(), + FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG) + ); + SET_MEMORY_GUARD_THRESHOLDS = linker.downcallHandle( lib.find("df_set_memory_guard_thresholds").orElseThrow(), FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG) @@ -719,6 +725,15 @@ public static void setMinTargetPartitions(int value) { } } + /** Sets the initial target_partitions for coordinator-reduce sessions. */ + public static void setReduceTargetPartitions(int value) { + try { + SET_REDUCE_TARGET_PARTITIONS.invokeExact((long) value); + } catch (Throwable t) { + logger.debug("Failed to set reduce target partitions", t); + } + } + /** Sets the memory guard thresholds (0.0–1.0): admission throttle, admission reject, execution spill, execution critical. */ public static void setMemoryGuardThresholds( double admissionThrottle, @@ -1040,9 +1055,20 @@ public static long executeLocalPlan(long sessionPtr, byte[] substrait, long cont } } + /** + * Positive sentinel returned by {@code df_sender_send} (via {@link #senderSend}) when the + * send was skipped because the consumer dropped the receiver before this batch could be sent + * — the benign "consumer finished first" case (e.g. a LimitExec satisfied its fetch). It + * rides the success half of the native return contract ({@code >= 0} success, {@code < 0} + * {@code -error_ptr}), so {@code checkResult} passes it through without throwing. + * MUST match {@code SENDER_SEND_RECEIVER_DROPPED} in {@code ffm.rs}. + */ + public static final long SENDER_SEND_RECEIVER_DROPPED = 1L; + /** * Pushes one Arrow C Data-exported batch (array + schema addresses) into the sender. The - * native side takes ownership of both FFI structs. + * native side takes ownership of both FFI structs. Returns {@code 0} on a normal send or + * {@link #SENDER_SEND_RECEIVER_DROPPED} if the consumer already dropped the receiver. */ public static long senderSend(long senderPtr, long arrayPtr, long schemaPtr) { NativeHandle.validatePointer(senderPtr, "sender"); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java index 1104ffce315d0..626956d97e062 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionReduceSinkTests.java @@ -75,6 +75,25 @@ public void testInputIdConstantMatchesDesign() { assertEquals("Single-input reduce uses the synthetic id 'input-0'", "input-0", DatafusionReduceSink.INPUT_ID); } + /** + * The benign "consumer finished first" case (a {@code LimitExec} satisfied its fetch and + * dropped the receiver while producers were still feeding) is signalled by {@code df_sender_send} + * returning the positive sentinel {@link NativeBridge#SENDER_SEND_RECEIVER_DROPPED} rather than a + * stringly-typed error — {@code feedToSender} keys off this code. The sentinel must be positive + * (so {@code checkResult} passes it through the success half of the return contract instead of + * treating it as an error pointer) and distinct from a normal send ({@code 0}). The Rust↔Java + * agreement on the value is enforced by the matching {@code SENDER_SEND_RECEIVER_DROPPED} in + * {@code ffm.rs}; the structural "only a dropped receiver maps here" guarantee is covered by the + * Rust {@code send_blocking_reports_receiver_dropped} unit test. + */ + public void testReceiverDroppedSentinelIsPositiveAndDistinctFromSuccess() { + assertTrue( + "sentinel must be positive so checkResult treats it as success, not an error pointer", + NativeBridge.SENDER_SEND_RECEIVER_DROPPED > 0 + ); + assertNotEquals("sentinel must be distinct from a normal send (0)", 0L, NativeBridge.SENDER_SEND_RECEIVER_DROPPED); + } + /** * End-to-end feed + drain: feeds three Arrow batches (values 1..9) into a real * {@link DatafusionReduceSink} running a {@code SELECT SUM(x) FROM "input-0"} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java index 47fb48682eb78..25cac6239a22d 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java @@ -148,14 +148,20 @@ public void executeFragmentStreamingAsync( ) { try { executor.execute(() -> { - LOGGER.debug("[FragmentExecution] shard={} task={}", shard.shardId(), task.getId()); + LOGGER.debug("[shard-producer] shard={} starting stream iteration", request.getShardId()); try (FragmentResources ctx = executeFragmentStreaming(request, shard, task)) { Iterator it = ctx.stream().iterator(); + int batches = 0; while (it.hasNext()) { responseHandler.onBatch(it.next()); + batches++; } + LOGGER.debug("[shard-producer] shard={} stream complete, {} batches produced", request.getShardId(), batches); responseHandler.onComplete(); } catch (Exception e) { + LOGGER.debug("[shard-producer] shard={} stream failed: {} cause={}", + request.getShardId(), e.getMessage(), + e.getCause() != null ? e.getCause().toString() : "none", e); responseHandler.onFailure(e); } }); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java index 7f74733eabe23..5df5ce355d0fa 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java @@ -95,10 +95,7 @@ private static void registerStreamingFragmentHandler( shard, (AnalyticsShardTask) task, channelResponseHandler(channel), - ContextAwareExecutor.wrap( - transportService.getThreadPool().executor(ThreadPool.Names.SEARCH), - transportService.getThreadPool() - ) + transportService.getThreadPool().executor(ThreadPool.Names.SEARCH) ); } ); @@ -130,10 +127,7 @@ private static void registerFetchByRowIdsHandler( shard, (AnalyticsShardTask) task, channelResponseHandler(channel), - ContextAwareExecutor.wrap( - transportService.getThreadPool().executor(ThreadPool.Names.SEARCH), - transportService.getThreadPool() - ) + transportService.getThreadPool().executor(ThreadPool.Names.SEARCH) ); } ); @@ -234,16 +228,21 @@ public String executor() { @Override public void handleStreamResponse(StreamTransportResponse stream) { try { - FragmentExecutionArrowResponse current; - FragmentExecutionArrowResponse last = null; - while ((current = stream.nextResponse()) != null) { - if (last != null) { - listener.onStreamResponse(last, false); + FragmentExecutionArrowResponse last = stream.nextResponse(); + while (last != null) { + FragmentExecutionArrowResponse next = stream.nextResponse(); + boolean isLast = next == null; + boolean keepReading = listener.onStreamResponse(last, isLast); + if (!keepReading) { + if (next != null) { + if (next.getRoot() != null) { + next.getRoot().close(); + } + stream.cancel("reduce input satisfied (downstream consumer finished)", null); + } + return; } - last = current; - } - if (last != null) { - listener.onStreamResponse(last, true); + last = next; } } catch (Exception e) { listener.onFailure(e); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/StreamingResponseListener.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/StreamingResponseListener.java index 686cdb1319cbe..f48118775d89d 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/StreamingResponseListener.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/StreamingResponseListener.java @@ -33,8 +33,11 @@ public interface StreamingResponseListener { * * @param response the response batch * @param isLast {@code true} if this is the final batch (terminal success event) + * @return {@code true} to keep draining this stream; {@code false} if the consumer is already + * satisfied (e.g. a downstream LimitExec finished) and the caller should cancel the + * stream and stop — the listener has already settled its terminal event in that case. */ - void onStreamResponse(Resp response, boolean isLast); + boolean onStreamResponse(Resp response, boolean isLast); /** * Called when the request fails. Terminal failure event. diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java index ab64bf3cbc86b..743ebf4f5a62e 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java @@ -417,20 +417,21 @@ private static final class GatherListener implements StreamingResponseListener children, Consumer { + closeChildInput(childId); Exception cause = child.getFailure(); failWithCause( cause != null diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/PassThroughStageExecution.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/PassThroughStageExecution.java index 1cc06a055dc13..ad2cb03cae84e 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/PassThroughStageExecution.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/PassThroughStageExecution.java @@ -29,14 +29,16 @@ */ public final class PassThroughStageExecution extends AbstractStageExecution implements SinkProvidingStageExecution { - private final RowProducingSink ownedSink; + private final ExchangeSink ownedSink; + private final ExchangeSource ownedSource; public PassThroughStageExecution(Stage stage, QueryContext config, ExchangeSink sink) { super(stage, config.queryId(), config.operationListeners(), config.parentTask()); - if ((sink instanceof RowProducingSink) == false) { - throw new IllegalArgumentException("PassThroughStageExecution requires a RowProducingSink"); + if ((sink instanceof ExchangeSource) == false) { + throw new IllegalArgumentException("PassThroughStageExecution requires a sink that also implements ExchangeSource"); } - this.ownedSink = (RowProducingSink) sink; + this.ownedSink = sink; + this.ownedSource = (ExchangeSource) sink; this.runner = new LocalTaskRunner(config.schedulerExecutor()); } @@ -52,6 +54,6 @@ public ExchangeSink inputSink(int childStageId) { @Override public ExchangeSource outputSource() { - return ownedSink; + return ownedSource; } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecution.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecution.java index 5b8cf68ca75a1..76a754648f903 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecution.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecution.java @@ -19,7 +19,6 @@ import org.opensearch.analytics.exec.stage.StageTaskId; import org.opensearch.analytics.planner.dag.InputSinkDecorator; import org.opensearch.analytics.planner.dag.Stage; -import org.opensearch.analytics.spi.CancellableExchangeSink; import org.opensearch.analytics.spi.ExchangeSink; import org.opensearch.analytics.spi.MultiInputExchangeSink; import org.opensearch.analytics.spi.ReducingExchangeSink; @@ -64,6 +63,8 @@ public boolean schedulesEagerly() { @Override public void closeChildInput(int childStageId) { + logger.debug("[reduce-stage] closeChildInput: stageId={} childStageId={} isMultiInput={}", + getStageId(), childStageId, backendSink instanceof MultiInputExchangeSink); if (backendSink instanceof MultiInputExchangeSink multi) { multi.sinkForChild(childStageId).close(); } @@ -97,28 +98,29 @@ public ExchangeSource outputSource() { @Override protected List materializeTasks() { return List.of(new LocalStageTask(new StageTaskId(getStageId(), 0), listener -> { + logger.debug("[reduce-stage] reduce task dispatched, stageId={}", getStageId()); reduceExecutor.execute(() -> { try { backendSink.reduce(listener); + logger.debug("[reduce-stage] reduce() returned normally, stageId={}", getStageId()); } catch (Exception e) { + logger.debug("[reduce-stage] reduce() threw, stageId={}: {}", getStageId(), e.getMessage()); listener.onFailure(e); } }); })); } + @Override + public boolean failWithCause(Exception cause) { + logger.debug("[reduce-stage] failWithCause: stageId={} cause={}", getStageId(), + cause != null ? cause.getMessage() : "null"); + return super.failWithCause(cause); + } + @Override protected void onTerminalTransition(State terminal) { - if (terminal == State.CANCELLED || terminal == State.FAILED) { - if (backendSink instanceof CancellableExchangeSink cancellable) { - logger.warn("[ReduceStageExecution] stage {} terminal={}, firing cancellable.cancel()", getStageId(), terminal); - try { - cancellable.cancel(); - } catch (Exception e) { - logger.warn("[ReduceStageExecution] cancel() threw for stage " + getStageId(), e); - } - } - } + logger.debug("[reduce-stage] onTerminalTransition: stageId={} terminal={}", getStageId(), terminal); try { backendSink.close(); } catch (Exception ignore) {} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecution.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecution.java index a8dfcd91c6c47..35100a6a5e46a 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecution.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecution.java @@ -126,15 +126,15 @@ public ExchangeSource outputSource() { StreamingResponseListener responseListenerFor(int sourceOrdinal, ActionListener listener) { return new StreamingResponseListener<>() { @Override - public void onStreamResponse(FragmentExecutionArrowResponse response, boolean isLast) { + public boolean onStreamResponse(FragmentExecutionArrowResponse response, boolean isLast) { VectorSchemaRoot vsr = response.getRoot(); if (getState().isTerminal()) { if (vsr != null) vsr.close(); - return; + return false; // stage already settled — stop draining, let the caller cancel the stream } if (vsr == null) { if (isLast) listener.onResponse(null); - return; + return true; } try { outputSink.feed(vsr, sourceOrdinal); @@ -147,10 +147,19 @@ public void onStreamResponse(FragmentExecutionArrowResponse response, boolean is wrapped.addSuppressed(closeFailure); } listener.onFailure(wrapped); - return; + return false; } metrics.addRowsProcessed(vsr.getRowCount()); + // Downstream consumer satisfied (e.g. a LimitExec above the reduce finished and dropped + // this input's receiver). Settle this task as success and tell the caller to cancel the + // stream so this shard stops scanning instead of feeding batches that will be discarded. + // Each input reacts independently on its own stream. + if (outputSink.isConsumerDone()) { + listener.onResponse(null); + return false; + } if (isLast) listener.onResponse(null); + return true; } @Override diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecutionTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecutionTests.java index 4a306698722f1..a8392f8c43bd3 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecutionTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecutionTests.java @@ -147,6 +147,32 @@ public void testArrowResponseFedToSinkOnHappyPath() { sink.close(); } + /** + * When the downstream consumer is satisfied (e.g. a LimitExec above the reduce finished and + * dropped this input's receiver), the shard listener must stop reading: it feeds the in-hand + * batch, settles its task as success, and returns {@code false} so the transport drain loop + * cancels the stream instead of scanning to exhaustion. The non-last flag proves we stop early. + */ + public void testStreamStopsAndTaskSucceedsWhenConsumerDone() { + AtomicReference> capturedListener = new AtomicReference<>(); + CapturingSink sink = new CapturingSink(); + sink.consumerDone = true; // consumer already satisfied before this batch arrives + + ShardFragmentStageExecution exec = buildExecution(sink, capturedListener); + scheduleAndDispatch(exec); + assertNotNull("listener should have been captured by dispatch", capturedListener.get()); + + VectorSchemaRoot root = createTestBatch(3); + FragmentExecutionArrowResponse response = new FragmentExecutionArrowResponse(root); + // isLast=false: there are more batches upstream, but the consumer is done — we stop anyway. + boolean keepReading = capturedListener.get().onStreamResponse(response, false); + + assertFalse("listener must signal stop when the consumer is done", keepReading); + assertEquals("the in-hand batch is still fed before stopping", 1, sink.fed.size()); + assertEquals("task completes as SUCCESS (not cancelled/failed)", StageExecution.State.SUCCEEDED, exec.getState()); + sink.close(); + } + /** * Mirrors {@code QueryExecution.scheduleStage} for unit-test purposes — calls * start() to materialise + transition, then iterates the stage's tasks via its @@ -469,12 +495,18 @@ private ClusterService mockClusterService() { private static final class CapturingSink implements ExchangeSink { final List fed = new ArrayList<>(); boolean closed = false; + volatile boolean consumerDone = false; @Override public void feed(VectorSchemaRoot batch) { fed.add(batch); } + @Override + public boolean isConsumerDone() { + return consumerDone; + } + @Override public void close() { closed = true; diff --git a/sandbox/qa/analytics-engine-rest/build.gradle b/sandbox/qa/analytics-engine-rest/build.gradle index 2e1623c806f3c..b640d67f51b96 100644 --- a/sandbox/qa/analytics-engine-rest/build.gradle +++ b/sandbox/qa/analytics-engine-rest/build.gradle @@ -120,7 +120,6 @@ integTest { systemProperty 'tests.security.manager', 'false' exclude '**/CoordinatorReduceMemtableIT.class' exclude '**/StreamingCoordinatorReduceIT.class' - exclude '**/QueryCacheIT.class' // Note: parallel forks against the same 2-node testCluster slow the suite down // (cluster is the bottleneck, not JVM startup) and cause cross-fork data races diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/AliasIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/AliasIT.java index 016e2b4be7be9..7fb4871414656 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/AliasIT.java +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/AliasIT.java @@ -32,8 +32,12 @@ public class AliasIT extends AnalyticsRestTestCase { private static boolean dataProvisioned = false; +<<<<<<< HEAD + private void ensureDataProvisioned() throws IOException { +======= @Override protected void onBeforeQuery() throws IOException { +>>>>>>> upstream/main if (dataProvisioned == false) { DatasetProvisioner.provision(client(), CALCS_A); DatasetProvisioner.provision(client(), CALCS_B); @@ -47,6 +51,10 @@ protected void onBeforeQuery() throws IOException { * Verifies the alias fans out — a missing-fan-out bug would return 17. */ public void testAliasSpansAllBackingIndices() throws IOException { +<<<<<<< HEAD + ensureDataProvisioned(); +======= +>>>>>>> upstream/main long count = singleCount("source=" + COMPAT_ALIAS + " | stats count() as c"); assertEquals("alias fan-out: 17 + 17", 34L, count); } @@ -56,6 +64,10 @@ public void testAliasSpansAllBackingIndices() throws IOException { * passes through as a singleton list). */ public void testConcreteIndexStillResolvesAsBefore() throws IOException { +<<<<<<< HEAD + ensureDataProvisioned(); +======= +>>>>>>> upstream/main long count = singleCount("source=" + CALCS_A.indexName + " | stats count() as c"); assertEquals("single concrete index", 17L, count); } @@ -66,6 +78,10 @@ public void testConcreteIndexStillResolvesAsBefore() throws IOException { * the user can fix the mapping rather than guess. */ public void testSchemaMismatchAliasIsRejected() throws IOException { +<<<<<<< HEAD + ensureDataProvisioned(); +======= +>>>>>>> upstream/main // Provision a third index whose `key` field is a long (calcs has it as keyword) and // alias it together with calcs_a. The planner should reject when the query references // either index. @@ -131,6 +147,10 @@ public void testAliasSkipsClosedBackingIndex() throws IOException { * rows than the user asked for. Better to error with a clear message. */ public void testFilterAliasIsRejected() throws IOException { +<<<<<<< HEAD + ensureDataProvisioned(); +======= +>>>>>>> upstream/main String filterAlias = "calcs_filter_only"; // PUT alias with a term filter against the existing calcs_a index. Request put = new Request("POST", "/_aliases"); @@ -151,7 +171,11 @@ public void testFilterAliasIsRejected() throws IOException { private long singleCount(String ppl) throws IOException { Map body = executePpl(ppl); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List> rows = (List>) body.get("rows"); +======= List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertNotNull("missing 'rows' for: " + ppl, rows); assertEquals("single count row expected: " + ppl, 1, rows.size()); Object cell = rows.get(0).get(0); @@ -159,9 +183,21 @@ private long singleCount(String ppl) throws IOException { return ((Number) cell).longValue(); } +<<<<<<< HEAD + private Map executePpl(String ppl) throws IOException { + Request request = new Request("POST", "/_analytics/ppl"); + request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); + Response response = client().performRequest(request); + return assertOkAndParse(response, "PPL: " + ppl); + } + + private String executePplExpectingFailure(String ppl) throws IOException { + Request request = new Request("POST", "/_analytics/ppl"); +======= private String executePplExpectingFailure(String ppl) throws IOException { Request request = new Request("POST", "/_plugins/_ppl"); +>>>>>>> upstream/main request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); try { Response response = client().performRequest(request); diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/DataStreamIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/DataStreamIT.java index da2b20df29837..fa0afdee36a36 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/DataStreamIT.java +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/DataStreamIT.java @@ -202,7 +202,11 @@ public void testDataStreamGroupByAcrossBackings() throws IOException { Map body = executePpl("source=" + STREAM + " | stats count() as c by category | sort category"); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List> rows = (List>) body.get("rows"); +======= List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertNotNull(rows); assertEquals("two distinct categories", 2, rows.size()); // After sort by category: alpha (3), beta (3). @@ -304,7 +308,11 @@ private long singleCount(String ppl) throws IOException { private long singleLongAgg(String ppl) throws IOException { Map body = executePpl(ppl); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List> rows = (List>) body.get("rows"); +======= List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertNotNull("missing 'rows' for: " + ppl, rows); assertEquals("single row expected: " + ppl, 1, rows.size()); Object cell = rows.get(0).get(0); @@ -312,9 +320,21 @@ private long singleLongAgg(String ppl) throws IOException { return ((Number) cell).longValue(); } +<<<<<<< HEAD + private Map executePpl(String ppl) throws IOException { + Request request = new Request("POST", "/_analytics/ppl"); + request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); + Response response = client().performRequest(request); + return assertOkAndParse(response, "PPL: " + ppl); + } + + private String executePplExpectingFailure(String ppl) throws IOException { + Request request = new Request("POST", "/_analytics/ppl"); +======= private String executePplExpectingFailure(String ppl) throws IOException { Request request = new Request("POST", "/_plugins/_ppl"); +>>>>>>> upstream/main request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); try { Response response = client().performRequest(request); diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/FilterDelegationIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/FilterDelegationIT.java index 35ef34ecc21fb..419fff8d9c001 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/FilterDelegationIT.java +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/FilterDelegationIT.java @@ -240,7 +240,11 @@ public void testAndCorrectnessAndPerfAndNative() throws Exception { createIndex(); indexDocs(); String ppl = "source = " + INDEX_NAME + " | where match(message, 'hello') and tag = 'hello' and value = 5 | stats count() as cnt"; +<<<<<<< HEAD + Map result = executePPL(ppl); +======= Map result = executePplViaShim(ppl); +>>>>>>> upstream/main @SuppressWarnings("unchecked") List> rows = (List>) result.get("rows"); assertEquals(10L, ((Number) rows.get(0).get(0)).longValue()); @@ -254,7 +258,11 @@ public void testAndOnlyPerfAndNative() throws Exception { createIndex(); indexDocs(); String ppl = "source = " + INDEX_NAME + " | where tag = 'hello' and value = 5 | stats count() as cnt"; +<<<<<<< HEAD + Map result = executePPL(ppl); +======= Map result = executePplViaShim(ppl); +>>>>>>> upstream/main @SuppressWarnings("unchecked") List> rows = (List>) result.get("rows"); assertEquals(10L, ((Number) rows.get(0).get(0)).longValue()); @@ -270,7 +278,11 @@ public void testOrCorrectnessWithPerfAndNative() throws Exception { createIndex(); indexDocs(); String ppl = "source = " + INDEX_NAME + " | where match(message, 'hello') or tag = 'hello' and value = 5 | stats count() as cnt"; +<<<<<<< HEAD + Map result = executePPL(ppl); +======= Map result = executePplViaShim(ppl); +>>>>>>> upstream/main @SuppressWarnings("unchecked") List> rows = (List>) result.get("rows"); assertEquals(10L, ((Number) rows.get(0).get(0)).longValue()); @@ -286,7 +298,11 @@ public void testOrOfTwoAndArms() throws Exception { indexDocs(); String ppl = "source = " + INDEX_NAME + " | where (match(message, 'hello') and tag = 'hello') or (tag = 'goodbye' and value = 3) | stats count() as cnt"; +<<<<<<< HEAD + Map result = executePPL(ppl); +======= Map result = executePplViaShim(ppl); +>>>>>>> upstream/main @SuppressWarnings("unchecked") List> rows = (List>) result.get("rows"); assertEquals(20L, ((Number) rows.get(0).get(0)).longValue()); diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/IndexPatternUnionIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/IndexPatternUnionIT.java index 3d79cf9c4660c..ef18204a84175 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/IndexPatternUnionIT.java +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/IndexPatternUnionIT.java @@ -97,7 +97,11 @@ public void testUnionPreservesRequestedColumnOrder() throws IOException { ensureProvisioned(); Map body = executePpl("source=" + ALIAS + " | fields name, age, alias"); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List columns = (List) body.get("columns"); +======= List columns = extractColumnNames(body); +>>>>>>> upstream/main assertEquals("union output column order must match requested fields", List.of("name", "age", "alias"), columns); } @@ -114,7 +118,11 @@ public void testDottedIndexNameRegistersAndScans() throws IOException { bulk(dotted, "{\"v\":1}\n{\"v\":2}\n"); Map body = executePpl("source=" + dotted + " | stats count() as c"); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List> rows = (List>) body.get("rows"); +======= List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertEquals("single count row", 1, rows.size()); assertEquals("dotted index must scan its 2 rows", 2L, ((Number) rows.get(0).get(0)).longValue()); } @@ -124,7 +132,11 @@ public void testAliasFansOutAcrossDifferingIndices() throws IOException { ensureProvisioned(); Map body = executePpl("source=" + ALIAS + " | stats count() as c"); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List> rows = (List>) body.get("rows"); +======= List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertEquals("single count row", 1, rows.size()); assertEquals("union_a (2) + union_b (1)", 3L, ((Number) rows.get(0).get(0)).longValue()); } @@ -137,7 +149,11 @@ private void assertUnionRows(String ppl) throws IOException { assertTrue("schema must carry all union columns", col.containsKey("name") && col.containsKey("age") && col.containsKey("alias")); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List> rows = (List>) body.get("rows"); +======= List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertEquals("union row count", 3, rows.size()); Map> byName = new HashMap<>(); @@ -164,7 +180,11 @@ private void assertUnionRows(String ppl) throws IOException { /** Maps each column name in the PPL response to its position in the row arrays. */ @SuppressWarnings("unchecked") private static Map columnIndex(Map body) { +<<<<<<< HEAD + List columns = (List) body.get("columns"); +======= List columns = extractColumnNames(body); +>>>>>>> upstream/main assertNotNull("response must carry columns", columns); Map col = new HashMap<>(); for (int i = 0; i < columns.size(); i++) { @@ -175,6 +195,14 @@ private static Map columnIndex(Map body) { // ── helpers ────────────────────────────────────────────────────────── +<<<<<<< HEAD + private Map executePpl(String ppl) throws IOException { + Request request = new Request("POST", "/_analytics/ppl"); + request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); + return assertOkAndParse(client().performRequest(request), "PPL: " + ppl); + } +======= +>>>>>>> upstream/main private void createParquetIndex(String name, String mappingJson) throws IOException { Request create = new Request("PUT", "/" + name); diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MultiIndexQueryShapesIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MultiIndexQueryShapesIT.java index b34662bdce7aa..b8d1c0b0f865a 100644 --- a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MultiIndexQueryShapesIT.java +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/MultiIndexQueryShapesIT.java @@ -75,7 +75,11 @@ public void testAliasAggregation() throws IOException { ensureProvisioned(); Map body = executePpl("source=" + PQONLY_ALIAS + " | stats sum(status) as total"); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List> rows = (List>) body.get("rows"); +======= List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertEquals(1, rows.size()); long total = ((Number) rows.get(0).get(0)).longValue(); assertEquals("sum(200+500+200+200+404)", 1504L, total); @@ -85,10 +89,17 @@ public void testAliasNullFillForDifferingFields() throws IOException { ensureProvisioned(); Map body = executePpl("source=" + PQONLY_ALIAS + " | fields message, source"); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List columns = (List) body.get("columns"); + assertTrue("must have source column", columns.contains("source")); + @SuppressWarnings("unchecked") + List> rows = (List>) body.get("rows"); +======= List columns = extractColumnNames(body); assertTrue("must have source column", columns.contains("source")); @SuppressWarnings("unchecked") List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertEquals("total rows: 3 + 2 = 5", 5, rows.size()); int sourceCol = columns.indexOf("source"); int nullCount = 0; @@ -132,7 +143,11 @@ public void testDelegationAliasAggregation() throws IOException { ensureProvisioned(); Map body = executePpl("source=" + PQLUC_ALIAS + " | stats sum(status) as total"); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List> rows = (List>) body.get("rows"); +======= List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertEquals(1, rows.size()); long total = ((Number) rows.get(0).get(0)).longValue(); assertEquals("sum(200+500+200+200+404)", 1504L, total); @@ -142,10 +157,17 @@ public void testDelegationAliasNullFill() throws IOException { ensureProvisioned(); Map body = executePpl("source=" + PQLUC_ALIAS + " | fields message, source"); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List columns = (List) body.get("columns"); + assertTrue("must have source column", columns.contains("source")); + @SuppressWarnings("unchecked") + List> rows = (List>) body.get("rows"); +======= List columns = extractColumnNames(body); assertTrue("must have source column", columns.contains("source")); @SuppressWarnings("unchecked") List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertEquals(5, rows.size()); int sourceCol = columns.indexOf("source"); int nullCount = 0; @@ -190,7 +212,11 @@ public void testMultiShardAggregationAcrossUnion() throws IOException { ensureMshardProvisioned(); Map body = executePpl("source=" + MSHARD_ALIAS + " | stats sum(val) as total"); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List> rows = (List>) body.get("rows"); +======= List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertEquals(1, rows.size()); long total = ((Number) rows.get(0).get(0)).longValue(); // sum(0..9) + sum(10..19) = 45 + 145 = 190 @@ -201,11 +227,19 @@ public void testMultiShardNullFillFieldsPresent() throws IOException { ensureMshardProvisioned(); Map body = executePpl("source=" + MSHARD_ALIAS + " | fields val, tag, extra"); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List columns = (List) body.get("columns"); + assertTrue("must have tag", columns.contains("tag")); + assertTrue("must have extra", columns.contains("extra")); + @SuppressWarnings("unchecked") + List> rows = (List>) body.get("rows"); +======= List columns = extractColumnNames(body); assertTrue("must have tag", columns.contains("tag")); assertTrue("must have extra", columns.contains("extra")); @SuppressWarnings("unchecked") List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertEquals("20 total rows", 20, rows.size()); int tagCol = columns.indexOf("tag"); int extraCol = columns.indexOf("extra"); @@ -249,10 +283,17 @@ public void testTypeCoverageNullFillAcrossTypes() throws IOException { ensureTypesProvisioned(); Map body = executePpl("source=" + TYPES_ALIAS + " | fields id, label, active, score, count"); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List columns = (List) body.get("columns"); + assertTrue("must have all union columns", columns.containsAll(List.of("id", "label", "active", "score", "count"))); + @SuppressWarnings("unchecked") + List> rows = (List>) body.get("rows"); +======= List columns = extractColumnNames(body); assertTrue("must have all union columns", columns.containsAll(List.of("id", "label", "active", "score", "count"))); @SuppressWarnings("unchecked") List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertEquals("4 total rows", 4, rows.size()); int labelCol = columns.indexOf("label"); int scoreCol = columns.indexOf("score"); @@ -269,7 +310,11 @@ public void testTypeCoverageAggregateOnSharedField() throws IOException { ensureTypesProvisioned(); Map body = executePpl("source=" + TYPES_ALIAS + " | stats sum(id) as total"); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List> rows = (List>) body.get("rows"); +======= List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertEquals(1, rows.size()); long total = ((Number) rows.get(0).get(0)).longValue(); assertEquals("sum(1+2+3+4)", 10L, total); @@ -294,11 +339,19 @@ public void testDynamicMappingUnionAcrossIndices() throws IOException { Map body = executePpl("source=" + dynAlias + " | fields id, city, country"); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List columns = (List) body.get("columns"); + assertTrue("must have city", columns.contains("city")); + assertTrue("must have country", columns.contains("country")); + @SuppressWarnings("unchecked") + List> rows = (List>) body.get("rows"); +======= List columns = extractColumnNames(body); assertTrue("must have city", columns.contains("city")); assertTrue("must have country", columns.contains("country")); @SuppressWarnings("unchecked") List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertEquals(4, rows.size()); int cityCol = columns.indexOf("city"); int countryCol = columns.indexOf("country"); @@ -342,7 +395,11 @@ public void testAliasTypeMismatchIsRejected() throws IOException { } private String executePplExpectingFailure(String ppl) throws IOException { +<<<<<<< HEAD + Request request = new Request("POST", "/_analytics/ppl"); +======= Request request = new Request("POST", "/_plugins/_ppl"); +>>>>>>> upstream/main request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); try { Response response = client().performRequest(request); @@ -362,7 +419,11 @@ private static void assertContains(String haystack, String needle) { private long singleCount(String ppl) throws IOException { Map body = executePpl(ppl); @SuppressWarnings("unchecked") +<<<<<<< HEAD + List> rows = (List>) body.get("rows"); +======= List> rows = (List>) body.get("datarows"); +>>>>>>> upstream/main assertNotNull("missing 'rows' for: " + ppl, rows); assertEquals("single count row expected: " + ppl, 1, rows.size()); Object cell = rows.get(0).get(0); @@ -370,6 +431,15 @@ private long singleCount(String ppl) throws IOException { return ((Number) cell).longValue(); } +<<<<<<< HEAD + private Map executePpl(String ppl) throws IOException { + Request request = new Request("POST", "/_analytics/ppl"); + request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); + Response response = client().performRequest(request); + return assertOkAndParse(response, "PPL: " + ppl); + } +======= +>>>>>>> upstream/main private void createParquetIndex(String name, String mappingJson) throws IOException { createIndexWithSettings(name, mappingJson, false, 1);