This is an automated email from the ASF dual-hosted git repository.
JNSimba pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris-spark-connector.git
The following commit(s) were added to refs/heads/master by this push:
new 36e000b [Improve] Use Arrow Flight SQL as default read mode (#367)
36e000b is described below
commit 36e000bc7782a84467b41fdfd905bffa1395b83e
Author: wudi <[email protected]>
AuthorDate: Tue Aug 25 09:40:29 2026 +0800
[Improve] Use Arrow Flight SQL as default read mode (#367)
Use Arrow Flight SQL as the default read mode.
---
.../doris/spark/client/DorisFrontendClient.java | 96 +++++++++-------
.../spark/client/read/DorisReadModeResolver.java | 91 ++++++++++++++++
.../client/write/AbstractStreamLoadProcessor.java | 10 +-
.../apache/doris/spark/config/DorisOptions.java | 4 +-
.../org/apache/doris/spark/rdd/DorisRDD.scala | 4 +-
.../apache/doris/spark/sql/ScalaDorisRowRDD.scala | 3 +-
.../spark/client/DorisFrontendClientTest.java | 36 ++++++
.../client/read/DorisReadModeResolverTest.java | 121 +++++++++++++++++++++
.../spark/sql/DorisWriterFailoverITCase.scala | 30 ++++-
.../apache/doris/spark/sql/DorisWriterITCase.scala | 2 +-
.../doris/spark/read/AbstractDorisScan.scala | 6 +-
.../doris/spark/read/AbstractDorisScan.scala | 6 +-
12 files changed, 352 insertions(+), 57 deletions(-)
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/DorisFrontendClient.java
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/DorisFrontendClient.java
index f0c57e8..3d3fa0b 100644
---
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/DorisFrontendClient.java
+++
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/DorisFrontendClient.java
@@ -116,43 +116,8 @@ public class DorisFrontendClient implements Serializable {
private LoadBalanceList<Frontend> initFrontends(DorisConfig config) throws
Exception {
String frontendNodes = config.getValue(DorisOptions.DORIS_FENODES);
String[] frontendNodeArray = frontendNodes.split(",");
- List<Frontend> frontendList = null;
if (config.getValue(DorisOptions.DORIS_FE_AUTO_FETCH)) {
- Exception ex = null;
- for (String frontendNode : frontendNodeArray) {
- String[] nodeDetails = frontendNode.split(":");
- try {
- LoadBalanceList<Frontend> list = new LoadBalanceList<>(
- Collections.singletonList(new Frontend(nodeDetails[0],
- nodeDetails.length > 1 ?
Integer.parseInt(nodeDetails[1]) : -1)));
- frontendList = requestFrontends(list, (frontend, client)
-> {
- String url = URLs.getFrontEndNodes(frontend.getHost(),
frontend.getHttpPort(),
- isHttpsEnabled);
- HttpGet httpGet = new HttpGet(url);
- HttpUtils.setAuth(httpGet, username, password);
- JsonNode dataNode;
- try {
- HttpResponse response = client.execute(httpGet);
- dataNode = extractDataFromResponse(response, url);
- } catch (IOException e) {
- throw new RuntimeException("fetch fe failed", e);
- }
- ArrayNode columnNames = (ArrayNode)
dataNode.get("columnNames");
- ArrayNode rows = (ArrayNode) dataNode.get("rows");
- return parseFrontends(columnNames, rows);
- });
- } catch (Exception e) {
- LOG.warn("fetch fe request on {} failed, err: {}",
frontendNode, e.getMessage());
- ex = e;
- }
- }
- if (frontendList == null || frontendList.isEmpty()) {
- if (ex == null) {
- throw new DorisException("frontend init fetch failed,
empty frontend list");
- }
- throw new DorisException("frontend init fetch failed", ex);
- }
- return new LoadBalanceList<>(frontendList);
+ return new LoadBalanceList<>(fetchFrontends(frontendNodeArray));
} else {
int queryPort = config.contains(DorisOptions.DORIS_QUERY_PORT) ?
config.getValue(DorisOptions.DORIS_QUERY_PORT) : -1;
@@ -167,6 +132,63 @@ public class DorisFrontendClient implements Serializable {
}
}
+ private List<Frontend> fetchFrontends(String[] frontendNodeArray) throws
Exception {
+ List<Frontend> frontendList = null;
+ Exception ex = null;
+ for (String frontendNode : frontendNodeArray) {
+ String[] nodeDetails = frontendNode.split(":");
+ try {
+ LoadBalanceList<Frontend> list = new LoadBalanceList<>(
+ Collections.singletonList(new Frontend(nodeDetails[0],
+ nodeDetails.length > 1 ?
Integer.parseInt(nodeDetails[1]) : -1)));
+ frontendList = requestFrontends(list, (frontend, client) -> {
+ String url = URLs.getFrontEndNodes(frontend.getHost(),
frontend.getHttpPort(),
+ isHttpsEnabled);
+ HttpGet httpGet = new HttpGet(url);
+ HttpUtils.setAuth(httpGet, username, password);
+ JsonNode dataNode;
+ try {
+ HttpResponse response = client.execute(httpGet);
+ dataNode = extractDataFromResponse(response, url);
+ } catch (IOException e) {
+ throw new RuntimeException("fetch fe failed", e);
+ }
+ ArrayNode columnNames = (ArrayNode)
dataNode.get("columnNames");
+ ArrayNode rows = (ArrayNode) dataNode.get("rows");
+ return parseFrontends(columnNames, rows);
+ });
+ } catch (Exception e) {
+ LOG.warn("fetch fe request on {} failed, err: {}",
frontendNode, e.getMessage());
+ ex = e;
+ }
+ }
+ if (frontendList == null || frontendList.isEmpty()) {
+ if (ex == null) {
+ throw new DorisException("frontend init fetch failed, empty
frontend list");
+ }
+ throw new DorisException("frontend init fetch failed", ex);
+ }
+ return frontendList;
+ }
+
+ public int tryGetArrowFlightSqlPort() {
+ int flightSqlPort = findArrowFlightSqlPort(frontends.getList());
+ if (flightSqlPort > 0) {
+ return flightSqlPort;
+ }
+ try {
+ String frontendNodes = config.getValue(DorisOptions.DORIS_FENODES);
+ return
findArrowFlightSqlPort(fetchFrontends(frontendNodes.split(",")));
+ } catch (Exception e) {
+ LOG.warn("failed to get Arrow Flight SQL port, err: {}",
e.getMessage());
+ return -1;
+ }
+ }
+
+ static int findArrowFlightSqlPort(List<Frontend> frontends) {
+ return frontends.isEmpty() ? -1 : frontends.get(0).getFlightSqlPort();
+ }
+
public <T> T requestFrontends(BiFunction<Frontend, CloseableHttpClient, T>
reqFunc) throws Exception {
return requestFrontends(frontends, reqFunc);
}
@@ -326,7 +348,7 @@ public class DorisFrontendClient implements Serializable {
});
}
- private List<Frontend> parseFrontends(ArrayNode columnNames, ArrayNode
rows) {
+ List<Frontend> parseFrontends(ArrayNode columnNames, ArrayNode rows) {
int hostIdx = -1;
int httpPortIdx = -1;
int queryPortIdx = -1;
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/read/DorisReadModeResolver.java
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/read/DorisReadModeResolver.java
new file mode 100644
index 0000000..f4766d1
--- /dev/null
+++
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/read/DorisReadModeResolver.java
@@ -0,0 +1,91 @@
+// 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.spark.client.read;
+
+import org.apache.doris.spark.client.DorisFrontendClient;
+import org.apache.doris.spark.config.DorisConfig;
+import org.apache.doris.spark.config.DorisOptions;
+import org.apache.doris.spark.exception.DorisRuntimeException;
+import org.apache.doris.spark.exception.OptionRequiredException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Locale;
+
+public class DorisReadModeResolver {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(DorisReadModeResolver.class);
+ private static final String ARROW = "arrow";
+ private static final String THRIFT = "thrift";
+
+ private DorisReadModeResolver() {
+ }
+
+ public static String resolve(DorisConfig config) {
+ try {
+ return resolve(config, () -> discoverPort(new
DorisFrontendClient(config)));
+ } catch (OptionRequiredException e) {
+ throw new DorisRuntimeException("failed to resolve Doris read
mode", e);
+ }
+ }
+
+ private static int discoverPort(DorisFrontendClient frontendClient) throws
Exception {
+ try {
+ return frontendClient.tryGetArrowFlightSqlPort();
+ } finally {
+ frontendClient.close();
+ }
+ }
+
+ static String resolve(DorisConfig config, FlightSqlPortSupplier
portSupplier)
+ throws OptionRequiredException {
+ String readMode =
config.getValue(DorisOptions.READ_MODE).toLowerCase(Locale.ROOT);
+ if (THRIFT.equals(readMode)) {
+ return THRIFT;
+ }
+ if (!ARROW.equals(readMode)) {
+ throw new IllegalArgumentException("Unknown read mode: " +
readMode);
+ }
+
+ int flightSqlPort =
config.getValue(DorisOptions.DORIS_READ_FLIGHT_SQL_PORT);
+ if (flightSqlPort > 0) {
+ return ARROW;
+ }
+
+ try {
+ flightSqlPort = portSupplier.get();
+ } catch (Exception e) {
+ LOG.warn("failed to discover Arrow Flight SQL port, falling back
to Thrift", e);
+ return THRIFT;
+ }
+ if (flightSqlPort > 0) {
+ config.setProperty(DorisOptions.DORIS_READ_FLIGHT_SQL_PORT,
String.valueOf(flightSqlPort));
+ LOG.info("use Arrow Flight SQL to read data, port is {}",
flightSqlPort);
+ return ARROW;
+ }
+
+ LOG.warn("Arrow Flight SQL port {} is invalid or unavailable, falling
back to Thrift", flightSqlPort);
+ return THRIFT;
+ }
+
+ @FunctionalInterface
+ interface FlightSqlPortSupplier {
+ int get() throws Exception;
+ }
+}
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/AbstractStreamLoadProcessor.java
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/AbstractStreamLoadProcessor.java
index 6283f65..941fbfc 100644
---
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/AbstractStreamLoadProcessor.java
+++
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/client/write/AbstractStreamLoadProcessor.java
@@ -98,6 +98,7 @@ public abstract class AbstractStreamLoadProcessor<R> extends
DorisWriter<R> impl
private Future<StreamLoadResponse> requestFuture = null;
private volatile String currentLabel;
+ private volatile boolean closing;
private Exception unexpectedException = null;
public AbstractStreamLoadProcessor(DorisConfig config) throws Exception {
@@ -429,9 +430,11 @@ public abstract class AbstractStreamLoadProcessor<R>
extends DorisWriter<R> impl
+ ", msg: " +
streamLoadResponse.getMessage());
}
} catch (Exception e) {
- logger.error("stream load exception", e);
- unexpectedException = e;
- currentThread.interrupt();
+ if (!closing) {
+ logger.error("stream load exception", e);
+ unexpectedException = e;
+ currentThread.interrupt();
+ }
throw e;
}
return streamLoadResponse;
@@ -440,6 +443,7 @@ public abstract class AbstractStreamLoadProcessor<R>
extends DorisWriter<R> impl
@Override
public void close() throws IOException {
+ closing = true;
createNewBatch = true;
isFirstRecordOfBatch = true;
unexpectedException = null;
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/config/DorisOptions.java
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/config/DorisOptions.java
index 10e9972..3d1b3a2 100644
---
a/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/config/DorisOptions.java
+++
b/spark-doris-connector/spark-doris-connector-base/src/main/java/org/apache/doris/spark/config/DorisOptions.java
@@ -126,11 +126,11 @@ public class DorisOptions {
public static final ConfigOption<String> LOAD_MODE =
ConfigOptions.name("doris.sink.mode").stringType().defaultValue("stream_load").withDescription("Write
mode, supports stream_load, copy_into and tvf.");
- public static final ConfigOption<String> READ_MODE =
ConfigOptions.name("doris.read.mode").stringType().defaultValue("thrift").withDescription("");
+ public static final ConfigOption<String> READ_MODE =
ConfigOptions.name("doris.read.mode").stringType().defaultValue("arrow").withDescription("");
public static final ConfigOption<String> DORIS_READ_FLIGHT_SQL_PREFIX =
ConfigOptions.name("doris.read.arrow-flight-sql.prefix").stringType().defaultValue("ApplicationName=Spark
ArrowFlightSQL Query").withDescription("");
- public static final ConfigOption<Integer> DORIS_READ_FLIGHT_SQL_PORT =
ConfigOptions.name("doris.read.arrow-flight-sql.port").intType().withoutDefaultValue().withDescription("");
+ public static final ConfigOption<Integer> DORIS_READ_FLIGHT_SQL_PORT =
ConfigOptions.name("doris.read.arrow-flight-sql.port").intType().defaultValue(-1).withDescription("");
public static final ConfigOption<String> DORIS_SINK_LABEL_PREFIX =
ConfigOptions.name("doris.sink.label.prefix").stringType().defaultValue("spark-doris").withDescription("Label
prefix used by Doris sink writes.");
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/main/scala/org/apache/doris/spark/rdd/DorisRDD.scala
b/spark-doris-connector/spark-doris-connector-base/src/main/scala/org/apache/doris/spark/rdd/DorisRDD.scala
index d6653ef..86e217d 100644
---
a/spark-doris-connector/spark-doris-connector-base/src/main/scala/org/apache/doris/spark/rdd/DorisRDD.scala
+++
b/spark-doris-connector/spark-doris-connector-base/src/main/scala/org/apache/doris/spark/rdd/DorisRDD.scala
@@ -18,7 +18,7 @@
package org.apache.doris.spark.rdd
import org.apache.doris.spark.client.entity.DorisReaderPartition
-import org.apache.doris.spark.client.read.{DorisFlightSqlReader,
DorisThriftReader}
+import org.apache.doris.spark.client.read.{DorisFlightSqlReader,
DorisReadModeResolver, DorisThriftReader}
import org.apache.doris.spark.config.{DorisConfig, DorisOptions}
import org.apache.spark.{Partition, SparkContext, TaskContext}
@@ -39,7 +39,7 @@ private[spark] class ScalaDorisRDDIterator[T](
extends AbstractDorisRDDIterator[T](context, partition) {
override def initReader(config: DorisConfig): Unit = {
- config.getValue(DorisOptions.READ_MODE).toLowerCase match {
+ DorisReadModeResolver.resolve(config) match {
case "thrift" =>
config.setProperty(DorisOptions.DORIS_VALUE_READER_CLASS,
classOf[DorisThriftReader].getName)
case "arrow" =>
config.setProperty(DorisOptions.DORIS_VALUE_READER_CLASS,
classOf[DorisFlightSqlReader].getName)
case rm: String => throw new IllegalArgumentException("Unknown read
mode: " + rm)
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/main/scala/org/apache/doris/spark/sql/ScalaDorisRowRDD.scala
b/spark-doris-connector/spark-doris-connector-base/src/main/scala/org/apache/doris/spark/sql/ScalaDorisRowRDD.scala
index 0e6038d..1e59612 100644
---
a/spark-doris-connector/spark-doris-connector-base/src/main/scala/org/apache/doris/spark/sql/ScalaDorisRowRDD.scala
+++
b/spark-doris-connector/spark-doris-connector-base/src/main/scala/org/apache/doris/spark/sql/ScalaDorisRowRDD.scala
@@ -18,6 +18,7 @@
package org.apache.doris.spark.sql
import org.apache.doris.spark.client.entity.DorisReaderPartition
+import org.apache.doris.spark.client.read.DorisReadModeResolver
import org.apache.doris.spark.config.{DorisConfig, DorisOptions}
import org.apache.doris.spark.rdd.{AbstractDorisRDD, AbstractDorisRDDIterator,
DorisPartition}
import org.apache.spark.sql.Row
@@ -38,7 +39,7 @@ private[spark] class ScalaDorisRowRDDIterator(context:
TaskContext,
extends AbstractDorisRDDIterator[Row](context, partition) {
override def initReader(config: DorisConfig): Unit = {
- config.getValue(DorisOptions.READ_MODE).toLowerCase match {
+ DorisReadModeResolver.resolve(config) match {
case "thrift" =>
config.setProperty(DorisOptions.DORIS_VALUE_READER_CLASS,
classOf[DorisRowThriftReader].getName)
case "arrow" =>
config.setProperty(DorisOptions.DORIS_VALUE_READER_CLASS,
classOf[DorisRowFlightSqlReader].getName)
case rm: String => throw new IllegalArgumentException("Unknown read
mode: " + rm)
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/DorisFrontendClientTest.java
b/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/DorisFrontendClientTest.java
index 5a5b6b6..bf41a9b 100644
---
a/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/DorisFrontendClientTest.java
+++
b/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/DorisFrontendClientTest.java
@@ -18,14 +18,19 @@
package org.apache.doris.spark.client;
import org.apache.doris.spark.client.entity.Backend;
+import org.apache.doris.spark.client.entity.Frontend;
import org.apache.doris.spark.testutil.HttpsTestServer;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpGet;
import org.junit.Assert;
import org.junit.Test;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+
import javax.net.ssl.SSLHandshakeException;
+import java.util.Arrays;
import java.util.List;
public class DorisFrontendClientTest {
@@ -130,6 +135,37 @@ public class DorisFrontendClientTest {
}
}
+ @Test
+ public void findArrowFlightSqlPortReturnsFirstFrontendPort() {
+ List<Frontend> frontends = Arrays.asList(
+ new Frontend("192.168.1.1", 8030, 9030, -1),
+ new Frontend("192.168.1.2", 8030, 9030, 9040));
+
+ Assert.assertEquals(-1,
DorisFrontendClient.findArrowFlightSqlPort(frontends));
+ }
+
+ @Test
+ public void findArrowFlightSqlPortReturnsInvalidPortWhenMissing() {
+ List<Frontend> frontends = Arrays.asList(
+ new Frontend("192.168.1.1", 8030, 9030, -1),
+ new Frontend("192.168.1.2", 8030, 9030, 0));
+
+ Assert.assertEquals(-1,
DorisFrontendClient.findArrowFlightSqlPort(frontends));
+ }
+
+ @Test
+ public void parseFrontendsIncludesArrowFlightSqlPort() throws Exception {
+ ObjectMapper mapper = new ObjectMapper();
+ ArrayNode columnNames = (ArrayNode) mapper.readTree(
+
"[\"Host\",\"HttpPort\",\"QueryPort\",\"ArrowFlightSqlPort\"]");
+ ArrayNode rows = (ArrayNode) mapper.readTree(
+ "[[\"192.168.1.1\",\"8030\",\"9030\",\"9040\"]]");
+
+ List<Frontend> frontends = new
DorisFrontendClient().parseFrontends(columnNames, rows);
+
+ Assert.assertEquals(9040,
DorisFrontendClient.findArrowFlightSqlPort(frontends));
+ }
+
private static boolean hasCause(Throwable failure, Class<? extends
Throwable> causeType) {
for (Throwable current = failure; current != null; current =
current.getCause()) {
if (causeType.isInstance(current)) {
diff --git
a/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/read/DorisReadModeResolverTest.java
b/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/read/DorisReadModeResolverTest.java
new file mode 100644
index 0000000..96d5f70
--- /dev/null
+++
b/spark-doris-connector/spark-doris-connector-base/src/test/java/org/apache/doris/spark/client/read/DorisReadModeResolverTest.java
@@ -0,0 +1,121 @@
+// 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.spark.client.read;
+
+import org.apache.doris.spark.config.DorisConfig;
+import org.apache.doris.spark.config.DorisOptions;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+public class DorisReadModeResolverTest {
+
+ @Test
+ public void readOptionsDefaultToArrowWithInvalidFlightSqlPort() {
+ Assert.assertEquals("arrow", DorisOptions.READ_MODE.getDefaultValue());
+ Assert.assertEquals(Integer.valueOf(-1),
DorisOptions.DORIS_READ_FLIGHT_SQL_PORT.getDefaultValue());
+ }
+
+ @Test
+ public void defaultArrowModeDiscoversFlightSqlPort() throws Exception {
+ DorisConfig config = createConfig(null, null);
+
+ String mode = DorisReadModeResolver.resolve(config, () -> 9040);
+
+ Assert.assertEquals("arrow", mode);
+ Assert.assertEquals(Integer.valueOf(9040),
config.getValue(DorisOptions.DORIS_READ_FLIGHT_SQL_PORT));
+ }
+
+ @Test
+ public void explicitThriftSkipsPortDiscovery() throws Exception {
+ DorisConfig config = createConfig("thrift", null);
+ AtomicBoolean discovered = new AtomicBoolean(false);
+
+ String mode = DorisReadModeResolver.resolve(config, () -> {
+ discovered.set(true);
+ return 9040;
+ });
+
+ Assert.assertEquals("thrift", mode);
+ Assert.assertFalse(discovered.get());
+ }
+
+ @Test
+ public void configuredPositivePortKeepsArrowMode() throws Exception {
+ DorisConfig config = createConfig("arrow", 9040);
+ AtomicBoolean discovered = new AtomicBoolean(false);
+
+ String mode = DorisReadModeResolver.resolve(config, () -> {
+ discovered.set(true);
+ return 9050;
+ });
+
+ Assert.assertEquals("arrow", mode);
+ Assert.assertFalse(discovered.get());
+ }
+
+ @Test
+ public void discoveredPositivePortKeepsArrowModeAndUpdatesConfig() throws
Exception {
+ DorisConfig config = createConfig("arrow", null);
+
+ String mode = DorisReadModeResolver.resolve(config, () -> 9040);
+
+ Assert.assertEquals("arrow", mode);
+ Assert.assertEquals(Integer.valueOf(9040),
config.getValue(DorisOptions.DORIS_READ_FLIGHT_SQL_PORT));
+ }
+
+ @Test
+ public void invalidDiscoveredPortFallsBackToThrift() throws Exception {
+ DorisConfig config = createConfig("arrow", null);
+
+ String mode = DorisReadModeResolver.resolve(config, () -> -1);
+
+ Assert.assertEquals("thrift", mode);
+ }
+
+ @Test
+ public void unknownReadModeIsRejected() throws Exception {
+ DorisConfig config = createConfig("unknown", null);
+
+ try {
+ DorisReadModeResolver.resolve(config, () -> 9040);
+ Assert.fail("Expected unknown read mode to be rejected");
+ } catch (IllegalArgumentException e) {
+ Assert.assertEquals("Unknown read mode: unknown", e.getMessage());
+ }
+ }
+
+ private DorisConfig createConfig(String readMode, Integer flightSqlPort)
throws Exception {
+ Map<String, String> options = new HashMap<>();
+ options.put("doris.fenodes", "127.0.0.1:8030");
+ options.put("doris.table.identifier", "db.tbl");
+ options.put("doris.user", "root");
+ options.put("doris.password", "");
+ if (readMode != null) {
+ options.put("doris.read.mode", readMode);
+ }
+ if (flightSqlPort != null) {
+ options.put("doris.read.arrow-flight-sql.port",
flightSqlPort.toString());
+ }
+ return DorisConfig.fromMap(options, false);
+ }
+
+}
diff --git
a/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/DorisWriterFailoverITCase.scala
b/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/DorisWriterFailoverITCase.scala
index b1edd80..475d7bd 100644
---
a/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/DorisWriterFailoverITCase.scala
+++
b/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/DorisWriterFailoverITCase.scala
@@ -183,15 +183,35 @@ class DorisWriterFailoverITCase extends
AbstractContainerTestBase {
session.stop()
}
- val actual = ContainerUtils.executeSQLStatement(
- getDorisQueryConnection,
- LOG,
- String.format("select * from %s.%s", DATABASE,
TABLE_WRITE_TBL_TASK_RETRY),
- 2)
val expected = util.Arrays.asList("doris,cn", "spark,us", "catalog,uk");
+ var actual = util.Collections.emptyList[String]()
+ val connection = getDorisQueryConnection
+ try {
+ waitForCondition("task retry rows to become visible") {
+ actual = ContainerUtils.executeSQLStatement(
+ connection,
+ LOG,
+ String.format("select * from %s.%s", DATABASE,
TABLE_WRITE_TBL_TASK_RETRY),
+ 2)
+ actual.size() >= expected.size()
+ }
+ } finally {
+ connection.close()
+ }
checkResultInAnyOrder("testFailoverForTaskRetry", expected.toArray,
actual.toArray)
}
+ private def waitForCondition(description: String)(condition: => Boolean):
Unit = {
+ val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(60)
+ while (System.nanoTime() < deadline) {
+ if (condition) {
+ return
+ }
+ TimeUnit.MILLISECONDS.sleep(100)
+ }
+ throw new AssertionError(s"Timed out waiting for $description")
+ }
+
private def waitForCondition(future: Future[_], description:
String)(condition: => Boolean): Unit = {
val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(60)
while (System.nanoTime() < deadline) {
diff --git
a/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/DorisWriterITCase.scala
b/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/DorisWriterITCase.scala
index 3ebf670..4290352 100644
---
a/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/DorisWriterITCase.scala
+++
b/spark-doris-connector/spark-doris-connector-it/src/test/java/org/apache/doris/spark/sql/DorisWriterITCase.scala
@@ -171,7 +171,7 @@ class DorisWriterITCase extends AbstractContainerTestBase {
("doris_csv", 1),
("spark_csv", 2)
)).toDF("name", "age")
- df.write
+ df.coalesce(1).write
.format("doris")
.option("doris.fenodes", getFenodes)
.option("doris.table.identifier", DATABASE + "." + TABLE_GROUP_COMMIT)
diff --git
a/spark-doris-connector/spark-doris-connector-spark-3-base/src/main/scala/org/apache/doris/spark/read/AbstractDorisScan.scala
b/spark-doris-connector/spark-doris-connector-spark-3-base/src/main/scala/org/apache/doris/spark/read/AbstractDorisScan.scala
index 4afcc01..3d89d50 100644
---
a/spark-doris-connector/spark-doris-connector-spark-3-base/src/main/scala/org/apache/doris/spark/read/AbstractDorisScan.scala
+++
b/spark-doris-connector/spark-doris-connector-spark-3-base/src/main/scala/org/apache/doris/spark/read/AbstractDorisScan.scala
@@ -18,8 +18,8 @@
package org.apache.doris.spark.read
import org.apache.doris.spark.client.entity.{Backend, DorisReaderPartition}
-import org.apache.doris.spark.client.read.ReaderPartitionGenerator
-import org.apache.doris.spark.config.{DorisConfig, DorisOptions}
+import org.apache.doris.spark.client.read.{DorisReadModeResolver,
ReaderPartitionGenerator}
+import org.apache.doris.spark.config.DorisConfig
import org.apache.spark.internal.Logging
import org.apache.spark.sql.connector.read.{Batch, InputPartition,
PartitionReaderFactory, Scan}
import org.apache.spark.sql.internal.SQLConf
@@ -29,7 +29,7 @@ import scala.language.implicitConversions
abstract class AbstractDorisScan(config: DorisConfig, schema: StructType)
extends Scan with Batch with Logging {
- private val scanMode =
ScanMode.valueOf(config.getValue(DorisOptions.READ_MODE).toUpperCase)
+ private val scanMode =
ScanMode.valueOf(DorisReadModeResolver.resolve(config).toUpperCase)
override def readSchema(): StructType = schema
diff --git
a/spark-doris-connector/spark-doris-connector-spark-4-base/src/main/scala/org/apache/doris/spark/read/AbstractDorisScan.scala
b/spark-doris-connector/spark-doris-connector-spark-4-base/src/main/scala/org/apache/doris/spark/read/AbstractDorisScan.scala
index 4afcc01..3d89d50 100644
---
a/spark-doris-connector/spark-doris-connector-spark-4-base/src/main/scala/org/apache/doris/spark/read/AbstractDorisScan.scala
+++
b/spark-doris-connector/spark-doris-connector-spark-4-base/src/main/scala/org/apache/doris/spark/read/AbstractDorisScan.scala
@@ -18,8 +18,8 @@
package org.apache.doris.spark.read
import org.apache.doris.spark.client.entity.{Backend, DorisReaderPartition}
-import org.apache.doris.spark.client.read.ReaderPartitionGenerator
-import org.apache.doris.spark.config.{DorisConfig, DorisOptions}
+import org.apache.doris.spark.client.read.{DorisReadModeResolver,
ReaderPartitionGenerator}
+import org.apache.doris.spark.config.DorisConfig
import org.apache.spark.internal.Logging
import org.apache.spark.sql.connector.read.{Batch, InputPartition,
PartitionReaderFactory, Scan}
import org.apache.spark.sql.internal.SQLConf
@@ -29,7 +29,7 @@ import scala.language.implicitConversions
abstract class AbstractDorisScan(config: DorisConfig, schema: StructType)
extends Scan with Batch with Logging {
- private val scanMode =
ScanMode.valueOf(config.getValue(DorisOptions.READ_MODE).toUpperCase)
+ private val scanMode =
ScanMode.valueOf(DorisReadModeResolver.resolve(config).toUpperCase)
override def readSchema(): StructType = schema
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
