This is an automated email from the ASF dual-hosted git repository.
Gabriel39 pushed a commit to branch branch-4.1
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/branch-4.1 by this push:
new d050e7dbd9f [fix](arrow-flight) Deduplicate parallel result endpoints
by backend (#68642)
d050e7dbd9f is described below
commit d050e7dbd9f7cd2bb86a3f7b295ceae3439e3b65
Author: Gabriel <[email protected]>
AuthorDate: Tue Sep 29 21:57:10 2026 +0800
[fix](arrow-flight) Deduplicate parallel result endpoints by backend
(#68642)
## Problem
With parallel result sinks, Nereids publishes one Flight endpoint per
pipeline instance even though instances on the same backend share a
query-id result buffer. Six result backends with eight instances each
therefore publish 48 endpoints for only six independent result streams.
## Fix
Publish one endpoint per backend in query-id mode, preserving first-seen
order and the existing instance-id behavior in non-parallel mode.
Backend IDs distinguish result owners even when Flight addresses are
shared.
Add five FE unit tests and a Flight regression covering unique tickets
and complete results for sequential/concurrent consumption with parallel
result sinks enabled and disabled.
## Validation
- Reproduced the original bug with three failing unit tests (including
48 endpoints instead of 6); all five tests pass after the fix.
- Compiled the changed FE class and tests with `javac --release 8`
against cached dependencies.
- FE Checkstyle and Groovy regression compilation checked locally.
- Full cluster regression is pending CI.
---
.../org/apache/doris/qe/NereidsCoordinator.java | 8 ++
.../qe/NereidsCoordinatorFlightEndpointsTest.java | 118 ++++++++++++++++++
.../test_flight_parallel_partitions.groovy | 134 +++++++++++++++++++++
3 files changed, 260 insertions(+)
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/qe/NereidsCoordinator.java
b/fe/fe-core/src/main/java/org/apache/doris/qe/NereidsCoordinator.java
index 638128326c3..146df08108c 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/qe/NereidsCoordinator.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/qe/NereidsCoordinator.java
@@ -74,8 +74,10 @@ import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.Comparator;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
+import java.util.Set;
import java.util.stream.Collectors;
/** NereidsCoordinator */
@@ -490,6 +492,7 @@ public class NereidsCoordinator extends Coordinator {
if (dataSink instanceof ResultSink || dataSink instanceof
ResultFileSink) {
if (connectContext != null &&
!connectContext.isReturnResultFromLocal()) {
Preconditions.checkState(connectContext.getConnectType().equals(ConnectType.ARROW_FLIGHT_SQL));
+ Set<Long> resultBackendIds = new HashSet<>();
for (AssignedJob instance : topPlan.getInstanceJobs()) {
BackendWorker worker = (BackendWorker)
instance.getAssignedWorker();
Backend backend = worker.getBackend();
@@ -498,6 +501,11 @@ public class NereidsCoordinator extends Coordinator {
}
TUniqueId finstId;
if
(connectContext.getSessionVariable().enableParallelResultSink()) {
+ // Parallel instances on a BE share one query-id
buffer, so their tickets
+ // cannot be published as independently consumable
result partitions.
+ if (!resultBackendIds.add(worker.id())) {
+ continue;
+ }
finstId = getQueryId();
} else {
finstId = instance.instanceId();
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/qe/NereidsCoordinatorFlightEndpointsTest.java
b/fe/fe-core/src/test/java/org/apache/doris/qe/NereidsCoordinatorFlightEndpointsTest.java
new file mode 100644
index 00000000000..6c6bb362190
--- /dev/null
+++
b/fe/fe-core/src/test/java/org/apache/doris/qe/NereidsCoordinatorFlightEndpointsTest.java
@@ -0,0 +1,118 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.qe;
+
+import org.apache.doris.nereids.trees.plans.distribute.PipelineDistributedPlan;
+import org.apache.doris.nereids.trees.plans.distribute.worker.BackendWorker;
+import org.apache.doris.nereids.trees.plans.distribute.worker.job.AssignedJob;
+import org.apache.doris.planner.ResultSink;
+import org.apache.doris.service.arrowflight.results.FlightSqlEndpointsLocation;
+import org.apache.doris.system.Backend;
+import org.apache.doris.thrift.TNetworkAddress;
+import org.apache.doris.thrift.TUniqueId;
+
+import com.google.common.collect.ImmutableList;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.List;
+
+public class NereidsCoordinatorFlightEndpointsTest {
+ private static final TUniqueId QUERY_ID = new TUniqueId(1, 2);
+
+ @Test
+ public void parallelInstancesShareOneEndpointPerBackend() throws Exception
{
+ List<FlightSqlEndpointsLocation> endpoints = endpoints(true, false, 1,
8);
+ Assert.assertEquals(1, endpoints.size());
+ Assert.assertEquals(QUERY_ID, endpoints.get(0).getFinstId());
+ }
+
+ @Test
+ public void parallelEndpointsRetainAllResultBackendsInOrder() throws
Exception {
+ List<FlightSqlEndpointsLocation> endpoints = endpoints(true, false, 6,
8);
+ Assert.assertEquals(6, endpoints.size());
+ for (int i = 0; i < endpoints.size(); i++) {
+ Assert.assertEquals(QUERY_ID, endpoints.get(i).getFinstId());
+ Assert.assertEquals(new TNetworkAddress("127.0.0.1", 8100 + i),
+ endpoints.get(i).getResultInternalServiceAddr());
+ }
+ }
+
+ @Test
+ public void sharedFlightAddressDoesNotMergeDifferentBackends() throws
Exception {
+ List<FlightSqlEndpointsLocation> endpoints = endpoints(true, false, 2,
2);
+ Assert.assertEquals(2, endpoints.size());
+ Assert.assertEquals(endpoints.get(0).getResultFlightServerAddr(),
+ endpoints.get(1).getResultFlightServerAddr());
+ Assert.assertNotEquals(endpoints.get(0).getResultInternalServiceAddr(),
+ endpoints.get(1).getResultInternalServiceAddr());
+ }
+
+ @Test
+ public void nonParallelEndpointsKeepDistinctInstanceIds() throws Exception
{
+ List<FlightSqlEndpointsLocation> endpoints = endpoints(false, false,
1, 2);
+ Assert.assertEquals(2, endpoints.size());
+ Assert.assertEquals(new TUniqueId(2, 0),
endpoints.get(0).getFinstId());
+ Assert.assertEquals(new TUniqueId(2, 1),
endpoints.get(1).getFinstId());
+ }
+
+ @Test
+ public void localResultDoesNotPublishFlightEndpoints() throws Exception {
+ Assert.assertTrue(endpoints(true, true, 2, 8).isEmpty());
+ }
+
+ private List<FlightSqlEndpointsLocation> endpoints(boolean parallel,
boolean local,
+ int backendCount, int instancesPerBackend) throws Exception {
+ ConnectContext context = new ConnectContext();
+ context.connectType = ConnectContext.ConnectType.ARROW_FLIGHT_SQL;
+ context.setReturnResultFromLocal(local);
+ context.getSessionVariable().setEnableParallelResultSink(parallel);
+ CoordinatorContext coordinatorContext =
Mockito.mock(CoordinatorContext.class);
+ setContextField(coordinatorContext, "connectContext", context);
+ setContextField(coordinatorContext, "dataSink",
Mockito.mock(ResultSink.class));
+ PipelineDistributedPlan plan =
Mockito.mock(PipelineDistributedPlan.class, Mockito.RETURNS_DEEP_STUBS);
+ List<AssignedJob> jobs = new ArrayList<>();
+ for (int i = 0; i < backendCount; i++) {
+ Backend backend = new Backend(i + 1, "127.0.0.1", 9000 + i);
+ backend.setBrpcPort(8100 + i);
+ // Backend identity must remain distinct even when Flight
locations are shared.
+ backend.setArrowFlightSqlPort(8050);
+ for (int j = 0; j < instancesPerBackend; j++) {
+ AssignedJob job = Mockito.mock(AssignedJob.class);
+ Mockito.when(job.getAssignedWorker()).thenReturn(new
BackendWorker(0, backend));
+ Mockito.when(job.instanceId()).thenReturn(new TUniqueId(2, i *
instancesPerBackend + j));
+ jobs.add(job);
+ }
+ }
+
Mockito.when(plan.getInstanceJobs()).thenReturn(ImmutableList.copyOf(jobs));
+
Mockito.when(plan.getFragmentJob().getFragment().getOutputExprs()).thenReturn(new
ArrayList<>());
+ NereidsCoordinator coordinator =
Mockito.mock(NereidsCoordinator.class, Mockito.CALLS_REAL_METHODS);
+ Mockito.doReturn(QUERY_ID).when(coordinator).getQueryId();
+ coordinator.processTopSink(coordinatorContext, plan);
+ return context.getFlightSqlEndpointsLocations();
+ }
+
+ private void setContextField(CoordinatorContext context, String name,
Object value) throws Exception {
+ Field field = CoordinatorContext.class.getDeclaredField(name);
+ field.setAccessible(true);
+ field.set(context, value);
+ }
+}
diff --git
a/regression-test/suites/arrow_flight_sql_p0/test_flight_parallel_partitions.groovy
b/regression-test/suites/arrow_flight_sql_p0/test_flight_parallel_partitions.groovy
new file mode 100644
index 00000000000..005341df0ee
--- /dev/null
+++
b/regression-test/suites/arrow_flight_sql_p0/test_flight_parallel_partitions.groovy
@@ -0,0 +1,134 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+import org.apache.arrow.driver.jdbc.shaded.com.google.protobuf.Any
+import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.CallOptions
+import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.FlightClient
+import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.Location
+import
org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.sql.FlightSqlClient
+import
org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.sql.impl.FlightSql
+import
org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.memory.RootAllocator
+
+import java.util.concurrent.Callable
+import java.util.concurrent.Executors
+import java.util.concurrent.TimeUnit
+
+suite("test_flight_parallel_partitions", "arrow_flight_sql") {
+ def frontend = jdbc_sql_return_maparray("SHOW FRONTENDS").find {
+ it.IsMaster.toString().equalsIgnoreCase("true") &&
it.Alive.toString().equalsIgnoreCase("true")
+ }
+ assertNotNull(frontend)
+ assertTrue(frontend.ArrowFlightSqlPort.toString().toInteger() > 0)
+ def database = jdbc_sql("SELECT DATABASE()")[0][0]
+ def table = "${database}.flight_parallel_partition_source"
+ def allocator = new RootAllocator(Long.MAX_VALUE)
+ def feClient = FlightClient.builder(allocator,
+ Location.forGrpcInsecure(frontend.Host.toString(),
frontend.ArrowFlightSqlPort.toString().toInteger())).build()
+ def client = new FlightSqlClient(feClient)
+ def auth
+ def readEndpoint = { endpoint, boolean collectRows ->
+ def rows = []
+ FlightClient.builder(allocator,
endpoint.locations[0]).build().withCloseable { beClient ->
+ beClient.getStream(endpoint.ticket, auth, CallOptions.timeout(30,
TimeUnit.SECONDS)).withCloseable { stream ->
+ while (stream.next()) {
+ if (collectRows) {
+ def vector = stream.root.getVector(0)
+ for (int i = 0; i < stream.root.rowCount; i++) {
+ rows.add(((Number)
vector.getObject(i)).longValue())
+ }
+ }
+ }
+ }
+ }
+ rows
+ }
+ def executeSetting = { String statement ->
+ client.execute(statement, auth).endpoints.each { readEndpoint(it,
false) }
+ }
+ try {
+ auth =
feClient.authenticateBasicToken(context.config.otherConfigs.get("extArrowFlightSqlUser"),
+
context.config.otherConfigs.get("extArrowFlightSqlPassword")).get()
+ executeSetting("SET enable_sql_cache=false")
+ executeSetting("SET enable_nereids_distribute_planner=true")
+ executeSetting("SET parallel_pipeline_task_num=8")
+ executeSetting("SET query_timeout=60")
+ jdbc_sql("DROP TABLE IF EXISTS ${table}")
+ jdbc_sql("CREATE TABLE ${table} (id BIGINT NOT NULL) DISTRIBUTED BY
HASH(id) BUCKETS 60 " +
+ "PROPERTIES(\"replication_num\"=\"1\")")
+ jdbc_sql("INSERT INTO ${table} SELECT number FROM
numbers(\"number\"=\"60\")")
+ def resultBackendCount = jdbc_sql_return_maparray("SHOW TABLETS FROM
${table}")
+ .collect { it.BackendId }.unique().size()
+ [true, false].each { parallel ->
+ executeSetting("SET enable_parallel_result_sink=${parallel}")
+ [false, true].each { concurrent ->
+ // Each execution owns fresh tickets; consuming a ticket does
not create a replayable partition.
+ def info = client.execute("SELECT id * 1000 + n AS sequence_id
FROM ${table} " +
+ "LATERAL VIEW explode_numbers(1000) expanded AS n",
auth)
+ assertTrue(!info.endpoints.isEmpty())
+ def tickets = info.endpoints.collect {
Base64.encoder.encodeToString(it.ticket.bytes) }
+ assertEquals(tickets.size(), tickets.toSet().size(),
"Duplicate Flight result tickets")
+ if (parallel) {
+ def resultAddresses = info.endpoints.collect { endpoint ->
+ def fields =
Any.parseFrom(endpoint.ticket.bytes).unpack(FlightSql.TicketStatementQuery.class)
+ .statementHandle.toStringUtf8().split("&")
+ "${fields[1]}:${fields[2]}".toString()
+ }
+ // Instance parallelism must not publish multiple readers
for the same BE result buffer.
+ assertEquals(info.endpoints.size(),
resultAddresses.toSet().size())
+ if (resultBackendCount > 1) {
+ assertTrue(info.endpoints.size() > 1, "Expected
multiple result backends")
+ }
+ } else {
+ assertEquals(1, info.endpoints.size())
+ }
+ def rows = []
+ if (concurrent) {
+ def executor = Executors.newFixedThreadPool(Math.min(8,
info.endpoints.size()))
+ try {
+ def futures = info.endpoints.collect { endpoint ->
+ executor.submit({ readEndpoint(endpoint, true) }
as Callable)
+ }
+ futures.each { rows.addAll(it.get(60,
TimeUnit.SECONDS)) }
+ } finally {
+ executor.shutdownNow()
+ assertTrue(executor.awaitTermination(35,
TimeUnit.SECONDS))
+ }
+ } else {
+ info.endpoints.each { rows.addAll(readEndpoint(it, true)) }
+ }
+ assertEquals(60000, rows.size())
+ assertEquals((0L..<60000L).toList(), rows.sort())
+ }
+ }
+ } finally {
+ try {
+ if (auth != null) {
+ feClient.closeSession(new
org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.CloseSessionRequest(),
auth)
+ }
+ } finally {
+ try {
+ client.close()
+ } finally {
+ try {
+ allocator.close()
+ } finally {
+ jdbc_sql("DROP TABLE IF EXISTS ${table}")
+ }
+ }
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
