From 64536a7bfd0f4adc2ee2ebfd7e0e3ce70ab20f24 Mon Sep 17 00:00:00 2001 From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com> Date: Fri, 3 Jan 2025 15:51:55 +0800 Subject: [PATCH] [BugFix] fix HttpResultWriter ill-formatted json string (backport #54590) (#54673) Co-authored-by: Kevin Cai --- be/src/runtime/http_result_writer.cpp | 8 +- be/test/CMakeLists.txt | 1 + be/test/runtime/http_result_writer_test.cpp | 108 ++++++++++++++++++++ 3 files changed, 113 insertions(+), 4 deletions(-) create mode 100644 be/test/runtime/http_result_writer_test.cpp diff --git a/be/src/runtime/http_result_writer.cpp b/be/src/runtime/http_result_writer.cpp index d1fbd096f1d3c..2a17531da73ff 100644 --- a/be/src/runtime/http_result_writer.cpp +++ b/be/src/runtime/http_result_writer.cpp @@ -53,14 +53,14 @@ void HttpResultWriter::_init_profile() { // transform one row into json format Status HttpResultWriter::_transform_row_to_json(const Columns& result_columns, int idx) { - int num_columns = result_columns.size(); + size_t num_columns = result_columns.size(); _row_str.append("{\"data\":["); - for (auto& result_column : result_columns) { + for (size_t i = 0; i < num_columns; ++i) { std::string row; - ASSIGN_OR_RETURN(row, cast_type_to_json_str(result_column, idx)); + ASSIGN_OR_RETURN(row, cast_type_to_json_str(result_columns[i], idx)); _row_str.append(row); - if (result_column != result_columns[num_columns - 1]) { + if (i != num_columns - 1) { _row_str.append(","); } } diff --git a/be/test/CMakeLists.txt b/be/test/CMakeLists.txt index e353113934536..4e489244e95f9 100644 --- a/be/test/CMakeLists.txt +++ b/be/test/CMakeLists.txt @@ -358,6 +358,7 @@ set(EXEC_FILES ./runtime/decimalv3_test.cpp ./runtime/external_scan_context_mgr_test.cpp ./runtime/fragment_mgr_test.cpp + ./runtime/http_result_writer_test.cpp ./runtime/int128_arithmetic_ops_test.cpp ./runtime/kafka_consumer_pipe_test.cpp ./runtime/local_tablets_channel_test.cpp diff --git a/be/test/runtime/http_result_writer_test.cpp b/be/test/runtime/http_result_writer_test.cpp new file mode 100644 index 0000000000000..a39396d698124 --- /dev/null +++ b/be/test/runtime/http_result_writer_test.cpp @@ -0,0 +1,108 @@ +// Copyright 2021-present StarRocks, Inc. All rights reserved. +// +// Licensed 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 +// +// https://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. + +#include "runtime/http_result_writer.h" + +#include + +#include + +#include "column/vectorized_fwd.h" +#include "exprs/column_ref.h" +#include "runtime/buffer_control_block.h" +#include "storage/chunk_helper.h" +#include "testutil/assert.h" +#include "util/uid_util.h" + +namespace starrocks { + +static ChunkPtr create_chunk(SchemaPtr schema, int col0value, int col1value) { + std::vector c0v{col0value}; + std::vector c1v{col1value}; + auto c0 = Int32Column::create(); + auto c1 = Int32Column::create(); + c0->append_numbers(c0v.data(), c0v.size() * sizeof(int)); + c1->append_numbers(c1v.data(), c1v.size() * sizeof(int)); + auto chunk = std::shared_ptr(new Chunk({c0, c1}, schema)); + chunk->set_slot_id_to_index(0, 0); + chunk->set_slot_id_to_index(1, 1); + return chunk; +} + +TEST(HttpResultWriterTest, BasicJsonFormat) { + // Mock all the necessary data structures + RuntimeState dummy_state; + SchemaPtr schema(new Schema()); + auto c0_field = std::make_shared(0, "c0", TYPE_INT, true); + auto c1_field = std::make_shared(1, "c1", TYPE_INT, true); + schema->append(c0_field); + schema->append(c1_field); + auto c0ref = std::make_unique(TypeDescriptor(TYPE_INT), 0); + auto c1ref = std::make_unique(TypeDescriptor(TYPE_INT), 1); + std::vector> managed_expr_ctxs; + std::vector expr_ctxs; + // create c0/c1 ExprContext + managed_expr_ctxs.emplace_back(std::make_unique(c0ref.get())); + expr_ctxs.push_back(managed_expr_ctxs.back().get()); + managed_expr_ctxs.emplace_back(std::make_unique(c1ref.get())); + expr_ctxs.push_back(managed_expr_ctxs.back().get()); + ASSERT_OK(Expr::prepare(expr_ctxs, &dummy_state)); + ASSERT_OK(Expr::open(expr_ctxs, &dummy_state)); + + TUniqueId uuid(generate_uuid()); + RuntimeProfile dummy_profile{"dummy"}; + auto sinker = std::make_unique(uuid, 100 * 1024); + + // create the writer + HttpResultWriter writer(sinker.get(), expr_ctxs, &dummy_profile, TResultSinkFormatType::type::JSON); + writer.init(&dummy_state); + + // generate a chunk with [1,1], check the output json + { + auto chunk = create_chunk(schema, 1, 1); + auto result = writer.process_chunk(chunk.get()); + EXPECT_TRUE(result.ok()) << result.status(); + auto resultPtrs = std::move(result.value()); + EXPECT_EQ(1L, resultPtrs.size()); + EXPECT_EQ(1L, resultPtrs[0]->result_batch.rows.size()); + EXPECT_EQ("{\"data\":[1,1]}\n", resultPtrs[0]->result_batch.rows[0]); + } + // generate a chunk with [2,1], check the output json + { + auto chunk = create_chunk(schema, 2, 1); + auto result = writer.process_chunk(chunk.get()); + EXPECT_TRUE(result.ok()) << result.status(); + auto resultPtrs = std::move(result.value()); + EXPECT_EQ(1L, resultPtrs.size()); + EXPECT_EQ(1L, resultPtrs[0]->result_batch.rows.size()); + EXPECT_EQ("{\"data\":[2,1]}\n", resultPtrs[0]->result_batch.rows[0]); + } + // generate a chunk with reuse of the result columns + // `SELECT count(*) as A, count(*) as B FROM t` + { + std::vector c0v{10}; + auto c0 = Int32Column::create(); + c0->append_numbers(c0v.data(), c0v.size() * sizeof(int)); + auto chunk = std::shared_ptr(new Chunk({c0, c0}, schema)); + chunk->set_slot_id_to_index(0, 0); + chunk->set_slot_id_to_index(1, 1); + auto result = writer.process_chunk(chunk.get()); + EXPECT_TRUE(result.ok()) << result.status(); + auto resultPtrs = std::move(result.value()); + EXPECT_EQ(1L, resultPtrs.size()); + EXPECT_EQ(1L, resultPtrs[0]->result_batch.rows.size()); + EXPECT_EQ("{\"data\":[10,10]}\n", resultPtrs[0]->result_batch.rows[0]); + } +} +} // namespace starrocks