[feature](nereids) support multi_distinct_collect_list and multi_distinct_array_agg

be/src/exprs/aggregate/aggregate_function_array_agg.cpp

@@ -17,6 +17,7 @@
 
 #include "exprs/aggregate/aggregate_function_array_agg.h"
 
+#include "common/exception.h"
 #include "core/call_on_type_index.h"
 #include "exprs/aggregate/aggregate_function_collect.h"
 #include "exprs/aggregate/aggregate_function_simple_factory.h"
@@ -25,17 +26,37 @@
 namespace doris {
 
 template <PrimitiveType T>
-AggregateFunctionPtr do_create_agg_function_collect(const DataTypes& argument_types,
+AggregateFunctionPtr do_create_agg_function_collect(bool distinct, const std::string& name,
+ const DataTypes& argument_types,
  const bool result_is_nullable,
  const AggregateFunctionAttr& attr) {
+ if (distinct) {
+ if constexpr (T == INVALID_TYPE) {
+ throw Exception(ErrorCode::INTERNAL_ERROR,
+ "unexpected type for array_agg distinct, please check the input");
+ } else {
+ // array_agg keeps NULL elements, so multi_distinct_array_agg must preserve a single
+ // NULL when the input contains nulls. Feed the nullable column straight to the
+ // aggregate (create_ignore_nullable, no null-skipping wrapper) and let the ShowNull
+ // collect function record whether a null was seen; the Set still dedups the values.
+ if (argument_types[0]->is_nullable()) {
+ return creator_without_type::create_ignore_nullable<AggregateFunctionCollect<
+ AggregateFunctionCollectSetData<T, false>, false, true>>(
+ argument_types, result_is_nullable, attr, name);
+ }
+ return creator_without_type::create<
+ AggregateFunctionCollect<AggregateFunctionCollectSetData<T, false>, false>>(
+ argument_types, result_is_nullable, attr, name);
+ }
+ }
  if (argument_types[0]->is_nullable()) {
  return creator_without_type::create_ignore_nullable<
  AggregateFunctionArrayAgg<AggregateFunctionArrayAggData<T>>>(
  argument_types, result_is_nullable, attr);
  } else {
  return creator_without_type::create<
  AggregateFunctionCollect<AggregateFunctionCollectListData<T, false>, false>>(
- argument_types, result_is_nullable, attr);
+ argument_types, result_is_nullable, attr, name);
  }
 }
 
@@ -44,23 +65,25 @@ AggregateFunctionPtr create_aggregate_function_array_agg(const std::string& name
  const DataTypePtr& result_type,
  const bool result_is_nullable,
  const AggregateFunctionAttr& attr) {
+ bool distinct = name == "multi_distinct_array_agg";
  AggregateFunctionPtr agg_fn;
  auto call = [&](const auto& type) -> bool {
  using DispatcType = std::decay_t<decltype(type)>;
- agg_fn = do_create_agg_function_collect<DispatcType::PType>(argument_types,
+ agg_fn = do_create_agg_function_collect<DispatcType::PType>(distinct, name, argument_types,
  result_is_nullable, attr);
  return true;
  };
 
  if (!dispatch_switch_all(argument_types[0]->get_primitive_type(), call)) {
  // We do not care what the real type is.
- agg_fn = do_create_agg_function_collect<INVALID_TYPE>(argument_types, result_is_nullable,
- attr);
+ agg_fn = do_create_agg_function_collect<INVALID_TYPE>(distinct, name, argument_types,
+result_is_nullable, attr);
  }
  return agg_fn;
 }
 
 void register_aggregate_function_array_agg(AggregateFunctionSimpleFactory& factory) {
  factory.register_function_both("array_agg", create_aggregate_function_array_agg);
+ factory.register_function_both("multi_distinct_array_agg", create_aggregate_function_array_agg);
 }
 } // namespace doris

be/src/exprs/aggregate/aggregate_function_collect.cpp

@@ -53,6 +53,8 @@ void register_aggregate_function_collect_list(AggregateFunctionSimpleFactory& fa
  // notice: array_agg only differs from collect_list in that array_agg will show null elements in array
  factory.register_function_both("collect_list", create_aggregate_function_collect);
  factory.register_function_both("collect_set", create_aggregate_function_collect);
+ factory.register_function_both("multi_distinct_collect_list",
+ create_aggregate_function_collect);
  factory.register_alias("collect_list", "group_array");
  factory.register_alias("collect_set", "group_uniq_array");
 }

be/src/exprs/aggregate/aggregate_function_collect.h

@@ -53,6 +53,9 @@ struct AggregateFunctionCollectSetData {
  using Set = doris::flat_hash_set<ElementType>;
  Set data_set;
  Int64 max_size = -1;
+ // Only used by the ShowNull collect variant (multi_distinct_array_agg): the Set cannot store
+ // a null element, so we remember whether any null was seen and emit a single NULL on output.
+ bool has_null = false;
 
  AggregateFunctionCollectSetData(const DataTypes& argument_types) {}
 
@@ -109,7 +112,10 @@ struct AggregateFunctionCollectSetData {
  }
  }
 
- void reset() { data_set.clear(); }
+ void reset() {
+ data_set.clear();
+ has_null = false;
+ }
 };
 
 template <PrimitiveType T, bool HasLimit>
@@ -122,6 +128,9 @@ struct AggregateFunctionCollectSetData<T, HasLimit> {
  using Set = doris::flat_hash_set<ElementType>;
  Set data_set;
  Int64 max_size = -1;
+ // Only used by the ShowNull collect variant (multi_distinct_array_agg): the Set cannot store
+ // a null element, so we remember whether any null was seen and emit a single NULL on output.
+ bool has_null = false;
 
  AggregateFunctionCollectSetData(const DataTypes& argument_types) {}
 
@@ -177,7 +186,10 @@ struct AggregateFunctionCollectSetData<T, HasLimit> {
  }
  }
 
- void reset() { data_set.clear(); }
+ void reset() {
+ data_set.clear();
+ has_null = false;
+ }
 };
 
 template <PrimitiveType T, bool HasLimit>
@@ -391,9 +403,10 @@ struct AggregateFunctionCollectListData<T, HasLimit> {
  void insert_result_into(IColumn& to) const { to.insert_range_from(*column_data, 0, size()); }
 };
 
-template <typename Data, bool HasLimit>
+template <typename Data, bool HasLimit, bool ShowNull = false>
 class AggregateFunctionCollect final
- : public IAggregateFunctionDataHelper<Data, AggregateFunctionCollect<Data, HasLimit>, true>,
+ : public IAggregateFunctionDataHelper<
+ Data, AggregateFunctionCollect<Data, HasLimit, ShowNull>, true>,
  VarargsExpression,
  NotNullableAggregateFunction {
  static constexpr bool ENABLE_ARENA =
@@ -402,25 +415,36 @@ class AggregateFunctionCollect final
  std::is_same_v<Data, AggregateFunctionCollectSetData<TYPE_VARCHAR, HasLimit>>;
 
 public:
- AggregateFunctionCollect(const DataTypes& argument_types_)
- : IAggregateFunctionDataHelper<Data, AggregateFunctionCollect<Data, HasLimit>, true>(
- {argument_types_}),
+ AggregateFunctionCollect(std::string name, const DataTypes& argument_types_)
+ : IAggregateFunctionDataHelper<Data, AggregateFunctionCollect<Data, HasLimit, ShowNull>,
+ true>({argument_types_}),
+ _name(std::move(name)),
  return_type(std::make_shared<DataTypeArray>(make_nullable(argument_types_[0]))) {}
 
- std::string get_name() const override {
- if constexpr (std::is_same_v<AggregateFunctionCollectListData<Data::PType, HasLimit>,
- Data>) {
- return "collect_list";
- } else {
- return "collect_set";
- }
- }
+ std::string get_name() const override { return _name; }
 
  DataTypePtr get_return_type() const override { return return_type; }
 
  void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num,
  Arena& arena) const override {
  auto& data = this->data(place);
+ if constexpr (ShowNull) {
+ // The nullable column is fed in directly (no null-skipping wrapper) so that
+ // multi_distinct_array_agg can preserve a NULL element like array_agg does.
+ const auto& nullable =
+ assert_cast<const ColumnNullable&, TypeCheckOnRelease::DISABLE>(*columns[0]);
+ if (nullable.is_null_at(row_num)) {
+ data.has_null = true;
+ return;
+ }
+ const IColumn& nested = nullable.get_nested_column();
+ if constexpr (ENABLE_ARENA) {
+ data.add(nested, row_num, arena);
+ } else {
+ data.add(nested, row_num);
+ }
+ return;
+ }
  if constexpr (HasLimit) {
  if (data.max_size == -1) {
  data.max_size =
@@ -448,15 +472,26 @@ class AggregateFunctionCollect final
  } else {
  data.merge(rhs_data);
  }
+ if constexpr (ShowNull) {
+ data.has_null |= rhs_data.has_null;
+ }
  }
 
  void serialize(ConstAggregateDataPtr __restrict place, BufferWritable& buf) const override {
  this->data(place).write(buf);
+ if constexpr (ShowNull) {
+ buf.write_binary(this->data(place).has_null);
+ }
  }
 
  void deserialize(AggregateDataPtr __restrict place, BufferReadable& buf,
  Arena&) const override {
  this->data(place).read(buf);
+ if constexpr (ShowNull) {
+ bool has_null = false;
+ buf.read_binary(has_null);
+ this->data(place).has_null = has_null;
+ }
  }
 
  void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn& to) const override {
@@ -465,10 +500,17 @@ class AggregateFunctionCollect final
  auto* col_null = assert_cast<ColumnNullable*>(&to_nested_col);
  this->data(place).insert_result_into(col_null->get_nested_column());
  col_null->get_null_map_data().resize_fill(col_null->get_nested_column().size(), 0);
+ if constexpr (ShowNull) {
+ if (this->data(place).has_null) {
+ col_null->get_nested_column().insert_default();
+ col_null->get_null_map_data().push_back(1);
+ }
+ }
  to_arr.get_offsets().push_back(to_nested_col.size());
  }
 
 private:
+ std::string _name;
  DataTypePtr return_type;
 };
 

be/src/exprs/aggregate/aggregate_function_collect_impl.h

@@ -27,7 +27,8 @@
 namespace doris {
 
 template <PrimitiveType T, bool HasLimit>
-AggregateFunctionPtr do_create_agg_function_collect(bool distinct, const DataTypes& argument_types,
+AggregateFunctionPtr do_create_agg_function_collect(bool distinct, const std::string& name,
+ const DataTypes& argument_types,
  const bool result_is_nullable,
 
  const AggregateFunctionAttr& attr) {
@@ -38,12 +39,12 @@ AggregateFunctionPtr do_create_agg_function_collect(bool distinct, const DataTyp
  } else {
  return creator_without_type::create<AggregateFunctionCollect<
  AggregateFunctionCollectSetData<T, HasLimit>, HasLimit>>(
- argument_types, result_is_nullable, attr);
+ argument_types, result_is_nullable, attr, name);
  }
  } else {
  return creator_without_type::create<
  AggregateFunctionCollect<AggregateFunctionCollectListData<T, HasLimit>, HasLimit>>(
- argument_types, result_is_nullable, attr);
+ argument_types, result_is_nullable, attr, name);
  }
 }
 
@@ -53,20 +54,22 @@ AggregateFunctionPtr create_aggregate_function_collect_impl(const std::string& n
  const bool result_is_nullable,
 
  const AggregateFunctionAttr& attr) {
- bool distinct = name == "collect_set";
+ // multi_distinct_collect_list must dedup in the aggregate state: the Nereids multi-distinct
+ // plan does not push the distinct argument into the group-by key, so the Set path is required.
+ bool distinct = name == "collect_set" || name == "multi_distinct_collect_list";
 
  AggregateFunctionPtr agg_fn;
  auto call = [&](const auto& type) -> bool {
  using DispatcType = std::decay_t<decltype(type)>;
  agg_fn = do_create_agg_function_collect<DispatcType::PType, HasLimit>(
- distinct, argument_types, result_is_nullable, attr);
+ distinct, name, argument_types, result_is_nullable, attr);
  return true;
  };
 
  if (!dispatch_switch_all(argument_types[0]->get_primitive_type(), call)) {
  // We do not care what the real type is.
- agg_fn = do_create_agg_function_collect<INVALID_TYPE, HasLimit>(distinct, argument_types,
- result_is_nullable, attr);
+ agg_fn = do_create_agg_function_collect<INVALID_TYPE, HasLimit>(
+distinct, name, argument_types, result_is_nullable, attr);
  }
  return agg_fn;
 }