Use the AST to implement aggregate functions.

Author: dburkart

aggregate.go

@@ -2,8 +2,6 @@ package duckql
 
 import (
 	"reflect"
-	"strconv"
-	"strings"
 )
 
 type AggregateFunctionColumn struct {
@@ -185,37 +183,3 @@ func sumOfColumn(c *AggregateFunctionColumn, rows ResultRows) ResultRows {
 		sumRow,
 	}
 }
-
-func ParseAggregateFunction(text string) *AggregateFunctionColumn {
-	var column AggregateFunctionColumn
-	var current strings.Builder
-
-	for _, r := range text {
-		switch r {
-		case '(':
-			functionName := current.String()
-			if f, ok := functionMap[functionName]; ok {
-				column.Function = f
-			} else {
-				return nil
-			}
-
-			current.Reset()
-		case ')':
-			var err error
-			column.UnderlyingColumn, err = strconv.Unquote(current.String())
-			if err != nil {
-				column.UnderlyingColumn = current.String()
-			}
-			current.Reset()
-		default:
-			current.WriteRune(r)
-		}
-	}
-
-	if current.String() == "" {
-		return &column
-	}
-
-	return nil
-}

filter_backing.go

@@ -259,60 +259,124 @@ func (f *SliceFilter) Rows() ResultRows {
 	var r ResultRows
 
 	source := f.intermediate.Result()
-
 	source = source.Filter(f.filter)
 
+	if len(source.Rows) == 0 {
+		return r
+	}
+
 	// Transform our intermediate columns into a lookup table
 	lookup := make(map[string]int)
 	for idx, column := range source.Columns {
 		lookup[column] = idx
 	}
 
-	for _, row := range source.Rows {
-		var narrow ResultRow
+	// Find column positions to narrow
+	var narrowColumns []int
+	for _, column := range f.resultColumns {
+		if column.Star.Line > 0 {
+			for idx, _ := range source.Rows[0] {
+				narrowColumns = append(narrowColumns, idx)
+			}
+			break
+		}
 
-		for _, column := range f.resultColumns {
-			if column.Star.Line > 0 {
-				narrow = row
+		switch t := column.Expr.(type) {
+		case *sql.Ident:
+			index, ok := lookup[t.Name]
+			if !ok {
+				// FIXME: There should be a better way
+				parts := strings.Split(source.Columns[0], ".")
+				if len(parts) > 1 {
+					index = lookup[parts[0]+"."+t.Name]
+				}
 			}
 
-			switch t := column.Expr.(type) {
-			case *sql.Ident:
-				index, ok := lookup[t.Name]
+			narrowColumns = append(narrowColumns, index)
+		case *sql.QualifiedRef:
+			if t.Star.Line != 0 {
+				// FIXME: Implement
+				continue
+			}
+
+			lh := t.Table.Name
+			rh := t.Column.Name
+
+			var index int
+			var ok bool
+			if source.Source != nil {
+				index, ok = lookup[rh]
+			} else {
+				index, ok = lookup[lh+"."+rh]
 				if !ok {
+					index = lookup[source.Aliases[lh]+"."+rh]
+				}
+			}
+
+			narrowColumns = append(narrowColumns, index)
+		default:
+			narrowColumns = append(narrowColumns, 0)
+		}
+	}
+
+	// Find aggregations
+	var aggregations []AggregateFunctionColumn
+	for idx, column := range f.resultColumns {
+		switch t := column.Expr.(type) {
+		case *sql.Call:
+			var underlying string
+			if t.Star.Line == 0 {
+				if len(t.Args) != 1 {
+					panic("unexpected number of args to function")
+				}
+
+				switch arg := t.Args[0].(type) {
+				case *sql.Ident:
+					// Validate?
+					index, ok := lookup[arg.Name]
+					if ok {
+						underlying = arg.Name
+						break
+					}
+
 					// FIXME: There should be a better way
 					parts := strings.Split(source.Columns[0], ".")
 					if len(parts) > 1 {
-						index = lookup[parts[0]+"."+t.Name]
+						index, ok = lookup[parts[0]+"."+arg.Name]
+						if ok {
+							underlying = parts[0] + "." + arg.Name
+						}
 					}
-				}
 
-				narrow = append(narrow, row[index])
-			case *sql.QualifiedRef:
-				if t.Star.Line != 0 {
-					// FIXME: Implement
-					continue
+					narrowColumns[idx] = index
 				}
+			} else {
+				narrowColumns[idx] = 0
+			}
 
-				lh := t.Table.Name
-				rh := t.Column.Name
-
-				var index int
-				var ok bool
-				if source.Source != nil {
-					index, ok = lookup[rh]
-				} else {
-					index, ok = lookup[lh+"."+rh]
-					if !ok {
-						index = lookup[source.Aliases[lh]+"."+rh]
-					}
-				}
+			aggregations = append(aggregations, AggregateFunctionColumn{
+				UnderlyingColumn: underlying,
+				ResultPosition:   idx,
+				Function:         functionMap[t.Name.Name],
+			})
+		}
+	}
 
-				narrow = append(narrow, row[index])
-			}
+	for _, row := range source.Rows {
+		if len(row) == len(narrowColumns) {
+			r = append(r, row)
+			continue
+		}
+
+		var newRow ResultRow
+		for _, column := range narrowColumns {
+			newRow = append(newRow, row[column])
 		}
+		r = append(r, newRow)
+	}
 
-		r = append(r, narrow)
+	for _, aggregation := range aggregations {
+		r = aggregation.Call(r)
 	}
 
 	f.resultColumns = []*sql.ResultColumn{}

sqlite_backing.go

@@ -18,7 +18,6 @@ type SQLiteBacking struct {
 
 // New creates a new SQLiteBacking with the given SQLite database connection
 func NewSQLiteBacking(db *gosql.DB, s *SQLizer) *SQLiteBacking {
-	s.HandleAggregateFunctions = false
 	return &SQLiteBacking{
 		sqlizer: s,
 		db:      db,

sqlizer.go

@@ -28,11 +28,9 @@ const (
 )
 
 type SQLizer struct {
-	Tables                   map[string]*Table
-	Permissions              uint
-	Backing                  BackingStore
-	AggregateFunctions       []*AggregateFunctionColumn
-	HandleAggregateFunctions bool
+	Tables      map[string]*Table
+	Permissions uint
+	Backing     BackingStore
 }
 
 func (s *SQLizer) SetPermissions(permissions uint) {
@@ -75,22 +73,7 @@ func (s *SQLizer) Execute(statement string) (ResultRows, error) {
 			return nil, err
 		}
 
-		rows := s.Backing.Rows()
-
-		if len(rows) > 0 && len(s.AggregateFunctions) > 0 && s.HandleAggregateFunctions {
-			var result ResultRows
-			for idx, aggregate := range s.AggregateFunctions {
-				r := aggregate.Call(rows)
-				if idx == 0 {
-					result = append(result, r[0])
-				} else {
-					result[0][aggregate.ResultPosition] = r[0][aggregate.ResultPosition]
-				}
-			}
-			rows = result
-		}
-
-		return rows, nil
+		return s.Backing.Rows(), nil
 	}
 
 	return nil, nil
@@ -204,8 +187,6 @@ func Initialize(structs ...any) *SQLizer {
 		sql.addStructTable(s)
 	}
 
-	sql.HandleAggregateFunctions = true
-
 	return &sql
 }
 

test/expectations/01_user_select_alias.txt

@@ -0,0 +1,14 @@
+.section = DDL
+---
+CREATE TABLE users
+(
+  id INTEGER,
+  name TEXT,
+  email TEXT
+)
+
+---
+.section = Result
+---
+John Doe
+Jane Smith

test/expectations/01_user_select_alias_no_alias.txt

@@ -0,0 +1,14 @@
+.section = DDL
+---
+CREATE TABLE users
+(
+  id INTEGER,
+  name TEXT,
+  email TEXT
+)
+
+---
+.section = Result
+---
+John Doe
+Jane Smith

test/expectations/04_account_join_default_alias.txt

@@ -0,0 +1,20 @@
+.section = DDL
+---
+CREATE TABLE accounts
+(
+  id INTEGER,
+  username TEXT,
+  email TEXT,
+  organization_id INTEGER
+)
+
+CREATE TABLE organizations
+(
+  id INTEGER,
+  name TEXT
+)
+
+---
+.section = Result
+---
+user1|Initech

validate.go

@@ -35,32 +35,16 @@ func (v *Validator) Visit(n sql.Node) (sql.Visitor, sql.Node, error) {
 		case *sql.Ident:
 			v.columns = append(v.columns, e.Name)
 		case *sql.Call:
-			var aggregate AggregateFunctionColumn
-			var star sql.Pos
+			var underlyingColumn string
 
 			// FIXME
 			if len(e.Args) > 0 {
-				aggregate.UnderlyingColumn = e.Args[0].(*sql.Ident).Name
+				underlyingColumn = e.Args[0].(*sql.Ident).Name
 			} else {
-				aggregate.UnderlyingColumn = "*"
-				star.Line = 1
+				underlyingColumn = "*"
 			}
 
-			aggregate.ResultPosition = len(v.columns)
-			aggregate.Function = functionMap[e.Name.Name]
-
-			v.s.AggregateFunctions = append(v.s.AggregateFunctions, &aggregate)
-			// FIXME
-			if v.s.HandleAggregateFunctions {
-				n = &sql.ResultColumn{
-					Star: star,
-					Expr: &sql.Ident{
-						Name:   aggregate.UnderlyingColumn,
-						Quoted: false,
-					},
-				}
-			}
-			v.columns = append(v.columns, aggregate.UnderlyingColumn)
+			v.columns = append(v.columns, underlyingColumn)
 		}
 
 	case *sql.QualifiedTableName: