From 7f61626c8dfb72c6440e5be79cf652911ac30f93 Mon Sep 17 00:00:00 2001 From: Xinyi Zou Date: Mon, 22 Apr 2024 10:46:30 +0800 Subject: [PATCH] [fix](arrow_flight_sql) Fix ArrowSchema column alias (#33490) run: select TABLE_SCHEMA as a, sum(TABLE_ROWS) as b from tables group by TABLE_SCHEMA limit 2; old output: TABLE_SCHEMA Nullable(Int64)_1 0 regression_test_mv_p0_sum_count 9 1 regression_test_query_p0_sql_functions_string_functions 70414 now output: a b 0 regression_test_mv_p0_sum_count 9 1 regression_test_query_p0_sql_functions_string_functions 70414 --- be/src/util/arrow/row_batch.cpp | 4 ++-- be/src/vec/exprs/vexpr.h | 1 + be/src/vec/exprs/vslot_ref.cpp | 10 +++++++++- be/src/vec/exprs/vslot_ref.h | 2 ++ .../main/java/org/apache/doris/analysis/SlotRef.java | 1 + .../java/org/apache/doris/service/FlightSqlJDBC.java | 1 - gensrc/thrift/Exprs.thrift | 1 + 7 files changed, 16 insertions(+), 4 deletions(-) diff --git a/be/src/util/arrow/row_batch.cpp b/be/src/util/arrow/row_batch.cpp index 439572f75e..8b97c7a76d 100644 --- a/be/src/util/arrow/row_batch.cpp +++ b/be/src/util/arrow/row_batch.cpp @@ -196,8 +196,8 @@ Status convert_expr_ctxs_arrow_schema(const vectorized::VExprContextSPtrs& outpu std::shared_ptr arrow_type; 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().empty() - ? root_expr->expr_name() + auto field_name = root_expr->is_slot_ref() && !root_expr->expr_label().empty() + ? root_expr->expr_label() : fmt::format("{}_{}", root_expr->data_type()->get_name(), i); fields.push_back( std::make_shared(field_name, arrow_type, root_expr->is_nullable())); diff --git a/be/src/vec/exprs/vexpr.h b/be/src/vec/exprs/vexpr.h index 57bb4a1cf6..8ad62100e0 100644 --- a/be/src/vec/exprs/vexpr.h +++ b/be/src/vec/exprs/vexpr.h @@ -84,6 +84,7 @@ public: virtual ~VExpr() = default; virtual const std::string& expr_name() const = 0; + virtual std::string expr_label() { return ""; } /// Initializes this expr instance for execution. This does not include initializing /// state in the VExprContext; 'context' should only be used to register a diff --git a/be/src/vec/exprs/vslot_ref.cpp b/be/src/vec/exprs/vslot_ref.cpp index b683a1fb15..de9a10137f 100644 --- a/be/src/vec/exprs/vslot_ref.cpp +++ b/be/src/vec/exprs/vslot_ref.cpp @@ -38,7 +38,11 @@ class VExprContext; namespace doris::vectorized { VSlotRef::VSlotRef(const doris::TExprNode& node) - : VExpr(node), _slot_id(node.slot_ref.slot_id), _column_id(-1), _column_name(nullptr) {} + : VExpr(node), + _slot_id(node.slot_ref.slot_id), + _column_id(-1), + _column_name(nullptr), + _column_label(node.label) {} VSlotRef::VSlotRef(const SlotDescriptor* desc) : VExpr(desc->type(), true, desc->is_nullable()), @@ -99,6 +103,10 @@ Status VSlotRef::execute(VExprContext* context, Block* block, int* result_column const std::string& VSlotRef::expr_name() const { return *_column_name; } +std::string VSlotRef::expr_label() { + return _column_label; +} + std::string VSlotRef::debug_string() const { std::stringstream out; out << "SlotRef(slot_id=" << _slot_id << VExpr::debug_string() << ")"; diff --git a/be/src/vec/exprs/vslot_ref.h b/be/src/vec/exprs/vslot_ref.h index c30ac64041..f96728eda5 100644 --- a/be/src/vec/exprs/vslot_ref.h +++ b/be/src/vec/exprs/vslot_ref.h @@ -44,6 +44,7 @@ public: Status execute(VExprContext* context, Block* block, int* result_column_id) override; const std::string& expr_name() const override; + std::string expr_label() override; std::string debug_string() const override; bool is_constant() const override { return false; } @@ -55,6 +56,7 @@ private: int _slot_id; int _column_id; const std::string* _column_name = nullptr; + const std::string _column_label; }; } // namespace vectorized } // namespace doris diff --git a/fe/fe-core/src/main/java/org/apache/doris/analysis/SlotRef.java b/fe/fe-core/src/main/java/org/apache/doris/analysis/SlotRef.java index 3eb5422a12..52927badba 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/analysis/SlotRef.java +++ b/fe/fe-core/src/main/java/org/apache/doris/analysis/SlotRef.java @@ -365,6 +365,7 @@ public class SlotRef extends Expr { msg.slot_ref = new TSlotRef(desc.getId().asInt(), desc.getParent().getId().asInt()); msg.slot_ref.setColUniqueId(desc.getUniqueId()); msg.setOutputColumn(outputColumn); + msg.setLabel(label); } @Override diff --git a/fe/fe-core/src/test/java/org/apache/doris/service/FlightSqlJDBC.java b/fe/fe-core/src/test/java/org/apache/doris/service/FlightSqlJDBC.java index 083d9537ea..a3c294bd7a 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/service/FlightSqlJDBC.java +++ b/fe/fe-core/src/test/java/org/apache/doris/service/FlightSqlJDBC.java @@ -42,7 +42,6 @@ public class FlightSqlJDBC { conn = DriverManager.getConnection(DB_URL, USER, PASS); stmt = conn.createStatement(); - stmt.executeQuery("set dry_run_query=true"); stmt.executeQuery("use information_schema;"); String sql = "show tables;"; diff --git a/gensrc/thrift/Exprs.thrift b/gensrc/thrift/Exprs.thrift index 6b7abc2ad9..5d46e4b38d 100644 --- a/gensrc/thrift/Exprs.thrift +++ b/gensrc/thrift/Exprs.thrift @@ -265,6 +265,7 @@ struct TExprNode { 33: optional TMatchPredicate match_predicate 34: optional TIPv4Literal ipv4_literal 35: optional TIPv6Literal ipv6_literal + 36: optional string label // alias name, a/b in `select xxx as a, count(1) as b` } // A flattened representation of a tree of Expr nodes, obtained by depth-first