diff --git a/be/src/exprs/function/array/function_array_flatten.cpp b/be/src/exprs/function/array/function_array_flatten.cpp index d5f5f25315642b..263b6368bfe6e8 100644 --- a/be/src/exprs/function/array/function_array_flatten.cpp +++ b/be/src/exprs/function/array/function_array_flatten.cpp @@ -57,7 +57,7 @@ class FunctionArrayFlatten : public IFunction { block.get_by_position(arguments[0]).column->convert_to_full_column_if_const(); auto* src_column_array_ptr = assert_cast(remove_nullable(src_column).get()); - const ColumnArray* nested_src_column_array_ptr = src_column_array_ptr; + ColumnPtr nested_src_column_array = remove_nullable(src_column); DataTypePtr src_data_type = block.get_by_position(arguments[0]).type; auto* src_data_type_array = @@ -69,22 +69,42 @@ class FunctionArrayFlatten : public IFunction { .data(); while (src_data_type_array->get_nested_type()->get_primitive_type() == TYPE_ARRAY) { - nested_src_column_array_ptr = assert_cast( - remove_nullable(src_column_array_ptr->get_data_ptr()).get()); + const auto* nullable = + check_and_get_column(src_column_array_ptr->get_data()); + auto nested_array = remove_nullable(src_column_array_ptr->get_data_ptr()); + const auto& array = assert_cast(*nested_array); + + // IF can leave [101, 102] beneath a NULL subarray. Remove that range before + // composing offsets, so [NULL, [9]] flattens to [9] at every nesting level. + if (nullable && nullable->has_null()) { + auto data = array.get_data().clone_empty(); + auto nested_offsets = ColumnArray::ColumnOffsets::create(); + auto& nested_offsets_data = nested_offsets->get_data(); + nested_offsets_data.reserve(array.size()); + for (size_t i = 0; i < array.size(); ++i) { + // NULL subarrays may own nonempty ranges; their children are invisible. + if (!nullable->is_null_at(i)) { + data->insert_range_from(array.get_data(), array.offset_at(i), + array.size_at(i)); + } + nested_offsets_data.push_back(data->size()); + } + nested_array = ColumnArray::create(std::move(data), std::move(nested_offsets)); + } + const auto& nested_array_offsets = + assert_cast(*nested_array).get_offsets(); for (size_t i = 0; i < input_rows_count; ++i) { - offsets[i] = nested_src_column_array_ptr->get_offsets()[offsets[i] - 1]; + offsets[i] = nested_array_offsets[offsets[i] - 1]; } - src_column_array_ptr = nested_src_column_array_ptr; + nested_src_column_array = std::move(nested_array); + src_column_array_ptr = assert_cast(nested_src_column_array.get()); src_data_type_array = assert_cast( remove_nullable(src_data_type_array->get_nested_type()).get()); } - block.replace_by_position( - result, ColumnArray::create(assert_cast( - nested_src_column_array_ptr->get_data()) - .clone(), - std::move(result_column_offsets))); + block.replace_by_position(result, ColumnArray::create(src_column_array_ptr->get_data_ptr(), + std::move(result_column_offsets))); return Status::OK(); } }; diff --git a/be/test/exprs/function/function_array_flatten_test.cpp b/be/test/exprs/function/function_array_flatten_test.cpp new file mode 100644 index 00000000000000..4f811acb09c112 --- /dev/null +++ b/be/test/exprs/function/function_array_flatten_test.cpp @@ -0,0 +1,153 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include + +#include +#include +#include + +#include "common/status.h" +#include "core/assert_cast.h" +#include "core/block/block.h" +#include "core/column/column_array.h" +#include "core/column/column_const.h" +#include "core/column/column_nullable.h" +#include "core/column/column_vector.h" +#include "core/data_type/data_type_array.h" +#include "core/data_type/data_type_nullable.h" +#include "core/data_type/data_type_number.h" +#include "exprs/function/simple_function_factory.h" +#include "testutil/function_utils.h" + +namespace doris { +namespace { + +ColumnPtr nullable_column(ColumnPtr data, std::initializer_list nulls) { + auto null_map = ColumnUInt8::create(); + null_map->get_data().assign(nulls.begin(), nulls.end()); + return ColumnNullable::create(std::move(data), std::move(null_map)); +} + +ColumnPtr array_column(ColumnPtr data, std::initializer_list offsets) { + auto offsets_column = ColumnArray::ColumnOffsets::create(); + offsets_column->get_data().assign(offsets.begin(), offsets.end()); + return ColumnArray::create(std::move(data), std::move(offsets_column)); +} + +ColumnPtr int_arrays(std::initializer_list values, std::initializer_list nulls, + std::initializer_list offsets) { + auto data = ColumnInt32::create(); + data->get_data().assign(values.begin(), values.end()); + return array_column(nullable_column(std::move(data), nulls), offsets); +} + +DataTypePtr array_type(size_t depth) { + DataTypePtr type = std::make_shared(); + for (size_t i = 0; i < depth; ++i) { + type = std::make_shared(make_nullable(type)); + } + return type; +} + +void expect_flatten(ColumnPtr source, DataTypePtr source_type, ColumnPtr expected) { + Block block; + const auto rows = source->size(); + auto result_type = array_type(1); + if (source_type->is_nullable()) { + result_type = make_nullable(result_type); + } + block.insert({std::move(source), source_type, "source"}); + auto function = SimpleFunctionFactory::instance().get_function( + "array_flatten", block.get_columns_with_type_and_name(), result_type); + ASSERT_NE(function, nullptr); + + FunctionUtils fn_utils(result_type, {source_type}, false); + auto* context = fn_utils.get_fn_ctx(); + ASSERT_TRUE(function->open(context, FunctionContext::FunctionStateScope::FRAGMENT_LOCAL).ok()); + ASSERT_TRUE(function->open(context, FunctionContext::FunctionStateScope::THREAD_LOCAL).ok()); + block.insert({nullptr, result_type, "result"}); + auto status = function->execute(context, block, {0}, 1, rows); + EXPECT_TRUE(function->close(context, FunctionContext::FunctionStateScope::THREAD_LOCAL).ok()); + EXPECT_TRUE(function->close(context, FunctionContext::FunctionStateScope::FRAGMENT_LOCAL).ok()); + ASSERT_TRUE(status.ok()) << status; + + auto result = block.get_by_position(1).column->convert_to_full_column_if_const(); + ASSERT_EQ(result->size(), expected->size()); + for (size_t i = 0; i < expected->size(); ++i) { + EXPECT_EQ(result->compare_at(i, i, *expected, 1), 0) << "row " << i; + } + const auto& result_array = assert_cast(*remove_nullable(result)); + EXPECT_EQ(result_array.get_offsets().back(), result_array.get_data().size()); +} + +} // namespace + +TEST(FunctionArrayFlattenTest, SkipsHiddenSubarrayElementsAcrossRows) { + // [NULL, [9]], [[201], [9]]; the first NULL still owns [101, 102]. + auto children = int_arrays({101, 102, 9, 201, 9}, {0, 0, 0, 0, 0}, {2, 3, 4, 5}); + auto source = array_column(nullable_column(std::move(children), {1, 0, 0, 0}), {2, 4}); + expect_flatten(std::move(source), array_type(2), int_arrays({9, 201, 9}, {0, 0, 0}, {1, 3})); +} + +TEST(FunctionArrayFlattenTest, SkipsNullParentsAndNullDescendants) { + // [NULL, [NULL, [9]]]; NULL parents and descendants both hide nonempty data. + auto children = int_arrays({101, 102, 9}, {0, 0, 0}, {1, 2, 3}); + auto parents = array_column(nullable_column(std::move(children), {0, 1, 0}), {1, 3}); + auto source = array_column(nullable_column(std::move(parents), {1, 0}), {2}); + expect_flatten(std::move(source), array_type(3), int_arrays({9}, {0}, {1})); +} + +TEST(FunctionArrayFlattenTest, HandlesEmptyRowsAndAllNullChildren) { + // [], [NULL, NULL], [[], [NULL, 9]], [] with hidden payload in both NULL arrays. + auto children = int_arrays({101, 102, 0, 9}, {0, 0, 1, 0}, {1, 2, 2, 4}); + auto source = array_column(nullable_column(std::move(children), {1, 1, 0, 0}), {0, 2, 4, 4}); + expect_flatten(std::move(source), array_type(2), int_arrays({0, 9}, {1, 0}, {0, 0, 2, 2})); +} + +TEST(FunctionArrayFlattenTest, PreservesOuterNullsThroughFunctionFramework) { + auto children = int_arrays({101, 9, 201, 9}, {0, 0, 0, 0}, {1, 2, 3, 4}); + auto arrays = array_column(nullable_column(std::move(children), {1, 0, 0, 0}), {2, 4}); + auto source = nullable_column(std::move(arrays), {1, 0}); + auto expected = nullable_column(int_arrays({201, 9}, {0, 0}, {0, 2}), {1, 0}); + expect_flatten(std::move(source), make_nullable(array_type(2)), std::move(expected)); +} + +TEST(FunctionArrayFlattenTest, HandlesConstantNestedArrayWithHiddenElements) { + auto children = int_arrays({101, 102, 9}, {0, 0, 0}, {2, 3}); + auto array = array_column(nullable_column(std::move(children), {1, 0}), {2}); + auto source = ColumnConst::create(std::move(array), 3); + expect_flatten(std::move(source), array_type(2), int_arrays({9, 9, 9}, {0, 0, 0}, {1, 2, 3})); +} + +TEST(FunctionArrayFlattenTest, HandlesOnlyNullInputThroughFunctionFramework) { + auto children = int_arrays({101, 9}, {0, 0}, {1, 2}); + auto arrays = array_column(nullable_column(std::move(children), {0, 0}), {1, 2}); + auto source = nullable_column(std::move(arrays), {1, 1}); + auto expected = nullable_column(int_arrays({}, {}, {0, 0}), {1, 1}); + expect_flatten(std::move(source), make_nullable(array_type(2)), std::move(expected)); +} + +TEST(FunctionArrayFlattenTest, PreservesArraysWithoutNullChildren) { + auto children = int_arrays({1, 0, 9}, {0, 1, 0}, {0, 2, 3}); + auto source = array_column(nullable_column(std::move(children), {0, 0, 0}), {0, 2, 3}); + expect_flatten(std::move(source), array_type(2), int_arrays({1, 0, 9}, {0, 1, 0}, {0, 2, 3})); + expect_flatten(int_arrays({1, 0, 9}, {0, 1, 0}, {0, 2, 3}), array_type(1), + int_arrays({1, 0, 9}, {0, 1, 0}, {0, 2, 3})); +} + +} // namespace doris diff --git a/regression-test/data/query_p0/sql_functions/array_functions/array_flatten.out b/regression-test/data/query_p0/sql_functions/array_functions/array_flatten.out index fa71438783db9f..af293f9ce8d80c 100644 --- a/regression-test/data/query_p0/sql_functions/array_functions/array_flatten.out +++ b/regression-test/data/query_p0/sql_functions/array_functions/array_flatten.out @@ -4,3 +4,11 @@ 2 \N \N [] \N ["b", null] 3 [1, 2, null] [null] [null, 2] [null, null, 3] [null, "aaaab", "ccc"] +-- !null_subarray -- +1 [null, [9]] [9] [9] +2 [[201], [9]] [201, 9] [9] + +-- !null_subarray_nested -- +1 [9] +2 [201, 9] + diff --git a/regression-test/suites/query_p0/sql_functions/array_functions/array_flatten.groovy b/regression-test/suites/query_p0/sql_functions/array_functions/array_flatten.groovy index 82d4a39a565204..06eaf999462cad 100644 --- a/regression-test/suites/query_p0/sql_functions/array_functions/array_flatten.groovy +++ b/regression-test/suites/query_p0/sql_functions/array_functions/array_flatten.groovy @@ -43,4 +43,33 @@ suite("array_flatten") { qt_test """ select k1, array_flatten(a1), array_flatten(aaa1), array_flatten(aa3), array_flatten(aa5), array_flatten(aa14) from t_array_flatten order by k1; """ + + sql """DROP TABLE IF EXISTS t_array_flatten_null_subarray""" + sql """ + CREATE TABLE t_array_flatten_null_subarray ( + `id` int, + `a` array + ) ENGINE=OLAP + DUPLICATE KEY(`id`) + DISTRIBUTED BY HASH(`id`) BUCKETS 1 + PROPERTIES ( + "replication_allocation" = "tag.location.default: 1", + "storage_format" = "V2" + ) + """ + sql """INSERT INTO t_array_flatten_null_subarray VALUES (1, [101, 102]), (2, [201])""" + + order_qt_null_subarray """ + SELECT id, + array(if(id = 1, NULL, a), [9]) AS visible_array, + array_flatten(array(if(id = 1, NULL, a), [9])) AS flat_array, + array_flatten(CAST([NULL, [9]] AS ARRAY>)) AS literal_flat + FROM t_array_flatten_null_subarray + """ + + order_qt_null_subarray_nested """ + SELECT id, + array_flatten(array(array(if(id = 1, NULL, a), [9]))) AS flat_array + FROM t_array_flatten_null_subarray + """ }