Skip to content

Commit

Permalink
[Test](function) not unpack when else column is const null in IF func…
Browse files Browse the repository at this point in the history 
…tion (apache#22419)
  • Loading branch information
@zhangstar333
zhangstar333 authored Aug 7, 2023
1 parent 1847e44 commit af8774c
Showing 1 changed file with 29 additions and 40 deletions.
69 changes: 29 additions & 40 deletions be/src/vec/functions/if.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -157,11 +157,13 @@ class FunctionIf : public IFunction {
}

DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
DataTypePtr type = nullptr;
get_least_supertype(DataTypes {arguments[1], arguments[2]}, &type);
DCHECK_NE(type, nullptr) << " arguments[1]: " << arguments[1]->get_name()
<< " arguments[2]: " << arguments[2]->get_name();
return type;
// if return type is custom, one of nullable return type will be nullable
bool nullable = arguments[1]->is_nullable() || arguments[2]->is_nullable();
if (nullable) {
return make_nullable(arguments[1]);
} else {
return arguments[1];
}
}

static ColumnPtr materialize_column_if_const(const ColumnPtr& column) {
Expand Down Expand Up @@ -293,7 +295,9 @@ class FunctionIf : public IFunction {
bool then_is_null = arg_then.column->only_null();
bool else_is_null = arg_else.column->only_null();

if (!then_is_null && !else_is_null) return false;
if (!then_is_null && !else_is_null) {
return false;
}

if (then_is_null && else_is_null) {
block.get_by_position(result).column =
Expand All @@ -309,24 +313,24 @@ class FunctionIf : public IFunction {
/// If then is NULL, we create Nullable column with null mask OR-ed with condition.
if (then_is_null) {
if (cond_col) {
if (is_column_nullable(*arg_else.column)) {
if (is_column_nullable(*arg_else.column)) { // if(cond, null, nullable)
auto arg_else_column = arg_else.column;
auto result_column = (*std::move(arg_else_column)).mutate();
assert_cast<ColumnNullable&>(*result_column)
.apply_null_map(assert_cast<const ColumnUInt8&>(*arg_cond.column));
block.replace_by_position(result, std::move(result_column));
} else {
} else { // if(cond, null, not_nullable)
block.replace_by_position(
result,
ColumnNullable::create(materialize_column_if_const(arg_else.column),
arg_cond.column));
}
} else if (cond_const_col) {
if (cond_const_col->get_value<UInt8>()) {
if (cond_const_col->get_value<UInt8>()) { // if(true, null, else)
block.get_by_position(result).column =
block.get_by_position(result).type->create_column()->clone_resized(
input_rows_count);
} else {
} else { // if(false, null, else)
block.get_by_position(result).column =
make_nullable_column_if_not(arg_else.column);
}
Expand All @@ -353,24 +357,24 @@ class FunctionIf : public IFunction {
negated_null_map_data[i] = !null_map_data[i];
}

if (is_column_nullable(*arg_then.column)) {
if (is_column_nullable(*arg_then.column)) { // if(cond, nullable, NULL)
auto arg_then_column = arg_then.column;
auto result_column = (*std::move(arg_then_column)).mutate();
assert_cast<ColumnNullable&>(*result_column)
.apply_negated_null_map(
assert_cast<const ColumnUInt8&>(*arg_cond.column));
block.replace_by_position(result, std::move(result_column));
} else {
} else { // if(cond, not_nullable, NULL)
block.replace_by_position(
result,
ColumnNullable::create(materialize_column_if_const(arg_then.column),
std::move(negated_null_map)));
}
} else if (cond_const_col) {
if (cond_const_col->get_value<UInt8>()) {
if (cond_const_col->get_value<UInt8>()) { // if(true, then, NULL)
block.get_by_position(result).column =
make_nullable_column_if_not(arg_then.column);
} else {
} else { // if(false, then, NULL)
block.get_by_position(result).column =
block.get_by_position(result).type->create_column()->clone_resized(
input_rows_count);
Expand All @@ -395,7 +399,9 @@ class FunctionIf : public IFunction {
auto* then_is_nullable = check_and_get_column<ColumnNullable>(*arg_then.column);
auto* else_is_nullable = check_and_get_column<ColumnNullable>(*arg_else.column);

if (!then_is_nullable && !else_is_nullable) return false;
if (!then_is_nullable && !else_is_nullable) {
return false;
}

/** Calculate null mask of result and nested column separately.
*/
Expand Down Expand Up @@ -439,6 +445,7 @@ class FunctionIf : public IFunction {
}

bool execute_for_null_condition(FunctionContext* context, Block& block,
const ColumnNumbers& arguments,
const ColumnWithTypeAndName& arg_cond,
const ColumnWithTypeAndName& arg_then,
const ColumnWithTypeAndName& arg_else, size_t result) {
Expand All @@ -459,19 +466,13 @@ class FunctionIf : public IFunction {
((ColumnVector<UInt8>&)(nullable->get_nested_column())).get_data().data();
auto rows = nullable->size();
for (size_t i = 0; i < rows; i++) {
nested_bool_data[i] = null_map[i] ? false : nested_bool_data[i];
nested_bool_data[i] = null_map[i] ? 0 : nested_bool_data[i];
}
auto column_size = block.columns();
block.insert({nullable->get_nested_column_ptr(), remove_nullable(arg_cond.type),
arg_cond.name});

Block temporary_block {{nullable->get_nested_column_ptr(),
remove_nullable(arg_cond.type), arg_cond.name},
arg_then,
arg_else,
block.get_by_position(result)};

execute_impl(context, temporary_block, {0, 1, 2}, 3, rows);

block.get_by_position(result).column =
std::move(temporary_block.get_by_position(3).column);
execute_impl(context, block, {column_size, arguments[1], arguments[2]}, result, rows);
return true;
}
return false;
Expand All @@ -493,21 +494,9 @@ class FunctionIf : public IFunction {
cond_column.column = materialize_column_if_const(cond_column.column);
const ColumnWithTypeAndName& arg_cond = block.get_by_position(arguments[0]);

if (auto* then_is_const = check_and_get_column<ColumnConst>(*arg_then.column)) {
if (check_and_get_column<ColumnNullable>(then_is_const->get_data_column())) {
ColumnWithTypeAndName& then_column = block.get_by_position(arguments[1]);
then_column.column = materialize_column_if_const(then_column.column);
}
}
if (auto* else_is_const = check_and_get_column<ColumnConst>(*arg_else.column)) {
if (check_and_get_column<ColumnNullable>(else_is_const->get_data_column())) {
ColumnWithTypeAndName& else_column = block.get_by_position(arguments[2]);
else_column.column = materialize_column_if_const(else_column.column);
}
}

Status ret = Status::OK();
if (execute_for_null_condition(context, block, arg_cond, arg_then, arg_else, result) ||
if (execute_for_null_condition(context, block, arguments, arg_cond, arg_then, arg_else,
result) ||
execute_for_null_then_else(context, block, arg_cond, arg_then, arg_else, result,
input_rows_count, ret) ||
execute_for_nullable_then_else(context, block, arg_cond, arg_then, arg_else, result,
Expand Down

0 comments on commit af8774c

Please sign in to comment.