diff --git a/be/benchmark/benchmark_arrow_validation.hpp b/be/benchmark/benchmark_arrow_validation.hpp new file mode 100644 index 00000000000000..f5a2da45fb4b05 --- /dev/null +++ b/be/benchmark/benchmark_arrow_validation.hpp @@ -0,0 +1,382 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "common/config.h" +#include "core/assert_cast.h" +#include "core/column/column_array.h" +#include "core/column/column_map.h" +#include "core/column/column_nullable.h" +#include "core/column/column_string.h" +#include "core/column/column_vector.h" +#include "core/data_type/data_type_array.h" +#include "core/data_type/data_type_map.h" +#include "core/data_type/data_type_number.h" +#include "core/data_type/data_type_string.h" +#include "core/data_type/primitive_type.h" +#include "util/unaligned.h" + +namespace doris { +namespace { + +enum class ArrowValidationBenchMode { OLD, CHECK_DISABLED, CHECK_ENABLED }; + +void throw_if_not_ok(const arrow::Status& status) { + if (!status.ok()) { + throw std::runtime_error(status.ToString()); + } +} + +void throw_if_not_ok(const Status& status) { + if (!status.ok()) { + throw std::runtime_error(status.to_string()); + } +} + +void keep_column_size(size_t size) { + benchmark::DoNotOptimize(size); +} + +template +double measure_seconds(Func&& func) { + const auto start = std::chrono::steady_clock::now(); + std::forward(func)(); + const auto finish = std::chrono::steady_clock::now(); + return std::chrono::duration(finish - start).count(); +} + +std::shared_ptr make_string_array(int64_t rows) { + arrow::StringBuilder builder; + throw_if_not_ok(builder.Reserve(rows)); + for (int64_t i = 0; i < rows; ++i) { + const std::string value = "arrow_validation_" + std::to_string(i % 1000); + throw_if_not_ok(builder.Append(value)); + } + std::shared_ptr array; + throw_if_not_ok(builder.Finish(&array)); + return std::static_pointer_cast(array); +} + +std::shared_ptr make_int64_array(int64_t rows) { + arrow::Int64Builder builder; + throw_if_not_ok(builder.Reserve(rows)); + for (int64_t i = 0; i < rows; ++i) { + throw_if_not_ok(builder.Append(i)); + } + std::shared_ptr array; + throw_if_not_ok(builder.Finish(&array)); + return std::static_pointer_cast(array); +} + +std::shared_ptr make_boolean_array(int64_t rows) { + arrow::BooleanBuilder builder; + throw_if_not_ok(builder.Reserve(rows)); + for (int64_t i = 0; i < rows; ++i) { + throw_if_not_ok(builder.Append((i & 1) == 0)); + } + std::shared_ptr array; + throw_if_not_ok(builder.Finish(&array)); + return std::static_pointer_cast(array); +} + +std::shared_ptr make_string_list_array(int64_t rows) { + auto value_builder = std::make_shared(); + arrow::ListBuilder builder(arrow::default_memory_pool(), value_builder); + throw_if_not_ok(builder.Reserve(rows)); + throw_if_not_ok(value_builder->Reserve(rows * 3)); + for (int64_t i = 0; i < rows; ++i) { + throw_if_not_ok(builder.Append()); + for (int64_t j = 0; j < 3; ++j) { + const std::string value = "item_" + std::to_string((i + j) % 1000); + throw_if_not_ok(value_builder->Append(value)); + } + } + std::shared_ptr array; + throw_if_not_ok(builder.Finish(&array)); + return std::static_pointer_cast(array); +} + +std::shared_ptr make_string_int_map_array(int64_t rows) { + auto keys = make_string_array(rows * 2); + auto items = make_int64_array(rows * 2); + std::vector offsets(rows + 1); + for (int64_t i = 0; i <= rows; ++i) { + offsets[i] = static_cast(i * 2); + } + auto offsets_buffer = arrow::Buffer::FromVector(std::move(offsets)); + auto map_type = arrow::map(arrow::utf8(), arrow::int64()); + return std::make_shared(map_type, rows, offsets_buffer, keys, items); +} + +void read_string_old(ColumnString& column, const arrow::BinaryArray& array, int64_t start, + int64_t end) { + std::shared_ptr buffer = array.value_data(); + const uint8_t* offsets_data = array.value_offsets()->data(); + const size_t offset_size = sizeof(int32_t); + + for (auto offset_i = start; offset_i < end; ++offset_i) { + if (!array.IsNull(offset_i)) { + auto start_offset = unaligned_load(offsets_data + offset_i * offset_size); + auto end_offset = unaligned_load(offsets_data + (offset_i + 1) * offset_size); + int32_t length = end_offset - start_offset; + const auto* raw_data = buffer->data() + start_offset; + column.insert_data(reinterpret_cast(raw_data), length); + } else { + column.insert_default(); + } + } +} + +void read_int64_old(ColumnInt64& column, const arrow::Array& array, int64_t start, int64_t end) { + auto& col_data = column.get_data(); + std::shared_ptr buffer = array.data()->buffers[1]; + const auto* raw_data = reinterpret_cast(buffer->data()) + start; + col_data.insert(raw_data, raw_data + end - start); +} + +void read_boolean_old(ColumnUInt8& column, const arrow::BooleanArray& array, int64_t start, + int64_t end) { + auto& col_data = column.get_data(); + for (int64_t bool_i = start; bool_i != end; ++bool_i) { + col_data.emplace_back(array.Value(bool_i)); + } +} + +void read_array_string_old(ColumnArray& column, const arrow::ListArray& array, int64_t start, + int64_t end) { + auto& offsets_data = column.get_offsets(); + auto arrow_offsets_array = array.offsets(); + auto* arrow_offsets = dynamic_cast(arrow_offsets_array.get()); + auto prev_size = offsets_data.back(); + const auto* base_offsets_ptr = reinterpret_cast(arrow_offsets->raw_values()); + const size_t offset_element_size = sizeof(int32_t); + const uint8_t* start_offset_ptr = base_offsets_ptr + start * offset_element_size; + const uint8_t* end_offset_ptr = base_offsets_ptr + end * offset_element_size; + auto arrow_nested_start_offset = unaligned_load(start_offset_ptr); + auto arrow_nested_end_offset = unaligned_load(end_offset_ptr); + + for (auto i = start + 1; i < end + 1; ++i) { + const uint8_t* current_offset_ptr = base_offsets_ptr + i * offset_element_size; + auto current_offset = unaligned_load(current_offset_ptr); + offsets_data.emplace_back(prev_size + current_offset - arrow_nested_start_offset); + } + auto& nested_nullable = assert_cast(column.get_data()); + read_string_old(assert_cast(nested_nullable.get_nested_column()), + static_cast(*array.values()), + arrow_nested_start_offset, arrow_nested_end_offset); + auto& null_map = nested_nullable.get_null_map_data(); + null_map.resize_fill(null_map.size() + arrow_nested_end_offset - arrow_nested_start_offset, 0); +} + +void read_map_string_int_old(ColumnMap& column, const arrow::MapArray& array, int64_t start, + int64_t end) { + auto& offsets_data = column.get_offsets(); + auto arrow_offsets_array = array.offsets(); + auto* arrow_offsets = dynamic_cast(arrow_offsets_array.get()); + auto prev_size = offsets_data.back(); + const auto* base_offsets_ptr = reinterpret_cast(arrow_offsets->raw_values()); + const size_t offset_element_size = sizeof(int32_t); + const uint8_t* start_offset_ptr = base_offsets_ptr + start * offset_element_size; + const uint8_t* end_offset_ptr = base_offsets_ptr + end * offset_element_size; + auto arrow_nested_start_offset = unaligned_load(start_offset_ptr); + auto arrow_nested_end_offset = unaligned_load(end_offset_ptr); + for (int64_t i = start + 1; i < end + 1; ++i) { + const uint8_t* current_offset_ptr = base_offsets_ptr + i * offset_element_size; + auto current_offset = unaligned_load(current_offset_ptr); + offsets_data.emplace_back(prev_size + current_offset - arrow_nested_start_offset); + } + read_string_old(assert_cast(column.get_keys()), + static_cast(*array.keys()), + arrow_nested_start_offset, arrow_nested_end_offset); + read_int64_old(assert_cast(column.get_values()), *array.items(), + arrow_nested_start_offset, arrow_nested_end_offset); +} + +template +void run_serde_benchmark(benchmark::State& state, const std::shared_ptr& type, + const std::shared_ptr& array) { + const bool old_config = config::enable_arrow_input_validation; + if constexpr (mode == ArrowValidationBenchMode::CHECK_DISABLED) { + config::enable_arrow_input_validation = false; + } else if constexpr (mode == ArrowValidationBenchMode::CHECK_ENABLED) { + config::enable_arrow_input_validation = true; + } + + for (auto _ : state) { + const double seconds = measure_seconds([&] { + auto column = type->create_column(); + throw_if_not_ok(type->get_serde()->read_column_from_arrow( + *column, array.get(), 0, array->length(), cctz::utc_time_zone())); + keep_column_size(column->size()); + }); + state.SetIterationTime(seconds); + } + config::enable_arrow_input_validation = old_config; +} + +template +void BM_ArrowValidation_String(benchmark::State& state) { + auto array = make_string_array(state.range(0)); + auto type = std::make_shared(); + + if constexpr (mode == ArrowValidationBenchMode::OLD) { + for (auto _ : state) { + const double seconds = measure_seconds([&] { + auto column = ColumnString::create(); + read_string_old(*column, *array, 0, array->length()); + keep_column_size(column->size()); + }); + state.SetIterationTime(seconds); + } + } else { + run_serde_benchmark(state, type, array); + } +} + +template +void BM_ArrowValidation_Int64(benchmark::State& state) { + auto array = make_int64_array(state.range(0)); + auto type = std::make_shared(); + + if constexpr (mode == ArrowValidationBenchMode::OLD) { + for (auto _ : state) { + const double seconds = measure_seconds([&] { + auto column = ColumnInt64::create(); + read_int64_old(*column, *array, 0, array->length()); + keep_column_size(column->size()); + }); + state.SetIterationTime(seconds); + } + } else { + run_serde_benchmark(state, type, array); + } +} + +template +void BM_ArrowValidation_Boolean(benchmark::State& state) { + auto array = make_boolean_array(state.range(0)); + auto type = std::make_shared(); + + if constexpr (mode == ArrowValidationBenchMode::OLD) { + for (auto _ : state) { + const double seconds = measure_seconds([&] { + auto column = ColumnUInt8::create(); + read_boolean_old(*column, *array, 0, array->length()); + keep_column_size(column->size()); + }); + state.SetIterationTime(seconds); + } + } else { + run_serde_benchmark(state, type, array); + } +} + +template +void BM_ArrowValidation_ArrayString(benchmark::State& state) { + auto array = make_string_list_array(state.range(0)); + auto type = std::make_shared(std::make_shared()); + + if constexpr (mode == ArrowValidationBenchMode::OLD) { + for (auto _ : state) { + const double seconds = measure_seconds([&] { + auto column = type->create_column(); + read_array_string_old(assert_cast(*column), *array, 0, + array->length()); + keep_column_size(column->size()); + }); + state.SetIterationTime(seconds); + } + } else { + run_serde_benchmark(state, type, array); + } +} + +template +void BM_ArrowValidation_MapStringInt(benchmark::State& state) { + auto array = make_string_int_map_array(state.range(0)); + auto type = std::make_shared(std::make_shared(), + std::make_shared()); + + if constexpr (mode == ArrowValidationBenchMode::OLD) { + for (auto _ : state) { + const double seconds = measure_seconds([&] { + auto column = type->create_column(); + read_map_string_int_old(assert_cast(*column), *array, 0, + array->length()); + keep_column_size(column->size()); + }); + state.SetIterationTime(seconds); + } + } else { + run_serde_benchmark(state, type, array); + } +} + +#define REGISTER_ARROW_VALIDATION_BENCH(FUNC, MODE, NAME) \ + BENCHMARK_TEMPLATE(FUNC, ArrowValidationBenchMode::MODE) \ + ->Name("BM_ArrowValidation_" #NAME "_" #MODE "/rows:4096") \ + ->Arg(4096) \ + ->UseManualTime(); \ + BENCHMARK_TEMPLATE(FUNC, ArrowValidationBenchMode::MODE) \ + ->Name("BM_ArrowValidation_" #NAME "_" #MODE "/rows:65536") \ + ->Arg(65536) \ + ->UseManualTime() + +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_String, OLD, String); +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_String, CHECK_DISABLED, String); +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_String, CHECK_ENABLED, String); + +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_Int64, OLD, Int64); +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_Int64, CHECK_DISABLED, Int64); +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_Int64, CHECK_ENABLED, Int64); + +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_Boolean, OLD, Boolean); +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_Boolean, CHECK_DISABLED, Boolean); +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_Boolean, CHECK_ENABLED, Boolean); + +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_ArrayString, OLD, ArrayString); +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_ArrayString, CHECK_DISABLED, ArrayString); +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_ArrayString, CHECK_ENABLED, ArrayString); + +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_MapStringInt, OLD, MapStringInt); +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_MapStringInt, CHECK_DISABLED, MapStringInt); +REGISTER_ARROW_VALIDATION_BENCH(BM_ArrowValidation_MapStringInt, CHECK_ENABLED, MapStringInt); + +#undef REGISTER_ARROW_VALIDATION_BENCH + +} // namespace +} // namespace doris diff --git a/be/benchmark/benchmark_main.cpp b/be/benchmark/benchmark_main.cpp index caf5459c46af51..ff57d89c802d77 100644 --- a/be/benchmark/benchmark_main.cpp +++ b/be/benchmark/benchmark_main.cpp @@ -17,6 +17,7 @@ #include +#include "benchmark_arrow_validation.hpp" #include "benchmark_bit_pack.hpp" #include "benchmark_column_array_view.hpp" #include "benchmark_column_array_view_distance.hpp" diff --git a/be/src/common/config.cpp b/be/src/common/config.cpp index 59dffdf1ac34e3..1beaf56941f9d3 100644 --- a/be/src/common/config.cpp +++ b/be/src/common/config.cpp @@ -71,6 +71,9 @@ DEFINE_Int32(brpc_port, "8060"); DEFINE_Int32(arrow_flight_sql_port, "8050"); +// Validate Arrow input buffers in opted-in Arrow readers before converting them to Doris columns. +DEFINE_Bool(enable_arrow_input_validation, "true"); + DEFINE_Int32(cdc_client_port, "9096"); DEFINE_String(cdc_client_java_opts, ""); diff --git a/be/src/common/config.h b/be/src/common/config.h index 033b755409cc64..e7b43146d68c64 100644 --- a/be/src/common/config.h +++ b/be/src/common/config.h @@ -119,6 +119,9 @@ DECLARE_Int32(brpc_port); // Default -1, do not start arrow flight sql server. DECLARE_Int32(arrow_flight_sql_port); +// Validate Arrow input buffers in opted-in Arrow readers before converting them to Doris columns. +DECLARE_Bool(enable_arrow_input_validation); + // port for cdc client scan oltp cdc data DECLARE_Int32(cdc_client_port); diff --git a/be/src/core/block/block.cpp b/be/src/core/block/block.cpp index f5208a8e071cb8..f2de6f9cde9d86 100644 --- a/be/src/core/block/block.cpp +++ b/be/src/core/block/block.cpp @@ -341,6 +341,21 @@ Status Block::check_type_and_column() const { return Status::OK(); } +Status Block::check_column_and_type_not_null() const { + for (size_t i = 0; i != data.size(); ++i) { + const auto& elem = data[i]; + if (!elem.column) { + return Status::InternalError("Column in block is nullptr, column index: {}, name: {}", + i, elem.name); + } + if (!elem.type) { + return Status::InternalError("Type in block is nullptr, column index: {}, name: {}", i, + elem.name); + } + } + return Status::OK(); +} + size_t Block::rows() const { for (const auto& elem : data) { if (elem.column) { diff --git a/be/src/core/block/block.h b/be/src/core/block/block.h index 5545465cef86f7..ffd3adc672f4e1 100644 --- a/be/src/core/block/block.h +++ b/be/src/core/block/block.h @@ -186,6 +186,8 @@ class Block { Status check_type_and_column() const; + Status check_column_and_type_not_null() const; + /// Approximate number of bytes used by column data in memory. /// This reflects the actual data footprint (e.g. string contents, numeric arrays) /// and is the metric used by adaptive batch size byte budgets. diff --git a/be/src/core/column/column.h b/be/src/core/column/column.h index 9e25b42d0d141e..9547ff1a7ca3da 100644 --- a/be/src/core/column/column.h +++ b/be/src/core/column/column.h @@ -446,7 +446,8 @@ class IColumn : public COW { * // nullable -> predict_column * // string (dictionary) -> column_dictionary */ - virtual Status filter_by_selector(const uint16_t* sel, size_t sel_size, IColumn* col_ptr) { + virtual Status filter_by_selector(const uint16_t* sel, size_t sel_size, + IColumn* col_ptr) const { throw doris::Exception(ErrorCode::NOT_IMPLEMENTED_ERROR, "Method filter_by_selector is not supported for {}, only " "column_nullable, column_dictionary and predict_column support", diff --git a/be/src/core/column/column_decimal.h b/be/src/core/column/column_decimal.h index 55e429f3183749..58116b942a6036 100644 --- a/be/src/core/column/column_decimal.h +++ b/be/src/core/column/column_decimal.h @@ -149,7 +149,8 @@ class ColumnDecimal final : public COWHelper> { memset(data.data() + old_size, 0, length * sizeof(data[0])); } - Status filter_by_selector(const uint16_t* sel, size_t sel_size, IColumn* col_ptr) override { + Status filter_by_selector(const uint16_t* sel, size_t sel_size, + IColumn* col_ptr) const override { Self* output = assert_cast(col_ptr); auto& res_data = output->get_data(); DCHECK(res_data.empty()) diff --git a/be/src/core/column/column_dictionary.h b/be/src/core/column/column_dictionary.h index 3f6987d6baa370..5ec603225cf144 100644 --- a/be/src/core/column/column_dictionary.h +++ b/be/src/core/column/column_dictionary.h @@ -153,19 +153,23 @@ class ColumnDictI32 final : public COWHelper { "permute not supported in ColumnDictionary"); } - Status filter_by_selector(const uint16_t* sel, size_t sel_size, IColumn* col_ptr) override { + Status filter_by_selector(const uint16_t* sel, size_t sel_size, + IColumn* col_ptr) const override { auto* res_col = assert_cast(col_ptr); - _strings.resize(sel_size); + if (sel_size == 0) { + return Status::OK(); + } + std::vector strings(sel_size); size_t length = 0; for (size_t i = 0; i != sel_size; ++i) { - auto& value = _dict.get_value(_codes[sel[i]]); - _strings[i].data = value.data; - _strings[i].size = value.size; + const auto& value = _dict.get_value(_codes[sel[i]]); + strings[i].data = value.data; + strings[i].size = value.size; length += value.size; } res_col->get_offsets().reserve(sel_size + res_col->get_offsets().size()); res_col->get_chars().reserve(length + res_col->get_chars().size()); - res_col->insert_many_strings_without_reserve(_strings.data(), sel_size); + res_col->insert_many_strings_without_reserve(strings.data(), sel_size); return Status::OK(); } @@ -455,7 +459,6 @@ class ColumnDictI32 final : public COWHelper { Dictionary _dict; Container _codes; std::pair _rowset_segment_id; - std::vector _strings; }; } // namespace doris diff --git a/be/src/core/column/column_nullable.cpp b/be/src/core/column/column_nullable.cpp index 7a560dde4b85e1..01d0fc072ed41d 100644 --- a/be/src/core/column/column_nullable.cpp +++ b/be/src/core/column/column_nullable.cpp @@ -471,17 +471,15 @@ size_t ColumnNullable::filter(const Filter& filter) { return data_result_size; } -Status ColumnNullable::filter_by_selector(const uint16_t* sel, size_t sel_size, IColumn* col_ptr) { +Status ColumnNullable::filter_by_selector(const uint16_t* sel, size_t sel_size, + IColumn* col_ptr) const { auto* nullable_col_ptr = assert_cast(col_ptr); - // Access the nested column via const path to avoid assert_mutable_ref (which requires - // exclusive ownership). The output col_ptr was just created, so its nested column is exclusive. - auto nest_col_raw = const_cast( - static_cast(nullable_col_ptr->_nested_column).get()); + auto nested_column = nullable_col_ptr->get_nested_column_ptr(); /// `get_null_map_data` will set `_need_update_has_null` to true auto& res_nullmap = nullable_col_ptr->get_null_map_data(); - RETURN_IF_ERROR(get_nested_column().filter_by_selector(sel, sel_size, nest_col_raw)); + RETURN_IF_ERROR(get_nested_column().filter_by_selector(sel, sel_size, nested_column.get())); DCHECK(res_nullmap.empty()); res_nullmap.resize(sel_size); auto& cur_nullmap = get_null_map_column().get_data(); diff --git a/be/src/core/column/column_nullable.h b/be/src/core/column/column_nullable.h index 34ee0dffebc9e2..3e7afb904e6c81 100644 --- a/be/src/core/column/column_nullable.h +++ b/be/src/core/column/column_nullable.h @@ -206,7 +206,8 @@ class ColumnNullable final : public COWHelper { size_t filter(const Filter& filter) override; - Status filter_by_selector(const uint16_t* sel, size_t sel_size, IColumn* col_ptr) override; + Status filter_by_selector(const uint16_t* sel, size_t sel_size, + IColumn* col_ptr) const override; MutableColumnPtr permute(const Permutation& perm, size_t limit) const override; // ColumnPtr index(const IColumn & indexes, size_t limit) const override; int compare_at(size_t n, size_t m, const IColumn& rhs_, int null_direction_hint) const override; diff --git a/be/src/core/column/column_string.cpp b/be/src/core/column/column_string.cpp index bb921b93d1a92e..c8b5339ac32828 100644 --- a/be/src/core/column/column_string.cpp +++ b/be/src/core/column/column_string.cpp @@ -372,7 +372,8 @@ size_t ColumnStr::filter(const IColumn::Filter& filter) { } template -Status ColumnStr::filter_by_selector(const uint16_t* sel, size_t sel_size, IColumn* col_ptr) { +Status ColumnStr::filter_by_selector(const uint16_t* sel, size_t sel_size, + IColumn* col_ptr) const { if constexpr (std::is_same_v) { auto* col = static_cast*>(col_ptr); Chars& res_chars = col->chars; diff --git a/be/src/core/column/column_string.h b/be/src/core/column/column_string.h index 5b748ae18e9b98..8972a4ca210ceb 100644 --- a/be/src/core/column/column_string.h +++ b/be/src/core/column/column_string.h @@ -482,7 +482,8 @@ class ColumnStr final : public COWHelper> { ColumnPtr filter(const IColumn::Filter& filt, ssize_t result_size_hint) const override; size_t filter(const IColumn::Filter& filter) override; - Status filter_by_selector(const uint16_t* sel, size_t sel_size, IColumn* col_ptr) override; + Status filter_by_selector(const uint16_t* sel, size_t sel_size, + IColumn* col_ptr) const override; MutableColumnPtr permute(const IColumn::Permutation& perm, size_t limit) const override; diff --git a/be/src/core/column/column_vector.h b/be/src/core/column/column_vector.h index 22a807f6929124..5112ad82a42cae 100644 --- a/be/src/core/column/column_vector.h +++ b/be/src/core/column/column_vector.h @@ -271,7 +271,8 @@ class ColumnVector final : public COWHelper> { void insert_value(const value_type value) { data.push_back(value); } - Status filter_by_selector(const uint16_t* sel, size_t sel_size, IColumn* col_ptr) override { + Status filter_by_selector(const uint16_t* sel, size_t sel_size, + IColumn* col_ptr) const override { Self* output = assert_cast(col_ptr); auto& res_data = output->get_data(); DCHECK(res_data.empty()) diff --git a/be/src/core/data_type_serde/arrow_validation.h b/be/src/core/data_type_serde/arrow_validation.h new file mode 100644 index 00000000000000..9d8b09b0cf198a --- /dev/null +++ b/be/src/core/data_type_serde/arrow_validation.h @@ -0,0 +1,320 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "common/compiler_util.h" +#include "common/exception.h" +#include "util/unaligned.h" + +namespace doris { +namespace arrow_validation_detail { + +inline std::string arrow_type_name(const arrow::Array& array) { + return array.type() ? array.type()->name() : "unknown"; +} + +inline void throw_invalid_arrow(std::string_view arrow_type, std::string_view message) { + throw Exception(ErrorCode::INVALID_ARGUMENT, "Invalid Arrow {}: {}", arrow_type, message); +} + +inline void throw_invalid_arrow(const arrow::Array& array, std::string_view message) { + throw_invalid_arrow(arrow_type_name(array), message); +} + +template +inline void throw_invalid_arrow(std::string_view arrow_type, std::string_view format, + Args&&... args) { + throw Exception(ErrorCode::INVALID_ARGUMENT, "Invalid Arrow {}: {}", arrow_type, + fmt::format(std::string(format), std::forward(args)...)); +} + +template +inline void throw_invalid_arrow(const arrow::Array& array, std::string_view format, + Args&&... args) { + throw_invalid_arrow(arrow_type_name(array), format, std::forward(args)...); +} + +inline void check_arrow_length_and_offset(const arrow::Array& array) { + if (UNLIKELY(array.length() < 0 || array.offset() < 0)) { + throw_invalid_arrow(array, "negative length or offset: length={}, offset={}", + array.length(), array.offset()); + } +} + +inline void check_arrow_no_offset(const arrow::Array& array) { + check_arrow_length_and_offset(array); + if (UNLIKELY(array.offset() != 0)) { + throw_invalid_arrow(array, "non-zero array offset is not supported: offset={}", + array.offset()); + } +} + +inline void check_add_overflow(size_t left, size_t right, const arrow::Array& array, + std::string_view item) { + if (UNLIKELY(left > std::numeric_limits::max() - right)) { + throw_invalid_arrow(array, "{} size overflow", item); + } +} + +inline std::shared_ptr get_int32_offsets_array(const arrow::Array& array) { + auto offsets_array = static_cast(array).offsets(); + auto offsets = std::dynamic_pointer_cast(offsets_array); + if (UNLIKELY(!offsets)) { + throw_invalid_arrow(array, "offsets array is not Int32Array"); + } + return offsets; +} + +} // namespace arrow_validation_detail + +inline void check_arrow_no_offset(const arrow::Array& array) { + arrow_validation_detail::check_arrow_no_offset(array); +} + +// Validate the caller's requested read range before any Arrow buffer access. +// This rejects negative or out-of-array start/end values up front. +inline void check_arrow_array_range(const arrow::Array& array, int64_t start, int64_t end) { + arrow_validation_detail::check_arrow_no_offset(array); + if (UNLIKELY(start < 0 || end < start || end > array.length())) { + arrow_validation_detail::throw_invalid_arrow( + array, "read range is invalid: start={}, end={}, length={}", start, end, + array.length()); + } +} + +// Validate buffers[0] when a validity bitmap exists. This must run before +// IsNull()/IsValid()/null_count(), because those APIs may scan the bitmap. +inline void check_arrow_validity_bitmap(const arrow::Array& array) { + arrow_validation_detail::check_arrow_length_and_offset(array); + const auto& buffers = array.data()->buffers; + if (buffers.empty() || !buffers[0]) { + if (UNLIKELY(array.data()->null_count > 0)) { + arrow_validation_detail::throw_invalid_arrow( + array, "validity bitmap is missing but null_count={}", + array.data()->null_count.load()); + } + return; + } + + const size_t offset = static_cast(array.offset()); + const size_t length = static_cast(array.length()); + arrow_validation_detail::check_add_overflow(offset, length, array, "validity bitmap"); + const size_t count = offset + length; + const size_t required = count / 8 + (count % 8 != 0 ? 1 : 0); + const size_t available = static_cast(buffers[0]->size()); + if (UNLIKELY(available < required)) { + arrow_validation_detail::throw_invalid_arrow( + array, "validity bitmap too small: {} bytes available, {} required", available, + required); + } +} + +// Validate buffers[1] for fixed-width arrays before raw_values(), Value(), or +// direct buffer memcpy. elem_size is the physical byte width of one value. +inline void check_arrow_fixed_width_buffer(const arrow::Array& array, size_t elem_size) { + check_arrow_validity_bitmap(array); + const auto& buffers = array.data()->buffers; + if (UNLIKELY(buffers.size() <= 1 || !buffers[1])) { + if (array.length() == 0) { + return; + } + arrow_validation_detail::throw_invalid_arrow(array, "data buffer is missing"); + } + + const size_t offset = static_cast(array.offset()); + const size_t length = static_cast(array.length()); + arrow_validation_detail::check_add_overflow(offset, length, array, "data buffer"); + const size_t count = offset + length; + if (UNLIKELY(elem_size != 0 && count > std::numeric_limits::max() / elem_size)) { + arrow_validation_detail::throw_invalid_arrow(array, "data buffer size overflow"); + } + const size_t required = elem_size * count; + const size_t available = static_cast(buffers[1]->size()); + if (UNLIKELY(available < required)) { + arrow_validation_detail::throw_invalid_arrow( + array, "data buffer too small: {} bytes available, {} required", available, + required); + } +} + +// Validate the bit-packed boolean data bitmap before BooleanArray::Value(). +inline void check_arrow_boolean_buffer(const arrow::Array& array) { + check_arrow_validity_bitmap(array); + const auto& buffers = array.data()->buffers; + if (UNLIKELY(buffers.size() <= 1 || !buffers[1])) { + if (array.length() == 0) { + return; + } + arrow_validation_detail::throw_invalid_arrow(array, "data bitmap is missing"); + } + + const size_t offset = static_cast(array.offset()); + const size_t length = static_cast(array.length()); + arrow_validation_detail::check_add_overflow(offset, length, array, "data bitmap"); + const size_t count = offset + length; + const size_t required = count / 8 + (count % 8 != 0 ? 1 : 0); + const size_t available = static_cast(buffers[1]->size()); + if (UNLIKELY(available < required)) { + arrow_validation_detail::throw_invalid_arrow( + array, "data bitmap too small: {} bytes available, {} required", available, + required); + } +} + +// Validate one variable-width value range after reading offsets and before +// forming value_data() + offset. +inline void check_arrow_value_range(const arrow::Array& array, int64_t offset, int64_t length, + size_t buffer_size) { + if (UNLIKELY(offset < 0 || length < 0)) { + arrow_validation_detail::throw_invalid_arrow( + array, "value range has negative offset or length: offset={}, length={}", offset, + length); + } + const size_t safe_offset = static_cast(offset); + const size_t safe_length = static_cast(length); + if (UNLIKELY(safe_offset > buffer_size || safe_length > buffer_size - safe_offset)) { + arrow_validation_detail::throw_invalid_arrow( + array, "value range exceeds data buffer: offset={}, length={}, buffer_size={}", + offset, length, buffer_size); + } +} + +namespace arrow_validation_detail { + +// Offsets buffers may come from external Arrow producers through Buffer::Wrap or FFI and are not +// guaranteed to be aligned to int32_t. Do not use Int32Array::Value() here because it performs a +// typed raw_values()[i] load and can trigger UBSan on misaligned buffers. Keep this validation path +// consistent with the array/map readers below, which load offsets through unaligned_load(). +inline int32_t read_int32_offset(const arrow::Int32Array& offsets, int64_t index) { + const auto* data = reinterpret_cast(offsets.raw_values()); + return unaligned_load(data + index * sizeof(int32_t)); +} + +inline int64_t check_arrow_offsets_range(const arrow::Int32Array& offsets, int64_t start, + int64_t end) { + check_arrow_array_range(offsets, 0, offsets.length()); + check_arrow_fixed_width_buffer(offsets, sizeof(int32_t)); + if (UNLIKELY(start < 0 || end < start || end >= offsets.length())) { + arrow_validation_detail::throw_invalid_arrow( + offsets, "offsets read range is invalid: start={}, end={}, offsets_length={}", + start, end, offsets.length()); + } + + int64_t previous_offset = read_int32_offset(offsets, start); + if (UNLIKELY(previous_offset < 0)) { + arrow_validation_detail::throw_invalid_arrow( + offsets, "offsets contain negative value: offset[{}]={}", start, previous_offset); + } + for (int64_t i = start + 1; i <= end; ++i) { + const int64_t current_offset = read_int32_offset(offsets, i); + if (UNLIKELY(current_offset < previous_offset)) { + arrow_validation_detail::throw_invalid_arrow( + offsets, + "offsets are not monotonically non-decreasing: offset[{}]={} < offset[{}]={}", + i, current_offset, i - 1, previous_offset); + } + previous_offset = current_offset; + } + return previous_offset; +} + +} // namespace arrow_validation_detail + +// Validate List offsets before reading offsets or recursing into values. +inline void check_arrow_list_offsets(const arrow::ListArray& array, int64_t start, int64_t end) { + check_arrow_array_range(array, start, end); + const auto offsets = arrow_validation_detail::get_int32_offsets_array(array); + const int64_t last_offset = + arrow_validation_detail::check_arrow_offsets_range(*offsets, start, end); + const int64_t values_length = array.values() ? array.values()->length() : 0; + if (UNLIKELY(last_offset > values_length)) { + arrow_validation_detail::throw_invalid_arrow( + array, "offsets exceed values length: last_offset={}, values_length={}", + last_offset, values_length); + } +} + +// Validate Map offsets before reading offsets or recursing into keys/items. +inline void check_arrow_map_offsets(const arrow::MapArray& array, int64_t start, int64_t end) { + check_arrow_array_range(array, start, end); + const auto offsets = arrow_validation_detail::get_int32_offsets_array(array); + const int64_t last_offset = + arrow_validation_detail::check_arrow_offsets_range(*offsets, start, end); + const int64_t keys_length = array.keys() ? array.keys()->length() : 0; + if (UNLIKELY(last_offset > keys_length)) { + arrow_validation_detail::throw_invalid_arrow( + array, "offsets exceed keys length: last_offset={}, keys_length={}", last_offset, + keys_length); + } + const int64_t items_length = array.items() ? array.items()->length() : 0; + if (UNLIKELY(last_offset > items_length)) { + arrow_validation_detail::throw_invalid_arrow( + array, "offsets exceed items length: last_offset={}, items_length={}", last_offset, + items_length); + } +} + +// Validate String/Binary offsets buffer before value_offset(), value_length(), +// raw_value_offsets(), or manual offset reads. Variable-width Arrow arrays need +// offset + length + 1 offset entries. +template +inline void check_arrow_binary_offsets_buffer(const ArrowBinaryArray& array) { + check_arrow_validity_bitmap(array); + const auto& buffers = array.data()->buffers; + if (UNLIKELY(buffers.size() <= 1 || !buffers[1])) { + arrow_validation_detail::throw_invalid_arrow(array, "offsets buffer is missing"); + } + + const size_t offset = static_cast(array.offset()); + const size_t length = static_cast(array.length()); + if (UNLIKELY(offset > std::numeric_limits::max() - length)) { + arrow_validation_detail::throw_invalid_arrow(array, "offsets entry count overflow"); + } + const size_t count = offset + length; + if (UNLIKELY(count == std::numeric_limits::max())) { + arrow_validation_detail::throw_invalid_arrow(array, "offsets entry count overflow"); + } + const size_t count_plus_one = count + 1; + if (UNLIKELY(count_plus_one > std::numeric_limits::max() / + sizeof(typename ArrowBinaryArray::offset_type))) { + arrow_validation_detail::throw_invalid_arrow(array, "offsets buffer size overflow"); + } + + const size_t required = count_plus_one * sizeof(typename ArrowBinaryArray::offset_type); + const size_t available = static_cast(buffers[1]->size()); + if (UNLIKELY(available < required)) { + arrow_validation_detail::throw_invalid_arrow( + array, "offsets buffer too small: {} bytes available, {} required", available, + required); + } +} + +} // namespace doris diff --git a/be/src/core/data_type_serde/data_type_array_serde.cpp b/be/src/core/data_type_serde/data_type_array_serde.cpp index a97c317889979a..55c77c56bee6ed 100644 --- a/be/src/core/data_type_serde/data_type_array_serde.cpp +++ b/be/src/core/data_type_serde/data_type_array_serde.cpp @@ -19,6 +19,7 @@ #include +#include "common/config.h" #include "common/status.h" #include "core/assert_cast.h" #include "core/column/column.h" @@ -27,6 +28,7 @@ #include "core/data_type/data_type.h" #include "core/data_type/data_type_array.h" #include "core/data_type/get_least_supertype.h" +#include "core/data_type_serde/arrow_validation.h" #include "core/data_type_serde/complex_type_deserialize_util.h" #include "core/string_ref.h" #include "exprs/function/function_helpers.h" @@ -324,6 +326,9 @@ Status DataTypeArraySerDe::read_column_from_arrow(IColumn& column, const arrow:: const auto* concrete_array = dynamic_cast(arrow_array); auto arrow_offsets_array = concrete_array->offsets(); auto* arrow_offsets = dynamic_cast(arrow_offsets_array.get()); + if (config::enable_arrow_input_validation) { + check_arrow_list_offsets(*concrete_array, start, end); + } auto prev_size = offsets_data.back(); const auto* base_offsets_ptr = reinterpret_cast(arrow_offsets->raw_values()); const size_t offset_element_size = sizeof(int32_t); diff --git a/be/src/core/data_type_serde/data_type_date_or_datetime_serde.cpp b/be/src/core/data_type_serde/data_type_date_or_datetime_serde.cpp index c7c9eb19630c88..d99f8340a5626a 100644 --- a/be/src/core/data_type_serde/data_type_date_or_datetime_serde.cpp +++ b/be/src/core/data_type_serde/data_type_date_or_datetime_serde.cpp @@ -20,10 +20,12 @@ #include #include +#include "common/config.h" #include "common/status.h" #include "core/column/column_const.h" #include "core/data_type/data_type_decimal.h" #include "core/data_type/data_type_number.h" +#include "core/data_type_serde/arrow_validation.h" #include "core/value/vdatetime_value.h" #include "exprs/function/cast/cast_base.h" #include "exprs/function/cast/cast_to_date_or_datetime_impl.hpp" @@ -192,6 +194,9 @@ Status DataTypeDateSerDe::_read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const { + if (config::enable_arrow_input_validation) { + check_arrow_no_offset(*arrow_array); + } auto& col_data = static_cast&>(column).get_data(); int64_t divisor = 1; int64_t multiplier = 1; diff --git a/be/src/core/data_type_serde/data_type_datetimev2_serde.cpp b/be/src/core/data_type_serde/data_type_datetimev2_serde.cpp index 999046503fb9c8..5b372cce23bd65 100644 --- a/be/src/core/data_type_serde/data_type_datetimev2_serde.cpp +++ b/be/src/core/data_type_serde/data_type_datetimev2_serde.cpp @@ -23,11 +23,13 @@ #include // IWYU pragma: keep #include +#include "common/config.h" #include "common/status.h" #include "core/column/column_const.h" #include "core/data_type/data_type_decimal.h" #include "core/data_type/data_type_number.h" #include "core/data_type/primitive_type.h" +#include "core/data_type_serde/arrow_validation.h" #include "core/data_type_serde/decoded_column_view.h" #include "core/types.h" #include "core/value/vdatetime_value.h" @@ -488,6 +490,9 @@ Status DataTypeDateTimeV2SerDe::read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const { + if (config::enable_arrow_input_validation) { + check_arrow_no_offset(*arrow_array); + } auto& col_data = static_cast(column).get_data(); int64_t divisor = 1; if (arrow_array->type()->id() == arrow::Type::TIMESTAMP) { diff --git a/be/src/core/data_type_serde/data_type_datev2_serde.cpp b/be/src/core/data_type_serde/data_type_datev2_serde.cpp index 04fc37a9bb375e..a7151ff020c4db 100644 --- a/be/src/core/data_type_serde/data_type_datev2_serde.cpp +++ b/be/src/core/data_type_serde/data_type_datev2_serde.cpp @@ -23,10 +23,12 @@ #include +#include "common/config.h" #include "core/column/column_const.h" #include "core/data_type/data_type_decimal.h" #include "core/data_type/data_type_number.h" #include "core/data_type/define_primitive_type.h" +#include "core/data_type_serde/arrow_validation.h" #include "core/data_type_serde/decoded_column_view.h" #include "core/types.h" #include "core/value/vdatetime_value.h" @@ -111,6 +113,9 @@ Status DataTypeDateV2SerDe::write_column_to_arrow(const IColumn& column, const N Status DataTypeDateV2SerDe::read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const { + if (config::enable_arrow_input_validation) { + check_arrow_no_offset(*arrow_array); + } auto& col_data = static_cast(column).get_data(); const auto* concrete_array = dynamic_cast(arrow_array); const auto* base_ptr = reinterpret_cast(concrete_array->raw_values()); diff --git a/be/src/core/data_type_serde/data_type_decimal_serde.cpp b/be/src/core/data_type_serde/data_type_decimal_serde.cpp index 6393e9965fc3b1..e547c67b8735b2 100644 --- a/be/src/core/data_type_serde/data_type_decimal_serde.cpp +++ b/be/src/core/data_type_serde/data_type_decimal_serde.cpp @@ -26,12 +26,14 @@ #include "arrow/type.h" #include "common/cast_set.h" +#include "common/config.h" #include "common/consts.h" #include "core/column/column.h" #include "core/column/column_decimal.h" #include "core/data_type/data_type_decimal.h" #include "core/data_type/define_primitive_type.h" #include "core/data_type/storage_field_type.h" +#include "core/data_type_serde/arrow_validation.h" #include "core/data_type_serde/decoded_column_view.h" #include "core/types.h" #include "exec/common/arithmetic_overflow.h" @@ -455,6 +457,9 @@ Status DataTypeDecimalSerDe::read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const { + if (config::enable_arrow_input_validation) { + check_arrow_no_offset(*arrow_array); + } auto& column_data = static_cast&>(column).get_data(); // Decimal for decimalv2 // Decimal for deicmalv3 diff --git a/be/src/core/data_type_serde/data_type_ipv4_serde.cpp b/be/src/core/data_type_serde/data_type_ipv4_serde.cpp index fdb091fe835d4e..096715a1c6f49b 100644 --- a/be/src/core/data_type_serde/data_type_ipv4_serde.cpp +++ b/be/src/core/data_type_serde/data_type_ipv4_serde.cpp @@ -19,7 +19,9 @@ #include +#include "common/config.h" #include "core/column/column_const.h" +#include "core/data_type_serde/arrow_validation.h" #include "core/types.h" #include "exprs/function/cast/cast_to_ip.h" #include "exprs/function/cast/cast_to_string.h" @@ -114,6 +116,9 @@ Status DataTypeIPv4SerDe::write_column_to_arrow(const IColumn& column, const Nul Status DataTypeIPv4SerDe::read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const { + if (config::enable_arrow_input_validation) { + check_arrow_no_offset(*arrow_array); + } auto& col_data = assert_cast(column).get_data(); int64_t row_count = end - start; /// buffers[0] is a null bitmap and buffers[1] are actual values diff --git a/be/src/core/data_type_serde/data_type_ipv6_serde.cpp b/be/src/core/data_type_serde/data_type_ipv6_serde.cpp index febf802d0e05ff..d53577cbefccc5 100644 --- a/be/src/core/data_type_serde/data_type_ipv6_serde.cpp +++ b/be/src/core/data_type_serde/data_type_ipv6_serde.cpp @@ -22,7 +22,9 @@ #include #include +#include "common/config.h" #include "core/column/column_const.h" +#include "core/data_type_serde/arrow_validation.h" #include "core/types.h" #include "exprs/function/cast/cast_to_ip.h" #include "exprs/function/cast/cast_to_string.h" @@ -143,6 +145,9 @@ Status DataTypeIPv6SerDe::write_column_to_arrow(const IColumn& column, const Nul Status DataTypeIPv6SerDe::read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const { + if (config::enable_arrow_input_validation) { + check_arrow_no_offset(*arrow_array); + } auto& col_data = assert_cast(column).get_data(); const auto* concrete_array = assert_cast(arrow_array); std::shared_ptr buffer = concrete_array->value_data(); diff --git a/be/src/core/data_type_serde/data_type_jsonb_serde.cpp b/be/src/core/data_type_serde/data_type_jsonb_serde.cpp index 7dc0e4cfd308c4..c728532eafb977 100644 --- a/be/src/core/data_type_serde/data_type_jsonb_serde.cpp +++ b/be/src/core/data_type_serde/data_type_jsonb_serde.cpp @@ -26,8 +26,10 @@ #include #include "arrow/array/builder_binary.h" +#include "common/config.h" #include "common/exception.h" #include "common/status.h" +#include "core/data_type_serde/arrow_validation.h" #include "core/value/jsonb_value.h" #include "exprs/json_functions.h" #include "util/jsonb_parser_simd.h" @@ -121,6 +123,9 @@ Status DataTypeJsonbSerDe::write_column_to_arrow(const IColumn& column, const Nu Status DataTypeJsonbSerDe::read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const { + if (config::enable_arrow_input_validation) { + check_arrow_no_offset(*arrow_array); + } if (arrow_array->type_id() == arrow::Type::STRING || arrow_array->type_id() == arrow::Type::BINARY) { const auto* concrete_array = dynamic_cast(arrow_array); diff --git a/be/src/core/data_type_serde/data_type_map_serde.cpp b/be/src/core/data_type_serde/data_type_map_serde.cpp index 5f00a887ff0092..1a2e0643d999b6 100644 --- a/be/src/core/data_type_serde/data_type_map_serde.cpp +++ b/be/src/core/data_type_serde/data_type_map_serde.cpp @@ -18,11 +18,13 @@ #include "core/data_type_serde/data_type_map_serde.h" #include "arrow/array/builder_nested.h" +#include "common/config.h" #include "common/exception.h" #include "common/status.h" #include "core/column/column.h" #include "core/column/column_const.h" #include "core/column/column_map.h" +#include "core/data_type_serde/arrow_validation.h" #include "core/data_type_serde/complex_type_deserialize_util.h" #include "core/string_ref.h" #include "util/jsonb_document.h" @@ -389,6 +391,9 @@ Status DataTypeMapSerDe::read_column_from_arrow(IColumn& column, const arrow::Ar const auto* concrete_map = dynamic_cast(arrow_array); auto arrow_offsets_array = concrete_map->offsets(); auto* arrow_offsets = dynamic_cast(arrow_offsets_array.get()); + if (config::enable_arrow_input_validation) { + check_arrow_map_offsets(*concrete_map, start, end); + } auto prev_size = offsets_data.back(); const auto* base_offsets_ptr = reinterpret_cast(arrow_offsets->raw_values()); diff --git a/be/src/core/data_type_serde/data_type_nullable_serde.cpp b/be/src/core/data_type_serde/data_type_nullable_serde.cpp index b805467bf15a9b..d157fed7730d49 100644 --- a/be/src/core/data_type_serde/data_type_nullable_serde.cpp +++ b/be/src/core/data_type_serde/data_type_nullable_serde.cpp @@ -24,11 +24,13 @@ #include #include +#include "common/config.h" #include "core/assert_cast.h" #include "core/column/column.h" #include "core/column/column_const.h" #include "core/column/column_nullable.h" #include "core/column/column_vector.h" +#include "core/data_type_serde/arrow_validation.h" #include "core/data_type_serde/data_type_serde.h" #include "core/data_type_serde/data_type_string_serde.h" #include "core/data_type_serde/decoded_column_view.h" @@ -344,6 +346,10 @@ Status DataTypeNullableSerDe::read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const { + if (config::enable_arrow_input_validation) { + check_arrow_array_range(*arrow_array, start, end); + check_arrow_validity_bitmap(*arrow_array); + } auto& col = reinterpret_cast(column); NullMap& map_data = col.get_null_map_data(); for (auto i = start; i < end; ++i) { diff --git a/be/src/core/data_type_serde/data_type_number_serde.cpp b/be/src/core/data_type_serde/data_type_number_serde.cpp index 7a4db5c4b4aac5..af6ae9d7f43a72 100644 --- a/be/src/core/data_type_serde/data_type_number_serde.cpp +++ b/be/src/core/data_type_serde/data_type_number_serde.cpp @@ -23,12 +23,14 @@ #include #include +#include "common/config.h" #include "common/exception.h" #include "common/status.h" #include "core/column/column_nullable.h" #include "core/data_type/define_primitive_type.h" #include "core/data_type/primitive_type.h" #include "core/data_type/storage_field_type.h" +#include "core/data_type_serde/arrow_validation.h" #include "core/data_type_serde/data_type_serde.h" #include "core/data_type_serde/decoded_column_view.h" #include "core/packed_int128.h" @@ -404,13 +406,19 @@ Status DataTypeNumberSerDe::read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const { + if (config::enable_arrow_input_validation) { + check_arrow_array_range(*arrow_array, start, end); + } auto row_count = end - start; auto& col_data = static_cast(column).get_data(); // now uint8 for bool if constexpr (T == TYPE_BOOLEAN) { const auto* concrete_array = dynamic_cast(arrow_array); - for (size_t bool_i = 0; bool_i != static_cast(concrete_array->length()); ++bool_i) { + if (config::enable_arrow_input_validation) { + check_arrow_boolean_buffer(*concrete_array); + } + for (int64_t bool_i = start; bool_i != end; ++bool_i) { col_data.emplace_back(concrete_array->Value(bool_i)); } return Status::OK(); @@ -419,7 +427,11 @@ Status DataTypeNumberSerDe::read_column_from_arrow(IColumn& column, // only for largeint(int128) type if (arrow_array->type_id() == arrow::Type::STRING) { const auto* concrete_array = dynamic_cast(arrow_array); + if (config::enable_arrow_input_validation) { + check_arrow_binary_offsets_buffer(*concrete_array); + } std::shared_ptr buffer = concrete_array->value_data(); + const size_t buffer_size = buffer ? static_cast(buffer->size()) : 0; CastParameters params; const auto* offsets_data = concrete_array->value_offsets()->data(); @@ -431,13 +443,17 @@ Status DataTypeNumberSerDe::read_column_from_arrow(IColumn& column, memcpy(&start_offset, offsets_data + offset_i * offset_size, offset_size); memcpy(&end_offset, offsets_data + (offset_i + 1) * offset_size, offset_size); - const auto* raw_data = buffer->data() + start_offset; const auto raw_data_len = end_offset - start_offset; - + if (config::enable_arrow_input_validation) { + check_arrow_value_range(*concrete_array, start_offset, raw_data_len, + buffer_size); + } if (raw_data_len == 0) { col_data.emplace_back( typename PrimitiveTypeTraits::CppType()); // Int128() is NULL } else { + const auto* raw_data = + reinterpret_cast(buffer->data() + start_offset); if constexpr (T == TYPE_DATETIMEV2 || T == TYPE_TIMESTAMPTZ) { StringRef str_ref(raw_data, raw_data_len); UInt64 val = 0; @@ -488,6 +504,10 @@ Status DataTypeNumberSerDe::read_column_from_arrow(IColumn& column, } /// buffers[0] is a null bitmap and buffers[1] are actual values + if (config::enable_arrow_input_validation) { + check_arrow_fixed_width_buffer(*arrow_array, + sizeof(typename PrimitiveTypeTraits::CppType)); + } std::shared_ptr buffer = arrow_array->data()->buffers[1]; // Handle empty array case: buffer can be null when row_count is 0. diff --git a/be/src/core/data_type_serde/data_type_string_serde.cpp b/be/src/core/data_type_serde/data_type_string_serde.cpp index 7f9a82c23cc6ac..80cf74634f0a16 100644 --- a/be/src/core/data_type_serde/data_type_string_serde.cpp +++ b/be/src/core/data_type_serde/data_type_string_serde.cpp @@ -20,8 +20,10 @@ #include #include +#include "common/config.h" #include "core/column/column_string.h" #include "core/data_type/define_primitive_type.h" +#include "core/data_type_serde/arrow_validation.h" #include "core/data_type_serde/decoded_column_view.h" #include "util/jsonb_document_cast.h" #include "util/jsonb_utils.h" @@ -406,7 +408,12 @@ Status DataTypeStringSerDeBase::read_column_from_arrow( if (arrow_array->type_id() == arrow::Type::STRING || arrow_array->type_id() == arrow::Type::BINARY) { const auto* concrete_array = dynamic_cast(arrow_array); + if (config::enable_arrow_input_validation) { + check_arrow_array_range(*concrete_array, start, end); + check_arrow_binary_offsets_buffer(*concrete_array); + } std::shared_ptr buffer = concrete_array->value_data(); + const size_t buffer_size = buffer ? static_cast(buffer->size()) : 0; const uint8_t* offsets_data = concrete_array->value_offsets()->data(); const size_t offset_size = sizeof(int32_t); @@ -418,16 +425,23 @@ Status DataTypeStringSerDeBase::read_column_from_arrow( memcpy(&end_offset, offsets_data + (offset_i + 1) * offset_size, offset_size); int32_t length = end_offset - start_offset; - const auto* raw_data = buffer->data() + start_offset; - - assert_cast(column).insert_data( - reinterpret_cast(raw_data), length); + if (config::enable_arrow_input_validation) { + check_arrow_value_range(*concrete_array, start_offset, length, buffer_size); + } + // insert_data() does not read the input pointer when length is zero. + const auto* raw_data = reinterpret_cast(buffer->data() + start_offset); + assert_cast(column).insert_data(raw_data, length); } else { assert_cast(column).insert_default(); } } } else if (arrow_array->type_id() == arrow::Type::FIXED_SIZE_BINARY) { const auto* concrete_array = dynamic_cast(arrow_array); + if (config::enable_arrow_input_validation) { + check_arrow_array_range(*concrete_array, start, end); + check_arrow_fixed_width_buffer(*concrete_array, + static_cast(concrete_array->byte_width())); + } uint32_t width = concrete_array->byte_width(); for (auto offset_i = start; offset_i < end; ++offset_i) { @@ -441,13 +455,24 @@ Status DataTypeStringSerDeBase::read_column_from_arrow( } else if (arrow_array->type_id() == arrow::Type::LARGE_STRING || arrow_array->type_id() == arrow::Type::LARGE_BINARY) { const auto* concrete_array = dynamic_cast(arrow_array); + if (config::enable_arrow_input_validation) { + check_arrow_array_range(*concrete_array, start, end); + check_arrow_binary_offsets_buffer(*concrete_array); + } std::shared_ptr buffer = concrete_array->value_data(); + const size_t buffer_size = buffer ? static_cast(buffer->size()) : 0; for (auto offset_i = start; offset_i < end; ++offset_i) { if (!concrete_array->IsNull(offset_i)) { - const auto* raw_data = buffer->data() + concrete_array->value_offset(offset_i); - assert_cast(column).insert_data( - (char*)raw_data, concrete_array->value_length(offset_i)); + const auto value_offset = concrete_array->value_offset(offset_i); + const auto value_length = concrete_array->value_length(offset_i); + if (config::enable_arrow_input_validation) { + check_arrow_value_range(*concrete_array, value_offset, value_length, + buffer_size); + } + // insert_data() does not read the input pointer when length is zero. + const auto* raw_data = reinterpret_cast(buffer->data() + value_offset); + assert_cast(column).insert_data(raw_data, value_length); } else { assert_cast(column).insert_default(); } diff --git a/be/src/core/data_type_serde/data_type_struct_serde.cpp b/be/src/core/data_type_serde/data_type_struct_serde.cpp index ee0643bc6aabed..e03aa9cb963f1f 100644 --- a/be/src/core/data_type_serde/data_type_struct_serde.cpp +++ b/be/src/core/data_type_serde/data_type_struct_serde.cpp @@ -20,10 +20,12 @@ #include #include "arrow/array/builder_nested.h" +#include "common/config.h" #include "common/status.h" #include "core/column/column.h" #include "core/column/column_const.h" #include "core/column/column_struct.h" +#include "core/data_type_serde/arrow_validation.h" #include "core/data_type_serde/complex_type_deserialize_util.h" #include "core/data_type_serde/data_type_serde.h" #include "core/string_ref.h" @@ -418,6 +420,9 @@ Status DataTypeStructSerDe::write_column_to_arrow(const IColumn& column, const N Status DataTypeStructSerDe::read_column_from_arrow(IColumn& column, const arrow::Array* arrow_array, int64_t start, int64_t end, const cctz::time_zone& ctz) const { + if (config::enable_arrow_input_validation) { + check_arrow_no_offset(*arrow_array); + } auto& struct_column = static_cast(column); const auto* concrete_struct = dynamic_cast(arrow_array); DCHECK_EQ(struct_column.tuple_size(), concrete_struct->num_fields()); diff --git a/be/src/exec/operator/operator.h b/be/src/exec/operator/operator.h index e8ee719546717b..8c63b9d8a719c0 100644 --- a/be/src/exec/operator/operator.h +++ b/be/src/exec/operator/operator.h @@ -634,6 +634,7 @@ class DataSinkOperatorXBase : public OperatorBase { } [[nodiscard]] Status sink(RuntimeState* state, Block* block, bool eos) { + RETURN_IF_ERROR(block->check_column_and_type_not_null()); RETURN_IF_ERROR(block->check_type_and_column()); return sink_impl(state, block, eos); } @@ -881,6 +882,7 @@ class OperatorXBase : public OperatorBase { Status terminate(RuntimeState* state) override; [[nodiscard]] Status get_block(RuntimeState* state, Block* block, bool* eos) { RETURN_IF_ERROR(get_block_impl(state, block, eos)); + RETURN_IF_ERROR(block->check_column_and_type_not_null()); RETURN_IF_ERROR(block->check_type_and_column()); return Status::OK(); } diff --git a/be/src/exprs/function/cast/cast_to_decimal.h b/be/src/exprs/function/cast/cast_to_decimal.h index 17e9af68723435..7a65d07a5dbcbe 100644 --- a/be/src/exprs/function/cast/cast_to_decimal.h +++ b/be/src/exprs/function/cast/cast_to_decimal.h @@ -781,6 +781,17 @@ class CastToImpl : public CastToBase { params.is_strict = (CastMode == CastModeType::StrictMode); size_t size = vec_from.size(); + if constexpr (IsDataTypeDecimalV3) { + if (to_scale == 0 && !narrow_integral) { + for (size_t i = 0; i < size; i++) { + vec_to_data[i].value = + static_cast(vec_from_data[i]); + } + block.get_by_position(result).column = std::move(col_to); + return Status::OK(); + } + } + RETURN_IF_ERROR(std::visit( [&](auto multiply_may_overflow, auto narrow_integral) { for (size_t i = 0; i < size; i++) { diff --git a/be/test/core/column/column_dictionary_test.cpp b/be/test/core/column/column_dictionary_test.cpp index e802eae7541c9f..389485e2181272 100644 --- a/be/test/core/column/column_dictionary_test.cpp +++ b/be/test/core/column/column_dictionary_test.cpp @@ -248,8 +248,9 @@ TEST_F(ColumnDictionaryTest, permute) { } TEST_F(ColumnDictionaryTest, filter_by_selector) { auto test_func = [&](const auto& source_column) { - auto src_size = source_column->size(); - const auto& codes_data = source_column->get_data(); + const auto& source = *source_column; + auto src_size = source.size(); + const auto& codes_data = source.get_data(); EXPECT_TRUE(src_size <= UINT16_MAX); auto target_column = ColumnString::create(); @@ -262,13 +263,12 @@ TEST_F(ColumnDictionaryTest, filter_by_selector) { size_t sel_size = src_size / 2; indices.resize(sel_size); - auto status = - source_column->filter_by_selector(indices.data(), sel_size, target_column.get()); + auto status = source.filter_by_selector(indices.data(), sel_size, target_column.get()); EXPECT_TRUE(status.ok()); EXPECT_EQ(target_column->size(), sel_size); for (size_t i = 0; i != sel_size; ++i) { auto real_data = target_column->get_data_at(i); - auto expect_data = source_column->get_value(codes_data[indices[i]]); + auto expect_data = source.get_value(codes_data[indices[i]]); if (real_data != expect_data) { std::cout << "index: " << i << ", real_data: " << real_data.to_string() << "\nexpect_data: " << expect_data.to_string() << std::endl; @@ -364,4 +364,4 @@ std::vector ColumnDictionaryTest::dict_array; ColumnDictI32::MutablePtr ColumnDictionaryTest::column_dict_char; ColumnDictI32::MutablePtr ColumnDictionaryTest::column_dict_varchar; ColumnDictI32::MutablePtr ColumnDictionaryTest::column_dict_str; -} // namespace doris \ No newline at end of file +} // namespace doris diff --git a/be/test/core/column/column_nullable_test.cpp b/be/test/core/column/column_nullable_test.cpp index 96db41296eff92..77f167c9ea8f2e 100644 --- a/be/test/core/column/column_nullable_test.cpp +++ b/be/test/core/column/column_nullable_test.cpp @@ -113,7 +113,8 @@ TEST(ColumnNullableTest, PredicateTest) { EXPECT_FALSE(null_dst->has_null()); uint16_t selector[] = {5, 8}; // both null - EXPECT_EQ(nullable_pred->filter_by_selector(selector, 2, null_dst.get()), Status::OK()); + const IColumn& nullable_src = *nullable_pred; + EXPECT_EQ(nullable_src.filter_by_selector(selector, 2, null_dst.get()), Status::OK()); // filter_by_selector must announce to update has_null to make below right. EXPECT_TRUE(null_dst->has_null()); } diff --git a/be/test/core/column/column_string_test.cpp b/be/test/core/column/column_string_test.cpp index 3581be59893244..dcd57487c84a03 100644 --- a/be/test/core/column/column_string_test.cpp +++ b/be/test/core/column/column_string_test.cpp @@ -993,8 +993,8 @@ TEST_F(ColumnStringTest, filter_by_selector) { } std::cout << std::endl; - auto status = - source_column->filter_by_selector(indices.data(), sel_size, target_column.get()); + const auto& source = *source_column; + auto status = source.filter_by_selector(indices.data(), sel_size, target_column.get()); EXPECT_TRUE(status.ok()); EXPECT_EQ(target_column->size(), sel_size); for (size_t i = 0; i != sel_size; ++i) { diff --git a/be/test/core/data_type/block_check_type.cpp b/be/test/core/data_type/block_check_type.cpp index 2756d8714c78d9..6915aabb92f357 100644 --- a/be/test/core/data_type/block_check_type.cpp +++ b/be/test/core/data_type/block_check_type.cpp @@ -39,4 +39,22 @@ TEST(BlockCheckType, test1) { EXPECT_FALSE(st.ok()); std::cout << st.msg() << std::endl; } -} // namespace doris \ No newline at end of file + +TEST(BlockCheckType, CheckColumnAndTypeNotNull) { + auto block = Block { + ColumnHelper::create_column_with_name({1, 2, 3, 4}), + ColumnHelper::create_column_with_name({1, 2, 3, 4}), + }; + + EXPECT_TRUE(block.check_column_and_type_not_null()); + + block.get_by_position(0).column = nullptr; + auto st = block.check_column_and_type_not_null(); + EXPECT_FALSE(st.ok()); + + block.get_by_position(0).column = ColumnHelper::create_column({1, 2, 3, 4}); + block.get_by_position(1).type = nullptr; + st = block.check_column_and_type_not_null(); + EXPECT_FALSE(st.ok()); +} +} // namespace doris diff --git a/be/test/core/data_type_serde/data_type_serde_arrow_validation_test.cpp b/be/test/core/data_type_serde/data_type_serde_arrow_validation_test.cpp new file mode 100644 index 00000000000000..fdce7af8f89b88 --- /dev/null +++ b/be/test/core/data_type_serde/data_type_serde_arrow_validation_test.cpp @@ -0,0 +1,275 @@ +// 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. + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "common/config.h" +#include "common/exception.h" +#include "core/column/column_array.h" +#include "core/column/column_map.h" +#include "core/column/column_nullable.h" +#include "core/column/column_string.h" +#include "core/data_type/primitive_type.h" +#include "core/data_type_serde/data_type_array_serde.h" +#include "core/data_type_serde/data_type_map_serde.h" +#include "core/data_type_serde/data_type_nullable_serde.h" +#include "core/data_type_serde/data_type_number_serde.h" +#include "core/data_type_serde/data_type_string_serde.h" + +namespace doris { +namespace { + +class ScopedArrowInputValidation { +public: + explicit ScopedArrowInputValidation(bool enabled) + : _old_value(config::enable_arrow_input_validation) { + config::enable_arrow_input_validation = enabled; + } + + ~ScopedArrowInputValidation() { config::enable_arrow_input_validation = _old_value; } + +private: + bool _old_value; +}; + +template +void expect_invalid_arrow(Func&& func, std::string_view message) { + bool thrown = false; + try { + std::forward(func)(); + } catch (const Exception& e) { + thrown = true; + EXPECT_EQ(e.code(), ErrorCode::INVALID_ARGUMENT) << e.to_string(); + } + EXPECT_TRUE(thrown) << message; +} + +std::shared_ptr wrap_offsets(const std::vector& offsets) { + return arrow::Buffer::Wrap(offsets); +} + +struct StringArrayHolder { + StringArrayHolder(std::vector offsets_, std::string_view values_) + : offsets(std::move(offsets_)), values(values_) { + auto value_buffer = arrow::Buffer::Wrap(values.data(), values.size()); + array = std::make_shared(offsets.size() - 1, wrap_offsets(offsets), + value_buffer); + } + + std::vector offsets; + std::string values; + std::shared_ptr array; +}; + +} // namespace + +TEST(DataTypeSerDeArrowValidationTest, RejectsShortStringOffsetsBuffer) { + ScopedArrowInputValidation validation(true); + + std::vector offsets = {0}; + std::string_view values = "abc"; + auto value_buffer = arrow::Buffer::Wrap(values.data(), values.size()); + auto array = std::make_shared(1, wrap_offsets(offsets), value_buffer); + auto column = ColumnString::create(); + DataTypeStringSerDe serde(TYPE_STRING); + + expect_invalid_arrow( + [&] { + static_cast(serde.read_column_from_arrow(*column, array.get(), 0, 1, + cctz::utc_time_zone())); + }, + "short string offsets buffer should be rejected"); +} + +TEST(DataTypeSerDeArrowValidationTest, RejectsStringValueRangeBeyondBuffer) { + ScopedArrowInputValidation validation(true); + + StringArrayHolder array({0, 8}, "abc"); + auto column = ColumnString::create(); + DataTypeStringSerDe serde(TYPE_STRING); + + expect_invalid_arrow( + [&] { + static_cast(serde.read_column_from_arrow(*column, array.array.get(), 0, 1, + cctz::utc_time_zone())); + }, + "string value range beyond data buffer should be rejected"); +} + +TEST(DataTypeSerDeArrowValidationTest, RejectsNonMonotonicStringOffsets) { + ScopedArrowInputValidation validation(true); + + StringArrayHolder array({3, 1}, "abcd"); + auto column = ColumnString::create(); + DataTypeStringSerDe serde(TYPE_STRING); + + expect_invalid_arrow( + [&] { + static_cast(serde.read_column_from_arrow(*column, array.array.get(), 0, 1, + cctz::utc_time_zone())); + }, + "non-monotonic string offsets should be rejected"); +} + +TEST(DataTypeSerDeArrowValidationTest, RejectsShortFixedWidthDataBuffer) { + ScopedArrowInputValidation validation(true); + + std::vector values = {1}; + auto data_buffer = arrow::Buffer::Wrap(values); + auto array = std::make_shared(2, data_buffer); + auto column = ColumnInt64::create(); + DataTypeNumberSerDe serde; + + expect_invalid_arrow( + [&] { + static_cast(serde.read_column_from_arrow(*column, array.get(), 0, 2, + cctz::utc_time_zone())); + }, + "short int64 data buffer should be rejected"); +} + +TEST(DataTypeSerDeArrowValidationTest, RejectsSlicedArrowArray) { + ScopedArrowInputValidation validation(true); + + std::vector values = {1, 2, 3}; + auto original = std::make_shared(3, arrow::Buffer::Wrap(values)); + auto sliced = original->Slice(1, 2); + auto column = ColumnInt64::create(); + DataTypeNumberSerDe serde; + + expect_invalid_arrow( + [&] { + static_cast(serde.read_column_from_arrow(*column, sliced.get(), 0, 2, + cctz::utc_time_zone())); + }, + "sliced Arrow array should be rejected"); +} + +TEST(DataTypeSerDeArrowValidationTest, RejectsShortBooleanDataBitmap) { + ScopedArrowInputValidation validation(true); + + std::vector bits = {0xFF}; + auto data_buffer = arrow::Buffer::Wrap(bits); + auto array = std::make_shared(9, data_buffer); + auto column = ColumnUInt8::create(); + DataTypeNumberSerDe serde; + + expect_invalid_arrow( + [&] { + static_cast(serde.read_column_from_arrow(*column, array.get(), 0, 9, + cctz::utc_time_zone())); + }, + "short boolean data bitmap should be rejected"); +} + +TEST(DataTypeSerDeArrowValidationTest, RejectsShortValidityBitmap) { + ScopedArrowInputValidation validation(true); + + std::vector validity = {0xFF}; + std::vector values(9, 1); + auto validity_buffer = arrow::Buffer::Wrap(validity); + auto data_buffer = arrow::Buffer::Wrap(values); + auto array = std::make_shared(9, data_buffer, validity_buffer); + auto column = ColumnNullable::create(ColumnInt64::create(), ColumnUInt8::create()); + auto nested_serde = std::make_shared>(); + DataTypeNullableSerDe serde(nested_serde); + + expect_invalid_arrow( + [&] { + static_cast(serde.read_column_from_arrow(*column, array.get(), 0, 9, + cctz::utc_time_zone())); + }, + "short validity bitmap should be rejected before IsNull"); +} + +TEST(DataTypeSerDeArrowValidationTest, RejectsMissingValidityBitmapWithNullCount) { + ScopedArrowInputValidation validation(true); + + std::vector values = {1, 2}; + auto data_buffer = arrow::Buffer::Wrap(values); + auto array = std::make_shared(2, data_buffer, + std::shared_ptr(), 1); + auto column = ColumnNullable::create(ColumnInt64::create(), ColumnUInt8::create()); + auto nested_serde = std::make_shared>(); + DataTypeNullableSerDe serde(nested_serde); + + expect_invalid_arrow( + [&] { + static_cast(serde.read_column_from_arrow(*column, array.get(), 0, 2, + cctz::utc_time_zone())); + }, + "missing validity bitmap with positive null_count should be rejected"); +} + +TEST(DataTypeSerDeArrowValidationTest, RejectsListOffsetsBeyondValuesLength) { + ScopedArrowInputValidation validation(true); + + StringArrayHolder values({0, 1}, "a"); + std::vector offsets = {0, 2}; + auto offsets_buffer = wrap_offsets(offsets); + auto array = std::make_shared(arrow::list(arrow::utf8()), 1, offsets_buffer, + values.array); + auto column = ColumnArray::create( + ColumnNullable::create(ColumnString::create(), ColumnUInt8::create()), + ColumnOffset64::create()); + auto nested_serde = std::make_shared( + std::make_shared(TYPE_STRING)); + DataTypeArraySerDe serde(nested_serde); + + expect_invalid_arrow( + [&] { + static_cast(serde.read_column_from_arrow(*column, array.get(), 0, 1, + cctz::utc_time_zone())); + }, + "list offsets beyond values length should be rejected"); +} + +TEST(DataTypeSerDeArrowValidationTest, RejectsMapOffsetsBeyondKeysLength) { + ScopedArrowInputValidation validation(true); + + StringArrayHolder keys({0, 1}, "k"); + std::vector item_values = {1}; + auto items = std::make_shared(1, arrow::Buffer::Wrap(item_values)); + std::vector offsets = {0, 2}; + auto offsets_buffer = wrap_offsets(offsets); + auto array = std::make_shared(arrow::map(arrow::utf8(), arrow::int64()), 1, + offsets_buffer, keys.array, items); + auto column = ColumnMap::create(ColumnString::create(), ColumnInt64::create(), + ColumnOffset64::create()); + DataTypeMapSerDe serde(std::make_shared(TYPE_STRING), + std::make_shared>()); + + expect_invalid_arrow( + [&] { + static_cast(serde.read_column_from_arrow(*column, array.get(), 0, 1, + cctz::utc_time_zone())); + }, + "map offsets beyond keys length should be rejected"); +} + +} // namespace doris