optimize array with mutiple arguments

This commit is contained in:
taiyang-li 2024-08-04 13:32:17 +08:00
parent e5d4eaa458
commit a498259ec8

View File

@ -1,10 +1,14 @@
#include <Functions/IFunction.h>
#include <Functions/FunctionFactory.h>
#include <Columns/ColumnArray.h>
#include <Columns/ColumnDecimal.h>
#include <Columns/ColumnString.h>
#include <DataTypes/DataTypeArray.h>
#include <DataTypes/getLeastSupertype.h>
#include <Columns/ColumnArray.h>
#include <Interpreters/castColumn.h>
#include <Functions/FunctionFactory.h>
#include <Functions/IFunction.h>
#include <Interpreters/Context.h>
#include <Interpreters/castColumn.h>
#include "Columns/ColumnFixedString.h"
#include "Columns/ColumnNullable.h"
namespace DB
@ -43,7 +47,7 @@ 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)
{
@ -81,42 +85,194 @@ public:
IColumn & out_data = out->getData();
IColumn::Offsets & out_offsets = out->getOffsets();
out_data.reserve(input_rows_count * num_elements);
/// Fill out_offsets
out_offsets.resize_exact(input_rows_count);
IColumn::Offset current_offset = 0;
if (num_elements == 1)
{
for (size_t i = 0; i < input_rows_count; ++i)
{
++current_offset;
out_offsets[i] = current_offset;
}
out_data.insertRangeFrom(*column_ptrs[0], 0, input_rows_count);
}
else
{
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
executeAny(column_ptrs, out_data, input_rows_count);
return out;
}
private:
bool executeAny(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
{
return false // NOLINT
|| 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);
}
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;
}
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 & concreate_out_data = assert_cast<ColumnString &>(out_data);
auto & out_chars = concreate_out_data.getChars();
auto & out_offsets = concreate_out_data.getOffsets();
out_chars.resize_exact(total_bytes);
out_offsets.resize_exact(input_rows_count * columns.size());
size_t curr_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[curr_out_offset], ref.data, ref.size);
out_chars[curr_out_offset + ref.size] = 0;
curr_out_offset += ref.size + 1;
out_offsets[base + col_i] = curr_out_offset;
}
}
return true;
}
bool executeFixedString(const ColumnRawPtrs & columns, IColumn & out_data, size_t input_rows_count) const
{
size_t total_bytes = 0;
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;
total_bytes += concrete_column->getChars().size();
concrete_columns[i] = concrete_column;
}
ColumnFixedString & concreate_out_data = assert_cast<ColumnFixedString &>(out_data);
auto & out_chars = concreate_out_data.getChars();
out_chars.resize_exact(total_bytes);
const size_t n = concreate_out_data.getN();
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 & concreate_out_data = assert_cast<ColumnNullable &>(out_data);
auto & out_null_map = concreate_out_data.getNullMapColumn();
auto & out_nested_column = concreate_out_data.getNestedColumn();
executeAny(null_maps, out_null_map, input_rows_count);
executeAny(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 * concreate_out_data = typeid_cast<ColumnTuple *>(&out_data);
if (!concreate_out_data)
return false;
const size_t tuple_size = concreate_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);
}
executeAny(elem_columns, concreate_out_data->getColumn(i), input_rows_count);
}
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;
};