[TMP FIX] Fix the invalid access to nulls in final avg (#42)

oap-project · Jan 31, 2023 · df31299 · df31299
1 parent 6e4af29
commit df31299
Showing 1 changed file with 11 additions and 6 deletions.
diff --git a/velox/functions/prestosql/aggregates/AverageAggregate.h b/velox/functions/prestosql/aggregates/AverageAggregate.h
@@ -164,18 +164,23 @@ class AverageAggregate : public exec::Aggregate {
     if (decodedPartial_.isConstantMapping()) {
       if (!decodedPartial_.isNullAt(0)) {
         auto decodedIndex = decodedPartial_.index(0);
-        auto count = baseCountVector->valueAt(decodedIndex);
-        auto sum = baseSumVector->valueAt(decodedIndex);
-        rows.applyToSelected([&](vector_size_t i) {
-          updateNonNullValue(groups[i], count, sum);
-        });
+        if (!baseSumVector->isNullAt(decodedIndex)) {
+          auto count = baseCountVector->valueAt(decodedIndex);
+          auto sum = baseSumVector->valueAt(decodedIndex);
+          rows.applyToSelected([&](vector_size_t i) {
+            updateNonNullValue(groups[i], count, sum);
+          });
+        }
       }
-    } else if (decodedPartial_.mayHaveNulls()) {
+    } else if (decodedPartial_.mayHaveNulls() || baseSumVector->mayHaveNulls()) {
       rows.applyToSelected([&](vector_size_t i) {
         if (decodedPartial_.isNullAt(i)) {
           return;
         }
         auto decodedIndex = decodedPartial_.index(i);
+        if (baseSumVector->isNullAt(decodedIndex)) {
+          return;
+        }
         updateNonNullValue(
             groups[i],
             baseCountVector->valueAt(decodedIndex),