Skip to content

Commit

Permalink
[Windowing] [Formatter] Reset format context in-between analytic func…
Browse files Browse the repository at this point in the history
…tion groups (#162)

* [Windowing] [Formatter] Reset format context in-between analytic function groups

* use last function order as scan order
  • Loading branch information
ohaibbq authored Mar 10, 2024
1 parent d5aff7d commit 5837a05
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 5 deletions.
21 changes: 16 additions & 5 deletions internal/formatter.go
Original file line number Diff line number Diff line change
Expand Up @@ -1081,7 +1081,10 @@ func (n *AnalyticScanNode) FormatSQL(ctx context.Context) (string, error) {
}
ctx = withAnalyticInputScan(ctx, formattedInput)
orderColumnNames := analyticOrderColumnNamesFromContext(ctx)
var scanOrderBy []*analyticOrderBy
for _, group := range n.node.FunctionGroupList() {
scanOrderBy = []*analyticOrderBy{}

if group.PartitionBy() != nil {
var partitionColumns []string
for _, columnRef := range group.PartitionBy().PartitionByList() {
Expand All @@ -1090,26 +1093,34 @@ func (n *AnalyticScanNode) FormatSQL(ctx context.Context) (string, error) {
partitionColumns,
colName,
)
orderColumnNames.values = append(orderColumnNames.values, &analyticOrderBy{
order := &analyticOrderBy{
column: colName,
isAsc: true,
})
}
orderColumnNames.values = append(orderColumnNames.values, order)
scanOrderBy = append(scanOrderBy, order)
}
ctx = withAnalyticPartitionColumnNames(ctx, partitionColumns)
}
if group.OrderBy() != nil {
for _, item := range group.OrderBy().OrderByItemList() {
colName := uniqueColumnName(ctx, item.ColumnRef().Column())
formattedColName := fmt.Sprintf("`%s`", colName)
orderColumnNames.values = append(orderColumnNames.values, &analyticOrderBy{
order := &analyticOrderBy{
column: formattedColName,
isAsc: !item.IsDescending(),
})
}
orderColumnNames.values = append(orderColumnNames.values, order)
scanOrderBy = append(scanOrderBy, order)
}
}
if _, err := newNode(group).FormatSQL(ctx); err != nil {
return "", err
}

// Reset context after each analytic function group
orderColumnNames.values = []*analyticOrderBy{}
ctx = withAnalyticPartitionColumnNames(ctx, nil)
}
columns := []string{}
columnMap := columnRefMap(ctx)
Expand All @@ -1126,7 +1137,7 @@ func (n *AnalyticScanNode) FormatSQL(ctx context.Context) (string, error) {
}
}
var orderColumnFormattedNames []string
for _, col := range orderColumnNames.values {
for _, col := range scanOrderBy {
if col.isAsc {
orderColumnFormattedNames = append(
orderColumnFormattedNames,
Expand Down
22 changes: 22 additions & 0 deletions query_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2103,6 +2103,28 @@ FROM finishers`,
{"Suzy Slane", createTimestampFormatFromString("2016-10-18 03:06:24+00"), "F35-39", "Desiree Berry"},
},
},
// Regression test for https://github.com/goccy/go-zetasqlite/issues/160
{
name: "window partitions are distinct from each other",
query: `
WITH inventory AS (
SELECT 'banana' AS item, 'fruit' AS kind, 2 AS purchases
UNION ALL SELECT 'onion', 'vegetable', 3
UNION ALL SELECT 'orange', 'fruit', 4
ORDER BY item ASC
)
SELECT
item,
purchases,
LEAD(item) OVER (PARTITION BY kind ORDER BY item ASC) AS next_in_kind,
LAG(item) OVER (ORDER BY purchases ASC) AS next_best_seller
FROM inventory`,
expectedRows: [][]interface{}{
{"banana", int64(2), "orange", nil},
{"onion", int64(3), nil, "banana"},
{"orange", int64(4), nil, "onion"},
},
},
{
name: "lag with option",
query: `
Expand Down

0 comments on commit 5837a05

Please sign in to comment.