diff --git a/be/src/format_v2/lance/lance_reader_helper.cpp b/be/src/format_v2/lance/lance_reader_helper.cpp index 8599b7c3850e0d..2b14381437e1a6 100644 --- a/be/src/format_v2/lance/lance_reader_helper.cpp +++ b/be/src/format_v2/lance/lance_reader_helper.cpp @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -272,26 +273,6 @@ Status arrow_field_to_doris_type(const std::shared_ptr& field, } } -// Determine whether a field subtree contains values that require Lance normalization. -Status field_requires_lance_normalization(const std::shared_ptr& field, - bool* requires_normalization) { - DORIS_CHECK(field != nullptr); - DORIS_CHECK(requires_normalization != nullptr); - - LanceExtensionKind extension_kind; - std::shared_ptr storage_type; - RETURN_IF_ERROR(get_lance_extension(field, &extension_kind, &storage_type)); - bool required = extension_kind == LanceExtensionKind::BFLOAT16 || - field->type()->id() == arrow::Type::EXTENSION; - for (const auto& child : storage_type->fields()) { - bool child_required = false; - RETURN_IF_ERROR(field_requires_lance_normalization(child, &child_required)); - required |= child_required; - } - *requires_normalization = required; - return Status::OK(); -} - // Widen little-endian Lance BFloat16 values to Arrow Float32 without precision loss. Status convert_bfloat16_array(const std::shared_ptr& array, arrow::MemoryPool* memory_pool, @@ -373,89 +354,6 @@ Status set_lance_nested_type(std::string_view field_name, return Status::OK(); } -// Check whether an Arrow type tree contains a registered extension wrapper. -bool type_contains_registered_extension(const std::shared_ptr& type) { - if (type->id() == arrow::Type::EXTENSION) { - return true; - } - for (const auto& field : type->fields()) { - if (type_contains_registered_extension(field->type())) { - return true; - } - } - return false; -} - -// Remove registered ExtensionArray wrappers only along extension-bearing branches. -Status unwrap_lance_extension_arrays(const std::shared_ptr& expected_type, - const std::shared_ptr& array, - std::shared_ptr* unwrapped) { - DORIS_CHECK(expected_type != nullptr); - DORIS_CHECK(array != nullptr); - DORIS_CHECK(unwrapped != nullptr); - - auto storage_array = array; - auto expected_storage_type = expected_type; - if (expected_type->id() == arrow::Type::EXTENSION) { - const auto extension_type = std::dynamic_pointer_cast(expected_type); - if (extension_type == nullptr) { - return Status::InvalidArgument("invalid expected Arrow extension type {}", - expected_type->ToString()); - } - expected_storage_type = extension_type->storage_type(); - } - if (array->type_id() == arrow::Type::EXTENSION) { - const auto extension_array = std::dynamic_pointer_cast(array); - if (extension_array == nullptr) { - return Status::InvalidArgument("invalid Arrow extension array: {}", - array->type()->ToString()); - } - storage_array = extension_array->storage(); - } - - const auto& child_data = storage_array->data()->child_data; - const auto& child_fields = expected_storage_type->fields(); - if (child_data.empty()) { - *unwrapped = std::move(storage_array); - return Status::OK(); - } - if (child_fields.size() != child_data.size()) { - return Status::InvalidArgument( - "Arrow array type {} has {} child fields but its data has {} children", - storage_array->type()->ToString(), child_fields.size(), child_data.size()); - } - - std::shared_ptr unwrapped_data; - arrow::FieldVector unwrapped_fields; - for (size_t child_idx = 0; child_idx < child_data.size(); ++child_idx) { - if (!type_contains_registered_extension(child_fields[child_idx]->type())) { - continue; - } - auto child_array = arrow::MakeArray(child_data[child_idx]); - std::shared_ptr unwrapped_child; - RETURN_IF_ERROR(unwrap_lance_extension_arrays(child_fields[child_idx]->type(), child_array, - &unwrapped_child)); - if (unwrapped_child.get() == child_array.get()) { - continue; - } - if (unwrapped_data == nullptr) { - unwrapped_data = storage_array->data()->Copy(); - unwrapped_fields = storage_array->type()->fields(); - } - unwrapped_data->child_data[child_idx] = unwrapped_child->data(); - unwrapped_fields[child_idx] = - unwrapped_fields[child_idx]->WithType(unwrapped_child->type()); - } - if (unwrapped_data == nullptr) { - *unwrapped = std::move(storage_array); - return Status::OK(); - } - RETURN_IF_ERROR( - set_lance_nested_type("", storage_array->type(), unwrapped_fields, &unwrapped_data)); - *unwrapped = arrow::MakeArray(std::move(unwrapped_data)); - return Status::OK(); -} - // Materialize the visible range into an offset-zero Arrow array for Doris SerDes. Status compact_lance_array(const std::shared_ptr& array, arrow::MemoryPool* memory_pool, @@ -575,78 +473,86 @@ Status compact_lance_array_if_needed(const std::shared_ptr& array, } // namespace -// Normalize Lance extensions and materialize sliced arrays for Doris Arrow SerDes. -Status normalize_lance_arrow_array(const std::shared_ptr& field, - const std::shared_ptr& array, - arrow::MemoryPool* memory_pool, - std::shared_ptr* normalized) { +Status LanceArrowArrayNormalizer::create(const std::shared_ptr& field, + LanceArrowArrayNormalizer* normalizer) { DORIS_CHECK(field != nullptr); - DORIS_CHECK(array != nullptr); - DORIS_CHECK(memory_pool != nullptr); - DORIS_CHECK(normalized != nullptr); + DORIS_CHECK(normalizer != nullptr); LanceExtensionKind extension_kind; std::shared_ptr storage_type; RETURN_IF_ERROR(get_lance_extension(field, &extension_kind, &storage_type)); + LanceArrowArrayNormalizer result; + result._field_name = field->name(); + result._storage_type = std::move(storage_type); + result._unwrap_registered_extension = field->type()->id() == arrow::Type::EXTENSION; + result._convert_bfloat16 = extension_kind == LanceExtensionKind::BFLOAT16; + result._requires_special_handling = + result._unwrap_registered_extension || result._convert_bfloat16; + result._child_normalizers.reserve(result._storage_type->num_fields()); + for (const auto& child_field : result._storage_type->fields()) { + LanceArrowArrayNormalizer child_normalizer; + RETURN_IF_ERROR(create(child_field, &child_normalizer)); + result._requires_special_handling |= child_normalizer._requires_special_handling; + result._child_normalizers.emplace_back(std::move(child_normalizer)); + } + *normalizer = std::move(result); + return Status::OK(); +} + +Status LanceArrowArrayNormalizer::normalize_for_doris( + const std::shared_ptr& array, arrow::MemoryPool* memory_pool, + std::shared_ptr* normalized) const { + DORIS_CHECK(array != nullptr); + DORIS_CHECK(memory_pool != nullptr); + DORIS_CHECK(normalized != nullptr); + + // The schema binding already resolved extension metadata and nested special types. For common + // arrays, runtime work is limited to checking whether the visible slice needs compaction. + if (!_requires_special_handling) { + return compact_lance_array_if_needed(array, memory_pool, normalized); + } + auto storage_array = array; - if (type_contains_registered_extension(field->type())) { - std::shared_ptr unwrapped_array; - RETURN_IF_ERROR(unwrap_lance_extension_arrays(field->type(), array, &unwrapped_array)); - storage_array = std::move(unwrapped_array); + if (_unwrap_registered_extension && array->type_id() == arrow::Type::EXTENSION) { + const auto extension_array = std::dynamic_pointer_cast(array); + if (extension_array == nullptr) { + return Status::InvalidArgument("invalid Arrow extension array: {}", + array->type()->ToString()); + } + storage_array = extension_array->storage(); } - if (storage_array->type_id() != storage_type->id()) { + // Extension and BFloat16 arrays are first converted to their physical storage representation. + if (storage_array->type_id() != _storage_type->id()) { return Status::InvalidArgument( - "Lance field '{}' storage type {} does not match array type {}", field->name(), - storage_type->ToString(), storage_array->type()->ToString()); + "Lance field '{}' storage type {} does not match array type {}", _field_name, + _storage_type->ToString(), storage_array->type()->ToString()); } - if (extension_kind == LanceExtensionKind::BFLOAT16) { + if (_convert_bfloat16) { return convert_bfloat16_array(storage_array, memory_pool, normalized); } - std::shared_ptr compacted_array; - RETURN_IF_ERROR(compact_lance_array_if_needed(storage_array, memory_pool, &compacted_array)); - storage_array = std::move(compacted_array); - - const auto& child_fields = storage_type->fields(); + // Normalize extension children before compacting the parent. Arrow cannot create a builder + // for list/struct types whose child is still an ExtensionType (Arrow 24 returns + // NotImplemented), while the rebuilt physical type can be compacted normally. const auto& child_data = storage_array->data()->child_data; - if (child_fields.empty()) { - *normalized = std::move(storage_array); - return Status::OK(); - } - if (child_fields.size() != child_data.size()) { + if (_child_normalizers.size() != child_data.size()) { return Status::InvalidArgument( "Lance field '{}' has {} child fields but its Arrow array has {} children", - field->name(), child_fields.size(), child_data.size()); - } - - bool requires_normalization = false; - for (const auto& child_field : child_fields) { - bool child_required = false; - RETURN_IF_ERROR(field_requires_lance_normalization(child_field, &child_required)); - if (child_required) { - requires_normalization = true; - break; - } - } - if (!requires_normalization) { - *normalized = std::move(storage_array); - return Status::OK(); + _field_name, _child_normalizers.size(), child_data.size()); } arrow::FieldVector normalized_fields; std::shared_ptr normalized_data; - for (size_t child_idx = 0; child_idx < child_fields.size(); ++child_idx) { - bool child_required = false; - RETURN_IF_ERROR( - field_requires_lance_normalization(child_fields[child_idx], &child_required)); - if (!child_required) { + for (size_t child_idx = 0; child_idx < _child_normalizers.size(); ++child_idx) { + const auto& child_normalizer = _child_normalizers[child_idx]; + if (!child_normalizer._requires_special_handling) { continue; } - auto child_array = arrow::MakeArray(storage_array->data()->child_data[child_idx]); + auto child_array = arrow::MakeArray(child_data[child_idx]); std::shared_ptr normalized_child; - RETURN_IF_ERROR(normalize_lance_arrow_array(child_fields[child_idx], child_array, - memory_pool, &normalized_child)); + RETURN_IF_ERROR( + child_normalizer.normalize_for_doris(child_array, memory_pool, &normalized_child)); if (normalized_child.get() == child_array.get()) { continue; } @@ -658,15 +564,15 @@ Status normalize_lance_arrow_array(const std::shared_ptr& field, normalized_fields[child_idx] = normalized_fields[child_idx]->WithType(normalized_child->type()); } - if (normalized_data == nullptr) { - *normalized = std::move(storage_array); - return Status::OK(); + if (normalized_data != nullptr) { + RETURN_IF_ERROR(set_lance_nested_type(_field_name, storage_array->type(), normalized_fields, + &normalized_data)); + storage_array = arrow::MakeArray(std::move(normalized_data)); } - RETURN_IF_ERROR(set_lance_nested_type(field->name(), storage_array->type(), normalized_fields, - &normalized_data)); - *normalized = arrow::MakeArray(std::move(normalized_data)); - return Status::OK(); + // Compact only after child types are physical, so MakeBuilder never sees a nested + // ExtensionType. This also preserves the existing zero-copy path when no slice is present. + return compact_lance_array_if_needed(storage_array, memory_pool, normalized); } #ifdef BE_TEST @@ -675,12 +581,33 @@ Status normalize_lance_arrow_array_for_test(const std::shared_ptr& const std::shared_ptr& array, std::shared_ptr* normalized, arrow::MemoryPool* memory_pool) { - return normalize_lance_arrow_array( - field, array, memory_pool != nullptr ? memory_pool : arrow::default_memory_pool(), - normalized); + LanceArrowArrayNormalizer normalizer; + RETURN_IF_ERROR(LanceArrowArrayNormalizer::create(field, &normalizer)); + return normalizer.normalize_for_doris( + array, memory_pool != nullptr ? memory_pool : arrow::default_memory_pool(), normalized); } #endif +Status import_lance_dataset_schema(LanceDataset* dataset, std::shared_ptr* schema) { + DORIS_CHECK(dataset != nullptr); + DORIS_CHECK(schema != nullptr); + + ArrowSchema arrow_schema {}; + if (lance_dataset_schema(dataset, &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()); + } + *schema = std::move(imported_schema).ValueUnsafe(); + return Status::OK(); +} + void LanceDatasetDeleter::operator()(LanceDataset* dataset) const { lance_dataset_close(dataset); } diff --git a/be/src/format_v2/lance/lance_reader_helper.h b/be/src/format_v2/lance/lance_reader_helper.h index 8a8fb5a4f5990e..c2f67e201bedce 100644 --- a/be/src/format_v2/lance/lance_reader_helper.h +++ b/be/src/format_v2/lance/lance_reader_helper.h @@ -34,6 +34,7 @@ struct LanceScanner; namespace arrow { class Array; +class DataType; class Field; class MemoryPool; class Schema; @@ -60,17 +61,37 @@ struct LanceBatchDeleter { size_t lance_vector_element_width(TVectorElementType::type type); +// Import the physical Arrow schema owned by a Lance dataset. The caller owns the resulting +// shared schema and may cache it for operations that need physical Arrow types. +Status import_lance_dataset_schema(LanceDataset* dataset, std::shared_ptr* schema); + // Validate and convert the fragment and index-segment identifiers carried by the FE into the // unsigned and packed representations expected by lance-c. Status parse_fragment_ids(const TLanceFileDesc& lance_params, std::vector* fragment_ids); Status parse_index_segment_uuids(const TLanceFileDesc& lance_params, std::vector* segment_uuids, size_t* segment_count); -// Normalize Lance extension arrays into Arrow arrays supported by Doris. -Status normalize_lance_arrow_array(const std::shared_ptr& field, - const std::shared_ptr& array, - arrow::MemoryPool* memory_pool, - std::shared_ptr* normalized); +// Resolves one Arrow field once when a stream schema is bound. Runtime conversion then reuses this +// plan instead of rescanning extension metadata and nested fields for every array. +class LanceArrowArrayNormalizer { +public: + static Status create(const std::shared_ptr& field, + LanceArrowArrayNormalizer* normalizer); + + // Return an Arrow array whose physical layout and type can be consumed by Doris SerDes. The + // input is returned unchanged when no compaction or type adaptation is required. + Status normalize_for_doris(const std::shared_ptr& array, + arrow::MemoryPool* memory_pool, + std::shared_ptr* normalized) const; + +private: + std::string _field_name; + std::shared_ptr _storage_type; + std::vector _child_normalizers; + bool _unwrap_registered_extension = false; + bool _convert_bfloat16 = false; + bool _requires_special_handling = false; +}; #ifdef BE_TEST // Expose Lance Arrow normalization for allocation-sensitive unit tests. diff --git a/be/src/format_v2/lance/lance_record_batch_converter.cpp b/be/src/format_v2/lance/lance_record_batch_converter.cpp new file mode 100644 index 00000000000000..7f0bf848a08dcd --- /dev/null +++ b/be/src/format_v2/lance/lance_record_batch_converter.cpp @@ -0,0 +1,261 @@ +// 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/lance/lance_record_batch_converter.h" + +#include +#include +#include +#include + +#include + +#include "common/consts.h" +#include "common/exception.h" +#include "core/block/block.h" +#include "core/column/column_nullable.h" +#include "core/column/column_string.h" +#include "core/data_type/data_type_nullable.h" +#include "format_v2/lance/lance_reader_helper.h" +#include "runtime/exec_env.h" +#include "runtime/runtime_state.h" +#include "storage/utils.h" + +namespace doris::format::lance { +namespace { + +constexpr size_t IGNORED_OUTPUT_INDEX = std::numeric_limits::max(); + +} // namespace + +arrow::MemoryPool* LanceRecordBatchConverter::_query_arrow_memory_pool( + RuntimeState* runtime_state) { + DORIS_CHECK(runtime_state != nullptr); + if (runtime_state->exec_env() != nullptr) { + if (auto* memory_pool = runtime_state->exec_env()->arrow_memory_pool(); + memory_pool != nullptr) { + return memory_pool; + } + } + return arrow::default_memory_pool(); +} + +Status LanceRecordBatchConverter::init(RuntimeState* runtime_state, + const std::vector& projected_columns, + SearchKind search_kind) { + DORIS_CHECK(runtime_state != nullptr); + _search_kind = search_kind; + _timezone = runtime_state->timezone_obj(); + _memory_pool = _query_arrow_memory_pool(runtime_state); + _output_columns.clear(); + _output_columns.reserve(projected_columns.size()); + _output_index_by_input_name.clear(); + _output_index_by_input_name.reserve(projected_columns.size()); + _global_rowid_output_index.reset(); + reset_schema(); + + for (const auto& column : projected_columns) { + if (column.type == nullptr) { + return Status::InvalidArgument("Lance projected column '{}' has no type", column.name); + } + if (column.name.starts_with(BeConsts::GLOBAL_ROWID_COL)) { + if (_search_kind == SearchKind::NORMAL) { + return Status::NotSupported( + "Lance global row id is currently supported only for external search"); + } + if (_global_rowid_output_index.has_value()) { + return Status::InvalidArgument("duplicate Lance global row id projected column: {}", + column.name); + } + if (remove_nullable(column.type)->get_primitive_type() != TYPE_STRING) { + return Status::InvalidArgument( + "Lance global row id column '{}' must have Doris STRING type, but was {}", + column.name, column.type->get_name()); + } + _global_rowid_output_index = _output_columns.size(); + // Lance returns the native `_rowid`, while Doris exposes it through a generated + // global-rowid column. One input column cannot also satisfy a regular `_rowid` + // projection, so reject an existing binding instead of silently replacing it. + if (!_output_index_by_input_name + .emplace(LANCE_ROW_ID_COLUMN, *_global_rowid_output_index) + .second) { + return Status::InvalidArgument( + "Lance global row id cannot be combined with projected column '{}'", + LANCE_ROW_ID_COLUMN); + } + _output_columns.push_back({column.name, nullptr}); + continue; + } + const auto output_index = _output_columns.size(); + if (!_output_index_by_input_name.emplace(column.name, output_index).second) { + return Status::InvalidArgument("duplicate Lance projected column: {}", column.name); + } + if (_search_kind == SearchKind::VECTOR && column.name == LANCE_DISTANCE_COLUMN && + remove_nullable(column.type)->get_primitive_type() != TYPE_FLOAT) { + return Status::InvalidArgument( + "Lance vector search column '{}' must have Doris FLOAT type, but was {}", + LANCE_DISTANCE_COLUMN, column.type->get_name()); + } + if (_search_kind == SearchKind::FULL_TEXT && column.name == LANCE_SCORE_COLUMN && + remove_nullable(column.type)->get_primitive_type() != TYPE_FLOAT) { + return Status::InvalidArgument( + "Lance full-text search column '{}' must have Doris FLOAT type, but was {}", + LANCE_SCORE_COLUMN, column.type->get_name()); + } + _output_columns.push_back({column.name, column.type->get_serde()}); + } + return Status::OK(); +} + +Status LanceRecordBatchConverter::_bind_schema(const std::shared_ptr& schema) { + DORIS_CHECK(schema != nullptr); + if (_bound_schema == schema || + (_bound_schema != nullptr && _bound_schema->Equals(*schema, true))) { + return Status::OK(); + } + + std::vector arrow_column_to_output(static_cast(schema->num_fields()), + IGNORED_OUTPUT_INDEX); + std::vector array_normalizer_for_arrow_column( + static_cast(schema->num_fields())); + std::vector output_has_source_column(_output_columns.size(), false); + for (int arrow_index = 0; arrow_index < schema->num_fields(); ++arrow_index) { + const auto& field = schema->field(arrow_index); + const auto output = _output_index_by_input_name.find(field->name()); + if (output == _output_index_by_input_name.end()) { + const bool is_omitted_search_result = + (_search_kind == SearchKind::VECTOR && + field->name() == LANCE_DISTANCE_COLUMN) || + (_search_kind == SearchKind::FULL_TEXT && field->name() == LANCE_SCORE_COLUMN); + // Lance may return this generated result even when Doris did not project it. + if (is_omitted_search_result) { + continue; + } + return Status::InternalError("Lance returned unknown column '{}'", field->name()); + } + const auto output_index = output->second; + if (output_has_source_column[output_index]) { + return Status::InternalError("Lance returned duplicate column '{}'", field->name()); + } + output_has_source_column[output_index] = true; + arrow_column_to_output[arrow_index] = output_index; + RETURN_IF_ERROR(LanceArrowArrayNormalizer::create( + field, &array_normalizer_for_arrow_column[static_cast(arrow_index)])); + } + for (size_t output_index = 0; output_index < _output_columns.size(); ++output_index) { + if (!output_has_source_column[output_index]) { + return Status::InternalError("Lance did not return requested column '{}'", + _output_columns[output_index].name); + } + } + _bound_schema = schema; + _arrow_column_to_output_index = std::move(arrow_column_to_output); + _array_normalizer_for_arrow_column = std::move(array_normalizer_for_arrow_column); + return Status::OK(); +} + +Status LanceRecordBatchConverter::convert_record_batch_to_block( + const std::shared_ptr& record_batch, Block* block, + const std::optional& global_rowid_context, size_t* rows) { + DORIS_CHECK(record_batch != nullptr); + DORIS_CHECK(block != nullptr); + DORIS_CHECK(rows != nullptr); + RETURN_IF_ERROR(_bind_schema(record_batch->schema())); + + const auto row_count = static_cast(record_batch->num_rows()); + auto columns_guard = block->mutate_columns_scoped(); + auto& columns = columns_guard.mutable_columns(); + for (int arrow_index = 0; arrow_index < record_batch->num_columns(); ++arrow_index) { + const auto output_index = _arrow_column_to_output_index[static_cast(arrow_index)]; + if (output_index == IGNORED_OUTPUT_INDEX) { + continue; + } + const auto& output = _output_columns[output_index]; + const auto& arrow_column = record_batch->column(arrow_index); + if (is_global_rowid_output(output_index)) { + if (!global_rowid_context.has_value()) { + return Status::InvalidArgument( + "Lance global row id requested without global row id context"); + } + RETURN_IF_ERROR(_append_global_row_ids(arrow_column, columns[output_index], + *global_rowid_context)); + continue; + } + const auto& field = record_batch->schema()->field(arrow_index); + try { + if (arrow_column->type_id() == arrow::Type::NA) { + columns[output_index]->insert_many_defaults(row_count); + continue; + } + std::shared_ptr normalized_column; + RETURN_IF_ERROR( + _array_normalizer_for_arrow_column[static_cast(arrow_index)] + .normalize_for_doris(arrow_column, _memory_pool, &normalized_column)); + RETURN_IF_ERROR(output.serde->read_column_from_arrow( + *columns[output_index], normalized_column.get(), 0, row_count, _timezone)); + } catch (const Exception& e) { + return Status::InternalError("convert Lance Arrow column '{}' failed: {}", + field->name(), e.what()); + } + } + *rows = row_count; + return Status::OK(); +} + +void LanceRecordBatchConverter::reset_schema() { + _bound_schema.reset(); + _arrow_column_to_output_index.clear(); + _array_normalizer_for_arrow_column.clear(); +} + +Status LanceRecordBatchConverter::_append_global_row_ids( + const std::shared_ptr& row_ids, MutableColumnPtr& output_column, + const GlobalRowIdContext& context) const { + if (row_ids->type_id() != arrow::Type::UINT64) { + return Status::InternalError("Lance row id column must be Arrow UINT64, but was {}", + row_ids->type()->ToString()); + } + + ColumnString* data_column = nullptr; + ColumnUInt8::Container* null_map = nullptr; + if (auto* nullable = check_and_get_column(*output_column)) { + data_column = check_and_get_column(nullable->get_nested_column()); + null_map = &nullable->get_null_map_data(); + } else { + data_column = check_and_get_column(*output_column); + } + if (data_column == nullptr) { + return Status::InternalError("Lance global row id output column must be STRING"); + } + + const auto typed_row_ids = std::static_pointer_cast(row_ids); + if (typed_row_ids->null_count() != 0) { + return Status::InternalError("Lance returned null row id"); + } + const auto row_count = static_cast(typed_row_ids->length()); + if (null_map != nullptr) { + null_map->resize_fill(null_map->size() + row_count, 0); + } + for (size_t row = 0; row < row_count; ++row) { + const GlobalRowLoacationV2 location(ROW_VERSION::LANCE_DATASET_ROW_ID, context.backend_id, + context.file_id, typed_row_ids->Value(row)); + data_column->insert_data(reinterpret_cast(&location), sizeof(location)); + } + return Status::OK(); +} + +} // namespace doris::format::lance diff --git a/be/src/format_v2/lance/lance_record_batch_converter.h b/be/src/format_v2/lance/lance_record_batch_converter.h new file mode 100644 index 00000000000000..69ebd4618a9e30 --- /dev/null +++ b/be/src/format_v2/lance/lance_record_batch_converter.h @@ -0,0 +1,91 @@ +// 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 "common/status.h" +#include "core/data_type_serde/data_type_serde.h" +#include "format_v2/lance/lance_reader_helper.h" +#include "format_v2/table_reader.h" + +namespace arrow { +class Array; +class MemoryPool; +class RecordBatch; +class Schema; +} // namespace arrow + +namespace doris { +class RuntimeState; +} + +namespace doris::format::lance { + +enum class SearchKind { NORMAL, VECTOR, FULL_TEXT }; + +// Bridges one Lance Arrow record-batch stream into an existing Doris Block. It owns the stable +// output-column metadata and lazily binds an incoming Arrow schema once, so ordinary batches only +// perform positional conversion. Each Arrow stream has an independent schema binding. +class LanceRecordBatchConverter { +public: + Status init(RuntimeState* runtime_state, const std::vector& projected_columns, + SearchKind search_kind); + + Status convert_record_batch_to_block( + const std::shared_ptr& record_batch, Block* block, + const std::optional& global_rowid_context, size_t* rows); + + bool requires_global_rowid() const { return _global_rowid_output_index.has_value(); } + bool is_global_rowid_output(size_t output_index) const { + return _global_rowid_output_index.has_value() && + *_global_rowid_output_index == output_index; + } + void reset_schema(); + +private: + struct OutputColumn { + std::string name; + DataTypeSerDeSPtr serde; + }; + + static arrow::MemoryPool* _query_arrow_memory_pool(RuntimeState* runtime_state); + Status _bind_schema(const std::shared_ptr& schema); + Status _append_global_row_ids(const std::shared_ptr& row_ids, + MutableColumnPtr& output_column, + const GlobalRowIdContext& context) const; + + SearchKind _search_kind = SearchKind::NORMAL; + cctz::time_zone _timezone; + arrow::MemoryPool* _memory_pool = nullptr; + std::vector _output_columns; + std::unordered_map _output_index_by_input_name; + std::optional _global_rowid_output_index; + std::shared_ptr _bound_schema; + std::vector _arrow_column_to_output_index; + std::vector _array_normalizer_for_arrow_column; +}; + +} // namespace doris::format::lance diff --git a/be/src/format_v2/lance/lance_runtime_filter_helper.cpp b/be/src/format_v2/lance/lance_runtime_filter_helper.cpp index 0072f5f81e7eda..d72a98bb8d947b 100644 --- a/be/src/format_v2/lance/lance_runtime_filter_helper.cpp +++ b/be/src/format_v2/lance/lance_runtime_filter_helper.cpp @@ -493,6 +493,12 @@ std::optional build_cache_key(const VExprContextSPtrs& conjuncts) { } // namespace +bool has_lance_runtime_filters(const VExprContextSPtrs& conjuncts) { + return std::ranges::any_of(conjuncts, [](const auto& conjunct) { + return get_runtime_filter(conjunct) != nullptr; + }); +} + std::shared_ptr get_or_create_lance_runtime_filter_sql( const VExprContextSPtrs& conjuncts, const arrow::Schema& physical_schema, ShardedKVCache* cache) { diff --git a/be/src/format_v2/lance/lance_runtime_filter_helper.h b/be/src/format_v2/lance/lance_runtime_filter_helper.h index 4dbf43e68da696..40529951ccc7c8 100644 --- a/be/src/format_v2/lance/lance_runtime_filter_helper.h +++ b/be/src/format_v2/lance/lance_runtime_filter_helper.h @@ -40,6 +40,10 @@ struct LanceRuntimeFilterSql { std::vector skipped_filter_ids; }; +// Return whether the conjunct list contains at least one Doris runtime filter. This lets callers +// avoid importing Lance's physical Arrow schema when no runtime-filter SQL can be produced. +bool has_lance_runtime_filters(const VExprContextSPtrs& conjuncts); + // Build one immutable SQL snapshot for all supported Doris runtime filters. When cache is non-null, // equivalent RF snapshots from parallel Lance readers in the same FileScanLocalState share the // conversion result. The returned snapshot also identifies RFs that cannot be represented exactly diff --git a/be/src/format_v2/table/lance_reader.cpp b/be/src/format_v2/table/lance_reader.cpp index dd431206e627a8..840338e4b333f3 100644 --- a/be/src/format_v2/table/lance_reader.cpp +++ b/be/src/format_v2/table/lance_reader.cpp @@ -28,57 +28,17 @@ #include #include #include -#include #include "common/config.h" -#include "common/consts.h" #include "common/logging.h" -#include "core/column/column_nullable.h" -#include "core/column/column_string.h" #include "exec/common/endian.h" #include "format_v2/lance/lance_reader_helper.h" #include "format_v2/lance/lance_runtime_filter_helper.h" #include "format_v2/lance/lance_session_manager.h" -#include "runtime/exec_env.h" #include "runtime/file_scan_profile.h" #include "runtime/runtime_state.h" -#include "storage/utils.h" namespace doris::format::lance { -namespace { - -Status import_dataset_schema(LanceDataset* dataset, std::shared_ptr* schema) { - DORIS_CHECK(dataset != nullptr); - DORIS_CHECK(schema != nullptr); - - ArrowSchema arrow_schema {}; - if (lance_dataset_schema(dataset, &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()); - } - *schema = std::move(imported_schema).ValueUnsafe(); - return Status::OK(); -} - -// Return the query-tracked Arrow pool, falling back only for standalone unit-test states. -arrow::MemoryPool* get_lance_arrow_memory_pool(RuntimeState* runtime_state) { - if (runtime_state != nullptr && runtime_state->exec_env() != nullptr) { - auto* memory_pool = runtime_state->exec_env()->arrow_memory_pool(); - if (memory_pool != nullptr) { - return memory_pool; - } - } - return arrow::default_memory_pool(); -} - -} // namespace LanceTableReader::~LanceTableReader() { static_cast(close()); @@ -110,7 +70,7 @@ Status LanceTableReader::fetch_schema(const TFileRangeDesc& range, } std::shared_ptr schema; - RETURN_IF_ERROR(import_dataset_schema(dataset.get(), &schema)); + RETURN_IF_ERROR(import_lance_dataset_schema(dataset.get(), &schema)); return convert_arrow_schema_to_doris(schema, column_names, column_types); } @@ -121,86 +81,19 @@ Status LanceTableReader::init(TableReadOptions&& options) { DORIS_CHECK(_scan_params != nullptr); RETURN_IF_ERROR(_resolve_search_kind()); - _ctz = _runtime_state->timezone_obj(); const auto& lance_scan_params = _scan_params->lance_scan_params; ADD_CHILD_TIMER_WITH_LEVEL(_scanner_profile, LANCE_READER_PROFILE, file_scan_profile::TABLE_READER, 1); _dataset_open_time = ADD_CHILD_TIMER_WITH_LEVEL(_scanner_profile, "LanceDatasetOpenTime", LANCE_READER_PROFILE, 1); - _scanner_configure_time = ADD_CHILD_TIMER_WITH_LEVEL( - _scanner_profile, "LanceScannerConfigureTime", LANCE_READER_PROFILE, 1); - _runtime_filter_sql_time = ADD_CHILD_TIMER_WITH_LEVEL( - _scanner_profile, "LanceRuntimeFilterSqlTime", LANCE_READER_PROFILE, 1); - _scanner_read_time = ADD_CHILD_TIMER_WITH_LEVEL(_scanner_profile, "LanceScannerReadTime", - LANCE_READER_PROFILE, 1); _arrow_to_doris_block_time = ADD_CHILD_TIMER_WITH_LEVEL( _scanner_profile, "LanceArrowToDorisBlockTime", LANCE_READER_PROFILE, 1); - _execution_iops = ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceExecutionIOOps", - TUnit::UNIT, LANCE_READER_PROFILE, 1); - _execution_requests = ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceExecutionIORequests", - TUnit::UNIT, LANCE_READER_PROFILE, 1); - _execution_bytes_read = ADD_CHILD_COUNTER_WITH_LEVEL( - _scanner_profile, "LanceExecutionIOBytesRead", TUnit::BYTES, LANCE_READER_PROFILE, 1); _data_cache_bytes_read_from_cache = ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceDataCacheBytesReadFromCache", TUnit::BYTES, LANCE_READER_PROFILE, 1); _data_cache_bytes_read_from_remote = ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceDataCacheBytesReadFromRemote", TUnit::BYTES, LANCE_READER_PROFILE, 1); - _index_partition_cache_miss_loads = - ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceIndexPartitionCacheMissLoads", - TUnit::UNIT, LANCE_READER_PROFILE, 1); - _index_comparisons = ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceIndexComparisons", - TUnit::UNIT, LANCE_READER_PROFILE, 1); - // These scan counts are emitted by Lance's FilteredRead execution node. For vector searches - // with an explicit fragment set, they normally describe the fragments, ranges, and rows read - // while applying the row-id prefilter. They are scan input counts, not ANN result counts. - _lance_count_metrics = { - {"fragments_scanned", - ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceFragmentsScanned", TUnit::UNIT, - LANCE_READER_PROFILE, 1)}, - {"ranges_scanned", - ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceRowOffsetRangesScanned", - TUnit::UNIT, LANCE_READER_PROFILE, 1)}, - {"rows_scanned", ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceRowsScanned", - TUnit::UNIT, LANCE_READER_PROFILE, 1)}, - {"partitions_ranked", - ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceIVFPartitionsRanked", TUnit::UNIT, - LANCE_READER_PROFILE, 1)}, - {"partitions_searched", - ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceIVFPartitionsSearched", - TUnit::UNIT, LANCE_READER_PROFILE, 1)}, - {"deltas_searched", - ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceVectorIndexSegmentsSearched", - TUnit::UNIT, LANCE_READER_PROFILE, 1)}, - {"scalar_segments_requested", - ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceScalarIndexSegmentsRequested", - TUnit::UNIT, LANCE_READER_PROFILE, 1)}, - {"scalar_segments_searched", - ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceScalarIndexSegmentsSearched", - TUnit::UNIT, LANCE_READER_PROFILE, 1)}, - {"scalar_segment_fallbacks", - ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceScalarIndexSegmentFallbacks", - TUnit::UNIT, LANCE_READER_PROFILE, 1)}, - {"scalar_segment_candidate_rows", - ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceScalarIndexCandidateRows", - TUnit::UNIT, LANCE_READER_PROFILE, 1)}, - }; - _lance_time_metrics = { - // This is wait time reported by the same Lance scan execution node described above, - // rather than Doris scanner scheduling wait time. - {"task_wait_time", ADD_CHILD_TIMER_WITH_LEVEL(_scanner_profile, "LanceTaskWaitTime", - LANCE_READER_PROFILE, 1)}, - {"find_partitions_elapsed", - ADD_CHILD_TIMER_WITH_LEVEL(_scanner_profile, "LanceIVFPartitionRankingTime", - LANCE_READER_PROFILE, 1)}, - {"scalar_segment_prepare_time", - ADD_CHILD_TIMER_WITH_LEVEL(_scanner_profile, "LanceScalarIndexSegmentPrepareTime", - LANCE_READER_PROFILE, 1)}, - {"scalar_segment_search_time", - ADD_CHILD_TIMER_WITH_LEVEL(_scanner_profile, "LanceScalarIndexSegmentSearchTime", - LANCE_READER_PROFILE, 1)}, - }; if (_search_kind != SearchKind::NORMAL) { RETURN_IF_ERROR(_validate_external_search_request()); const auto& request = lance_scan_params.external_search_request; @@ -236,15 +129,6 @@ Status LanceTableReader::init(TableReadOptions&& options) { _scanner_profile->add_info_string("LanceTopK", std::to_string(top_k)); _scanner_profile->add_info_string("LanceOffset", std::to_string(offset)); _scanner_profile->add_info_string("LanceTopKPlusOffset", std::to_string(top_k + offset)); - _planned_index_segment_count = - ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LancePlannedIndexSegmentCount", - TUnit::UNIT, LANCE_READER_PROFILE, 1); - _planned_indexed_fragment_count = - ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LancePlannedIndexedFragmentCount", - TUnit::UNIT, LANCE_READER_PROFILE, 1); - _planned_flat_search_fragment_count = ADD_CHILD_COUNTER_WITH_LEVEL( - _scanner_profile, "LancePlannedFlatSearchFragmentCount", TUnit::UNIT, - LANCE_READER_PROFILE, 1); } if (_scan_params->__isset.lance_scan_params && lance_scan_params.__isset.lance_substrait_filter) { @@ -254,51 +138,7 @@ Status LanceTableReader::init(TableReadOptions&& options) { std::to_string(lance_scan_params.lance_substrait_filter.size())); } - _output_name_to_idx.clear(); - _output_name_to_idx.reserve(_projected_columns.size()); - _global_rowid_output_idx.reset(); - 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 (column.name.starts_with(BeConsts::GLOBAL_ROWID_COL)) { - if (_search_kind == SearchKind::NORMAL) { - return Status::NotSupported( - "Lance global row id is currently supported only for external search"); - } - if (_global_rowid_output_idx.has_value()) { - return Status::InvalidArgument("duplicate Lance global row id projected column: {}", - column.name); - } - if (remove_nullable(column.type)->get_primitive_type() != TYPE_STRING) { - return Status::InvalidArgument( - "Lance global row id column '{}' must have Doris STRING type, but was {}", - column.name, column.type->get_name()); - } - _global_rowid_output_idx = idx; - continue; - } - if (!_output_name_to_idx.emplace(column.name, idx).second) { - return Status::InvalidArgument("duplicate Lance projected column: {}", column.name); - } - if (_search_kind == SearchKind::VECTOR && column.name == LANCE_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 {}", - LANCE_DISTANCE_COLUMN, column.type->get_name()); - } - } else if (_search_kind == SearchKind::FULL_TEXT && column.name == LANCE_SCORE_COLUMN) { - const auto score_type = remove_nullable(column.type); - if (score_type->get_primitive_type() != TYPE_FLOAT) { - return Status::InvalidArgument( - "Lance full-text search column '{}' must have Doris FLOAT type, but was {}", - LANCE_SCORE_COLUMN, column.type->get_name()); - } - } - } - return Status::OK(); + return _record_batch_converter.init(_runtime_state, _projected_columns, _search_kind); } Status LanceTableReader::prepare_split(const SplitReadOptions& options) { @@ -316,7 +156,7 @@ Status LanceTableReader::prepare_split(const SplitReadOptions& options) { if (_is_table_level_count_active()) { return Status::OK(); } - if (_global_rowid_output_idx.has_value() && !_global_rowid_context.has_value()) { + if (_record_batch_converter.requires_global_rowid() && !_global_rowid_context.has_value()) { return Status::InvalidArgument( "Lance global row id requested without global row id context"); } @@ -469,7 +309,8 @@ Status LanceTableReader::read_by_row_ids(const TFileRangeDesc& range, size_t rows = 0; { SCOPED_TIMER(_arrow_to_doris_block_time); - RETURN_IF_ERROR(_fill_block_from_record_batch(record_batch, block, &rows)); + RETURN_IF_ERROR(_record_batch_converter.convert_record_batch_to_block( + record_batch, block, _global_rowid_context, &rows)); } fetched_rows += rows; } @@ -748,10 +589,7 @@ Status LanceTableReader::_open_dataset(const DatasetKey& key) { static_cast(key.version), &raw_dataset)); dataset.reset(raw_dataset); } - std::shared_ptr schema; - RETURN_IF_ERROR(import_dataset_schema(dataset.get(), &schema)); _dataset = dataset.release(); - _dataset_schema = std::move(schema); return Status::OK(); } @@ -790,12 +628,98 @@ Status LanceTableReader::_prepare_fts_query_context() { return Status::OK(); } +void LanceTableReader::_init_scanner_profile() { + if (_scanner_configure_time != nullptr) { + return; + } + + _scanner_configure_time = ADD_CHILD_TIMER_WITH_LEVEL( + _scanner_profile, "LanceScannerConfigureTime", LANCE_READER_PROFILE, 1); + _runtime_filter_sql_time = ADD_CHILD_TIMER_WITH_LEVEL( + _scanner_profile, "LanceRuntimeFilterSqlTime", LANCE_READER_PROFILE, 1); + _scanner_read_time = ADD_CHILD_TIMER_WITH_LEVEL(_scanner_profile, "LanceScannerReadTime", + LANCE_READER_PROFILE, 1); + _execution_iops = ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceExecutionIOOps", + TUnit::UNIT, LANCE_READER_PROFILE, 1); + _execution_requests = ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceExecutionIORequests", + TUnit::UNIT, LANCE_READER_PROFILE, 1); + _execution_bytes_read = ADD_CHILD_COUNTER_WITH_LEVEL( + _scanner_profile, "LanceExecutionIOBytesRead", TUnit::BYTES, LANCE_READER_PROFILE, 1); + _index_partition_cache_miss_loads = + ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceIndexPartitionCacheMissLoads", + TUnit::UNIT, LANCE_READER_PROFILE, 1); + _index_comparisons = ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceIndexComparisons", + TUnit::UNIT, LANCE_READER_PROFILE, 1); + + // These scan counts are emitted by Lance's FilteredRead execution node. For vector searches + // with an explicit fragment set, they normally describe the fragments, ranges, and rows read + // while applying the row-id prefilter. They are scan input counts, not ANN result counts. + _lance_count_metrics = { + {"fragments_scanned", + ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceFragmentsScanned", TUnit::UNIT, + LANCE_READER_PROFILE, 1)}, + {"ranges_scanned", + ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceRowOffsetRangesScanned", + TUnit::UNIT, LANCE_READER_PROFILE, 1)}, + {"rows_scanned", ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceRowsScanned", + TUnit::UNIT, LANCE_READER_PROFILE, 1)}, + {"partitions_ranked", + ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceIVFPartitionsRanked", TUnit::UNIT, + LANCE_READER_PROFILE, 1)}, + {"partitions_searched", + ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceIVFPartitionsSearched", + TUnit::UNIT, LANCE_READER_PROFILE, 1)}, + {"deltas_searched", + ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceVectorIndexSegmentsSearched", + TUnit::UNIT, LANCE_READER_PROFILE, 1)}, + {"scalar_segments_requested", + ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceScalarIndexSegmentsRequested", + TUnit::UNIT, LANCE_READER_PROFILE, 1)}, + {"scalar_segments_searched", + ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceScalarIndexSegmentsSearched", + TUnit::UNIT, LANCE_READER_PROFILE, 1)}, + {"scalar_segment_fallbacks", + ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceScalarIndexSegmentFallbacks", + TUnit::UNIT, LANCE_READER_PROFILE, 1)}, + {"scalar_segment_candidate_rows", + ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LanceScalarIndexCandidateRows", + TUnit::UNIT, LANCE_READER_PROFILE, 1)}, + }; + _lance_time_metrics = { + // This is wait time reported by the same Lance scan execution node described above, + // rather than Doris scanner scheduling wait time. + {"task_wait_time", ADD_CHILD_TIMER_WITH_LEVEL(_scanner_profile, "LanceTaskWaitTime", + LANCE_READER_PROFILE, 1)}, + {"find_partitions_elapsed", + ADD_CHILD_TIMER_WITH_LEVEL(_scanner_profile, "LanceIVFPartitionRankingTime", + LANCE_READER_PROFILE, 1)}, + {"scalar_segment_prepare_time", + ADD_CHILD_TIMER_WITH_LEVEL(_scanner_profile, "LanceScalarIndexSegmentPrepareTime", + LANCE_READER_PROFILE, 1)}, + {"scalar_segment_search_time", + ADD_CHILD_TIMER_WITH_LEVEL(_scanner_profile, "LanceScalarIndexSegmentSearchTime", + LANCE_READER_PROFILE, 1)}, + }; + if (_search_kind != SearchKind::NORMAL) { + _planned_index_segment_count = + ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LancePlannedIndexSegmentCount", + TUnit::UNIT, LANCE_READER_PROFILE, 1); + _planned_indexed_fragment_count = + ADD_CHILD_COUNTER_WITH_LEVEL(_scanner_profile, "LancePlannedIndexedFragmentCount", + TUnit::UNIT, LANCE_READER_PROFILE, 1); + _planned_flat_search_fragment_count = ADD_CHILD_COUNTER_WITH_LEVEL( + _scanner_profile, "LancePlannedFlatSearchFragmentCount", TUnit::UNIT, + LANCE_READER_PROFILE, 1); + } +} + Status LanceTableReader::_open_scanner(const TFileRangeDesc& range) { + _init_scanner_profile(); SCOPED_TIMER(_scanner_configure_time); std::vector columns; columns.reserve(_projected_columns.size() + 1); for (size_t idx = 0; idx < _projected_columns.size(); ++idx) { - if (_global_rowid_output_idx == idx) { + if (_record_batch_converter.is_global_rowid_output(idx)) { continue; } const auto& column = _projected_columns[idx]; @@ -813,10 +737,14 @@ Status LanceTableReader::_open_scanner(const TFileRangeDesc& range) { std::string sql_filter; std::shared_ptr runtime_filter_sql; if (_search_kind == SearchKind::NORMAL) { - DORIS_CHECK(_dataset_schema != nullptr); - SCOPED_TIMER(_runtime_filter_sql_time); - runtime_filter_sql = get_or_create_lance_runtime_filter_sql(_conjuncts, *_dataset_schema, - _runtime_filter_cache); + if (has_lance_runtime_filters(_conjuncts)) { + if (_dataset_schema == nullptr) { + RETURN_IF_ERROR(import_lance_dataset_schema(_dataset, &_dataset_schema)); + } + SCOPED_TIMER(_runtime_filter_sql_time); + runtime_filter_sql = get_or_create_lance_runtime_filter_sql( + _conjuncts, *_dataset_schema, _runtime_filter_cache); + } } else { const auto& request = lance_scan_params.external_search_request; if (request.__isset.search_filter && @@ -824,9 +752,8 @@ Status LanceTableReader::_open_scanner(const TFileRangeDesc& range) { sql_filter = request.search_filter.payload; } } - LanceScanner* scanner = - lance_scanner_new(_dataset, columns.size() == 1 ? nullptr : columns.data(), - sql_filter.empty() ? nullptr : sql_filter.c_str()); + LanceScanner* scanner = lance_scanner_new(_dataset, columns.data(), + sql_filter.empty() ? nullptr : sql_filter.c_str()); if (scanner == nullptr) { return lance_error("create Lance scanner"); } @@ -839,7 +766,8 @@ Status LanceTableReader::_open_scanner(const TFileRangeDesc& range) { return lance_error("set Lance scanner statistics callback"); } - if (_global_rowid_output_idx.has_value() && lance_scanner_with_row_id(scanner, true) != 0) { + if (_record_batch_converter.requires_global_rowid() && + lance_scanner_with_row_id(scanner, true) != 0) { return lance_error("enable Lance row id output"); } @@ -1238,10 +1166,12 @@ void LanceTableReader::_close_dataset() { _dataset = nullptr; } _dataset_schema.reset(); + _record_batch_converter.reset_schema(); } void LanceTableReader::_collect_data_cache_statistics() { - if (_dataset == nullptr) { + if (_dataset == nullptr || (_data_cache_bytes_read_from_cache == nullptr && + _data_cache_bytes_read_from_remote == nullptr)) { return; } @@ -1292,111 +1222,8 @@ Status LanceTableReader::_fill_block_from_lance_batch(LanceBatch* batch, Block* result.status().message()); } - return _fill_block_from_record_batch(std::move(result).ValueUnsafe(), block, rows); -} - -Status LanceTableReader::_append_global_row_ids(const std::shared_ptr& row_ids, - MutableColumnPtr& output_column) const { - DORIS_CHECK(row_ids != nullptr); - DORIS_CHECK(_global_rowid_context.has_value()); - if (row_ids->type_id() != arrow::Type::UINT64) { - return Status::InternalError("Lance row id column must be Arrow UINT64, but was {}", - row_ids->type()->ToString()); - } - - ColumnString* data_column = nullptr; - ColumnUInt8::Container* null_map = nullptr; - if (auto* nullable = check_and_get_column(*output_column)) { - data_column = check_and_get_column(nullable->get_nested_column()); - null_map = &nullable->get_null_map_data(); - } else { - data_column = check_and_get_column(*output_column); - } - if (data_column == nullptr) { - return Status::InternalError("Lance global row id output column must be STRING"); - } - - const auto typed_row_ids = std::static_pointer_cast(row_ids); - if (typed_row_ids->null_count() != 0) { - return Status::InternalError("Lance returned null row id"); - } - const auto row_count = static_cast(typed_row_ids->length()); - if (null_map != nullptr) { - null_map->resize_fill(null_map->size() + row_count, 0); - } - const auto& context = *_global_rowid_context; - for (size_t row = 0; row < row_count; ++row) { - const GlobalRowLoacationV2 location(ROW_VERSION::LANCE_DATASET_ROW_ID, context.backend_id, - context.file_id, typed_row_ids->Value(row)); - data_column->insert_data(reinterpret_cast(&location), sizeof(location)); - } - return Status::OK(); -} - -Status LanceTableReader::_fill_block_from_record_batch( - const std::shared_ptr& record_batch, Block* block, size_t* rows) { - DORIS_CHECK(record_batch != nullptr); - DORIS_CHECK(block != nullptr); - DORIS_CHECK(rows != nullptr); - const auto row_count = static_cast(record_batch->num_rows()); - std::unordered_set materialized_columns; - materialized_columns.reserve(record_batch->num_columns()); - auto columns_guard = block->mutate_columns_scoped(); - auto& columns = columns_guard.mutable_columns(); - for (int arrow_idx = 0; arrow_idx < record_batch->num_columns(); ++arrow_idx) { - const auto& field = record_batch->schema()->field(arrow_idx); - if (field->name() == LANCE_ROW_ID_COLUMN && _global_rowid_output_idx.has_value()) { - const auto output_idx = *_global_rowid_output_idx; - const auto& output_name = _projected_columns[output_idx].name; - if (!materialized_columns.emplace(output_name).second) { - return Status::InternalError("Lance returned duplicate column '{}'", - LANCE_ROW_ID_COLUMN); - } - RETURN_IF_ERROR( - _append_global_row_ids(record_batch->column(arrow_idx), columns[output_idx])); - continue; - } - const auto output_it = _output_name_to_idx.find(field->name()); - if (output_it == _output_name_to_idx.end()) { - if ((_search_kind == SearchKind::VECTOR && field->name() == LANCE_DISTANCE_COLUMN) || - (_search_kind == SearchKind::FULL_TEXT && field->name() == LANCE_SCORE_COLUMN)) { - // Lance auto-projects the generated search result column. It is valid for Doris - // slot pruning to omit that optional result column. - continue; - } - return Status::InternalError("Lance returned unknown column '{}'", field->name()); - } - if (!materialized_columns.emplace(field->name()).second) { - return Status::InternalError("Lance returned duplicate column '{}'", field->name()); - } - const auto output_idx = output_it->second; - try { - const auto& arrow_column = record_batch->column(arrow_idx); - if (arrow_column->type_id() == arrow::Type::NA) { - columns[output_idx]->insert_many_defaults(row_count); - continue; - } - std::shared_ptr normalized_column; - RETURN_IF_ERROR(normalize_lance_arrow_array(field, arrow_column, - get_lance_arrow_memory_pool(_runtime_state), - &normalized_column)); - RETURN_IF_ERROR(columns_guard.get_datatype_by_position(output_idx) - ->get_serde() - ->read_column_from_arrow(*columns[output_idx], - normalized_column.get(), 0, row_count, - _ctz)); - } catch (const Exception& e) { - return Status::InternalError("convert Lance Arrow column '{}' failed: {}", - field->name(), e.what()); - } - } - for (const auto& column : _projected_columns) { - if (!materialized_columns.contains(column.name)) { - return Status::InternalError("Lance did not return requested column '{}'", column.name); - } - } - *rows = row_count; - return Status::OK(); + return _record_batch_converter.convert_record_batch_to_block( + std::move(result).ValueUnsafe(), block, _global_rowid_context, rows); } Status LanceTableReader::_dataset_key(const TFileRangeDesc& range, DatasetKey* key) const { diff --git a/be/src/format_v2/table/lance_reader.h b/be/src/format_v2/table/lance_reader.h index 944041bed3d950..88deea610b301a 100644 --- a/be/src/format_v2/table/lance_reader.h +++ b/be/src/format_v2/table/lance_reader.h @@ -17,8 +17,6 @@ #pragma once -#include - #include #include #include @@ -29,6 +27,7 @@ #include #include "common/status.h" +#include "format_v2/lance/lance_record_batch_converter.h" #include "format_v2/table_reader.h" #include "runtime/runtime_profile.h" @@ -41,12 +40,6 @@ namespace doris { class ShardedKVCache; } -namespace arrow { -class Array; -class RecordBatch; -class Schema; -} // namespace arrow - namespace doris::format::lance { // A FORMAT_LANCE table reader. Unlike file formats such as Parquet, a Lance split is not a @@ -87,6 +80,7 @@ class LanceTableReader final : public TableReader { Status _open_dataset(const DatasetKey& key); Status _prepare_fts_query_context(); Status _open_scanner(const TFileRangeDesc& range); + void _init_scanner_profile(); Status _configure_scan_options(LanceScanner* scanner) const; Status _configure_normal_scan(LanceScanner* scanner, const TLanceFileDesc& lance_params) const; Status _configure_vector_search(LanceScanner* scanner, @@ -102,10 +96,6 @@ class LanceTableReader final : public TableReader { void _close_scanner(); void _close_dataset(); Status _fill_block_from_lance_batch(LanceBatch* batch, Block* block, size_t* rows); - Status _fill_block_from_record_batch(const std::shared_ptr& record_batch, - Block* block, size_t* rows); - Status _append_global_row_ids(const std::shared_ptr& row_ids, - MutableColumnPtr& output_column) const; Status _dataset_key(const TFileRangeDesc& range, DatasetKey* key) const; LanceDataset* _dataset = nullptr; @@ -113,9 +103,7 @@ class LanceTableReader final : public TableReader { LanceScanner* _scanner = nullptr; ShardedKVCache* _runtime_filter_cache = nullptr; std::optional _opened_dataset_key; - std::unordered_map _output_name_to_idx; - std::optional _global_rowid_output_idx; - cctz::time_zone _ctz; + LanceRecordBatchConverter _record_batch_converter; size_t _scanner_batch_size = 0; RuntimeProfile::Counter* _planned_index_segment_count = nullptr; RuntimeProfile::Counter* _planned_indexed_fragment_count = nullptr; @@ -137,7 +125,6 @@ class LanceTableReader final : public TableReader { std::unordered_map _lance_count_metrics; std::unordered_map _lance_time_metrics; LanceFtsQueryContext* _fts_query_context = nullptr; - enum class SearchKind { NORMAL, VECTOR, FULL_TEXT }; SearchKind _search_kind = SearchKind::NORMAL; bool _eof = false; }; diff --git a/be/test/format_v2/lance/lance_nested_null_test.cpp b/be/test/format_v2/lance/lance_nested_null_test.cpp index 8d1115687df943..c36b936de6527a 100644 --- a/be/test/format_v2/lance/lance_nested_null_test.cpp +++ b/be/test/format_v2/lance/lance_nested_null_test.cpp @@ -57,8 +57,9 @@ class LanceNestedNullTest : public testing::Test { void read(const std::shared_ptr& array, const DataTypePtr& type, MutableColumnPtr* column) { std::shared_ptr normalized; - ASSERT_TRUE(normalize_lance_arrow_array(arrow::field("value", array->type()), array, - arrow::default_memory_pool(), &normalized) + ASSERT_TRUE(normalize_lance_arrow_array_for_test(arrow::field("value", array->type()), + array, &normalized, + arrow::default_memory_pool()) .ok()); *column = type->create_column(); ASSERT_TRUE(type->get_serde() diff --git a/be/test/format_v2/table/lance_reader_test.cpp b/be/test/format_v2/table/lance_reader_test.cpp index 6baccb3bd9dffd..b0e1b2bb831bf1 100644 --- a/be/test/format_v2/table/lance_reader_test.cpp +++ b/be/test/format_v2/table/lance_reader_test.cpp @@ -299,6 +299,14 @@ Status init_reader(LanceTableReader* reader, const Columns& projected_columns, }); } +Status convert_record_batch_for_test(RuntimeState* runtime_state, const Columns& projected_columns, + const std::shared_ptr& record_batch, + Block* block, size_t* rows) { + LanceRecordBatchConverter converter; + RETURN_IF_ERROR(converter.init(runtime_state, projected_columns, SearchKind::NORMAL)); + return converter.convert_record_batch_to_block(record_batch, block, std::nullopt, rows); +} + Status prepare_range(LanceTableReader* reader, TFileRangeDesc range, std::optional global_rowid_context = std::nullopt) { // Assign after value initialization so adding optional split state cannot break this fixture's @@ -2302,7 +2310,7 @@ TEST(LanceTableReaderTypeTest, ReadsAdditionalArrowAndLanceTypes) { Block block; add_output_columns(&block, columns); size_t rows = 0; - ASSERT_TRUE(reader._fill_block_from_record_batch(record_batch, &block, &rows).ok()); + ASSERT_TRUE(convert_record_batch_for_test(&state, columns, record_batch, &block, &rows).ok()); ASSERT_EQ(2, rows); const auto& null_values = assert_cast(*block.get_by_position(0).column); @@ -2377,7 +2385,7 @@ TEST(LanceTableReaderTypeTest, ReadsSlicedDurationAndJsonValues) { Block block; add_output_columns(&block, columns); size_t rows = 0; - ASSERT_TRUE(reader._fill_block_from_record_batch(record_batch, &block, &rows).ok()); + ASSERT_TRUE(convert_record_batch_for_test(&state, columns, record_batch, &block, &rows).ok()); ASSERT_EQ(2, rows); const auto& duration = assert_cast(*block.get_by_position(0).column); @@ -2444,7 +2452,7 @@ TEST(LanceTableReaderTypeTest, ReadsRegisteredJsonNestedInSlicedList) { Block block; add_output_columns(&block, columns); size_t rows = 0; - ASSERT_TRUE(reader._fill_block_from_record_batch(record_batch, &block, &rows).ok()); + ASSERT_TRUE(convert_record_batch_for_test(&state, columns, record_batch, &block, &rows).ok()); ASSERT_EQ(1, rows); const auto& nullable_list = @@ -2693,7 +2701,7 @@ TEST(LanceTableReaderTypeTest, NormalizesVisibleBFloat16ValuesInSlicedMap) { Block block; add_output_columns(&block, columns); size_t rows = 0; - ASSERT_TRUE(reader._fill_block_from_record_batch(record_batch, &block, &rows).ok()); + ASSERT_TRUE(convert_record_batch_for_test(&state, columns, record_batch, &block, &rows).ok()); ASSERT_EQ(1, rows); const auto& nullable_map = assert_cast(*block.get_by_position(0).column); @@ -2740,7 +2748,7 @@ TEST(LanceTableReaderTypeTest, ReadsLanceJsonLargeBinaryValues) { Block block; add_output_columns(&block, columns); size_t rows = 0; - ASSERT_TRUE(reader._fill_block_from_record_batch(record_batch, &block, &rows).ok()); + ASSERT_TRUE(convert_record_batch_for_test(&state, columns, record_batch, &block, &rows).ok()); ASSERT_EQ(2, rows); const auto& json_values = assert_cast(*block.get_by_position(0).column); EXPECT_EQ((ColumnUInt8::Container {0, 1}), json_values.get_null_map_data()); @@ -2772,7 +2780,7 @@ TEST(LanceTableReaderTypeTest, ReadsDurationBoundaryValues) { Block block; add_output_columns(&block, columns); size_t rows = 0; - ASSERT_TRUE(reader._fill_block_from_record_batch(record_batch, &block, &rows).ok()); + ASSERT_TRUE(convert_record_batch_for_test(&state, columns, record_batch, &block, &rows).ok()); ASSERT_EQ(5, rows); const auto& durations = assert_cast(*block.get_by_position(0).column); const auto& duration_values = assert_cast(durations.get_nested_column());