diff --git a/be/src/common/config.cpp b/be/src/common/config.cpp index 9ec587bd100ec5..2c89946f647978 100644 --- a/be/src/common/config.cpp +++ b/be/src/common/config.cpp @@ -1081,6 +1081,8 @@ DEFINE_mInt32(merged_hdfs_min_io_size, "8192"); // OrcReader DEFINE_mInt32(orc_natural_read_size_mb, "8"); +DEFINE_Validator(orc_natural_read_size_mb, + [](const int config) -> bool { return config > 0 && config <= 1024; }); // Perform the always_true check at intervals determined by runtime_filter_sampling_frequency DEFINE_mInt32(runtime_filter_sampling_frequency, "32"); DEFINE_mInt32(execution_max_rpc_timeout_sec, "3600"); diff --git a/be/src/format_v2/orc/orc_file_input_stream.cpp b/be/src/format_v2/orc/orc_file_input_stream.cpp new file mode 100644 index 00000000000000..a5144d7b2017b0 --- /dev/null +++ b/be/src/format_v2/orc/orc_file_input_stream.cpp @@ -0,0 +1,377 @@ +// 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 "format_v2/orc/orc_file_input_stream.h" + +#include +#include +#include + +#include "common/status.h" +#include "core/custom_allocator.h" +#include "io/fs/buffered_reader.h" +#include "io/fs/tracing_file_reader.h" +#include "io/io_common.h" +#include "orc/Exceptions.hh" +#include "runtime/runtime_profile.h" +#include "util/slice.h" + +namespace doris::format::orc { +namespace { + +struct OrcMergedRangeStatistics { + int64_t copy_time = 0; + int64_t read_time = 0; + int64_t request_io = 0; + int64_t merged_io = 0; + int64_t request_bytes = 0; + int64_t merged_bytes = 0; + int64_t apply_bytes = 0; + int64_t cluster_num = 1; +}; + +class OrcMergedRangeFileReader final : public io::FileReader { +public: + OrcMergedRangeFileReader(RuntimeProfile* profile, io::FileReaderSPtr file_reader, + io::PrefetchRange range) + : _profile(profile), + _file_reader(std::move(file_reader)), + _range(range), + _size(_file_reader->size()) { + _statistics.apply_bytes += _range.end_offset - _range.start_offset; + if (_profile != nullptr) { + const char* profile_name = "MergedSmallIO"; + ADD_TIMER_WITH_LEVEL(_profile, profile_name, 1); + _copy_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, "CopyTime", profile_name, 1); + _read_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, "ReadTime", profile_name, 1); + _request_io = ADD_CHILD_COUNTER_WITH_LEVEL(_profile, "RequestIO", TUnit::UNIT, + profile_name, 1); + _merged_io = ADD_CHILD_COUNTER_WITH_LEVEL(_profile, "MergedIO", TUnit::UNIT, + profile_name, 1); + _request_bytes = ADD_CHILD_COUNTER_WITH_LEVEL(_profile, "RequestBytes", TUnit::BYTES, + profile_name, 1); + _merged_bytes = ADD_CHILD_COUNTER_WITH_LEVEL(_profile, "MergedBytes", TUnit::BYTES, + profile_name, 1); + _apply_bytes = ADD_CHILD_COUNTER_WITH_LEVEL(_profile, "ApplyBytes", TUnit::BYTES, + profile_name, 1); + _over_read_bytes = ADD_CHILD_COUNTER_WITH_LEVEL(_profile, "OverReadBytes", TUnit::BYTES, + profile_name, 1); + _cluster_num = ADD_CHILD_COUNTER_WITH_LEVEL(_profile, "ClusterNum", TUnit::UNIT, + profile_name, 1); + } + } + + Status close() override { + _closed = true; + return Status::OK(); + } + + const io::Path& path() const override { return _file_reader->path(); } + size_t size() const override { return _size; } + bool closed() const override { return _closed; } + int64_t mtime() const override { return _file_reader->mtime(); } + +protected: + Status read_at_impl(size_t offset, Slice result, size_t* bytes_read, + const io::IOContext* io_ctx) override { + ++_statistics.request_io; + _statistics.request_bytes += static_cast(result.size); + *bytes_read = 0; + if (result.size == 0) { + return Status::OK(); + } + if (offset < _range.start_offset || offset + result.size > _range.end_offset) { + return Status::IOError("ORC stripe read [{}, {}) is outside merged range [{}, {})", + offset, offset + result.size, _range.start_offset, + _range.end_offset); + } + + RETURN_IF_ERROR(_load(io_ctx)); + { + SCOPED_RAW_TIMER(&_statistics.copy_time); + std::memcpy(result.data, _cache.get() + offset - _range.start_offset, result.size); + } + *bytes_read = result.size; + return Status::OK(); + } + + void _collect_profile_before_close() override { + if (_profile == nullptr) { + return; + } + COUNTER_UPDATE(_copy_time, _statistics.copy_time); + COUNTER_UPDATE(_read_time, _statistics.read_time); + COUNTER_UPDATE(_request_io, _statistics.request_io); + COUNTER_UPDATE(_merged_io, _statistics.merged_io); + COUNTER_UPDATE(_request_bytes, _statistics.request_bytes); + COUNTER_UPDATE(_merged_bytes, _statistics.merged_bytes); + COUNTER_UPDATE(_apply_bytes, _statistics.apply_bytes); + COUNTER_UPDATE(_over_read_bytes, + std::max(_statistics.merged_bytes - _statistics.request_bytes, 0)); + COUNTER_UPDATE(_cluster_num, _statistics.cluster_num); + } + +private: + Status _load(const io::IOContext* io_ctx) { + if (_loaded) { + return Status::OK(); + } + + const size_t range_size = _range.end_offset - _range.start_offset; + _cache = make_unique_buffer(range_size); + size_t total_read = 0; + { + SCOPED_RAW_TIMER(&_statistics.read_time); + while (total_read < range_size) { + size_t loop_read = 0; + RETURN_IF_ERROR(_file_reader->read_at( + _range.start_offset + total_read, + Slice(_cache.get() + total_read, range_size - total_read), &loop_read, + io_ctx)); + ++_statistics.merged_io; + _statistics.merged_bytes += static_cast(loop_read); + if (loop_read == 0) { + return Status::IOError("Short read for ORC merged range [{}, {})", + _range.start_offset, _range.end_offset); + } + total_read += loop_read; + } + } + _loaded = true; + return Status::OK(); + } + + RuntimeProfile* _profile = nullptr; + io::FileReaderSPtr _file_reader; + io::PrefetchRange _range; + size_t _size = 0; + bool _closed = false; + bool _loaded = false; + DorisUniqueBufferPtr _cache; + OrcMergedRangeStatistics _statistics; + + RuntimeProfile::Counter* _copy_time = nullptr; + RuntimeProfile::Counter* _read_time = nullptr; + RuntimeProfile::Counter* _request_io = nullptr; + RuntimeProfile::Counter* _merged_io = nullptr; + RuntimeProfile::Counter* _request_bytes = nullptr; + RuntimeProfile::Counter* _merged_bytes = nullptr; + RuntimeProfile::Counter* _apply_bytes = nullptr; + RuntimeProfile::Counter* _over_read_bytes = nullptr; + RuntimeProfile::Counter* _cluster_num = nullptr; +}; + +class OrcStripeInputStream final : public ::orc::InputStream { +public: + OrcStripeInputStream(std::string file_name, io::FileReaderSPtr file_reader, + const io::IOContext* io_ctx, uint64_t natural_read_size) + : _file_name(std::move(file_name)), + _file_reader(std::move(file_reader)), + _io_ctx(io_ctx), + _natural_read_size(natural_read_size) {} + + uint64_t getLength() const override { return _file_reader->size(); } + uint64_t getNaturalReadSize() const override { return _natural_read_size; } + + void read(void* buf, uint64_t length, uint64_t offset) override { + uint64_t bytes_read = 0; + auto* out = static_cast(buf); + while (bytes_read < length) { + if (_io_ctx != nullptr && _io_ctx->should_stop) { + throw ::orc::ParseError("stop"); + } + size_t loop_read = 0; + Status st = _file_reader->read_at( + static_cast(offset + bytes_read), + Slice(out + bytes_read, static_cast(length - bytes_read)), &loop_read, + _io_ctx); + if (!st.ok()) { + throw ::orc::ParseError("Failed to read " + _file_name + ": " + + st.to_string_no_stack()); + } + if (loop_read == 0) { + break; + } + bytes_read += loop_read; + } + if (bytes_read != length) { + throw ::orc::ParseError("Short read from " + _file_name); + } + } + + const std::string& getName() const override { return _file_name; } + +private: + std::string _file_name; + io::FileReaderSPtr _file_reader; + const io::IOContext* _io_ctx = nullptr; + uint64_t _natural_read_size = 0; +}; + +struct StripeStreamRange { + ::orc::StreamId stream_id; + io::PrefetchRange range; +}; + +} // namespace + +OrcFileInputStream::OrcFileInputStream(std::string file_name, io::FileReaderSPtr file_reader, + const io::IOContext* io_ctx, RuntimeProfile* profile, + OrcFileInputStreamOptions options) + : _file_name(std::move(file_name)), + _file_reader(std::move(file_reader)), + _default_reader(io_ctx != nullptr && io_ctx->file_reader_stats != nullptr + ? std::make_shared( + _file_reader, io_ctx->file_reader_stats) + : _file_reader), + _io_ctx(io_ctx), + _profile(profile), + _options(options) { + DORIS_CHECK_GT(_options.natural_read_size, 0); + DORIS_CHECK_GE(_options.once_max_read_bytes, 0); + DORIS_CHECK_GE(_options.max_merge_distance_bytes, 0); +} + +OrcFileInputStream::~OrcFileInputStream() { + _flush_active_clusters(); +} + +uint64_t OrcFileInputStream::getLength() const { + return _default_reader->size(); +} + +uint64_t OrcFileInputStream::getNaturalReadSize() const { + return _options.natural_read_size; +} + +void OrcFileInputStream::read(void* buf, uint64_t length, uint64_t offset) { + OrcStripeInputStream(_file_name, _default_reader, _io_ctx, _options.natural_read_size) + .read(buf, length, offset); +} + +const std::string& OrcFileInputStream::getName() const { + return _file_name; +} + +void OrcFileInputStream::beforeReadStripe( + std::unique_ptr<::orc::StripeInformation> current_stripe_information, + const std::vector& selected_columns, + std::unordered_map<::orc::StreamId, std::shared_ptr<::orc::InputStream>>& streams) { + _flush_active_clusters(); + _active_stripe_streams.clear(); + + std::vector small_streams; + std::vector<::orc::StreamId> direct_streams; + uint64_t offset = current_stripe_information->getOffset(); + for (uint64_t stream_index = 0; stream_index < current_stripe_information->getNumberOfStreams(); + ++stream_index) { + auto stream = current_stripe_information->getStreamInformation(stream_index); + const uint64_t column_id = stream->getColumnId(); + if (column_id >= selected_columns.size()) { + throw ::orc::ParseError( + fmt::format("Invalid ORC stream column id {} in {}, selected column count {}", + column_id, _file_name, selected_columns.size())); + } + const uint64_t length = stream->getLength(); + if (selected_columns[column_id]) { + ::orc::StreamId stream_id(column_id, stream->getKind()); + if (length == 0 || std::cmp_greater(length, _options.once_max_read_bytes)) { + direct_streams.push_back(stream_id); + } else { + small_streams.push_back({stream_id, io::PrefetchRange(offset, offset + length)}); + } + } + offset += length; + } + + for (const auto& stream_id : direct_streams) { + _add_direct_stream(stream_id, streams); + } + if (small_streams.empty()) { + return; + } + + std::sort(small_streams.begin(), small_streams.end(), + [](const StripeStreamRange& left, const StripeStreamRange& right) { + return left.range.start_offset < right.range.start_offset; + }); + std::vector small_ranges; + small_ranges.reserve(small_streams.size()); + for (const auto& stream : small_streams) { + small_ranges.push_back(stream.range); + } + const auto merged_ranges = io::PrefetchRange::merge_adjacent_seq_ranges( + small_ranges, _options.max_merge_distance_bytes, _options.once_max_read_bytes); + + size_t stream_index = 0; + for (const auto& merged_range : merged_ranges) { + std::vector> cluster_streams; + while (stream_index < small_streams.size() && + small_streams[stream_index].range.start_offset < merged_range.end_offset) { + DORIS_CHECK_LE(small_streams[stream_index].range.end_offset, merged_range.end_offset); + cluster_streams.emplace_back(small_streams[stream_index].stream_id, + small_streams[stream_index].range); + ++stream_index; + } + DORIS_CHECK(!cluster_streams.empty()); + if (cluster_streams.size() == 1) { + _add_direct_stream(cluster_streams.front().first, streams); + } else { + _add_clustered_streams(cluster_streams, merged_range, streams); + } + } + DORIS_CHECK_EQ(stream_index, small_streams.size()); +} + +void OrcFileInputStream::_flush_active_clusters() { + for (const auto& reader : _active_cluster_readers) { + reader->collect_profile_before_close(); + } + _active_cluster_readers.clear(); +} + +void OrcFileInputStream::_add_direct_stream( + const ::orc::StreamId& stream_id, + std::unordered_map<::orc::StreamId, std::shared_ptr<::orc::InputStream>>& streams) { + auto stream = std::make_shared(_file_name, _default_reader, _io_ctx, + _options.natural_read_size); + streams.emplace(stream_id, stream); + _active_stripe_streams.push_back(std::move(stream)); +} + +void OrcFileInputStream::_add_clustered_streams( + const std::vector>& cluster_streams, + const io::PrefetchRange& cluster_range, + std::unordered_map<::orc::StreamId, std::shared_ptr<::orc::InputStream>>& streams) { + io::FileReaderSPtr cluster_reader = + std::make_shared(_profile, _file_reader, cluster_range); + if (_io_ctx != nullptr && _io_ctx->file_reader_stats != nullptr) { + cluster_reader = std::make_shared(std::move(cluster_reader), + _io_ctx->file_reader_stats); + } + _active_cluster_readers.push_back(cluster_reader); + for (const auto& [stream_id, range] : cluster_streams) { + DORIS_CHECK_GE(range.start_offset, cluster_range.start_offset); + DORIS_CHECK_LE(range.end_offset, cluster_range.end_offset); + auto stream = std::make_shared(_file_name, cluster_reader, _io_ctx, + _options.natural_read_size); + streams.emplace(stream_id, stream); + _active_stripe_streams.push_back(std::move(stream)); + } +} + +} // namespace doris::format::orc diff --git a/be/src/format_v2/orc/orc_file_input_stream.h b/be/src/format_v2/orc/orc_file_input_stream.h new file mode 100644 index 00000000000000..3ced728389004f --- /dev/null +++ b/be/src/format_v2/orc/orc_file_input_stream.h @@ -0,0 +1,84 @@ +// 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 "io/fs/buffered_reader.h" + +namespace doris { + +class RuntimeProfile; + +namespace io { +struct IOContext; +} + +namespace format::orc { + +struct OrcFileInputStreamOptions { + uint64_t natural_read_size = 8L * 1024L * 1024L; + int64_t once_max_read_bytes = 8L * 1024L * 1024L; + int64_t max_merge_distance_bytes = 1L * 1024L * 1024L; +}; + +class OrcFileInputStream final : public ::orc::InputStream { +public: + OrcFileInputStream(std::string file_name, io::FileReaderSPtr file_reader, + const io::IOContext* io_ctx, RuntimeProfile* profile, + OrcFileInputStreamOptions options); + ~OrcFileInputStream() override; + + uint64_t getLength() const override; + uint64_t getNaturalReadSize() const override; + void read(void* buf, uint64_t length, uint64_t offset) override; + const std::string& getName() const override; + + void beforeReadStripe(std::unique_ptr<::orc::StripeInformation> current_stripe_information, + const std::vector& selected_columns, + std::unordered_map<::orc::StreamId, std::shared_ptr<::orc::InputStream>>& + streams) override; + +private: + void _flush_active_clusters(); + void _add_direct_stream( + const ::orc::StreamId& stream_id, + std::unordered_map<::orc::StreamId, std::shared_ptr<::orc::InputStream>>& streams); + void _add_clustered_streams( + const std::vector>& cluster_streams, + const io::PrefetchRange& cluster_range, + std::unordered_map<::orc::StreamId, std::shared_ptr<::orc::InputStream>>& streams); + + std::string _file_name; + io::FileReaderSPtr _file_reader; + io::FileReaderSPtr _default_reader; + const io::IOContext* _io_ctx = nullptr; + RuntimeProfile* _profile = nullptr; + OrcFileInputStreamOptions _options; + + std::vector> _active_stripe_streams; + std::vector _active_cluster_readers; +}; + +} // namespace format::orc +} // namespace doris diff --git a/be/src/format_v2/orc/orc_reader.cpp b/be/src/format_v2/orc/orc_reader.cpp index 32b8aed2a888db..73306473987c39 100644 --- a/be/src/format_v2/orc/orc_reader.cpp +++ b/be/src/format_v2/orc/orc_reader.cpp @@ -44,6 +44,7 @@ #include #include "common/cast_set.h" +#include "common/config.h" #include "common/consts.h" #include "common/exception.h" #include "core/block/block.h" @@ -68,6 +69,7 @@ #include "exprs/vliteral.h" #include "exprs/vslot_ref.h" #include "format_v2/column_mapper.h" +#include "format_v2/orc/orc_file_input_stream.h" #include "format_v2/orc/orc_search_argument.h" #include "io/fs/file_reader.h" #include "runtime/exec_env.h" @@ -82,7 +84,6 @@ namespace doris::format::orc { namespace { constexpr uint64_t DEFAULT_ORC_READ_BATCH_SIZE = 4096; -constexpr uint64_t DEFAULT_ORC_NATURAL_READ_SIZE = 128 * 1024; constexpr int DECIMAL_PRECISION_FOR_HIVE11 = BeConsts::MAX_DECIMAL128_PRECISION; constexpr int DECIMAL_SCALE_FOR_HIVE11 = 10; constexpr const char* ORC_LIST_ELEMENT_NAME = "element"; @@ -140,54 +141,6 @@ Status set_orc_reader_timezone(const std::string& timezone, return Status::OK(); } -// Thin adapter from Doris FileReader to ORC's InputStream API. Keep IO policy, -// tracing, and retry behavior in the underlying FileReader. -class DorisOrcInputStream final : public ::orc::InputStream { -public: - DorisOrcInputStream(std::string file_name, io::FileReaderSPtr file_reader, - io::IOContext* io_ctx) - : _file_name(std::move(file_name)), - _file_reader(std::move(file_reader)), - _io_ctx(io_ctx) {} - - uint64_t getLength() const override { return _file_reader->size(); } - - uint64_t getNaturalReadSize() const override { return DEFAULT_ORC_NATURAL_READ_SIZE; } - - void read(void* buf, uint64_t length, uint64_t offset) override { - uint64_t bytes_read = 0; - auto* out = static_cast(buf); - while (bytes_read < length) { - if (_io_ctx != nullptr && _io_ctx->should_stop) { - throw ::orc::ParseError("stop"); - } - size_t loop_read = 0; - Status st = _file_reader->read_at( - static_cast(offset + bytes_read), - Slice(out + bytes_read, static_cast(length - bytes_read)), &loop_read, - _io_ctx); - if (!st.ok()) { - throw ::orc::ParseError("Failed to read " + _file_name + ": " + - st.to_string_no_stack()); - } - if (loop_read == 0) { - break; - } - bytes_read += loop_read; - } - if (bytes_read != length) { - throw ::orc::ParseError("Short read from " + _file_name); - } - } - - const std::string& getName() const override { return _file_name; } - -private: - std::string _file_name; - io::FileReaderSPtr _file_reader; - io::IOContext* _io_ctx = nullptr; -}; - // selected_rows is a source-row remap produced by ORC lazy callback: // predicate columns are decoded first, then surviving row ids drive follower decodes. size_t decode_row_count(size_t rows, const std::vector* selected_rows) { @@ -870,8 +823,21 @@ Status OrcReader::init(RuntimeState* state) { options.setMemoryPool(*ExecEnv::GetInstance()->orc_memory_pool()); options.setReaderMetrics(&_state->reader_metrics); - auto input_stream = std::make_unique(_file_description->path, - _tracing_file_reader, _io_ctx.get()); + OrcFileInputStreamOptions input_stream_options; + const auto natural_read_size_mb = config::orc_natural_read_size_mb; + if (natural_read_size_mb <= 0 || natural_read_size_mb > 1024) { + return Status::InvalidArgument( + "Invalid orc_natural_read_size_mb {}, valid range is [1, 1024]", + natural_read_size_mb); + } + input_stream_options.natural_read_size = static_cast(natural_read_size_mb) << 20; + if (state != nullptr) { + input_stream_options.once_max_read_bytes = state->query_options().orc_once_max_read_bytes; + input_stream_options.max_merge_distance_bytes = + state->query_options().orc_max_merge_distance_bytes; + } + auto input_stream = std::make_unique( + _file_description->path, _file_reader, _io_ctx.get(), _profile, input_stream_options); try { _state->reader = ::orc::createReader(std::move(input_stream), options); _state->root_type = &_state->reader->getType(); diff --git a/be/test/format_v2/orc/orc_file_input_stream_test.cpp b/be/test/format_v2/orc/orc_file_input_stream_test.cpp new file mode 100644 index 00000000000000..54c89505edc3f5 --- /dev/null +++ b/be/test/format_v2/orc/orc_file_input_stream_test.cpp @@ -0,0 +1,343 @@ +// 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 "format_v2/orc/orc_file_input_stream.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "runtime/runtime_profile.h" + +namespace doris::format::orc { +namespace { + +struct TestStream { + uint64_t column_id; + ::orc::StreamKind kind; + uint64_t length; +}; + +class RecordingFileReader final : public io::FileReader { +public: + explicit RecordingFileReader(size_t size) : _data(size) { + for (size_t offset = 0; offset < size; ++offset) { + _data[offset] = static_cast(offset % 251); + } + } + + Status close() override { + _closed = true; + return Status::OK(); + } + const io::Path& path() const override { return _path; } + size_t size() const override { return _data.size(); } + bool closed() const override { return _closed; } + int64_t mtime() const override { return 0; } + + const std::vector& reads() const { return _reads; } + char value_at(size_t offset) const { return _data[offset]; } + +protected: + Status read_at_impl(size_t offset, Slice result, size_t* bytes_read, + const io::IOContext* io_ctx) override { + _reads.emplace_back(offset, std::min(offset + result.size, _data.size())); + *bytes_read = std::min(result.size, _data.size() - offset); + std::memcpy(result.data, _data.data() + offset, *bytes_read); + return Status::OK(); + } + +private: + std::vector _data; + std::vector _reads; + bool _closed = false; + io::Path _path = "/tmp/orc_v2_input_stream"; +}; + +class TestStreamInformation final : public ::orc::StreamInformation { +public: + TestStreamInformation(TestStream stream, uint64_t offset) : _stream(stream), _offset(offset) {} + + ::orc::StreamKind getKind() const override { return _stream.kind; } + uint64_t getColumnId() const override { return _stream.column_id; } + uint64_t getOffset() const override { return _offset; } + uint64_t getLength() const override { return _stream.length; } + +private: + TestStream _stream; + uint64_t _offset; +}; + +class TestStripeInformation final : public ::orc::StripeInformation { +public: + TestStripeInformation(uint64_t offset, std::vector streams) + : _offset(offset), _streams(std::move(streams)) {} + + uint64_t getOffset() const override { return _offset; } + uint64_t getLength() const override { + uint64_t length = 0; + for (const auto& stream : _streams) { + length += stream.length; + } + return length; + } + uint64_t getIndexLength() const override { return 0; } + uint64_t getDataLength() const override { return getLength(); } + uint64_t getFooterLength() const override { return 0; } + uint64_t getNumberOfRows() const override { return 1; } + uint64_t getNumberOfStreams() const override { return _streams.size(); } + + std::unique_ptr<::orc::StreamInformation> getStreamInformation( + uint64_t stream_id) const override { + uint64_t offset = _offset; + for (uint64_t index = 0; index < stream_id; ++index) { + offset += _streams[index].length; + } + return std::make_unique(_streams[stream_id], offset); + } + + ::orc::ColumnEncodingKind getColumnEncoding(uint64_t col_id) const override { + return ::orc::ColumnEncodingKind_DIRECT; + } + uint64_t getDictionarySize(uint64_t col_id) const override { return 0; } + const std::string& getWriterTimezone() const override { return _timezone; } + +private: + uint64_t _offset; + std::vector _streams; + std::string _timezone = "UTC"; +}; + +using StripeStreamMap = std::unordered_map<::orc::StreamId, std::shared_ptr<::orc::InputStream>>; + +std::shared_ptr<::orc::InputStream> find_stream(const StripeStreamMap& streams, uint64_t column_id, + ::orc::StreamKind kind) { + auto it = streams.find(::orc::StreamId(column_id, kind)); + EXPECT_NE(it, streams.end()); + return it == streams.end() ? nullptr : it->second; +} + +std::vector selected_columns(std::initializer_list columns) { + std::vector selected(8, false); + for (uint64_t column : columns) { + selected[column] = true; + } + return selected; +} + +TEST(OrcFileInputStreamTest, MergesAdjacentSelectedStreams) { + auto reader = std::make_shared(256); + OrcFileInputStream input("test.orc", reader, nullptr, nullptr, + {.once_max_read_bytes = 16, .max_merge_distance_bytes = 0}); + StripeStreamMap streams; + input.beforeReadStripe(std::make_unique( + 100, std::vector {{1, ::orc::StreamKind_DATA, 4}, + {2, ::orc::StreamKind_DATA, 4}}), + selected_columns({1, 2}), streams); + + auto second = find_stream(streams, 2, ::orc::StreamKind_DATA); + ASSERT_NE(second, nullptr); + std::array second_data {}; + second->read(second_data.data(), second_data.size(), 104); + + auto first = find_stream(streams, 1, ::orc::StreamKind_DATA); + ASSERT_NE(first, nullptr); + std::array first_data {}; + first->read(first_data.data(), first_data.size(), 100); + + ASSERT_EQ(reader->reads().size(), 1); + EXPECT_EQ(reader->reads().front(), io::PrefetchRange(100, 108)); + EXPECT_EQ(second_data[0], reader->value_at(104)); + EXPECT_EQ(first_data[0], reader->value_at(100)); +} + +TEST(OrcFileInputStreamTest, SupportsRepeatedAndBackwardClusterReads) { + auto reader = std::make_shared(256); + OrcFileInputStream input("test.orc", reader, nullptr, nullptr, + {.once_max_read_bytes = 32, .max_merge_distance_bytes = 8}); + StripeStreamMap streams; + input.beforeReadStripe(std::make_unique( + 40, std::vector {{1, ::orc::StreamKind_DATA, 4}, + {0, ::orc::StreamKind_DATA, 3}, + {2, ::orc::StreamKind_LENGTH, 5}}), + selected_columns({1, 2}), streams); + + auto later = find_stream(streams, 2, ::orc::StreamKind_LENGTH); + auto earlier = find_stream(streams, 1, ::orc::StreamKind_DATA); + ASSERT_NE(later, nullptr); + ASSERT_NE(earlier, nullptr); + + std::array later_data {}; + later->read(later_data.data(), later_data.size(), 47); + std::array repeated_data {}; + later->read(repeated_data.data(), repeated_data.size(), 48); + std::array earlier_data {}; + earlier->read(earlier_data.data(), earlier_data.size(), 40); + + ASSERT_EQ(reader->reads().size(), 1); + EXPECT_EQ(reader->reads().front(), io::PrefetchRange(40, 52)); + EXPECT_EQ(repeated_data[0], reader->value_at(48)); + EXPECT_EQ(earlier_data[0], reader->value_at(40)); +} + +TEST(OrcFileInputStreamTest, RespectsMergeDistanceAndMaximumSpan) { + auto distance_reader = std::make_shared(256); + OrcFileInputStream distance_input("test.orc", distance_reader, nullptr, nullptr, + {.once_max_read_bytes = 32, .max_merge_distance_bytes = 2}); + StripeStreamMap distance_streams; + distance_input.beforeReadStripe( + std::make_unique( + 0, std::vector {{1, ::orc::StreamKind_DATA, 4}, + {0, ::orc::StreamKind_DATA, 3}, + {2, ::orc::StreamKind_DATA, 4}}), + selected_columns({1, 2}), distance_streams); + std::array buffer {}; + find_stream(distance_streams, 1, ::orc::StreamKind_DATA)->read(buffer.data(), buffer.size(), 0); + find_stream(distance_streams, 2, ::orc::StreamKind_DATA)->read(buffer.data(), buffer.size(), 7); + ASSERT_EQ(distance_reader->reads().size(), 2); + EXPECT_EQ(distance_reader->reads()[0], io::PrefetchRange(0, 4)); + EXPECT_EQ(distance_reader->reads()[1], io::PrefetchRange(7, 11)); + + auto span_reader = std::make_shared(256); + OrcFileInputStream span_input("test.orc", span_reader, nullptr, nullptr, + {.once_max_read_bytes = 8, .max_merge_distance_bytes = 0}); + StripeStreamMap span_streams; + span_input.beforeReadStripe( + std::make_unique( + 20, std::vector {{1, ::orc::StreamKind_DATA, 4}, + {2, ::orc::StreamKind_DATA, 5}}), + selected_columns({1, 2}), span_streams); + find_stream(span_streams, 1, ::orc::StreamKind_DATA)->read(buffer.data(), buffer.size(), 20); + std::array second_buffer {}; + find_stream(span_streams, 2, ::orc::StreamKind_DATA) + ->read(second_buffer.data(), second_buffer.size(), 24); + ASSERT_EQ(span_reader->reads().size(), 2); + EXPECT_EQ(span_reader->reads()[0], io::PrefetchRange(20, 24)); + EXPECT_EQ(span_reader->reads()[1], io::PrefetchRange(24, 29)); +} + +TEST(OrcFileInputStreamTest, KeepsLargeAndSingleStreamsDirect) { + auto reader = std::make_shared(256); + OrcFileInputStream input("test.orc", reader, nullptr, nullptr, + {.once_max_read_bytes = 8, .max_merge_distance_bytes = 8}); + StripeStreamMap streams; + input.beforeReadStripe(std::make_unique( + 60, std::vector {{1, ::orc::StreamKind_DATA, 9}, + {2, ::orc::StreamKind_DATA, 4}}), + selected_columns({1, 2}), streams); + + std::array large_data {}; + find_stream(streams, 1, ::orc::StreamKind_DATA)->read(large_data.data(), large_data.size(), 62); + std::array single_data {}; + find_stream(streams, 2, ::orc::StreamKind_DATA) + ->read(single_data.data(), single_data.size(), 69); + + ASSERT_EQ(reader->reads().size(), 2); + EXPECT_EQ(reader->reads()[0], io::PrefetchRange(62, 65)); + EXPECT_EQ(reader->reads()[1], io::PrefetchRange(69, 73)); +} + +TEST(OrcFileInputStreamTest, ZeroOnceMaxReadBytesKeepsStreamsDirect) { + auto reader = std::make_shared(256); + OrcFileInputStream input("test.orc", reader, nullptr, nullptr, + {.once_max_read_bytes = 0, .max_merge_distance_bytes = 16}); + StripeStreamMap streams; + input.beforeReadStripe(std::make_unique( + 120, std::vector {{1, ::orc::StreamKind_DATA, 4}, + {2, ::orc::StreamKind_DATA, 4}}), + selected_columns({1, 2}), streams); + + std::array first_data {}; + find_stream(streams, 1, ::orc::StreamKind_DATA) + ->read(first_data.data(), first_data.size(), 120); + std::array second_data {}; + find_stream(streams, 2, ::orc::StreamKind_DATA) + ->read(second_data.data(), second_data.size(), 124); + + ASSERT_EQ(reader->reads().size(), 2); + EXPECT_EQ(reader->reads()[0], io::PrefetchRange(120, 124)); + EXPECT_EQ(reader->reads()[1], io::PrefetchRange(124, 128)); +} + +TEST(OrcFileInputStreamTest, SkipsUnselectedStreams) { + auto reader = std::make_shared(256); + OrcFileInputStream input("test.orc", reader, nullptr, nullptr, + {.once_max_read_bytes = 16, .max_merge_distance_bytes = 16}); + StripeStreamMap streams; + input.beforeReadStripe(std::make_unique( + 80, std::vector {{1, ::orc::StreamKind_DATA, 4}, + {2, ::orc::StreamKind_DATA, 4}}), + selected_columns({2}), streams); + + EXPECT_EQ(streams.size(), 1); + EXPECT_EQ(streams.count(::orc::StreamId(1, ::orc::StreamKind_DATA)), 0); + EXPECT_EQ(streams.count(::orc::StreamId(2, ::orc::StreamKind_DATA)), 1); +} + +TEST(OrcFileInputStreamTest, RejectsInvalidStreamColumnId) { + auto reader = std::make_shared(256); + OrcFileInputStream input("test.orc", reader, nullptr, nullptr, + {.once_max_read_bytes = 16, .max_merge_distance_bytes = 16}); + StripeStreamMap streams; + + EXPECT_THROW(input.beforeReadStripe( + std::make_unique( + 80, std::vector {{9, ::orc::StreamKind_DATA, 4}}), + selected_columns({1, 2}), streams), + ::orc::ParseError); +} + +TEST(OrcFileInputStreamTest, PublishesClusterProfileExactlyOnce) { + auto reader = std::make_shared(256); + RuntimeProfile profile("orc_v2_input_stream"); + { + OrcFileInputStream input("test.orc", reader, nullptr, &profile, + {.once_max_read_bytes = 16, .max_merge_distance_bytes = 0}); + StripeStreamMap streams; + input.beforeReadStripe( + std::make_unique( + 100, std::vector {{1, ::orc::StreamKind_DATA, 4}, + {2, ::orc::StreamKind_DATA, 4}}), + selected_columns({1, 2}), streams); + + std::array data {}; + find_stream(streams, 2, ::orc::StreamKind_DATA)->read(data.data(), data.size(), 104); + find_stream(streams, 1, ::orc::StreamKind_DATA)->read(data.data(), data.size(), 100); + + StripeStreamMap next_streams; + input.beforeReadStripe( + std::make_unique(200, std::vector {}), + selected_columns({}), next_streams); + ASSERT_NE(profile.get_counter("RequestIO"), nullptr); + ASSERT_NE(profile.get_counter("MergedIO"), nullptr); + EXPECT_EQ(profile.get_counter("RequestIO")->value(), 2); + EXPECT_EQ(profile.get_counter("MergedIO")->value(), 1); + } + + EXPECT_EQ(profile.get_counter("RequestIO")->value(), 2); + EXPECT_EQ(profile.get_counter("MergedIO")->value(), 1); +} + +} // namespace +} // namespace doris::format::orc diff --git a/be/test/format_v2/orc/orc_reader_test.cpp b/be/test/format_v2/orc/orc_reader_test.cpp index 0e36498c41d9bf..2d4ec6c803154b 100644 --- a/be/test/format_v2/orc/orc_reader_test.cpp +++ b/be/test/format_v2/orc/orc_reader_test.cpp @@ -38,6 +38,7 @@ #include #include +#include "common/config.h" #include "core/assert_cast.h" #include "core/block/block.h" #include "core/column/column_array.h" @@ -69,6 +70,7 @@ #include "format_v2/expr/delete_predicate.h" #include "format_v2/file_reader.h" #include "gen_cpp/Types_types.h" +#include "io/fs/buffered_reader.h" #include "io/io_common.h" #include "runtime/runtime_profile.h" #include "runtime/runtime_state.h" @@ -235,6 +237,7 @@ constexpr int64_t COMPLEX_ROW_COUNT = 3; constexpr int64_t DEEP_NESTED_ROW_COUNT = 4; constexpr int64_t DEEP_NESTED_BATCH_CAPACITY = 16; constexpr int64_t NULL_ROW = 1; +constexpr int64_t PREFETCH_ROW_COUNT = 256; VExprSPtr function_expr(const std::string& function_name, const DataTypePtr& return_type, const std::vector& arg_types, @@ -2769,6 +2772,94 @@ void write_orc_file(const std::string& file_path) { out.write(memory_stream.getData(), static_cast(memory_stream.getLength())); } +void write_orc_prefetch_file(const std::string& file_path) { + auto type = std::unique_ptr<::orc::Type>( + ::orc::Type::buildTypeFromString("struct")); + + MemoryOutputStream memory_stream(4 * 1024 * 1024); + ::orc::WriterOptions options; + options.setCompression(::orc::CompressionKind_NONE); + options.setMemoryPool(::orc::getDefaultPool()); + options.setDictionaryKeySizeThreshold(0); + options.setStripeSize(64 * 1024 * 1024); + auto writer = ::orc::createWriter(*type, &memory_stream, options); + auto batch = writer->createRowBatch(PREFETCH_ROW_COUNT); + auto& struct_batch = dynamic_cast<::orc::StructVectorBatch&>(*batch); + auto& a_batch = dynamic_cast<::orc::LongVectorBatch&>(*struct_batch.fields[0]); + auto& b_batch = dynamic_cast<::orc::LongVectorBatch&>(*struct_batch.fields[1]); + auto& c_batch = dynamic_cast<::orc::LongVectorBatch&>(*struct_batch.fields[2]); + auto& payload_batch = dynamic_cast<::orc::StringVectorBatch&>(*struct_batch.fields[3]); + + std::vector payloads; + payloads.reserve(PREFETCH_ROW_COUNT); + for (int64_t row = 0; row < PREFETCH_ROW_COUNT; ++row) { + a_batch.data[row] = row; + b_batch.data[row] = row * 2; + c_batch.data[row] = row * 3; + payloads.push_back("payload-" + std::to_string(row)); + set_string_value(payload_batch, row, payloads.back()); + } + struct_batch.numElements = PREFETCH_ROW_COUNT; + a_batch.numElements = PREFETCH_ROW_COUNT; + b_batch.numElements = PREFETCH_ROW_COUNT; + c_batch.numElements = PREFETCH_ROW_COUNT; + payload_batch.numElements = PREFETCH_ROW_COUNT; + + writer->add(*batch); + writer->close(); + + std::ofstream out(file_path, std::ios::binary); + out.write(memory_stream.getData(), static_cast(memory_stream.getLength())); +} + +bool has_mergeable_orc_stream_cluster(const std::string& file_path, + int64_t max_merge_distance_bytes, + int64_t once_max_read_bytes) { + std::ifstream in(file_path, std::ios::binary | std::ios::ate); + const auto file_size = in.tellg(); + in.seekg(0); + std::vector data(static_cast(file_size)); + in.read(data.data(), file_size); + + ::orc::ReaderOptions options; + options.setMemoryPool(*::orc::getDefaultPool()); + auto input_stream = std::make_unique(data.data(), data.size()); + auto reader = ::orc::createReader(std::move(input_stream), options); + if (reader->getNumberOfStripes() != 1) { + return false; + } + + auto stripe = reader->getStripe(0); + std::vector small_ranges; + for (uint64_t stream_id = 0; stream_id < stripe->getNumberOfStreams(); ++stream_id) { + const auto stream = stripe->getStreamInformation(stream_id); + if (stream->getLength() > 0 && + stream->getLength() <= static_cast(once_max_read_bytes)) { + small_ranges.emplace_back(stream->getOffset(), + stream->getOffset() + stream->getLength()); + } + } + if (small_ranges.size() < 2) { + return false; + } + + const auto merged_ranges = io::PrefetchRange::merge_adjacent_seq_ranges( + small_ranges, max_merge_distance_bytes, once_max_read_bytes); + for (const auto& merged_range : merged_ranges) { + size_t member_count = 0; + for (const auto& range : small_ranges) { + if (range.start_offset >= merged_range.start_offset && + range.end_offset <= merged_range.end_offset) { + ++member_count; + } + } + if (member_count > 1) { + return true; + } + } + return false; +} + void write_primitive_orc_file(const std::string& file_path) { auto type = std::unique_ptr<::orc::Type>(::orc::Type::buildTypeFromString( "struct(); + io_ctx->file_reader_stats = &file_reader_stats; + auto reader = create_reader_for_path(file_path, &profile, io_ctx); + + TQueryOptions query_options; + query_options.__set_orc_once_max_read_bytes(ONCE_MAX_READ_BYTES); + query_options.__set_orc_max_merge_distance_bytes(MAX_MERGE_DISTANCE_BYTES); + RuntimeState state {query_options, TQueryGlobals()}; + ASSERT_TRUE(reader->init(&state).ok()); + + std::vector schema; + ASSERT_TRUE(reader->get_schema(&schema).ok()); + ASSERT_EQ(schema.size(), 4); + + auto request = std::make_shared(); + for (const auto& field : schema) { + request->non_predicate_columns.push_back(make_projection(field)); + } + ASSERT_TRUE(reader->open(request).ok()); + + size_t result_rows = 0; + std::vector first_values; + std::vector second_values; + std::vector third_values; + std::vector payload_values; + bool eof = false; + while (!eof) { + Block block = build_file_block(schema); + size_t rows = 0; + ASSERT_TRUE(reader->get_block(&block, &rows, &eof).ok()); + if (rows > 0) { + const auto& first_column = assert_cast( + assert_cast(*block.get_by_position(0).column) + .get_nested_column()); + const auto& second_column = assert_cast( + assert_cast(*block.get_by_position(1).column) + .get_nested_column()); + const auto& third_column = assert_cast( + assert_cast(*block.get_by_position(2).column) + .get_nested_column()); + const auto& payload_column = assert_cast( + assert_cast(*block.get_by_position(3).column) + .get_nested_column()); + for (size_t row = 0; row < rows; ++row) { + first_values.push_back(first_column.get_element(row)); + second_values.push_back(second_column.get_element(row)); + third_values.push_back(third_column.get_element(row)); + payload_values.push_back(payload_column.get_data_at(row).to_string()); + } + } + result_rows += rows; + } + EXPECT_EQ(result_rows, PREFETCH_ROW_COUNT); + ASSERT_EQ(first_values.size(), PREFETCH_ROW_COUNT); + ASSERT_EQ(second_values.size(), PREFETCH_ROW_COUNT); + ASSERT_EQ(third_values.size(), PREFETCH_ROW_COUNT); + ASSERT_EQ(payload_values.size(), PREFETCH_ROW_COUNT); + for (size_t row : std::array {0, PREFETCH_ROW_COUNT / 2, PREFETCH_ROW_COUNT - 1}) { + EXPECT_EQ(first_values[row], row); + EXPECT_EQ(second_values[row], row * 2); + EXPECT_EQ(third_values[row], row * 3); + EXPECT_EQ(payload_values[row], "payload-" + std::to_string(row)); + } + ASSERT_TRUE(reader->close().ok()); + + ASSERT_NE(profile.get_counter("RequestIO"), nullptr); + ASSERT_NE(profile.get_counter("MergedIO"), nullptr); + ASSERT_NE(profile.get_counter("ApplyBytes"), nullptr); + ASSERT_NE(profile.get_counter("ClusterNum"), nullptr); + ASSERT_NE(profile.get_counter("OverReadBytes"), nullptr); + EXPECT_GT(profile.get_counter("RequestIO")->value(), profile.get_counter("MergedIO")->value()); + EXPECT_GT(profile.get_counter("MergedIO")->value(), 0); + EXPECT_GT(profile.get_counter("ApplyBytes")->value(), 0); + EXPECT_GT(profile.get_counter("ClusterNum")->value(), 0); + EXPECT_GE(profile.get_counter("OverReadBytes")->value(), 0); +} + +TEST_F(NewOrcReaderTest, StripePrefetchCanBeDisabledByZeroOnceMaxReadBytes) { + static constexpr int64_t MAX_MERGE_DISTANCE_BYTES = 1L * 1024L * 1024L; + const auto file_path = (_test_dir / "stripe_prefetch_disabled.orc").string(); + write_orc_prefetch_file(file_path); + + RuntimeProfile profile("orc_v2_stripe_prefetch_disabled"); + auto reader = create_reader_for_path(file_path, &profile); + + TQueryOptions query_options; + query_options.__set_orc_once_max_read_bytes(0); + query_options.__set_orc_max_merge_distance_bytes(MAX_MERGE_DISTANCE_BYTES); + RuntimeState state {query_options, TQueryGlobals()}; + ASSERT_TRUE(reader->init(&state).ok()); + + std::vector schema; + ASSERT_TRUE(reader->get_schema(&schema).ok()); + ASSERT_EQ(schema.size(), 4); + + auto request = std::make_shared(); + request->non_predicate_columns = {make_projection(schema[0]), make_projection(schema[3])}; + ASSERT_TRUE(reader->open(request).ok()); + + Block block = build_file_block({schema[0], schema[3]}); + size_t rows = 0; + bool eof = false; + ASSERT_TRUE(reader->get_block(&block, &rows, &eof).ok()); + ASSERT_GT(rows, 0); + + const auto& first_column = assert_cast( + assert_cast(*block.get_by_position(0).column) + .get_nested_column()); + const auto& payload_column = assert_cast( + assert_cast(*block.get_by_position(1).column) + .get_nested_column()); + EXPECT_EQ(first_column.get_element(0), 0); + EXPECT_EQ(payload_column.get_data_at(0).to_string(), "payload-0"); + EXPECT_EQ(first_column.get_element(rows - 1), rows - 1); + EXPECT_EQ(payload_column.get_data_at(rows - 1).to_string(), + "payload-" + std::to_string(rows - 1)); + + size_t result_rows = rows; + while (!eof) { + Block next_block = build_file_block({schema[0], schema[3]}); + rows = 0; + ASSERT_TRUE(reader->get_block(&next_block, &rows, &eof).ok()); + result_rows += rows; + } + EXPECT_EQ(result_rows, PREFETCH_ROW_COUNT); + ASSERT_TRUE(reader->close().ok()); + + EXPECT_EQ(profile.get_counter("MergedIO"), nullptr); +} + +TEST_F(NewOrcReaderTest, RejectsInvalidNaturalReadSizeConfig) { + const auto old_natural_read_size_mb = config::orc_natural_read_size_mb; + config::orc_natural_read_size_mb = 0; + auto reader = create_reader(); + RuntimeState state {TQueryOptions(), TQueryGlobals()}; + const auto status = reader->init(&state); + config::orc_natural_read_size_mb = old_natural_read_size_mb; + + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("orc_natural_read_size_mb"), std::string::npos); +} + TEST_F(NewOrcReaderTest, GetBlockStopsWhenIoContextShouldStop) { auto io_ctx = std::make_shared(); auto reader = create_reader_for_path(_file_path, nullptr, io_ctx);