From a8709baa8efb49e0db67e849366c75755fac7f4b 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 02ffe0520..652dcf895 100644
--- a/http_handler.go
+++ b/http_handler.go
@@ -1375,7 +1375,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 5cad8d227..d78c2a8c0 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 bcfff76cc..319806e8d 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 497ce62b5..f8dc88cb9 100644
--- a/sql3/sql_defs_aggregate_test.go
+++ b/sql3/sql_defs_aggregate_test.go
@@ -5,7 +5,7 @@ import (
 	"github.com/featurebasedb/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",