This is an automated email from the ASF dual-hosted git repository. yiguolei pushed a commit to branch branch-4.2 in repository https://gitbox.apache.org/repos/asf/doris.git
commit 7d0c1bf425e1692e87ed84b737ff1c8644ed77e3 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]
