From 5ad8e934587b0263355510e3a7faa6922e7f4a9d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ra=C3=BAl=20Mar=C3=ADn?= <git@rmr.ninja>
Date: Fri, 27 Aug 2021 18:45:20 +0200
Subject: [PATCH] Speedup sumIf with Nullable types

---
 src/AggregateFunctions/AggregateFunctionIf.cpp | 16 ++++++++++++++++
 src/AggregateFunctions/AggregateFunctionSum.h  | 13 ++++++++-----
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/src/AggregateFunctions/AggregateFunctionIf.cpp b/src/AggregateFunctions/AggregateFunctionIf.cpp
index 89688ce1ffd..2ff1d3a4435 100644
--- a/src/AggregateFunctions/AggregateFunctionIf.cpp
+++ b/src/AggregateFunctions/AggregateFunctionIf.cpp
@@ -112,6 +112,22 @@ public:
         }
     }
 
+    void addBatchSinglePlace(size_t batch_size, AggregateDataPtr place, const IColumn ** columns, Arena * arena, ssize_t) const override
+    {
+        const ColumnNullable * column = assert_cast<const ColumnNullable *>(columns[0]);
+        const UInt8 * null_map = column->getNullMapData().data();
+        const IColumn * nested_column = &column->getNestedColumn();
+        const IColumn * filter_column = columns[num_arguments - 1];
+        if (const ColumnNullable * nullable_column = typeid_cast<const ColumnNullable *>(filter_column))
+            filter_column = &nullable_column->getNestedColumn();
+        const IColumn * column_param[] = {nested_column, filter_column};
+
+        this->nested_function->addBatchSinglePlaceNotNull(batch_size, this->nestedPlace(place), column_param, null_map, arena, 1);
+        if constexpr (result_is_nullable)
+            if (!memoryIsByte(null_map, batch_size, 1))
+                this->setFlag(place);
+    }
+
 #if USE_EMBEDDED_COMPILER
 
     void compileAdd(llvm::IRBuilderBase & builder, llvm::Value * aggregate_data_ptr, const DataTypes & arguments_types, const std::vector<llvm::Value *> & argument_values) const override
diff --git a/src/AggregateFunctions/AggregateFunctionSum.h b/src/AggregateFunctions/AggregateFunctionSum.h
index 77539240d97..7dd5c0eee8e 100644
--- a/src/AggregateFunctions/AggregateFunctionSum.h
+++ b/src/AggregateFunctions/AggregateFunctionSum.h
@@ -393,19 +393,22 @@ public:
     }
 
     void addBatchSinglePlaceNotNull(
-        size_t batch_size, AggregateDataPtr place, const IColumn ** columns, const UInt8 * null_map, Arena * arena, ssize_t if_argument_pos)
+        size_t batch_size, AggregateDataPtr place, const IColumn ** columns, const UInt8 * null_map, Arena *, ssize_t if_argument_pos)
         const override
     {
+        const auto & column = assert_cast<const ColVecType &>(*columns[0]);
         if (if_argument_pos >= 0)
         {
-            const auto & flags = assert_cast<const ColumnUInt8 &>(*columns[if_argument_pos]).getData();
+            /// Merge the 2 sets of flags (null and if) into a single one. This allows us to use parallelizable sums when available
+            UInt8 final_flags[batch_size];
+            memcpy(final_flags, assert_cast<const ColumnUInt8 &>(*columns[if_argument_pos]).getData().data(), batch_size * sizeof(UInt8));
             for (size_t i = 0; i < batch_size; ++i)
-                if (!null_map[i] && flags[i])
-                    add(place, columns, i, arena);
+                final_flags[i] = !null_map[i] && final_flags[i];
+
+            this->data(place).addManyConditional(column.getData().data(), final_flags, batch_size);
         }
         else
         {
-            const auto & column = assert_cast<const ColVecType &>(*columns[0]);
             this->data(place).addManyNotNull(column.getData().data(), null_map, batch_size);
         }
     }