Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Windowing] Rewrite window function implementation to use real SQLite windows #169

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open

[Windowing] Rewrite window function implementation to use real SQLite windows #169

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7f7bf52
[Windowing] Rewrite window function implementation to use real SQLite…
ohaibbq Feb 7, 2024
e02bc8f
remove wip
ohaibbq Feb 15, 2024
7ed3074
update fork
ohaibbq Feb 15, 2024
e98ab9d
lint
ohaibbq Feb 16, 2024
bc44157
more lint
ohaibbq Feb 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,3 +63,5 @@ require (
google.golang.org/grpc v1.54.0 // indirect
google.golang.org/protobuf v1.30.0 // indirect
)

replace github.com/mattn/go-sqlite3 => github.com/ohaibbq/go-sqlite3 v0.0.0-20240211011509-f8d4d3382d11
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This will be replaced with upstream once my PR makes it into the new release.

4 changes: 2 additions & 2 deletions go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,12 +100,12 @@ github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NB
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
github.com/mattn/go-sqlite3 v1.14.16 h1:yOQRA0RpS5PFz/oikGwBEqvAWhWg5ufRz4ETLjwpU1Y=
github.com/mattn/go-sqlite3 v1.14.16/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg=
github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs=
github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY=
github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3 h1:+n/aFZefKZp7spd8DFdX7uMikMLXX4oubIzJF4kv/wI=
github.com/minio/c2goasm v0.0.0-20190812172519-36a3d3bbc4f3/go.mod h1:RagcQ7I8IeTMnF8JTXieKnO4Z6JCsikNEzj0DwauVzE=
github.com/ohaibbq/go-sqlite3 v0.0.0-20240211011509-f8d4d3382d11 h1:GaOapuUZae9qDJokb4kKWLjolR38lBN/LyZtZap1q74=
github.com/ohaibbq/go-sqlite3 v0.0.0-20240211011509-f8d4d3382d11/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
github.com/pierrec/lz4/v4 v4.1.15 h1:MO0/ucJhngq7299dKLwIMtgTfbkoSPF6AoMYDd8Q4q0=
github.com/pierrec/lz4/v4 v4.1.15/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4=
github.com/pkg/errors v0.8.0 h1:WdK/asTD0HN+q6hsWO3/vpuAkAr+tw6aNJNDFFf0+qw=
Expand Down
13 changes: 0 additions & 13 deletions internal/context.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ type (
funcMapKey struct{}
analyticOrderColumnNamesKey struct{}
analyticPartitionColumnNamesKey struct{}
analyticInputScanKey struct{}
arraySubqueryColumnNameKey struct{}
currentTimeKey struct{}
tableNameToColumnListMapKey struct{}
Expand Down Expand Up @@ -117,18 +116,6 @@ func analyticPartitionColumnNamesFromContext(ctx context.Context) []string {
return value.([]string)
}

func withAnalyticInputScan(ctx context.Context, input string) context.Context {
return context.WithValue(ctx, analyticInputScanKey{}, input)
}

func analyticInputScanFromContext(ctx context.Context) string {
value := ctx.Value(analyticInputScanKey{})
if value == nil {
return ""
}
return value.(string)
}

type arraySubqueryColumnNames struct {
names []string
}
Expand Down
164 changes: 115 additions & 49 deletions internal/formatter.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -336,6 +336,21 @@ func (n *AggregateFunctionCallNode) FormatSQL(ctx context.Context) (string, erro
), nil
}

var windowFuncFixedRanges = map[string]string{
"zetasqlite_window_ntile": "ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING",
"zetasqlite_window_cume_dist": "GROUPS BETWEEN 1 FOLLOWING AND UNBOUNDED FOLLOWING",
"zetasqlite_window_dense_rank": "RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW",
"zetasqlite_window_rank": "GROUPS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW EXCLUDE TIES",
"zetasqlite_window_percent_rank": "GROUPS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING",
"zetasqlite_window_row_number": "ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW",
"zetasqlite_window_lag": "ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW",
"zetasqlite_window_lead": "ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING",
}

var windowFunctionsIgnoreNullsByDefault = map[string]bool{
"zetasqlite_window_percentile_disc": true,
}

func (n *AnalyticFunctionCallNode) FormatSQL(ctx context.Context) (string, error) {
if n.node == nil {
return "", nil
Expand All @@ -346,70 +361,122 @@ func (n *AnalyticFunctionCallNode) FormatSQL(ctx context.Context) (string, error
if err != nil {
return "", err
}
var opts []string
if n.node.Distinct() {
opts = append(opts, "zetasqlite_distinct()")
}
switch n.node.NullHandlingModifier() {
case ast.RespectNulls:
// do nothing
default:
opts = append(opts, "zetasqlite_ignore_nulls()")
}
args = append(args, opts...)
for _, column := range analyticPartitionColumnNamesFromContext(ctx) {
args = append(args, getWindowPartitionOptionFuncSQL(column))
funcMap := funcMapFromContext(ctx)

overClause := []string{}
partitionColumns := analyticPartitionColumnNamesFromContext(ctx)

if len(partitionColumns) > 0 {
overClause = append(overClause, "PARTITION BY")
columns := []string{}
for _, column := range partitionColumns {
columns = append(columns, fmt.Sprintf("%s COLLATE zetasqlite_collate", column))
}
overClause = append(overClause, strings.Join(columns, ", "))
}
for _, col := range orderColumns {
args = append(args, getWindowOrderByOptionFuncSQL(col.column, col.isAsc))

frame := n.node.WindowFrame()
frameSQL, found := windowFuncFixedRanges[funcName]
if found && frame != nil {
return "", fmt.Errorf("%s: window framing clause is not allowed for analytic function", n.node.BaseFunctionCallNode.Function().Name())
}
windowFrame := n.node.WindowFrame()
if windowFrame != nil {
args = append(args, getWindowFrameUnitOptionFuncSQL(windowFrame.FrameUnit()))
startSQL, err := n.getWindowBoundaryOptionFuncSQL(ctx, windowFrame.StartExpr(), true)
if !found {
frameSQL, err = n.getWindowBoundaryOptionFuncSQL(ctx, n.node.WindowFrame())
if err != nil {
return "", err
return "", nil
}
endSQL, err := n.getWindowBoundaryOptionFuncSQL(ctx, windowFrame.EndExpr(), false)
if err != nil {
return "", err
}

if len(orderColumns) > 0 {
overClause = append(overClause, "ORDER BY")
columns := []string{}
for _, column := range orderColumns {
dir := "ASC"
if !column.isAsc {
dir = "DESC"
}
columns = append(columns, fmt.Sprintf("%s COLLATE zetasqlite_collate %s", column.column, dir))
}
args = append(args, startSQL)
args = append(args, endSQL)
overClause = append(overClause, strings.Join(columns, ", "))
}
args = append(args, getWindowRowIDOptionFuncSQL())
input := analyticInputScanFromContext(ctx)
funcMap := funcMapFromContext(ctx)

overClause = append(overClause, frameSQL)

if n.node.Distinct() {
args = append(args, "zetasqlite_distinct()")
}

_, ignoreNullsByDefault := windowFunctionsIgnoreNullsByDefault[funcName]

switch n.node.NullHandlingModifier() {
case ast.IgnoreNulls:
args = append(args, "zetasqlite_ignore_nulls()")
case ast.DefaultNullHandling:
if ignoreNullsByDefault {
args = append(args, "zetasqlite_ignore_nulls()")
}
}

if spec, exists := funcMap[funcName]; exists {
return spec.CallSQL(ctx, n.node.BaseFunctionCallNode, args)
}
return fmt.Sprintf(
"( SELECT %s(%s) %s )",
"%s(%s) OVER (%s)",
funcName,
strings.Join(args, ","),
input,
strings.Join(overClause, " "),
), nil
}

func (n *AnalyticFunctionCallNode) getWindowBoundaryOptionFuncSQL(ctx context.Context, expr *ast.WindowFrameExprNode, isStart bool) (string, error) {
typ := expr.BoundaryType()
switch typ {
case ast.UnboundedPrecedingType, ast.CurrentRowType, ast.UnboundedFollowingType:
if isStart {
return getWindowBoundaryStartOptionFuncSQL(typ, ""), nil
}
return getWindowBoundaryEndOptionFuncSQL(typ, ""), nil
case ast.OffsetPrecedingType, ast.OffsetFollowingType:
literal, err := newNode(expr.Expression()).FormatSQL(ctx)
if err != nil {
return "", err
}
if isStart {
return getWindowBoundaryStartOptionFuncSQL(typ, literal), nil
func getWindowBoundarySQL(boundaryType ast.BoundaryType, literal string) string {
switch boundaryType {
case ast.UnboundedPrecedingType:
return "UNBOUNDED PRECEDING"
case ast.OffsetPrecedingType:
return fmt.Sprintf("%s PRECEDING", literal)
case ast.CurrentRowType:
return "CURRENT ROW"
case ast.OffsetFollowingType:
return fmt.Sprintf("%s FOLLOWING", literal)
case ast.UnboundedFollowingType:
return "UNBOUNDED FOLLOWING"
}
return ""
}

func (n *AnalyticFunctionCallNode) getWindowBoundaryOptionFuncSQL(ctx context.Context, node *ast.WindowFrameNode) (string, error) {
if node == nil {
return "ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", nil
}

frames := [2]*ast.WindowFrameExprNode{node.StartExpr(), node.EndExpr()}
sql := []string{}
for _, expr := range frames {

typ := expr.BoundaryType()
switch typ {
case ast.UnboundedPrecedingType, ast.CurrentRowType, ast.UnboundedFollowingType:
sql = append(sql, getWindowBoundarySQL(typ, ""))
case ast.OffsetPrecedingType, ast.OffsetFollowingType:
literal, err := newNode(expr.Expression()).FormatSQL(ctx)
if err != nil {
return "", err
}
sql = append(sql, getWindowBoundarySQL(typ, literal))
default:
return "", fmt.Errorf("unexpected boundary type %d", typ)
}
return getWindowBoundaryEndOptionFuncSQL(typ, literal), nil
}
return "", fmt.Errorf("unexpected boundary type %d", typ)
var unit string
switch node.FrameUnit() {
case ast.FrameUnitRows:
unit = "ROWS"
case ast.FrameUnitRange:
unit = "RANGE"
default:
return "", fmt.Errorf("unexpected frame unit %d", node.FrameUnit())
}
return fmt.Sprintf("%s BETWEEN %s AND %s", unit, sql[0], sql[1]), nil
}

func (n *ExtendedCastElementNode) FormatSQL(ctx context.Context) (string, error) {
Expand Down Expand Up @@ -1041,7 +1108,6 @@ func (n *AnalyticScanNode) FormatSQL(ctx context.Context) (string, error) {
if err != nil {
return "", err
}
ctx = withAnalyticInputScan(ctx, formattedInput)
orderColumnNames := analyticOrderColumnNamesFromContext(ctx)
for _, group := range n.node.FunctionGroupList() {
if group.PartitionBy() != nil {
Expand Down Expand Up @@ -1107,7 +1173,7 @@ func (n *AnalyticScanNode) FormatSQL(ctx context.Context) (string, error) {
}
orderColumnNames.values = []*analyticOrderBy{}
return fmt.Sprintf(
"SELECT %s FROM (SELECT *, ROW_NUMBER() OVER() AS `row_id` %s) %s",
"SELECT %s %s %s",
strings.Join(columns, ","),
formattedInput,
orderBy,
Expand Down
Loading
Loading