github-actions[bot] commented on code in PR #69644: URL: https://github.com/apache/doris/pull/69644#discussion_r4378604674
########## be/src/format_v2/table/es_reader.cpp: ########## @@ -0,0 +1,295 @@ +// 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/table/es_reader.h" + +#include <utility> + +#include "core/block/block.h" +#include "exec/es/es_scan_reader.h" +#include "exec/es/es_scroll_parser.h" +#include "exec/es/es_scroll_query.h" +#include "format_v2/column_mapper.h" +#include "format_v2/materialized_reader_util.h" +#include "runtime/descriptors.h" +#include "runtime/file_scan_profile.h" +#include "runtime/runtime_state.h" + +namespace doris::format::es { +namespace { + +Status validate_es_range(const TFileRangeDesc& range) { + if (!range.__isset.table_format_params || range.table_format_params.table_format_type != "es" || + !range.table_format_params.__isset.es_params) { + return Status::InvalidArgument("ES v2 reader requires es table format and es_params"); + } + return Status::OK(); +} + +class HttpScrollStream final : public EsScrollStream { +public: + HttpScrollStream(const std::string& host, const std::map<std::string, std::string>& properties, + bool docvalues) + : _reader(host, properties, docvalues) {} + Status open() override { return _reader.open(); } + Status next(bool* eof, std::unique_ptr<ScrollParser>& parser) override { + return _reader.get_next(eof, parser); + } + Status close() override { return _reader.close(); } + +private: + ESScanReader _reader; +}; + +} // namespace + +EsFileReader::EsFileReader(std::shared_ptr<io::FileSystemProperties>& system_properties, + std::unique_ptr<io::FileDescription>& file_description, + std::shared_ptr<io::IOContext> io_ctx, RuntimeProfile* profile, + TFileRangeDesc range, const TFileScanRangeParams& params, + std::vector<SlotDescriptor*> file_slot_descs, + EsScrollStreamFactory stream_factory) + : FileReader(system_properties, file_description, std::move(io_ctx), profile), + _range(std::move(range)), + _params(params), + _file_slots(std::move(file_slot_descs)), + _stream_factory(std::move(stream_factory)) {} + +EsFileReader::~EsFileReader() { + WARN_IF_ERROR(close(), "Failed to close ES scroll stream"); +} + +void EsFileReader::_init_profile() { + if (_profile == nullptr) { + return; + } + _io_time = file_scan_profile::ensure_hierarchy(_profile).io; + const char* parent = "EsFileReader"; + _total_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, parent, file_scan_profile::FILE_READER, 1); + _read_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, "EsReadTime", parent, 1); + _materialize_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, "EsMaterializeTime", parent, 1); + _filter_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, "EsFilterTime", parent, 1); + _batches_read = ADD_CHILD_COUNTER_WITH_LEVEL(_profile, "EsBatchesRead", TUnit::UNIT, parent, 1); + _rows_read = ADD_CHILD_COUNTER_WITH_LEVEL(_profile, "EsRowsRead", TUnit::UNIT, parent, 1); +} + +Status EsFileReader::init(RuntimeState* state) { + DORIS_CHECK(state != nullptr); + _state = state; + set_batch_size(state->batch_size()); + _init_profile(); + RETURN_IF_ERROR(validate_es_range(_range)); + _eof = false; + return Status::OK(); +} + +Status EsFileReader::get_schema(std::vector<ColumnDefinition>* schema) const { + DORIS_CHECK(schema != nullptr); + schema->clear(); + schema->reserve(_file_slots.size()); + for (size_t i = 0; i < _file_slots.size(); ++i) { + const auto* slot = _file_slots[i]; + schema->push_back({.identifier = Field::create_field<TYPE_INT>(cast_set<int32_t>(i)), + .local_id = cast_set<int32_t>(i), + .name = slot->col_name(), + .type = slot->type(), + .children = materialized_schema_children(slot->type())}); + } + return Status::OK(); +} + +std::unique_ptr<TableColumnMapper> EsFileReader::create_column_mapper( + TableColumnMapperOptions options) const { + return std::make_unique<MaterializedColumnMapper>(options); +} + +void EsFileReader::set_batch_size(size_t batch_size) { + DORIS_CHECK(batch_size > 0); + _batch_size = batch_size; +} + +Status EsFileReader::open(std::shared_ptr<FileScanRequest> request) { + SCOPED_TIMER(_total_time); + RETURN_IF_ERROR(FileReader::open(std::move(request))); + auto status = _open_stream(); + if (!status.ok()) { + WARN_IF_ERROR(close(), "Failed to clear ES scroll after open error"); + } + return status; +} + +Status EsFileReader::_open_stream() { + DORIS_CHECK(_stream == nullptr); + RETURN_IF_CANCELLED(_state); + if (_io_ctx != nullptr && _io_ctx->should_stop) { + _eof = true; + return Status::OK(); + } + std::vector<std::string> names; + for (const auto& [id, position] : _request->local_positions) { + DORIS_CHECK(id.value() >= 0 && id.value() < _file_slots.size()); + _read_slots.push_back(_file_slots[id.value()]); + _read_positions.push_back(position.value()); + names.push_back(_file_slots[id.value()]->col_name()); + } + auto properties = _params.es_properties; + for (const auto& [key, value] : _range.table_format_params.es_params) { + properties[key] = value; + } + properties[ESScanReader::KEY_BATCH_SIZE] = std::to_string(_batch_size); Review Comment: [P2] Preserve the small-LIMIT ES request optimization For a residual-free query such as `SELECT payload FROM es_table LIMIT 1`, the legacy `EsScanLocalState` sets `KEY_TERMINATE_AFTER`, producing a one-shot search with `size=1`. The V2 path never forwards that limit, so this property makes the request size the whole session batch (8160 by default), and `_read_block` materializes that batch before `Scanner` checks its limit. On a populated shard, source-backed wide payloads therefore require fetching and decoding thousands of documents, plus creating/deleting a scroll context, to return one row. Carry the limit into the ES request under the same no-residual-predicate guard, and add a request-size assertion to the LIMIT coverage. ########## fe/fe-core/src/main/java/org/apache/doris/datasource/es/source/EsScanNode.java: ########## @@ -169,8 +194,65 @@ protected void toThrift(TPlanNode msg) { esScanNode.setFieldsContext(table.fieldsContext()); } esScanNode.setProperties(properties); - msg.es_scan_node = esScanNode; - super.toThrift(msg); + return esScanNode; + } + + @Override + public List<TScanRangeLocations> getScanRangeLocations(long maxScanRangeLength) { + List<TScanRangeLocations> shardLocations = super.getScanRangeLocations(maxScanRangeLength); + if (!useFileScannerV2) { + return shardLocations; + } + TEsScanNode esNode = buildEsScanNode(); + TFileScanRangeParams params = new TFileScanRangeParams(); + params.setFormatType(TFileFormatType.FORMAT_ES_HTTP); + params.setFileType(TFileType.FILE_STREAM); + params.setSrcTupleId(-1); + params.setDestTupleId(desc.getId().asInt()); + params.setNumOfColumnsFromFile(desc.getSlots().size()); + params.setRequiredSlots(new ArrayList<>()); + for (SlotDescriptor slot : desc.getSlots()) { + params.addToRequiredSlots(new TFileScanSlotInfo().setSlotId(slot.getId().asInt()).setIsFileSlot(true)); + } + params.setEsProperties(esNode.getProperties()); + if (esNode.isSetDocvalueContext()) { + params.setEsDocvalueContext(esNode.getDocvalueContext()); + } + if (esNode.isSetFieldsContext()) { + params.setEsFieldsContext(esNode.getFieldsContext()); + } + List<TScanRangeLocations> fileLocations = new ArrayList<>(shardLocations.size()); + for (TScanRangeLocations location : shardLocations) { + TEsScanRange shard = location.getScanRange().getEsScanRange(); + Map<String, String> shardParams = new HashMap<>(); + shardParams.put("index", shard.getIndex()); + shardParams.put("shard_id", String.valueOf(shard.getShardId())); + if (shard.isSetType()) { + shardParams.put("type", shard.getType()); + } + shardParams.put("es_hosts", shard.getEsHosts().stream() + .map(host -> host.getHostname() + ":" + host.getPort()).collect(Collectors.joining(","))); + TTableFormatFileDesc tableFormat = new TTableFormatFileDesc(); + tableFormat.setTableFormatType("es"); + tableFormat.setEsParams(shardParams); + TFileRangeDesc range = new TFileRangeDesc(); + range.setPath("es://" + shard.getIndex() + "/" + shard.getShardId()); + range.setStartOffset(0); + range.setSize(-1); Review Comment: [P2] Keep unknown ES lengths from becoming negative scheduling weights With parallel scrolling and the same eligible BEs for each shard (for example, ES and Doris deployed on separate hosts), this `-1` is summed by `LoadBalanceScanWorkerSelector.getScanRangeSize` and added to `WorkerWorkload.scanBytes`. The first selected BE becomes `-1` while the others remain `0`, so every following shard chooses it again (`-2`, `-3`, ...), regardless of the shuffled location order. The default Nereids distribute planner therefore concentrates the ES scan on one BE instead of spreading shards across the cluster. Preserve a nonnegative scheduling weight for unknown-length ES ranges, or normalize the sentinel before workload accounting, and cover multi-BE shard assignment. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
