[Windowing] Handle multiple sort expressions (#10)

Recidiviz · Jan 29, 2024 · 3cd60af · 3cd60af
1 parent ed799e3
commit 3cd60af
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 23 deletions.
diff --git a/internal/function_window_option.go b/internal/function_window_option.go
@@ -413,32 +413,33 @@ func (s *WindowFuncAggregatedStatus) Done(cb func([]Value, int, int) error) erro
 	sortedValues := make([]*WindowOrderedValue, len(values))
 	copy(sortedValues, values)
 	if len(sortedValues) != 0 {
-		for orderBy := 0; orderBy < len(sortedValues[0].OrderBy); orderBy++ {
-			isAsc := sortedValues[0].OrderBy[orderBy].IsAsc
-			if isAsc {
-				sort.Slice(sortedValues, func(i, j int) bool {
-					if sortedValues[i].OrderBy[orderBy].Value == nil {
-						return true
-					}
-					if sortedValues[j].OrderBy[orderBy].Value == nil {
-						return false
-					}
-					cond, _ := sortedValues[i].OrderBy[orderBy].Value.LT(sortedValues[j].OrderBy[orderBy].Value)
+		sort.Slice(sortedValues, func(i, j int) bool {
+			for orderBy := 0; orderBy < len(sortedValues[0].OrderBy); orderBy++ {
+				iV := sortedValues[i].OrderBy[orderBy].Value
+				jV := sortedValues[j].OrderBy[orderBy].Value
+				isAsc := sortedValues[0].OrderBy[orderBy].IsAsc
+				if iV == nil {
+					return true
+				}
+				if jV == nil {
+					return false
+				}
+				isEqual, _ := iV.EQ(jV)
+				if isEqual {
+					// break tie with subsequent fields
+					continue
+				}
+				if isAsc {
+					cond, _ := iV.LT(jV)
 					return cond
-				})
-			} else {
-				sort.Slice(sortedValues, func(i, j int) bool {
-					if sortedValues[i].OrderBy[orderBy].Value == nil {
-						return true
-					}
-					if sortedValues[j].OrderBy[orderBy].Value == nil {
-						return false
-					}
-					cond, _ := sortedValues[i].OrderBy[orderBy].Value.GT(sortedValues[j].OrderBy[orderBy].Value)
+				} else {
+					cond, _ := iV.GT(jV)
 					return cond
-				})
+				}
 			}
-		}
+			return false
+		})
+
 	}
 	s.SortedValues = sortedValues
 	start, err := s.getIndexFromBoundary(s.Start)

diff --git a/query_test.go b/query_test.go
@@ -1616,6 +1616,36 @@ FROM finishers`,
 				{"Suzy Slane", "03:06:24", "F35-39", "Nobody"},
 			},
 		},
+		{
+			name: `window order by`,
+			query: `WITH toks AS (
+			SELECT DATE '2024-01-01' AS dt, 'c' AS letter
+			UNION ALL SELECT DATE '2024-02-01', 'b'
+			UNION ALL SELECT DATE '2024-02-01', 'c'
+			UNION ALL SELECT DATE '2024-03-01', 'a'
+)
+SELECT dt, letter, ROW_NUMBER() OVER (ORDER BY dt, letter) AS rn FROM toks
+`,
+			expectedRows: [][]interface{}{
+				{"2024-01-01", "c", int64(1)},
+				{"2024-02-01", "b", int64(2)},
+				{"2024-02-01", "c", int64(3)},
+				{"2024-03-01", "a", int64(4)},
+			},
+		},
+		{
+			name: `window order by handles nil`,
+			query: `WITH toks AS (
+			SELECT DATE '2024-01-01' AS dt, 'c' AS letter
+			UNION ALL SELECT DATE '2024-02-01', null
+)
+SELECT dt, letter, ROW_NUMBER() OVER (ORDER BY dt, letter) AS rn FROM toks
+`,
+			expectedRows: [][]interface{}{
+				{"2024-01-01", "c", int64(1)},
+				{"2024-02-01", nil, int64(2)},
+			},
+		},
 		{
 			name: `percentile_cont`,
 			query: `