[feature-wip](arrow-flight)(step5) Support JDBC and PreparedStatement and Fix Bug

be/src/common/config.cpp

@@ -61,6 +61,8 @@ DEFINE_Int32(brpc_port, "8060");
 
 DEFINE_Int32(arrow_flight_sql_port, "-1");
 
+DEFINE_mString(public_access_ip, "");
+
 // the number of bthreads for brpc, the default value is set to -1,
 // which means the number of bthreads is #cpu-cores
 DEFINE_Int32(brpc_num_threads, "256");

be/src/common/config.h

@@ -100,6 +100,11 @@ DECLARE_Int32(brpc_port);
 // Default -1, do not start arrow flight sql server.
 DECLARE_Int32(arrow_flight_sql_port);
 
+// If priority_networks is incorrect but cannot be modified, set public_access_ip as BE’s real IP.
+// For ADBC client fetch result, default is empty, the ADBC client uses the backend ip to fetch the result.
+// If ADBC client cannot access the backend ip, can set public_access_ip to modify the fetch result ip.
+DECLARE_mString(public_access_ip);
+
 // the number of bthreads for brpc, the default value is set to -1,
 // which means the number of bthreads is #cpu-cores
 DECLARE_Int32(brpc_num_threads);

be/src/service/internal_service.cpp

@@ -725,6 +725,9 @@ void PInternalServiceImpl::fetch_arrow_flight_schema(google::protobuf::RpcContro
  auto st = serialize_arrow_schema(&schema, &schema_str);
  if (st.ok()) {
  result->set_schema(std::move(schema_str));
+ if (config::public_access_ip != "") {
+ result->set_be_arrow_flight_ip(config::public_access_ip);
+ }
  }
  st.to_protobuf(result->mutable_status());
  });

be/src/util/arrow/row_batch.cpp

@@ -168,12 +168,13 @@ Status convert_to_arrow_schema(const RowDescriptor& row_desc,
 Status convert_expr_ctxs_arrow_schema(const vectorized::VExprContextSPtrs& output_vexpr_ctxs,
  std::shared_ptr<arrow::Schema>* result) {
  std::vector<std::shared_ptr<arrow::Field>> fields;
- for (auto expr_ctx : output_vexpr_ctxs) {
+ for (int i = 0; i < output_vexpr_ctxs.size(); i++) {
  std::shared_ptr<arrow::DataType> arrow_type;
- auto root_expr = expr_ctx->root();
+ auto root_expr = output_vexpr_ctxs.at(i)->root();
  RETURN_IF_ERROR(convert_to_arrow_type(root_expr->type(), &arrow_type));
- auto field_name = root_expr->is_slot_ref() ? root_expr->expr_name()
- : root_expr->data_type()->get_name();
+ auto field_name = root_expr->is_slot_ref() && !root_expr->expr_name().empty()
+ ? root_expr->expr_name()
+ : fmt::format("{}_{}", root_expr->data_type()->get_name(), i);
  fields.push_back(
  std::make_shared<arrow::Field>(field_name, arrow_type, root_expr->is_nullable()));
  }

fe/fe-core/pom.xml

@@ -699,6 +699,10 @@ under the License.
  <groupId>io.grpc</groupId>
  <artifactId>grpc-core</artifactId>
  </dependency>
+ <dependency>
+ <groupId>org.apache.arrow</groupId>
+ <artifactId>flight-sql-jdbc-driver</artifactId>
+ </dependency>
  <dependency>
  <groupId>io.grpc</groupId>
  <artifactId>grpc-context</artifactId>

fe/fe-core/src/main/java/org/apache/doris/planner/DistributedPlanner.java

@@ -268,7 +268,6 @@ private PlanFragment createMergeFragment(PlanFragment inputFragment)
  mergePlan.init(ctx.getRootAnalyzer());
  Preconditions.checkState(mergePlan.hasValidStats());
  PlanFragment fragment = new PlanFragment(ctx.getNextFragmentId(), mergePlan, DataPartition.UNPARTITIONED);
- fragment.setResultSinkType(ctx.getRootAnalyzer().getContext().getResultSinkType());
  inputFragment.setDestination(mergePlan);
  return fragment;
  }

fe/fe-core/src/main/java/org/apache/doris/planner/OriginalPlanner.java

@@ -257,6 +257,7 @@ public void createPlanFragments(StatementBase statement, Analyzer analyzer, TQue
  LOG.debug("substitute result Exprs {}", resExprs);
  rootFragment.setOutputExprs(resExprs);
  }
+ rootFragment.setResultSinkType(ConnectContext.get().getResultSinkType());
  LOG.debug("finalize plan fragments");
  for (PlanFragment fragment : fragments) {
  fragment.finalize(queryStmt);

fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java

@@ -112,6 +112,7 @@ public enum ConnectType {
  protected volatile long loginTime;
  // for arrow flight
  protected volatile String peerIdentity;
+ private final Map<String, String> preparedQuerys = new HashMap<>();
  private String runningQuery;
  private TNetworkAddress resultFlightServerAddr;
  private TNetworkAddress resultInternalServiceAddr;
@@ -611,6 +612,18 @@ public void resetLoginTime() {
  this.loginTime = System.currentTimeMillis();
  }
 
+ public void addPreparedQuery(String preparedStatementId, String preparedQuery) {
+ preparedQuerys.put(preparedStatementId, preparedQuery);
+ }
+
+ public String getPreparedQuery(String preparedStatementId) {
+ return preparedQuerys.get(preparedStatementId);
+ }
+
+ public void removePreparedQuery(String preparedStatementId) {
+ preparedQuerys.remove(preparedStatementId);
+ }
+
  public void setRunningQuery(String runningQuery) {
  this.runningQuery = runningQuery;
  }

fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java

@@ -1687,7 +1687,7 @@ private TNetworkAddress toArrowFlightHost(TNetworkAddress host) throws Exception
  if (backend.getArrowFlightSqlPort() < 0) {
  return null;
  }
- return new TNetworkAddress(backend.getHost(), backend.getArrowFlightSqlPort());
+ return backend.getArrowFlightAddress();
  }
 
  // estimate if this fragment contains UnionNode

fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java

@@ -605,9 +605,6 @@ public void finalizeQuery() {
  }
 
  private void handleQueryWithRetry(TUniqueId queryId) throws Exception {
- if (context.getConnectType() == ConnectType.ARROW_FLIGHT_SQL) {
- context.setReturnResultFromLocal(false);
- }
  // queue query here
  syncJournalIfNeeded();
  QueueOfferToken offerRet = null;
@@ -642,6 +639,9 @@ private void handleQueryWithRetry(TUniqueId queryId) throws Exception {
  DebugUtil.printId(queryId), i, DebugUtil.printId(newQueryId));
  context.setQueryId(newQueryId);
  }
+ if (context.getConnectType() == ConnectType.ARROW_FLIGHT_SQL) {
+ context.setReturnResultFromLocal(false);
+ }
  handleQueryStmt();
  break;
  } catch (RpcException e) {
@@ -2305,18 +2305,23 @@ private void handleLockTablesStmt() {
  }
 
  public void handleExplainStmt(String result, boolean isNereids) throws IOException {
- // TODO support arrow flight sql
  ShowResultSetMetaData metaData = ShowResultSetMetaData.builder()
  .addColumn(new Column("Explain String" + (isNereids ? "(Nereids Planner)" : "(Old Planner)"),
  ScalarType.createVarchar(20)))
  .build();
- sendMetaData(metaData);
+ if (context.getConnectType() == ConnectType.MYSQL) {
+ sendMetaData(metaData);
 
- // Send result set.
- for (String item : result.split("\n")) {
- serializer.reset();
- serializer.writeLenEncodedString(item);
- context.getMysqlChannel().sendOnePacket(serializer.toByteBuffer());
+ // Send result set.
+ for (String item : result.split("\n")) {
+ serializer.reset();
+ serializer.writeLenEncodedString(item);
+ context.getMysqlChannel().sendOnePacket(serializer.toByteBuffer());
+ }
+ } else if (context.getConnectType() == ConnectType.ARROW_FLIGHT_SQL) {
+ context.getFlightSqlChannel()
+ .addResult(DebugUtil.printId(context.queryId()), context.getRunningQuery(), metaData, result);
+ context.setReturnResultFromLocal(true);
  }
  context.getState().setEof();
  }