[FIX](function) fix size function for array map (#23920)

Issue Number: close #xxx
now we use select size(map(1, 2)); which will make be core
and we can make size function handle array & map column both
This commit is contained in:
amory
2023-09-06 14:32:06 +08:00
committed by GitHub
parent 29e1c59cca
commit dc28878f0e
11 changed files with 1070 additions and 181 deletions

View File

@ -146,64 +146,6 @@ public:
}
};
class FunctionMapSize : public IFunction {
public:
static constexpr auto name = "map_size";
static FunctionPtr create() { return std::make_shared<FunctionMapSize>(); }
/// Get function name.
String get_name() const override { return name; }
bool is_variadic() const override { return false; }
size_t get_number_of_arguments() const override { return 1; }
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
DataTypePtr datatype = arguments[0];
if (datatype->is_nullable()) {
datatype = assert_cast<const DataTypeNullable*>(datatype.get())->get_nested_type();
}
DCHECK(is_map(datatype)) << "first argument for function: " << name
<< " should be DataTypeMap";
return std::make_shared<DataTypeInt64>();
}
Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
size_t result, size_t input_rows_count) override {
const auto& [left_column, left_const] =
unpack_if_const(block.get_by_position(arguments[0]).column);
const ColumnMap* map_column = nullptr;
// const UInt8* map_null_map = nullptr;
if (left_column->is_nullable()) {
auto nullable_column = reinterpret_cast<const ColumnNullable*>(left_column.get());
map_column = check_and_get_column<ColumnMap>(nullable_column->get_nested_column());
// map_null_map = nullable_column->get_null_map_column().get_data().data();
} else {
map_column = check_and_get_column<ColumnMap>(*left_column.get());
}
if (!map_column) {
return Status::RuntimeError("unsupported types for function {}({})", get_name(),
block.get_by_position(arguments[0]).type->get_name());
}
auto dst_column = ColumnInt64::create(input_rows_count);
auto& dst_data = dst_column->get_data();
if (left_const) {
for (size_t i = 0; i < map_column->size(); i++) {
dst_data[i] = map_column->size_at(0);
}
} else {
for (size_t i = 0; i < map_column->size(); i++) {
dst_data[i] = map_column->size_at(i);
}
}
block.replace_by_position(result, std::move(dst_column));
return Status::OK();
}
};
template <bool is_key>
class FunctionMapContains : public IFunction {
public:
@ -354,7 +296,6 @@ public:
void register_function_map(SimpleFunctionFactory& factory) {
factory.register_function<FunctionMap>();
factory.register_function<FunctionMapSize>();
factory.register_function<FunctionMapContains<true>>();
factory.register_function<FunctionMapContains<false>>();
factory.register_function<FunctionMapEntries<true>>();