Skip to content

Commit

Permalink
1
Browse files Browse the repository at this point in the history
  • Loading branch information
@wyxxxcat
wyxxxcat committed Sep 23, 2024
1 parent 6bceb32 commit 0ca3187
Show file tree
Hide file tree
Showing 4 changed files with 61 additions and 19 deletions.
55 changes: 39 additions & 16 deletions be/src/vec/functions/function_string.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -389,7 +389,7 @@ class FunctionAutoPartitionName : public IFunction {
String get_name() const override { return name; }
size_t get_number_of_arguments() const override { return 0; }
bool is_variadic() const override { return true; }

bool use_default_implementation_for_nulls() const override { return false; }
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
return std::make_shared<DataTypeString>();
}
Expand All @@ -401,19 +401,36 @@ class FunctionAutoPartitionName : public IFunction {
return Status::InvalidArgument(
"auto_partition_name must contains at least two arguments");
}
auto const_null_map = ColumnUInt8::create(input_rows_count, 0);
auto null_map = ColumnUInt8::create(input_rows_count, 0);
std::vector<const ColumnString::Chars*> chars_list(argument_size);
std::vector<const ColumnString::Offsets*> offsets_list(argument_size);
std::vector<bool> is_const_args(argument_size);
std::vector<const ColumnUInt8::Container*> null_list(argument_size);
std::vector<ColumnPtr> argument_null_columns(argument_size);

std::vector<ColumnPtr> argument_columns(argument_size);
for (int i = 0; i < argument_size; ++i) {
argument_columns[i] =
block.get_by_position(arguments[i]).column->convert_to_full_column_if_const();
if (const auto* nullable =
check_and_get_column<const ColumnNullable>(*argument_columns[i])) {
null_list[i] = &nullable->get_null_map_data();
argument_null_columns[i] = nullable->get_null_map_column_ptr();
argument_columns[i] = nullable->get_nested_column_ptr();
} else {
null_list[i] = &const_null_map->get_data();
}

const auto& [col, is_const] =
unpack_if_const(block.get_by_position(arguments[i]).column);

const auto* col_str = assert_cast<const ColumnString*>(col.get());
const auto* col_str = assert_cast<const ColumnString*>(argument_columns[i].get());
chars_list[i] = &col_str->get_chars();
offsets_list[i] = &col_str->get_offsets();
is_const_args[i] = is_const;
}

auto res = ColumnString::create();
auto& res_data = res->get_chars();
auto& res_offset = res->get_offsets();
Expand All @@ -422,9 +439,9 @@ class FunctionAutoPartitionName : public IFunction {
const char* partition_type = chars_list[0]->raw_data();
// partition type is list|range
if (std::strncmp(partition_type, "list", 4) == 0) {
return _auto_partition_type_of_list(chars_list, offsets_list, is_const_args, res_data,
res_offset, input_rows_count, argument_size, block,
result, res);
return _auto_partition_type_of_list(chars_list, offsets_list, is_const_args, null_list,
res_data, res_offset, input_rows_count,
argument_size, block, result, res);
} else {
return _auto_partition_type_of_range(chars_list, offsets_list, is_const_args, res_data,
res_offset, input_rows_count, argument_size, block,
Expand Down Expand Up @@ -472,8 +489,9 @@ class FunctionAutoPartitionName : public IFunction {
}
Status _auto_partition_type_of_list(std::vector<const ColumnString::Chars*>& chars_list,
std::vector<const ColumnString::Offsets*>& offsets_list,
std::vector<bool>& is_const_args, auto& res_data,
auto& res_offset, size_t input_rows_count,
std::vector<bool>& is_const_args,
const std::vector<const ColumnUInt8::Container*>& null_list,
auto& res_data, auto& res_offset, size_t input_rows_count,
size_t argument_size, Block& block, size_t result,
auto& res) const {
int curr_len = 0;
Expand All @@ -484,16 +502,21 @@ class FunctionAutoPartitionName : public IFunction {
for (int col = 1; col < argument_size; col++) {
const auto& current_offsets = *offsets_list[col];
const auto& current_chars = *chars_list[col];
const auto& current_nullmap = *null_list[col];

auto idx = index_check_const(row, is_const_args[col]);
int size = current_offsets[idx] - current_offsets[idx - 1];
const char* raw_chars =
reinterpret_cast<const char*>(&current_chars[current_offsets[idx - 1]]);

// convert string to u16string in order to convert to unicode strings
const std::string raw_str(raw_chars, size);
auto u16string = _string_to_u16string(raw_str);
res_p += _string_to_unicode(u16string) + std::to_string(u16string.size());
if (current_nullmap[row]) {
res_p += 'X';
} else {
auto idx = index_check_const(row, is_const_args[col]);

int size = current_offsets[idx] - current_offsets[idx - 1];
const char* raw_chars =
reinterpret_cast<const char*>(&current_chars[current_offsets[idx - 1]]);
// convert string to u16string in order to convert to unicode strings
const std::string raw_str(raw_chars, size);
auto u16string = _string_to_u16string(raw_str);
res_p += _string_to_unicode(u16string) + std::to_string(u16string.size());
}
}

// check the name of length
Expand Down
5 changes: 2 additions & 3 deletions ...n/java/org/apache/doris/nereids/trees/expressions/functions/scalar/AutoPartitionName.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,9 @@
import org.apache.doris.catalog.FunctionSignature;
import org.apache.doris.nereids.exceptions.AnalysisException;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.functions.AlwaysNotNullable;
import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
import org.apache.doris.nereids.trees.expressions.functions.PropagateNullable;
import org.apache.doris.nereids.trees.expressions.literal.VarcharLiteral;
import org.apache.doris.nereids.trees.expressions.shape.UnaryExpression;
import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
import org.apache.doris.nereids.types.StringType;
import org.apache.doris.nereids.types.VarcharType;
Expand All @@ -40,7 +39,7 @@
* GenerateFunction.
*/
public class AutoPartitionName extends ScalarFunction
implements UnaryExpression, ExplicitlyCastableSignature, PropagateNullable {
implements ExplicitlyCastableSignature, AlwaysNotNullable {

public static final List<FunctionSignature> SIGNATURES = ImmutableList.of(
FunctionSignature.ret(VarcharType.SYSTEM_DEFAULT).varArgs(VarcharType.SYSTEM_DEFAULT),
Expand Down
16 changes: 16 additions & 0 deletions regression-test/data/nereids_function_p0/scalar_function/A.out
View file
Original file line number Diff line number Diff line change
Expand Up @@ -484,6 +484,16 @@ pchar116varchar119
pchar126varchar129
pchar136varchar139

-- !sql_auto_partition_name_list_column --
pchar116Xvarchar119
pchar126Xvarchar129
pchar136Xvarchar139

-- !sql_auto_partition_name_list_column --
pXstring17XX
pXstring27XX
pXstring37XX

-- !sql_auto_partition_name_list_literal_empty --
p0

Expand All @@ -496,6 +506,12 @@ p_2dhello6
-- !sql_auto_partition_name_list_literal_mixed --
p4023ffe5257e7cworld111112e2e2e2e20

-- !sql_auto_partition_name_list_literal_mixed --
plist4X11X

-- !sql_auto_partition_name_list_literal_null --
pXXX

-- !sql_auto_partition_name_range_literal_notnull --
p20221212000000

Expand Down
4 changes: 4 additions & 0 deletions regression-test/suites/nereids_function_p0/scalar_function/A.groovy
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,10 +80,14 @@ suite("nereids_scalar_fn_A") {
qt_sql_auto_partition_name_list_column_type_mixed "select auto_partition_name('list', kchrs1, kbool) from fn_test_not_nullable where id < 3 order by kchrs1"
qt_sql_auto_partition_name_list_column_literal_mixed "select auto_partition_name('list', kstr, 'hello') from fn_test_not_nullable where id < 3 order by kstr"
qt_sql_auto_partition_name_list_column "select auto_partition_name('list', kchrs1, kvchrs1) from fn_test_not_nullable where id < 3 order by kchrs1"
qt_sql_auto_partition_name_list_column "select auto_partition_name('list', kchrs1, null, kvchrs1) from fn_test where id < 3 order by kchrs1"
qt_sql_auto_partition_name_list_column "select auto_partition_name('list', null, kstr, null, null) from fn_test where id < 3 order by kstr"
qt_sql_auto_partition_name_list_literal_empty "select auto_partition_name('list', '')"
qt_sql_auto_partition_name_list_literal_mixed "select auto_partition_name('list', '你好', true, false)"
qt_sql_auto_partition_name_list_literal_mixed "select auto_partition_name('list', '-hello')"
qt_sql_auto_partition_name_list_literal_mixed "select auto_partition_name('list', '@#￥%~|world11111....')"
qt_sql_auto_partition_name_list_literal_mixed "select auto_partition_name('list', 'list', null, true, null)"
qt_sql_auto_partition_name_list_literal_null "select auto_partition_name('list', null, null, null);"
qt_sql_auto_partition_name_range_literal_notnull "select auto_partition_name('range', 'day', '2022-12-12 19:20:30')"
qt_sql_auto_partition_name_range_literal_notnull "select auto_partition_name('range', 'month', '2022-12-12 19:20:30')"
qt_sql_auto_partition_name_range_literal_notnull "select auto_partition_name('range', 'year', '2022-12-12 19:20:30')"
Expand Down

0 comments on commit 0ca3187

Please sign in to comment.