Skip to content

[fix](be) Reuse file block columns during v2 finalization #65589

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
xylaaaaa wants to merge 1 commit into
base: master
Choose a base branch
from
+138 −7
Open

[fix](be) Reuse file block columns during v2 finalization #65589

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 21 additions & 7 deletions be/src/format_v2/table_reader.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -655,15 +655,29 @@ class TableReader {
// Finalize file-local block to table/global schema block.
Status finalize_chunk(Block* block, const size_t rows) {
SCOPED_TIMER(_profile.finalize_timer);
size_t idx = 0;
std::vector<ColumnPtr> materialized_columns;
materialized_columns.reserve(_data_reader.column_mapper->mappings().size());
for (const auto& mapping : _data_reader.column_mapper->mappings()) {
ColumnPtr column;
RETURN_IF_ERROR(_materialize_mapping_column(mapping, &_data_reader.block_template, rows,
&column));
block->replace_by_position(idx, IColumn::mutate(std::move(column)));
idx++;
materialized_columns.push_back(std::move(column));
}
for (size_t idx = 0; idx < materialized_columns.size(); ++idx) {
block->replace_by_position(idx, std::move(materialized_columns[idx]));
}
// Table-format virtual columns can depend on auxiliary file-local columns such as Iceberg
// row positions, so materialize them before releasing the file block.
RETURN_IF_ERROR(materialize_virtual_columns(block));
// Projection results can alias columns owned by the file-local block. Release those
// owners before requesting mutable output columns so direct mappings transfer their
// payload instead of triggering a COW deep clone.
_data_reader.block_template.clear_column_data(
cast_set<int64_t>(_data_reader.file_block_layout.size()));
for (size_t idx = 0; idx < block->columns(); ++idx) {
block->replace_by_position(
idx, IColumn::mutate(std::move(block->get_by_position(idx).column)));
}
// Enforce CHAR/VARCHAR length declared by the table schema after all file-to-table
// materialization has finished.
RETURN_IF_ERROR(_truncate_char_or_varchar_columns(block));
Expand Down Expand Up @@ -1152,7 +1166,7 @@ class TableReader {
rows, st.to_string(), mapping.debug_string());
}
ColumnPtr result_column = current_block->get_by_position(res_id).column;
*column = _detach_column(std::move(result_column));
*column = std::move(result_column);
return Status::OK();
}
if (mapping.default_expr != nullptr) {
Expand All @@ -1162,7 +1176,7 @@ class TableReader {
mapping.default_expr, current_block, &result));
ColumnPtr result_column = result.column;
RETURN_IF_ERROR(_align_column_nullability(&result_column, mapping.table_type));
*column = _detach_column(std::move(result_column));
*column = std::move(result_column);
} else {
DORIS_CHECK(mapping.constant_index.has_value());
Block eval_block;
Expand All @@ -1173,12 +1187,12 @@ class TableReader {
mapping.default_expr, &eval_block, &result));
ColumnPtr result_column = result.column;
RETURN_IF_ERROR(_align_column_nullability(&result_column, mapping.table_type));
*column = _detach_column(std::move(result_column));
*column = std::move(result_column);
}
return Status::OK();
}
ColumnPtr result_column = mapping.table_type->create_column_const_with_default_value(rows);
*column = _detach_column(std::move(result_column));
*column = std::move(result_column);
return Status::OK();
}

Expand Down
117 changes: 117 additions & 0 deletions be/test/format_v2/table_reader_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1014,6 +1014,10 @@ struct FakeFileReaderState {
std::shared_ptr<FileScanRequest> last_request;
std::shared_ptr<ConditionCacheContext> condition_cache_ctx;
std::shared_ptr<io::IOContext> io_ctx;
const UInt8* string_payload = nullptr;
const UInt8* map_key_payload = nullptr;
const UInt8* map_value_payload = nullptr;
const ColumnArray::Offset64* map_offsets = nullptr;
};

class FakeFileReader final : public FileReader {
Expand Down Expand Up @@ -1074,6 +1078,7 @@ class FakeFileReader final : public FileReader {
auto column = ColumnString::create();
column->insert_data("one", 3);
column->insert_data("two", 3);
_state->string_payload = column->get_chars().data();
file_block->replace_by_position(block_position.value(),
make_not_null_nullable_column(std::move(column)));
} else if (file_column_id == LocalColumnId(2)) {
Expand All @@ -1094,6 +1099,30 @@ class FakeFileReader final : public FileReader {
file_block->replace_by_position(
block_position.value(),
make_not_null_nullable_column(std::move(struct_column)));
} else if (file_column_id == LocalColumnId(3)) {
auto keys = ColumnString::create();
keys->insert_data("first", 5);
keys->insert_data("second", 6);
keys->insert_data("third", 5);
_state->map_key_payload = keys->get_chars().data();

auto values = ColumnString::create();
values->insert_data("one", 3);
values->insert_data("two", 3);
values->insert_data("three", 5);
_state->map_value_payload = values->get_chars().data();

auto offsets = ColumnArray::ColumnOffsets::create();
offsets->insert_value(1);
offsets->insert_value(3);
_state->map_offsets = offsets->get_data().data();

auto map_column = ColumnMap::create(
make_not_null_nullable_column(std::move(keys)),
make_not_null_nullable_column(std::move(values)), std::move(offsets));
file_block->replace_by_position(
block_position.value(),
make_not_null_nullable_column(std::move(map_column)));
} else {
return Status::InvalidArgument("Unexpected fake file column id {}",
file_column_id.value());
Expand Down Expand Up @@ -1671,6 +1700,94 @@ TEST(TableReaderTest, SlotlessConjunctDisablesAggregatePushdown) {
ASSERT_TRUE(reader.close().ok());
}

TEST(TableReaderTest, DirectMappingTransfersFileColumnOwnership) {
std::vector<ColumnDefinition> file_schema;
file_schema.push_back(make_file_column(1, "value", std::make_shared<DataTypeString>()));

std::vector<ColumnDefinition> projected_columns;
projected_columns.push_back(make_table_column(1, "value", std::make_shared<DataTypeString>()));
set_name_identifiers(&projected_columns);

RuntimeState state {TQueryOptions(), TQueryGlobals()};
auto fake_state = std::make_shared<FakeFileReaderState>();
FakeTableReader reader(file_schema, fake_state);
ASSERT_TRUE(reader.init({
.projected_columns = projected_columns,
.conjuncts = {},
.format = FileFormat::PARQUET,
.scan_params = nullptr,
.io_ctx = nullptr,
.runtime_state = &state,
.scanner_profile = nullptr,
})
.ok());

SplitReadOptions split_options;
split_options.current_range.__set_path("fake-table-reader-input");
ASSERT_TRUE(reader.prepare_split(split_options).ok());

Block block = build_table_block(projected_columns);
bool eos = false;
ASSERT_TRUE(reader.get_block(&block, &eos).ok());
ASSERT_FALSE(eos);
ASSERT_NE(fake_state->string_payload, nullptr);

const auto& result = assert_cast<const ColumnString&>(expect_not_null_table_column(block, 0));
EXPECT_EQ(result.get_chars().data(), fake_state->string_payload);
ASSERT_TRUE(reader.close().ok());
}

TEST(TableReaderTest, DirectMapMappingTransfersNestedColumnOwnership) {
const auto string_type = std::make_shared<DataTypeString>();
const auto map_type = std::make_shared<DataTypeMap>(string_type, string_type);
auto file_map = make_file_column(3, "attributes", map_type);
file_map.children = {make_file_column(0, "key", string_type),
make_file_column(1, "value", string_type)};
std::vector<ColumnDefinition> file_schema = {file_map};

auto table_map = make_table_column(3, "attributes", map_type);
table_map.children = {make_table_column(0, "key", string_type),
make_table_column(1, "value", string_type)};
std::vector<ColumnDefinition> projected_columns = {table_map};
set_name_identifiers(&projected_columns);

RuntimeState state {TQueryOptions(), TQueryGlobals()};
auto fake_state = std::make_shared<FakeFileReaderState>();
FakeTableReader reader(file_schema, fake_state);
ASSERT_TRUE(reader.init({
.projected_columns = projected_columns,
.conjuncts = {},
.format = FileFormat::PARQUET,
.scan_params = nullptr,
.io_ctx = nullptr,
.runtime_state = &state,
.scanner_profile = nullptr,
})
.ok());

SplitReadOptions split_options;
split_options.current_range.__set_path("fake-table-reader-input");
ASSERT_TRUE(reader.prepare_split(split_options).ok());

Block block = build_table_block(projected_columns);
bool eos = false;
ASSERT_TRUE(reader.get_block(&block, &eos).ok());
ASSERT_FALSE(eos);
ASSERT_NE(fake_state->map_key_payload, nullptr);
ASSERT_NE(fake_state->map_value_payload, nullptr);
ASSERT_NE(fake_state->map_offsets, nullptr);

const auto& result = assert_cast<const ColumnMap&>(expect_not_null_table_column(block, 0));
const auto& keys = assert_cast<const ColumnString&>(
expect_not_null_nullable_nested_column(result.get_keys()));
const auto& values = assert_cast<const ColumnString&>(
expect_not_null_nullable_nested_column(result.get_values()));
EXPECT_EQ(keys.get_chars().data(), fake_state->map_key_payload);
EXPECT_EQ(values.get_chars().data(), fake_state->map_value_payload);
EXPECT_EQ(result.get_offsets().data(), fake_state->map_offsets);
ASSERT_TRUE(reader.close().ok());
}

TEST(TableReaderTest, AbortSplitClearsReaderAfterIgnorableNotFound) {
std::vector<ColumnDefinition> file_schema;
file_schema.push_back(make_file_column(0, "id", std::make_shared<DataTypeInt32>()));
Expand Down
Loading