Merge pull request #67707 from bigo-sg/opt_array_map

Optimize function `array()` and `map()`
This commit is contained in:
Robert Schulze 2024-09-09 14:37:50 +00:00 committed by GitHub
commit d749869f00
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 217 additions and 63 deletions

View File

@ -1,11 +1,15 @@
#include <Functions/IFunction.h>
#include <Functions/FunctionFactory.h>
#include <Columns/ColumnArray.h>
#include <Columns/ColumnDecimal.h>
#include <Columns/ColumnFixedString.h>
#include <Columns/ColumnNullable.h>
#include <Columns/ColumnString.h>
#include <Core/Settings.h>
#include <DataTypes/DataTypeArray.h>
#include <DataTypes/getLeastSupertype.h>
#include <Columns/ColumnArray.h>
#include <Core/Settings.h>
#include <Interpreters/castColumn.h>
#include <Functions/FunctionFactory.h>
#include <Functions/IFunction.h>
#include <Interpreters/Context.h>
#include <Interpreters/castColumn.h>
namespace DB
@ -44,11 +48,13 @@ public:
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count) const override
{
size_t num_elements = arguments.size();
const size_t num_elements = arguments.size();
if (num_elements == 0)
{
/// We should return constant empty array.
return result_type->createColumnConstWithDefaultValue(input_rows_count);
}
const DataTypePtr & elem_type = static_cast<const DataTypeArray &>(*result_type).getNestedType();
@ -60,7 +66,6 @@ public:
Columns columns_holder(num_elements);
ColumnRawPtrs column_ptrs(num_elements);
for (size_t i = 0; i < num_elements; ++i)
{
const auto & arg = arguments[i];
@ -77,35 +82,199 @@ public:
}
/// Create and fill the result array.
auto out = ColumnArray::create(elem_type->createColumn());
IColumn & out_data = out->getData();
IColumn::Offsets & out_offsets = out->getOffsets();
out_data.reserve(input_rows_count * num_elements);
out_offsets.resize(input_rows_count);
/// Fill out_offsets
out_offsets.resize_exact(input_rows_count);
IColumn::Offset current_offset = 0;
for (size_t i = 0; i < input_rows_count; ++i)
{
for (size_t j = 0; j < num_elements; ++j)
out_data.insertFrom(*column_ptrs[j], i);
current_offset += num_elements;
out_offsets[i] = current_offset;
}
/// Fill out_data
out_data.reserve(input_rows_count * num_elements);
if (num_elements == 1)
out_data.insertRangeFrom(*column_ptrs[0], 0, input_rows_count);
else
execute(column_ptrs, out_data, input_rows_count);
return out;
}
private:
bool execute(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
{
return executeNumber<UInt8>(columns, out_data, input_rows_count) || executeNumber<UInt16>(columns, out_data, input_rows_count)
|| executeNumber<UInt32>(columns, out_data, input_rows_count) || executeNumber<UInt64>(columns, out_data, input_rows_count)
|| executeNumber<UInt128>(columns, out_data, input_rows_count) || executeNumber<UInt256>(columns, out_data, input_rows_count)
|| executeNumber<Int8>(columns, out_data, input_rows_count) || executeNumber<Int16>(columns, out_data, input_rows_count)
|| executeNumber<Int32>(columns, out_data, input_rows_count) || executeNumber<Int64>(columns, out_data, input_rows_count)
|| executeNumber<Int128>(columns, out_data, input_rows_count) || executeNumber<Int256>(columns, out_data, input_rows_count)
|| executeNumber<Float32>(columns, out_data, input_rows_count) || executeNumber<Float64>(columns, out_data, input_rows_count)
|| executeNumber<Decimal32>(columns, out_data, input_rows_count)
|| executeNumber<Decimal64>(columns, out_data, input_rows_count)
|| executeNumber<Decimal128>(columns, out_data, input_rows_count)
|| executeNumber<Decimal256>(columns, out_data, input_rows_count)
|| executeNumber<DateTime64>(columns, out_data, input_rows_count) || executeString(columns, out_data, input_rows_count)
|| executeNullable(columns, out_data, input_rows_count) || executeTuple(columns, out_data, input_rows_count)
|| executeFixedString(columns, out_data, input_rows_count) || executeGeneric(columns, out_data, input_rows_count);
}
template <typename T>
bool executeNumber(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
{
using Container = ColumnVectorOrDecimal<T>::Container;
std::vector<const Container *> containers(columns.size(), nullptr);
for (size_t i = 0; i < columns.size(); ++i)
{
const ColumnVectorOrDecimal<T> * concrete_column = checkAndGetColumn<ColumnVectorOrDecimal<T>>(columns[i]);
if (!concrete_column)
return false;
containers[i] = &concrete_column->getData();
}
ColumnVectorOrDecimal<T> & concrete_out_data = assert_cast<ColumnVectorOrDecimal<T> &>(out_data);
Container & out_container = concrete_out_data.getData();
out_container.resize_exact(columns.size() * input_rows_count);
for (size_t row_i = 0; row_i < input_rows_count; ++row_i)
{
const size_t base = row_i * columns.size();
for (size_t col_i = 0; col_i < columns.size(); ++col_i)
out_container[base + col_i] = (*containers[col_i])[row_i];
}
return true;
}
bool executeString(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
{
size_t total_bytes = 0;
std::vector<const ColumnString *> concrete_columns(columns.size(), nullptr);
for (size_t i = 0; i < columns.size(); ++i)
{
const ColumnString * concrete_column = checkAndGetColumn<ColumnString>(columns[i]);
if (!concrete_column)
return false;
total_bytes += concrete_column->getChars().size();
concrete_columns[i] = concrete_column;
}
ColumnString & concrete_out_data = assert_cast<ColumnString &>(out_data);
auto & out_chars = concrete_out_data.getChars();
auto & out_offsets = concrete_out_data.getOffsets();
out_chars.resize_exact(total_bytes);
out_offsets.resize_exact(input_rows_count * columns.size());
size_t cur_out_offset = 0;
for (size_t row_i = 0; row_i < input_rows_count; ++row_i)
{
const size_t base = row_i * columns.size();
for (size_t col_i = 0; col_i < columns.size(); ++col_i)
{
StringRef ref = concrete_columns[col_i]->getDataAt(row_i);
memcpySmallAllowReadWriteOverflow15(&out_chars[cur_out_offset], ref.data, ref.size);
out_chars[cur_out_offset + ref.size] = 0;
cur_out_offset += ref.size + 1;
out_offsets[base + col_i] = cur_out_offset;
}
}
return true;
}
bool executeFixedString(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
{
std::vector<const ColumnFixedString *> concrete_columns(columns.size(), nullptr);
for (size_t i = 0; i < columns.size(); ++i)
{
const ColumnFixedString * concrete_column = checkAndGetColumn<ColumnFixedString>(columns[i]);
if (!concrete_column)
return false;
concrete_columns[i] = concrete_column;
}
ColumnFixedString & concrete_out_data = assert_cast<ColumnFixedString &>(out_data);
auto & out_chars = concrete_out_data.getChars();
const size_t n = concrete_out_data.getN();
size_t total_bytes = n * columns.size() * input_rows_count;
out_chars.resize_exact(total_bytes);
size_t curr_out_offset = 0;
for (size_t row_i = 0; row_i < input_rows_count; ++row_i)
{
for (size_t col_i = 0; col_i < columns.size(); ++col_i)
{
StringRef ref = concrete_columns[col_i]->getDataAt(row_i);
memcpySmallAllowReadWriteOverflow15(&out_chars[curr_out_offset], ref.data, n);
curr_out_offset += n;
}
}
return true;
}
bool executeNullable(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
{
ColumnRawPtrs null_maps(columns.size(), nullptr);
ColumnRawPtrs nested_columns(columns.size(), nullptr);
for (size_t i = 0; i < columns.size(); ++i)
{
const ColumnNullable * concrete_column = checkAndGetColumn<ColumnNullable>(columns[i]);
if (!concrete_column)
return false;
null_maps[i] = &concrete_column->getNullMapColumn();
nested_columns[i] = &concrete_column->getNestedColumn();
}
ColumnNullable & concrete_out_data = assert_cast<ColumnNullable &>(out_data);
auto & out_null_map = concrete_out_data.getNullMapColumn();
auto & out_nested_column = concrete_out_data.getNestedColumn();
execute(null_maps, out_null_map, input_rows_count);
execute(nested_columns, out_nested_column, input_rows_count);
return true;
}
bool executeTuple(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
{
ColumnTuple * concrete_out_data = typeid_cast<ColumnTuple *>(&out_data);
if (!concrete_out_data)
return false;
const size_t tuple_size = concrete_out_data->tupleSize();
for (size_t i = 0; i < tuple_size; ++i)
{
ColumnRawPtrs elem_columns(columns.size(), nullptr);
for (size_t j = 0; j < columns.size(); ++j)
{
const ColumnTuple * concrete_column = assert_cast<const ColumnTuple *>(columns[j]);
elem_columns[j] = &concrete_column->getColumn(i);
}
execute(elem_columns, concrete_out_data->getColumn(i), input_rows_count);
}
return true;
}
bool executeGeneric(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
{
for (size_t i = 0; i < input_rows_count; ++i)
for (const auto * column : columns)
out_data.insertFrom(*column, i);
return true;
}
String getName() const override
{
return name;
}
bool addField(DataTypePtr type_res, const Field & f, Array & arr) const;
bool use_variant_as_common_type = false;
};

View File

@ -2,6 +2,8 @@
#include <Columns/ColumnMap.h>
#include <Columns/ColumnNullable.h>
#include <Columns/ColumnsCommon.h>
#include <Columns/ColumnsNumber.h>
#include <Core/Settings.h>
#include <DataTypes/DataTypeArray.h>
#include <DataTypes/DataTypeMap.h>
#include <DataTypes/DataTypeTuple.h>
@ -13,7 +15,6 @@
#include <Interpreters/Context.h>
#include <Interpreters/castColumn.h>
#include <Common/HashTable/HashSet.h>
#include <Core/Settings.h>
namespace DB
@ -36,11 +37,18 @@ class FunctionMap : public IFunction
public:
static constexpr auto name = "map";
explicit FunctionMap(bool use_variant_as_common_type_) : use_variant_as_common_type(use_variant_as_common_type_) {}
explicit FunctionMap(ContextPtr context_)
: context(context_)
, use_variant_as_common_type(
context->getSettingsRef().allow_experimental_variant_type && context->getSettingsRef().use_variant_as_common_type)
, function_array(FunctionFactory::instance().get("array", context))
, function_map_from_arrays(FunctionFactory::instance().get("mapFromArrays", context))
{
}
static FunctionPtr create(ContextPtr context)
{
return std::make_shared<FunctionMap>(context->getSettingsRef().allow_experimental_variant_type && context->getSettingsRef().use_variant_as_common_type);
return std::make_shared<FunctionMap>(context);
}
String getName() const override
@ -101,62 +109,38 @@ public:
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count) const override
{
size_t num_elements = arguments.size();
if (num_elements == 0)
return result_type->createColumnConstWithDefaultValue(input_rows_count);
ColumnsWithTypeAndName key_args;
ColumnsWithTypeAndName value_args;
for (size_t i = 0; i < num_elements; i += 2)
{
key_args.emplace_back(arguments[i]);
value_args.emplace_back(arguments[i+1]);
}
const auto & result_type_map = static_cast<const DataTypeMap &>(*result_type);
const DataTypePtr & key_type = result_type_map.getKeyType();
const DataTypePtr & value_type = result_type_map.getValueType();
const DataTypePtr & key_array_type = std::make_shared<DataTypeArray>(key_type);
const DataTypePtr & value_array_type = std::make_shared<DataTypeArray>(value_type);
Columns columns_holder(num_elements);
ColumnRawPtrs column_ptrs(num_elements);
/// key_array = array(args[0], args[2]...)
ColumnPtr key_array = function_array->build(key_args)->execute(key_args, key_array_type, input_rows_count);
/// value_array = array(args[1], args[3]...)
ColumnPtr value_array = function_array->build(value_args)->execute(value_args, value_array_type, input_rows_count);
for (size_t i = 0; i < num_elements; ++i)
{
const auto & arg = arguments[i];
const auto to_type = i % 2 == 0 ? key_type : value_type;
ColumnPtr preprocessed_column = castColumn(arg, to_type);
preprocessed_column = preprocessed_column->convertToFullColumnIfConst();
columns_holder[i] = std::move(preprocessed_column);
column_ptrs[i] = columns_holder[i].get();
}
/// Create and fill the result map.
MutableColumnPtr keys_data = key_type->createColumn();
MutableColumnPtr values_data = value_type->createColumn();
MutableColumnPtr offsets = DataTypeNumber<IColumn::Offset>().createColumn();
size_t total_elements = input_rows_count * num_elements / 2;
keys_data->reserve(total_elements);
values_data->reserve(total_elements);
offsets->reserve(input_rows_count);
IColumn::Offset current_offset = 0;
for (size_t i = 0; i < input_rows_count; ++i)
{
for (size_t j = 0; j < num_elements; j += 2)
{
keys_data->insertFrom(*column_ptrs[j], i);
values_data->insertFrom(*column_ptrs[j + 1], i);
}
current_offset += num_elements / 2;
offsets->insert(current_offset);
}
auto nested_column = ColumnArray::create(
ColumnTuple::create(Columns{std::move(keys_data), std::move(values_data)}),
std::move(offsets));
return ColumnMap::create(nested_column);
/// result = mapFromArrays(key_array, value_array)
ColumnsWithTypeAndName map_args{{key_array, key_array_type, ""}, {value_array, value_array_type, ""}};
return function_map_from_arrays->build(map_args)->execute(map_args, result_type, input_rows_count);
}
private:
ContextPtr context;
bool use_variant_as_common_type = false;
FunctionOverloadResolverPtr function_array;
FunctionOverloadResolverPtr function_map_from_arrays;
};
/// mapFromArrays(keys, values) is a function that allows you to make key-value pair from a pair of arrays or maps
@ -173,6 +157,7 @@ public:
bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
bool useDefaultImplementationForNulls() const override { return false; }
bool useDefaultImplementationForConstants() const override { return true; }
bool useDefaultImplementationForLowCardinalityColumns() const override { return false; }
DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
{