zhangstar333 commented on code in PR #65730: URL: https://github.com/apache/doris/pull/65730#discussion_r3710132681
########## be/src/format_v2/table/lance_reader.cpp: ########## @@ -0,0 +1,912 @@ +// 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/lance_reader.h" + +#include <arrow/c/bridge.h> +#include <arrow/record_batch.h> +#include <arrow/type.h> +#include <lance/lance.h> + +#include <bit> +#include <cstring> +#include <limits> +#include <memory> + +#include "core/data_type/data_type_array.h" +#include "core/data_type/data_type_factory.hpp" +#include "core/data_type/data_type_map.h" +#include "core/data_type/data_type_struct.h" +#include "exec/common/endian.h" + +namespace doris::format::lance { +namespace { + +struct LanceDatasetDeleter { + void operator()(LanceDataset* dataset) const { lance_dataset_close(dataset); } +}; + +struct LanceScannerDeleter { + void operator()(LanceScanner* scanner) const { lance_scanner_close(scanner); } +}; + +struct LanceBatchDeleter { + void operator()(LanceBatch* batch) const { lance_batch_free(batch); } +}; + +constexpr std::string_view DISTANCE_COLUMN = "_distance"; + +size_t vector_element_width(TVectorElementType::type type) { + switch (type) { + case TVectorElementType::FLOAT16: + return sizeof(uint16_t); + case TVectorElementType::FLOAT32: + return sizeof(float); + case TVectorElementType::FLOAT64: + return sizeof(double); + case TVectorElementType::UINT8: + case TVectorElementType::INT8: + return sizeof(uint8_t); + } + return 0; +} + +std::string vector_metric_name(TVectorMetric::type metric) { + switch (metric) { + case TVectorMetric::DEFAULT: + return "default"; + case TVectorMetric::L2: + return "l2"; + case TVectorMetric::COSINE: + return "cosine"; + case TVectorMetric::DOT_PRODUCT: + return "dot"; + case TVectorMetric::HAMMING: + return "hamming"; + } + return "unknown"; +} + +int arrow_time_precision(arrow::TimeUnit::type unit) { + switch (unit) { + case arrow::TimeUnit::SECOND: + return 0; + case arrow::TimeUnit::MILLI: + return 3; + case arrow::TimeUnit::MICRO: + case arrow::TimeUnit::NANO: + return 6; + } + return 6; +} + +Status arrow_type_to_doris_type(const std::shared_ptr<arrow::DataType>& arrow_type, + DataTypePtr* doris_type) { + const auto nullable_primitive = [&](PrimitiveType type, int precision = 0, int scale = 0, + int len = -1) { + *doris_type = + DataTypeFactory::instance().create_data_type(type, true, precision, scale, len); + return Status::OK(); + }; + + switch (arrow_type->id()) { + case arrow::Type::BOOL: + return nullable_primitive(TYPE_BOOLEAN); + case arrow::Type::INT8: + return nullable_primitive(TYPE_TINYINT); + case arrow::Type::UINT8: + case arrow::Type::INT16: + return nullable_primitive(TYPE_SMALLINT); + case arrow::Type::UINT16: + case arrow::Type::INT32: + return nullable_primitive(TYPE_INT); + case arrow::Type::UINT32: + case arrow::Type::INT64: + return nullable_primitive(TYPE_BIGINT); + case arrow::Type::UINT64: + return nullable_primitive(TYPE_LARGEINT); + case arrow::Type::HALF_FLOAT: + case arrow::Type::FLOAT: + return nullable_primitive(TYPE_FLOAT); + case arrow::Type::DOUBLE: + return nullable_primitive(TYPE_DOUBLE); + case arrow::Type::STRING: + case arrow::Type::LARGE_STRING: + return nullable_primitive(TYPE_STRING); + case arrow::Type::BINARY: + case arrow::Type::LARGE_BINARY: + return nullable_primitive(TYPE_VARBINARY, 0, 0, std::numeric_limits<int32_t>::max()); + case arrow::Type::FIXED_SIZE_BINARY: { + const auto binary = std::static_pointer_cast<arrow::FixedSizeBinaryType>(arrow_type); + return nullable_primitive(TYPE_VARBINARY, 0, 0, binary->byte_width()); + } + case arrow::Type::DATE32: + case arrow::Type::DATE64: + return nullable_primitive(TYPE_DATEV2); + case arrow::Type::TIME32: + case arrow::Type::TIME64: { + const auto time = std::static_pointer_cast<arrow::TimeType>(arrow_type); + return nullable_primitive(TYPE_TIMEV2, 0, arrow_time_precision(time->unit())); + } + case arrow::Type::TIMESTAMP: { + const auto timestamp = std::static_pointer_cast<arrow::TimestampType>(arrow_type); + const auto doris_type = timestamp->timezone().empty() ? TYPE_DATETIMEV2 : TYPE_TIMESTAMPTZ; + return nullable_primitive(doris_type, 0, arrow_time_precision(timestamp->unit())); + } + case arrow::Type::DECIMAL128: + case arrow::Type::DECIMAL256: { + const auto decimal = std::static_pointer_cast<arrow::DecimalType>(arrow_type); + const int precision = decimal->precision(); + const PrimitiveType doris_decimal_type = precision <= 9 ? TYPE_DECIMAL32 + : precision <= 18 ? TYPE_DECIMAL64 + : precision <= 38 ? TYPE_DECIMAL128I + : TYPE_DECIMAL256; + return nullable_primitive(doris_decimal_type, precision, decimal->scale()); + } + case arrow::Type::LIST: + case arrow::Type::LARGE_LIST: + case arrow::Type::FIXED_SIZE_LIST: { + const auto list = std::static_pointer_cast<arrow::BaseListType>(arrow_type); + DataTypePtr value_type; + RETURN_IF_ERROR(arrow_type_to_doris_type(list->value_type(), &value_type)); + *doris_type = make_nullable(std::make_shared<DataTypeArray>(value_type)); + return Status::OK(); + } + case arrow::Type::MAP: { + const auto map = std::static_pointer_cast<arrow::MapType>(arrow_type); + DataTypePtr key_type; + DataTypePtr item_type; + RETURN_IF_ERROR(arrow_type_to_doris_type(map->key_type(), &key_type)); + RETURN_IF_ERROR(arrow_type_to_doris_type(map->item_type(), &item_type)); + *doris_type = make_nullable(std::make_shared<DataTypeMap>(key_type, item_type)); + return Status::OK(); + } + case arrow::Type::STRUCT: { + const auto struct_type = std::static_pointer_cast<arrow::StructType>(arrow_type); + DataTypes field_types; + Strings field_names; + field_types.reserve(struct_type->num_fields()); + field_names.reserve(struct_type->num_fields()); + for (const auto& field : struct_type->fields()) { + DataTypePtr field_type; + RETURN_IF_ERROR(arrow_type_to_doris_type(field->type(), &field_type)); + field_types.emplace_back(std::move(field_type)); + field_names.emplace_back(field->name()); + } + *doris_type = make_nullable(std::make_shared<DataTypeStruct>(field_types, field_names)); + return Status::OK(); + } + default: + return Status::NotSupported("unsupported Lance Arrow type: {}", arrow_type->ToString()); + } +} + +} // namespace + +LanceTableReader::~LanceTableReader() { + static_cast<void>(close()); +} + +Status LanceTableReader::fetch_schema(const TFileRangeDesc& range, + const TFileScanRangeParams& scan_params, + std::vector<std::string>* column_names, + std::vector<DataTypePtr>* column_types) const { + if (column_names == nullptr || column_types == nullptr) { + return Status::InvalidArgument("Lance schema output must not be null"); + } + RETURN_IF_ERROR(_validate_range(range)); + + const auto& params = range.table_format_params.lance_params; + const auto storage_options = _storage_options(&scan_params); + std::vector<const char*> storage_option_ptrs; + storage_option_ptrs.reserve(storage_options.size() + 1); + for (const auto& option : storage_options) { + storage_option_ptrs.emplace_back(option.c_str()); + } + storage_option_ptrs.emplace_back(nullptr); + + std::unique_ptr<LanceDataset, LanceDatasetDeleter> dataset( + lance_dataset_open(params.dataset_uri.c_str(), + storage_options.empty() ? nullptr : storage_option_ptrs.data(), + static_cast<uint64_t>(params.version))); + if (dataset == nullptr) { + return _lance_error("open Lance dataset for schema"); + } + + ArrowSchema arrow_schema {}; + if (lance_dataset_schema(dataset.get(), &arrow_schema) != 0) { + return _lance_error("get Lance dataset schema"); + } + auto imported_schema = arrow::ImportSchema(&arrow_schema); + if (!imported_schema.ok()) { + if (arrow_schema.release != nullptr) { + arrow_schema.release(&arrow_schema); + } + return Status::InternalError("import Lance Arrow schema failed: {}", + imported_schema.status().message()); + } + + const auto schema = std::move(imported_schema).ValueUnsafe(); + std::vector<std::string> parsed_names; + std::vector<DataTypePtr> parsed_types; + parsed_names.reserve(schema->num_fields()); + parsed_types.reserve(schema->num_fields()); + std::unordered_set<std::string> unique_names; + unique_names.reserve(schema->num_fields()); + for (const auto& field : schema->fields()) { + if (!unique_names.emplace(field->name()).second) { + return Status::InvalidArgument("duplicate Lance schema column: {}", field->name()); + } + DataTypePtr doris_type; + RETURN_IF_ERROR(arrow_type_to_doris_type(field->type(), &doris_type)); + parsed_names.emplace_back(field->name()); + parsed_types.emplace_back(std::move(doris_type)); + } + *column_names = std::move(parsed_names); + *column_types = std::move(parsed_types); + return Status::OK(); +} + +Status LanceTableReader::init(TableReadOptions&& options) { + RETURN_IF_ERROR(TableReader::init(std::move(options))); + DORIS_CHECK(_runtime_state != nullptr); + DORIS_CHECK(_scanner_profile != nullptr); + DORIS_CHECK(_scan_params != nullptr); + + _ctz = _runtime_state->timezone_obj(); + _vector_search = _scan_params->__isset.external_search_request; + _search_split_prepared = false; + if (_vector_search) { + RETURN_IF_ERROR(_validate_external_search_request()); + const auto& vector = _scan_params->external_search_request.query.vector; + _scanner_profile->add_info_string("ExternalSearchType", "VECTOR"); + _scanner_profile->add_info_string("LanceVectorColumn", vector.column); + _scanner_profile->add_info_string("LanceTopK", std::to_string(vector.top_k)); + _scanner_profile->add_info_string("LanceOffset", std::to_string(vector.offset)); + _scanner_profile->add_info_string("LanceMetric", vector.__isset.metric + ? vector_metric_name(vector.metric) + : "default"); + _scanner_profile->add_info_string("LanceVectorDimension", + std::to_string(vector.query_vector.dimension)); + } + if (_scan_params->__isset.lance_substrait_filter) { + _scanner_profile->add_info_string("LancePushdownFormat", "SUBSTRAIT"); + _scanner_profile->add_info_string( + "LanceSubstraitFilterBytes", + std::to_string(_scan_params->lance_substrait_filter.size())); + } + + _output_name_to_idx.clear(); + _output_name_to_idx.reserve(_projected_columns.size()); + for (size_t idx = 0; idx < _projected_columns.size(); ++idx) { + const auto& column = _projected_columns[idx]; + if (column.type == nullptr) { + return Status::InvalidArgument("Lance projected column '{}' has no type", column.name); + } + if (!_output_name_to_idx.emplace(column.name, idx).second) { + return Status::InvalidArgument("duplicate Lance projected column: {}", column.name); + } + if (_vector_search && column.name == DISTANCE_COLUMN) { + const auto distance_type = remove_nullable(column.type); + if (distance_type->get_primitive_type() != TYPE_FLOAT) { + return Status::InvalidArgument( + "Lance vector search column '{}' must have Doris FLOAT type, but was {}", + DISTANCE_COLUMN, column.type->get_name()); + } + } + } + return Status::OK(); +} + +Status LanceTableReader::prepare_split(const SplitReadOptions& options) { + RETURN_IF_ERROR(_validate_range(options.current_range)); + _close_scanner(); + _eof = false; + + RETURN_IF_ERROR(TableReader::prepare_split(options)); + // Lance does not currently provide metadata aggregate pushdown. Do not let a generic + // table-level count supplied by a future planner bypass fragment reads. + _remaining_table_level_count = -1; + if (current_split_pruned()) { + return Status::OK(); + } + + if (_vector_search) { + const auto& lance_params = options.current_range.table_format_params.lance_params; + if (lance_params.version <= 0) { + return Status::InvalidArgument( + "Lance vector search requires a fixed positive dataset version"); + } + if (lance_params.__isset.fragment_ids) { + return Status::InvalidArgument("Lance vector search split must not set fragment ids"); + } + if (_search_split_prepared) { + return Status::InvalidArgument( + "Lance vector search supports exactly one whole-dataset split"); + } + } + + const auto key = _dataset_key(options.current_range); + if (_dataset == nullptr) { + RETURN_IF_ERROR(_open_dataset(key)); + _opened_dataset_key = key; + } else if (!_opened_dataset_key.has_value() || *_opened_dataset_key != key) { + return Status::InvalidArgument( + "Lance reader cannot mix dataset snapshots or storage options in one scan"); + } + + RETURN_IF_ERROR(_prepare_conjuncts()); + RETURN_IF_ERROR(_open_scanner(options.current_range)); + _search_split_prepared = _vector_search; + return Status::OK(); +} + +Status LanceTableReader::get_block(Block* block, bool* eos) { + DORIS_CHECK(block != nullptr); + DORIS_CHECK(eos != nullptr); + DORIS_CHECK(block->columns() == _projected_columns.size()); + *eos = false; + + if (_eof) { + *eos = true; + return Status::OK(); + } + if (_scanner == nullptr) { + return Status::InternalError("Lance scanner is not initialized for the current split"); + } + + const auto target_rows = std::max<size_t>(1, _scanner_batch_size); + while (true) { + block->clear_column_data(_projected_columns.size()); + size_t raw_rows = 0; + while (raw_rows < target_rows) { + if (_io_ctx != nullptr && _io_ctx->should_stop) { + _eof = true; + _close_scanner(); + *eos = true; + return Status::OK(); + } + + LanceBatch* raw_batch = nullptr; + const int32_t scan_status = lance_scanner_next(_scanner, &raw_batch); Review Comment: TODO, now is hard to control it with using lance sdk -- 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]
