[fix](window_func) fix bug of agg function used in window function an…

…d add many test cases (apache#40678) ## Proposed changes Support usage in window function for many agg functions by adding `reset` method. Fix bug of `window_funnel` used in window funciton; Fix bug of wrong result of `orthogonal_bitmap_intersect`.
kaijchen · Sep 20, 2024 · 11dfd19 · 11dfd19
1 parent 51ba957
commit 11dfd19
Show file tree

Hide file tree

Showing 14 changed files with 1,440 additions and 53 deletions.
diff --git a/be/src/vec/aggregate_functions/aggregate_function.h b/be/src/vec/aggregate_functions/aggregate_function.h
@@ -170,9 +170,6 @@ class IAggregateFunction {
  virtual void deserialize_and_merge_from_column(AggregateDataPtr __restrict place,
  const IColumn& column, Arena* arena) const = 0;
 
- /// Returns true if a function requires Arena to handle own states (see add(), merge(), deserialize()).
- virtual bool allocates_memory_in_arena() const { return false; }
-
  /// Inserts results into a column.
  virtual void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn& to) const = 0;
 

diff --git a/be/src/vec/aggregate_functions/aggregate_function_binary.h b/be/src/vec/aggregate_functions/aggregate_function_binary.h
@@ -62,12 +62,12 @@ struct AggregateFunctionBinary
 
  String get_name() const override { return StatFunc::Data::name(); }
 
+ void reset(AggregateDataPtr __restrict place) const override { this->data(place).reset(); }
+
  DataTypePtr get_return_type() const override {
  return std::make_shared<DataTypeNumber<ResultType>>();
  }
 
- bool allocates_memory_in_arena() const override { return false; }
-
  void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num,
  Arena*) const override {
  this->data(place).add(

diff --git a/be/src/vec/aggregate_functions/aggregate_function_collect.h b/be/src/vec/aggregate_functions/aggregate_function_collect.h
@@ -626,8 +626,6 @@ class AggregateFunctionCollect
  return std::make_shared<DataTypeArray>(make_nullable(return_type));
  }
 
- bool allocates_memory_in_arena() const override { return ENABLE_ARENA; }
-
  void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num,
  Arena* arena) const override {
  auto& data = this->data(place);

diff --git a/be/src/vec/aggregate_functions/aggregate_function_corr.cpp b/be/src/vec/aggregate_functions/aggregate_function_corr.cpp
@@ -76,6 +76,15 @@ struct CorrMoment {
  }
 
  static String name() { return "corr"; }
+
+ void reset() {
+ m0 = {};
+ x1 = {};
+ y1 = {};
+ xy = {};
+ x2 = {};
+ y2 = {};
+ }
 };
 
 AggregateFunctionPtr create_aggregate_corr_function(const std::string& name,

diff --git a/be/src/vec/aggregate_functions/aggregate_function_distinct.h b/be/src/vec/aggregate_functions/aggregate_function_distinct.h
@@ -338,8 +338,6 @@ class AggregateFunctionDistinct
 
  DataTypePtr get_return_type() const override { return nested_func->get_return_type(); }
 
- bool allocates_memory_in_arena() const override { return true; }
-
  AggregateFunctionPtr transmit_to_stable() override {
  return AggregateFunctionPtr(new AggregateFunctionDistinct<Data, true>(
  nested_func, IAggregateFunction::argument_types));

diff --git a/be/src/vec/aggregate_functions/aggregate_function_foreach.h b/be/src/vec/aggregate_functions/aggregate_function_foreach.h
@@ -219,10 +219,6 @@ class AggregateFunctionForEach : public IAggregateFunctionDataHelper<AggregateFu
  offsets_to.push_back(offsets_to.back() + state.dynamic_array_size);
  }
 
- bool allocates_memory_in_arena() const override {
- return nested_function->allocates_memory_in_arena();
- }
-
  void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num,
  Arena* arena) const override {
  std::vector<const IColumn*> nested(num_arguments);

diff --git a/be/src/vec/aggregate_functions/aggregate_function_group_array_intersect.h b/be/src/vec/aggregate_functions/aggregate_function_group_array_intersect.h
@@ -93,6 +93,11 @@ struct AggregateFunctionGroupArrayIntersectData {
  Set value;
  bool init = false;
 
+ void reset() {
+ init = false;
+ value = std::make_unique<NullableNumericOrDateSetType>();
+ }
+
  void process_col_data(auto& column_data, size_t offset, size_t arr_size, bool& init, Set& set) {
  const bool is_column_data_nullable = column_data.is_nullable();
 
@@ -163,7 +168,7 @@ class AggregateFunctionGroupArrayIntersect
 
  DataTypePtr get_return_type() const override { return argument_type; }
 
- bool allocates_memory_in_arena() const override { return false; }
+ void reset(AggregateDataPtr __restrict place) const override { this->data(place).reset(); }
 
  void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num,
  Arena*) const override {
@@ -331,6 +336,11 @@ struct AggregateFunctionGroupArrayIntersectGenericData {
  : value(std::make_unique<NullableStringSet>()) {}
  Set value;
  bool init = false;
+
+ void reset() {
+ init = false;
+ value = std::make_unique<NullableStringSet>();
+ }
 };
 
 /** Template parameter with true value should be used for columns that store their elements in memory continuously.
@@ -357,7 +367,7 @@ class AggregateFunctionGroupArrayIntersectGeneric
 
  DataTypePtr get_return_type() const override { return input_data_type; }
 
- bool allocates_memory_in_arena() const override { return true; }
+ void reset(AggregateDataPtr __restrict place) const override { this->data(place).reset(); }
 
  void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num,
  Arena* arena) const override {

diff --git a/be/src/vec/aggregate_functions/aggregate_function_null.h b/be/src/vec/aggregate_functions/aggregate_function_null.h
@@ -177,10 +177,6 @@ class AggregateFunctionNullBaseInline : public IAggregateFunctionHelper<Derived>
  nested_function->insert_result_into(nested_place(place), to);
  }
  }
-
- bool allocates_memory_in_arena() const override {
- return nested_function->allocates_memory_in_arena();
- }
 };
 
 /** There are two cases: for single argument and variadic.
@@ -329,10 +325,6 @@ class AggregateFunctionNullVariadicInline final
  arena);
  }
 
- bool allocates_memory_in_arena() const override {
- return this->nested_function->allocates_memory_in_arena();
- }
-
 private:
  // The array length is fixed in the implementation of some aggregate functions.
  // Therefore we choose 256 as the appropriate maximum length limit.

diff --git a/be/src/vec/aggregate_functions/aggregate_function_orthogonal_bitmap.h b/be/src/vec/aggregate_functions/aggregate_function_orthogonal_bitmap.h
@@ -52,6 +52,11 @@ struct AggOrthBitmapBaseData {
 public:
  using ColVecData = std::conditional_t<IsNumber<T>, ColumnVector<T>, ColumnString>;
 
+ void reset() {
+ bitmap = {};
+ first_init = true;
+ }
+
  void add(const IColumn** columns, size_t row_num) {
  const auto& bitmap_col =
  assert_cast<const ColumnBitmap&, TypeCheckOnRelease::DISABLE>(*columns[0]);
@@ -99,6 +104,11 @@ struct AggOrthBitMapIntersect : public AggOrthBitmapBaseData<T> {
 
  static DataTypePtr get_return_type() { return std::make_shared<DataTypeBitMap>(); }
 
+ void reset() {
+ AggOrthBitmapBaseData<T>::reset();
+ result.reset();
+ }
+
  void merge(const AggOrthBitMapIntersect& rhs) {
  if (rhs.first_init) {
  return;
@@ -120,7 +130,8 @@ struct AggOrthBitMapIntersect : public AggOrthBitmapBaseData<T> {
 
  void get(IColumn& to) const {
  auto& column = assert_cast<ColumnBitmap&>(to);
- column.get_data().emplace_back(result);
+ column.get_data().emplace_back(result.empty() ? AggOrthBitmapBaseData<T>::bitmap.intersect()
+ : result);
  }
 
 private:
@@ -170,6 +181,11 @@ struct AggOrthBitMapIntersectCount : public AggOrthBitmapBaseData<T> {
 
  static DataTypePtr get_return_type() { return std::make_shared<DataTypeInt64>(); }
 
+ void reset() {
+ AggOrthBitmapBaseData<T>::reset();
+ result = 0;
+ }
+
  void merge(const AggOrthBitMapIntersectCount& rhs) {
  if (rhs.first_init) {
  return;
@@ -225,6 +241,11 @@ struct AggOrthBitmapExprCalBaseData {
  }
  }
 
+ void reset() {
+ bitmap_expr_cal = {};
+ first_init = true;
+ }
+
 protected:
  doris::BitmapExprCalculation bitmap_expr_cal;
  bool first_init = true;
@@ -263,6 +284,11 @@ struct AggOrthBitMapExprCal : public AggOrthBitmapExprCalBaseData<T> {
  ->bitmap_expr_cal.bitmap_calculate());
  }
 
+ void reset() {
+ AggOrthBitmapExprCalBaseData<T>::reset();
+ result.reset();
+ }
+
 private:
  BitmapValue result;
 };
@@ -299,6 +325,11 @@ struct AggOrthBitMapExprCalCount : public AggOrthBitmapExprCalBaseData<T> {
  ->bitmap_expr_cal.bitmap_calculate_count());
  }
 
+ void reset() {
+ AggOrthBitmapExprCalBaseData<T>::reset();
+ result = 0;
+ }
+
 private:
  int64_t result = 0;
 };
@@ -330,6 +361,11 @@ struct OrthBitmapUnionCountData {
  column.get_data().emplace_back(result ? result : value.cardinality());
  }
 
+ void reset() {
+ value.reset();
+ result = 0;
+ }
+
 private:
  BitmapValue value;
  int64_t result = 0;
@@ -347,6 +383,8 @@ class AggFunctionOrthBitmapFunc final
 
  DataTypePtr get_return_type() const override { return Impl::get_return_type(); }
 
+ void reset(AggregateDataPtr __restrict place) const override { this->data(place).reset(); }
+
  void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num,
  Arena*) const override {
  this->data(place).init_add_key(columns, row_num, _argument_size);

diff --git a/be/src/vec/aggregate_functions/aggregate_function_uniq.h b/be/src/vec/aggregate_functions/aggregate_function_uniq.h
@@ -76,6 +76,8 @@ struct AggregateFunctionUniqExactData {
  Set set;
 
  static String get_name() { return "multi_distinct"; }
+
+ void reset() { set.clear(); }
 };
 
 namespace detail {
@@ -115,6 +117,8 @@ class AggregateFunctionUniq final
 
  DataTypePtr get_return_type() const override { return std::make_shared<DataTypeInt64>(); }
 
+ void reset(AggregateDataPtr __restrict place) const override { this->data(place).reset(); }
+
  void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num,
  Arena*) const override {
  detail::OneAdder<T, Data>::add(this->data(place), *columns[0], row_num);

diff --git a/be/src/vec/aggregate_functions/aggregate_function_uniq_distribute_key.h b/be/src/vec/aggregate_functions/aggregate_function_uniq_distribute_key.h
@@ -68,6 +68,11 @@ struct AggregateFunctionUniqDistributeKeyData {
 
  Set set;
  UInt64 count = 0;
+
+ void reset() {
+ set.clear();
+ count = 0;
+ }
 };
 
 template <typename T, typename Data>
@@ -83,6 +88,8 @@ class AggregateFunctionUniqDistributeKey final
 
  DataTypePtr get_return_type() const override { return std::make_shared<DataTypeInt64>(); }
 
+ void reset(AggregateDataPtr __restrict place) const override { this->data(place).reset(); }
+
  void add(AggregateDataPtr __restrict place, const IColumn** columns, ssize_t row_num,
  Arena*) const override {
  detail::OneAdder<T, Data>::add(this->data(place), *columns[0], row_num);