fix panic in stats histogram union and intersect (#3116)

Author: jycor

sql/stats/filter.go

@@ -32,20 +32,21 @@ func Union(ctx *sql.Context, b1, b2 []sql.HistogramBucket, types []sql.Type) ([]
 			if err != nil {
 				return nil, err
 			}
-			switch cmp {
-			case 0:
-				if k == len(key1)-1 {
-					ret = append(ret, b1[i])
-					i++
-					j++
-				}
-				continue
-			case +1:
+			if cmp == +1 {
 				ret = append(ret, b2[j])
 				j++
-			case -1:
+				break
+			}
+			if cmp == -1 {
 				ret = append(ret, b1[i])
 				i++
+				break
+			}
+			// if keys are equal, merge buckets
+			if k == len(key1)-1 {
+				ret = append(ret, b1[i])
+				i++
+				j++
 			}
 		}
 	}
@@ -78,18 +79,19 @@ func Intersect(ctx *sql.Context, b1, b2 []sql.HistogramBucket, types []sql.Type)
 			if err != nil {
 				return nil, err
 			}
-			switch cmp {
-			case 0:
-				if k == len(key1)-1 {
-					ret = append(ret, b1[i])
-					i++
-					j++
-				}
-				continue
-			case +1:
+			if cmp == +1 {
 				j++
-			case -1:
+				break
+			}
+			if cmp == -1 {
 				i++
+				break
+			}
+			// if keys are equal, merge buckets
+			if k == len(key1)-1 {
+				ret = append(ret, b1[i])
+				i++
+				j++
 			}
 		}
 	}

sql/stats/filter_test.go

@@ -962,3 +962,121 @@ func TestUpdateCounts(t *testing.T) {
 		})
 	}
 }
+
+func TestHistogramUnion(t *testing.T) {
+	ctx := sql.NewEmptyContext()
+	tests := []struct {
+		types  []sql.Type
+		h1, h2 []sql.HistogramBucket
+		exp    []sql.HistogramBucket
+	}{
+		{
+			types: []sql.Type{types.Int64},
+			h1: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{9}},
+			},
+			h2: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1}},
+			},
+			exp: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1}},
+				&Bucket{BoundVal: []interface{}{9}},
+			},
+		},
+		{
+			types: []sql.Type{types.Int64, types.Int64},
+			h1: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1, 1}},
+				&Bucket{BoundVal: []interface{}{1, 9}},
+			},
+			h2: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{9, 9}},
+			},
+			exp: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1, 1}},
+				&Bucket{BoundVal: []interface{}{1, 9}},
+				&Bucket{BoundVal: []interface{}{9, 9}},
+			},
+		},
+		{
+			types: []sql.Type{types.Int64, types.Int64},
+			h1: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1, 9}},
+				&Bucket{BoundVal: []interface{}{9, 9}},
+			},
+			h2: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1, 1}},
+			},
+			exp: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1, 1}},
+				&Bucket{BoundVal: []interface{}{1, 9}},
+				&Bucket{BoundVal: []interface{}{9, 9}},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(fmt.Sprintf("union: %v and %v", tt.h1, tt.h2), func(t *testing.T) {
+			res, err := Union(ctx, tt.h1, tt.h2, tt.types)
+			require.NoError(t, err)
+			require.Equal(t, tt.exp, res)
+		})
+	}
+}
+
+func TestHistogramIntersect(t *testing.T) {
+	ctx := sql.NewEmptyContext()
+	tests := []struct {
+		types  []sql.Type
+		h1, h2 []sql.HistogramBucket
+		exp    []sql.HistogramBucket
+	}{
+		{
+			types: []sql.Type{types.Int64},
+			h1: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1}},
+			},
+			h2: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1}},
+			},
+			exp: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1}},
+			},
+		},
+		{
+			types: []sql.Type{types.Int64, types.Int64},
+			h1: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1, 1}},
+				&Bucket{BoundVal: []interface{}{1, 9}},
+			},
+			h2: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1, 1}},
+				&Bucket{BoundVal: []interface{}{9, 9}},
+			},
+			exp: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1, 1}},
+			},
+		},
+		{
+			types: []sql.Type{types.Int64, types.Int64},
+			h1: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1, 1}},
+				&Bucket{BoundVal: []interface{}{9, 9}},
+			},
+			h2: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1, 1}},
+				&Bucket{BoundVal: []interface{}{9, 9}},
+			},
+			exp: []sql.HistogramBucket{
+				&Bucket{BoundVal: []interface{}{1, 1}},
+				&Bucket{BoundVal: []interface{}{9, 9}},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(fmt.Sprintf("union: %v and %v", tt.h1, tt.h2), func(t *testing.T) {
+			res, err := Intersect(ctx, tt.h1, tt.h2, tt.types)
+			require.NoError(t, err)
+			require.Equal(t, tt.exp, res)
+		})
+	}
+}