From 91e3b8457af7735cd59dfb8593524d00edf74746 Mon Sep 17 00:00:00 2001
From: pokeeffe-molecula <85502298+pokeeffe-molecula@users.noreply.github.com>
Date: Tue, 20 Sep 2022 14:53:10 -0500
Subject: [PATCH] fb-1075 (#2221)

* handle multi field count correctly

COUNT() should ignore null values.
If the data type of the expression supports an existence bitmap for the underlying FeatureBase data type we will use it to eliminate nulls from the aggregate

* simplify aggregate for existence test

we can use a direct != null instead of an indirect not(=null), and
avoid relying on the probably-broken behavior in the executor that
tries to silently fix up Row(x=3) tests on BSI fields which wanted
Row(x==3).

Co-authored-by: Seebs <seebs@molecula.com>
---
 http_handler.go                 |  2 +-
 sql3/planner/expression.go      |  6 +++++-
 sql3/planner/oppqlaggregate.go  | 31 ++++++++++++++++++++++++++++---
 sql3/sql_defs_aggregate_test.go | 15 ++++++++++++++-
 4 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/http_handler.go b/http_handler.go
index 636b8d238..cf4caf5ba 100644
--- a/http_handler.go
+++ b/http_handler.go
@@ -1379,7 +1379,7 @@ func (h *Handler) handlePostSQL(w http.ResponseWriter, r *http.Request) {
 	// Write response back to client.
 	w.Header().Set("Content-Type", "application/json")
 
-	// the pandas data frame format in json per https://molecula.atlassian.net/wiki/spaces/MOLECULA/pages/540999700/Queries
+	// the pandas data frame format in json
 
 	// Opening bracket.
 	w.Write([]byte("{"))
diff --git a/sql3/planner/expression.go b/sql3/planner/expression.go
index 60a27208d..f7e6186e1 100644
--- a/sql3/planner/expression.go
+++ b/sql3/planner/expression.go
@@ -28,7 +28,11 @@ func coerceValue(sourceType parser.ExprDataType, targetType parser.ExprDataType,
 			return value, nil
 
 		case *parser.DataTypeID:
-			return value, nil
+			val, ok := value.(int64)
+			if !ok {
+				return nil, sql3.NewErrInternalf("unexpected value type '%T'", value)
+			}
+			return val, nil
 
 		case *parser.DataTypeDecimal:
 			val, ok := value.(int64)
diff --git a/sql3/planner/oppqlaggregate.go b/sql3/planner/oppqlaggregate.go
index 84e89b84a..8589ea7a8 100644
--- a/sql3/planner/oppqlaggregate.go
+++ b/sql3/planner/oppqlaggregate.go
@@ -140,14 +140,39 @@ func (i *pqlAggregateRowIter) Next(ctx context.Context) (types.Row, error) {
 
 		case types.AGGREGATE_COUNT:
 			if cond == nil {
-				cond = &pql.Call{Name: "All"}
+				// COUNT() should ignore null values
+				// if the data type of the expression supports an existence bitmap for
+				// the underlying FeatureBase data type use it to eliminate nulls from the aggregate
+				switch expr.dataType.(type) {
+				case *parser.DataTypeInt, *parser.DataTypeTimestamp, *parser.DataTypeDecimal:
+					cond = &pql.Call{
+						Name: "Row",
+						Args: map[string]interface{}{
+							expr.columnName: &pql.Condition{Op: pql.NEQ, Value: nil},
+						},
+					}
+				default:
+					cond = &pql.Call{Name: "All"}
+				}
 			}
-
 			call = &pql.Call{Name: "Count", Children: []*pql.Call{cond}}
 
 		case types.AGGREGATE_AVG:
 			if cond == nil {
-				cond = &pql.Call{Name: "All"}
+				// COUNT() should ignore null values
+				// if the data type of the expression supports an existence bitmap for
+				// the underlying FeatureBase data type use it to eliminate nulls from the aggregate
+				switch expr.dataType.(type) {
+				case *parser.DataTypeInt, *parser.DataTypeTimestamp, *parser.DataTypeDecimal:
+					cond = &pql.Call{
+						Name: "Row",
+						Args: map[string]interface{}{
+							expr.columnName: &pql.Condition{Op: pql.NEQ, Value: nil},
+						},
+					}
+				default:
+					cond = &pql.Call{Name: "All"}
+				}
 			}
 
 			call = &pql.Call{
diff --git a/sql3/sql_defs_aggregate_test.go b/sql3/sql_defs_aggregate_test.go
index d52b11007..d11910f23 100644
--- a/sql3/sql_defs_aggregate_test.go
+++ b/sql3/sql_defs_aggregate_test.go
@@ -5,7 +5,7 @@ import (
 	"github.com/molecula/featurebase/v3/sql3/parser"
 )
 
-//aggregate function tests
+// aggregate function tests
 var countTests = tableTest{
 	table: tbl(
 		"count_test",
@@ -50,6 +50,19 @@ var countTests = tableTest{
 			),
 			compare: compareExactUnordered,
 		},
+		{
+			sqls: sqls(
+				"SELECT COUNT(i1) as a, COUNT(i2) as b FROM count_test",
+			),
+			expHdrs: hdrs(
+				hdr("a", fldTypeInt),
+				hdr("b", fldTypeInt),
+			),
+			expRows: rows(
+				row(int64(6), int64(2)),
+			),
+			compare: compareExactUnordered,
+		},
 		{
 			sqls: sqls(
 				"SELECT COUNT(*) + 10 - 11 * 2 AS count_rows FROM count_test",