try to fix heap-use-after-free issues

This commit is contained in:
Yarik Briukhovetskyi 2024-09-13 13:34:36 +02:00 committed by GitHub
parent a41bb7e526
commit c3f2b9714b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1235,12 +1235,14 @@ bool KeyCondition::tryPrepareSetIndex(
size_t indexes_mapping_size = indexes_mapping.size();
assert(set_types_size == set_columns.size());
Columns transformed_set_columns = set_columns;
for (size_t indexes_mapping_index = 0; indexes_mapping_index < indexes_mapping_size; ++indexes_mapping_index)
{
const auto & key_column_type = data_types[indexes_mapping_index];
size_t set_element_index = indexes_mapping[indexes_mapping_index].tuple_index;
auto set_element_type = set_types[set_element_index];
auto set_column = set_columns[set_element_index];
ColumnPtr set_column = set_columns[set_element_index];
if (!set_transforming_chains[indexes_mapping_index].empty())
{
@ -1261,7 +1263,7 @@ bool KeyCondition::tryPrepareSetIndex(
if (canBeSafelyCasted(set_element_type, key_column_type))
{
set_columns[set_element_index] = castColumn({set_column, set_element_type, {}}, key_column_type);
transformed_set_columns[set_element_index] = castColumn({set_column, set_element_type, {}}, key_column_type);
continue;
}
@ -1270,6 +1272,9 @@ bool KeyCondition::tryPrepareSetIndex(
const NullMap * set_column_null_map = nullptr;
// Keep a reference to the original set_column to ensure the data remains valid
ColumnPtr original_set_column = set_column;
if (isNullableOrLowCardinalityNullable(set_element_type))
{
if (WhichDataType(set_element_type).isLowCardinality())
@ -1279,16 +1284,28 @@ bool KeyCondition::tryPrepareSetIndex(
}
set_element_type = removeNullable(set_element_type);
const auto & set_column_nullable = assert_cast<const ColumnNullable &>(*set_column);
const auto & null_map_data = set_column_nullable.getNullMapData();
// Obtain the nullable column without reassigning set_column immediately
const auto * set_column_nullable = typeid_cast<const ColumnNullable *>(set_column.get());
if (!set_column_nullable)
return false;
const NullMap & null_map_data = set_column_nullable->getNullMapData();
if (!null_map_data.empty())
set_column_null_map = &null_map_data;
set_column = set_column_nullable.getNestedColumnPtr();
ColumnPtr nested_column = set_column_nullable->getNestedColumnPtr();
// Reassign set_column after we have obtained necessary references
set_column = nested_column;
}
auto nullable_set_column = castColumnAccurateOrNull({set_column, set_element_type, {}}, key_column_type);
const auto & nullable_set_column_typed = assert_cast<const ColumnNullable &>(*nullable_set_column);
const auto & nullable_set_column_null_map = nullable_set_column_typed.getNullMapData();
ColumnPtr nullable_set_column = castColumnAccurateOrNull({set_column, set_element_type, {}}, key_column_type);
const auto * nullable_set_column_typed = typeid_cast<const ColumnNullable *>(nullable_set_column.get());
if (!nullable_set_column_typed)
return false;
const NullMap & nullable_set_column_null_map = nullable_set_column_typed->getNullMapData();
size_t nullable_set_column_null_map_size = nullable_set_column_null_map.size();
IColumn::Filter filter(nullable_set_column_null_map_size);
@ -1303,19 +1320,21 @@ bool KeyCondition::tryPrepareSetIndex(
filter[i] = !nullable_set_column_null_map[i];
}
set_column = nullable_set_column_typed.filter(filter, 0);
set_column = nullable_set_column_typed->filter(filter, 0);
}
else
{
for (size_t i = 0; i < nullable_set_column_null_map_size; ++i)
filter[i] = !nullable_set_column_null_map[i];
set_column = nullable_set_column_typed.getNestedColumn().filter(filter, 0);
set_column = nullable_set_column_typed->getNestedColumn().filter(filter, 0);
}
set_columns[set_element_index] = std::move(set_column);
transformed_set_columns[set_element_index] = std::move(set_column);
}
set_columns = std::move(transformed_set_columns);
out.set_index = std::make_shared<MergeTreeSetIndex>(set_columns, std::move(indexes_mapping));
/// When not all key columns are used or when there are multiple elements in