Merge pull request #3178 from dolthub/angela/window

Do not materialize any output for WindowPartitionIter if input is empty

Author: angelamayxie

enginetest/queries/logic_test_scripts.go

@@ -951,7 +951,6 @@ var SQLLogicSubqueryTests = []ScriptTest{
 				},
 			},
 			{
-				Skip:  true,
 				Query: "SELECT * FROM (SELECT c_id AS c_c_id, bill FROM c) sq1, LATERAL (SELECT row_number() OVER () AS rownum FROM o WHERE c_id = c_c_id) sq2 ORDER BY c_c_id, bill, rownum;",
 				Expected: []sql.Row{
 					{1, "CA", 1},

enginetest/queries/script_queries.go

@@ -11584,6 +11584,7 @@ select * from t1 except (
 				Query:   "select * from test01 where pk in (11)",
 				Expected: []sql.Row{
 					{"11"},
+					{"11-5"},
 					{"11d"},
 					{"11wha?"},
 				},
@@ -11628,6 +11629,64 @@ select * from t1 except (
 			},
 		},
 	},
+	{
+		// https://github.com/dolthub/dolt/issues/6899
+		Name: "window function tests",
+		SetUpScript: []string{
+			"CREATE TABLE c (c_id INT PRIMARY KEY, bill TEXT);",
+			"CREATE TABLE o (o_id INT PRIMARY KEY, c_id INT, ship TEXT);",
+			"INSERT INTO c VALUES (1, 'CA'), (2, 'TX'), (3, 'MA'), (4, 'TX'), (5, NULL), (6, 'FL');",
+			"INSERT INTO o VALUES (10, 1, 'CA'), (20, 1, 'CA'), (30, 1, 'CA'), (40, 2, 'CA'), (50, 2, 'TX'), (60, 2, NULL), (70, 4, 'WY'), (80, 4, NULL), (90, 6, 'WA');",
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query:    "select row_number() over () as rn from o where c_id=-999",
+				Expected: []sql.Row{},
+			},
+			{
+				// TODO: valid query in Postgres. https://github.com/dolthub/doltgresql/issues/1796
+				Dialect:  "mysql",
+				Query:    "select row_number() over () as rn from o where c_id=1",
+				Expected: []sql.Row{{1}, {2}, {3}},
+			},
+			{
+				Query:    "select rank() over() as rnk from o where c_id=-999",
+				Expected: []sql.Row{},
+			},
+			{
+				// TODO: valid query in Postgres. https://github.com/dolthub/doltgresql/issues/1796
+				Dialect: "mysql",
+				Query:   "select o_id, c_id, rank() over(order by o_id) as rnk from o where c_id=1",
+				Expected: []sql.Row{
+					{10, 1, uint64(1)},
+					{20, 1, uint64(2)},
+					{30, 1, uint64(3)},
+				},
+			},
+			{
+				Query:    "select dense_rank() over() as rnk from o where c_id=-999",
+				Expected: []sql.Row{},
+			},
+			{
+				// TODO: valid query in Postgres. But Postgres orders nil at the end. Maybe rewrite query to filter out
+				//  ship=null https://github.com/dolthub/doltgresql/issues/1796
+				Dialect: "mysql",
+				Query:   "select ship, dense_rank() over (order by ship) as drnk from o where c_id in (1, 2) order by ship",
+				Expected: []sql.Row{
+					{nil, uint64(1)},
+					{"CA", uint64(2)},
+					{"CA", uint64(2)},
+					{"CA", uint64(2)},
+					{"CA", uint64(2)},
+					{"TX", uint64(3)},
+				},
+			},
+			{
+				Query:    "select count(*) from o where c_id=-999",
+				Expected: []sql.Row{{0}},
+			},
+		},
+	},
 }
 
 var SpatialScriptTests = []ScriptTest{

sql/expression/function/aggregation/window_framer.go

@@ -598,7 +598,11 @@ func isNewOrderByValue(ctx *sql.Context, orderByExprs []sql.Expression, last sql
 	}
 
 	for i := range lastExp {
-		if lastExp[i] != thisExp[i] {
+		compare, err := orderByExprs[i].Type().Compare(ctx, lastExp[i], thisExp[i])
+		if err != nil {
+			return false, err
+		}
+		if compare != 0 {
 			return true, nil
 		}
 	}

sql/expression/function/aggregation/window_partition.go

@@ -232,9 +232,7 @@ func (i *WindowPartitionIter) initializePartitions(ctx *sql.Context) ([]sql.Wind
 // At this stage, result rows are appended with the original row index for resorting. The size of
 // [i.output] will be smaller than [i.input] if the outer sql.Node is a plan.GroupBy with fewer partitions than rows.
 func (i *WindowPartitionIter) materializeOutput(ctx *sql.Context) (sql.WindowBuffer, error) {
-	// handle nil input specially if no partition clause
-	// ex: COUNT(*) on nil rows returns 0, not nil
-	if len(i.input) == 0 && len(i.w.PartitionBy) > 0 {
+	if len(i.input) == 0 {
 		return nil, io.EOF
 	}
 

sql/expression/function/aggregation/window_partition_test.go

@@ -23,7 +23,6 @@ import (
 
 	"github.com/dolthub/go-mysql-server/memory"
 	"github.com/dolthub/go-mysql-server/sql"
-	"github.com/dolthub/go-mysql-server/sql/expression"
 	"github.com/dolthub/go-mysql-server/sql/types"
 )
 
@@ -216,7 +215,7 @@ func TestWindowPartition_MaterializeOutput(t *testing.T) {
 		require.ElementsMatch(t, expOutput, output)
 	})
 
-	t.Run("nil input with partition by", func(t *testing.T) {
+	t.Run("nil input", func(t *testing.T) {
 		ctx := sql.NewEmptyContext()
 		i := NewWindowPartitionIter(
 			&WindowPartition{
@@ -232,24 +231,6 @@ func TestWindowPartition_MaterializeOutput(t *testing.T) {
 		require.Equal(t, io.EOF, err)
 		require.ElementsMatch(t, nil, output)
 	})
-
-	t.Run("nil input no partition by", func(t *testing.T) {
-		ctx := sql.NewEmptyContext()
-		i := NewWindowPartitionIter(
-			&WindowPartition{
-				PartitionBy: nil,
-				Aggs: []*Aggregation{
-					NewAggregation(NewCountAgg(expression.NewGetField(0, types.Int64, "z", true)), NewGroupByFramer()),
-				},
-			})
-		i.input = []sql.Row{}
-		i.partitions = []sql.WindowInterval{{0, 0}}
-		i.outputOrdering = nil
-		output, err := i.materializeOutput(ctx)
-		require.NoError(t, err)
-		expOutput := []sql.Row{{int64(0), nil}}
-		require.ElementsMatch(t, expOutput, output)
-	})
 }
 
 func TestWindowPartition_SortAndFilterOutput(t *testing.T) {