瀏覽代碼

fix GROUP BY column types, DATE OIDs, and GROUP BY alias resolution

Danilo Fragoso 6 小時之前
父節點
當前提交
a8d5410617
共有 6 個文件被更改,包括 210 次插入25 次删除
  1. 二進制
      bin/pizzasql
  2. 12 9
      pkg/analyzer/analyzer.go
  3. 112 0
      pkg/executor/dashboard_repro_test.go
  4. 81 16
      pkg/executor/executor.go
  5. 4 0
      pkg/pgserver/connection.go
  6. 1 0
      pkg/pgserver/oid_test.go

二進制
bin/pizzasql


+ 12 - 9
pkg/analyzer/analyzer.go

@@ -140,14 +140,9 @@ func (a *Analyzer) analyzeSelect(stmt *parser.SelectStmt) error {
 	hasAggregate := false
 	hasGroupBy := len(stmt.GroupBy) > 0
 
-	// Analyze GROUP BY expressions first
-	for _, expr := range stmt.GroupBy {
-		if _, err := a.analyzeExpr(expr); err != nil {
-			return err
-		}
-	}
-
-	// Analyze SELECT columns and collect aliases for ORDER BY/HAVING reference
+	// Analyze SELECT columns and collect aliases for GROUP BY/ORDER BY/HAVING
+	// reference. SELECT aliases are registered before GROUP BY so a query may
+	// group by an output alias, matching SQLite.
 	selectAliases := make(map[string]*ExprInfo)
 	for _, col := range stmt.Columns {
 		if col.Star {
@@ -178,11 +173,19 @@ func (a *Analyzer) analyzeSelect(stmt *parser.SelectStmt) error {
 		}
 	}
 
-	// Register SELECT aliases as virtual columns for ORDER BY/HAVING reference
+	// Register SELECT aliases as virtual columns for GROUP BY/ORDER BY/HAVING
+	// reference
 	for alias, info := range selectAliases {
 		a.scope.DefineSelectAlias(alias, info.Type)
 	}
 
+	// Analyze GROUP BY expressions after aliases are in scope
+	for _, expr := range stmt.GroupBy {
+		if _, err := a.analyzeExpr(expr); err != nil {
+			return err
+		}
+	}
+
 	// Validate GROUP BY semantics
 	if hasAggregate && !hasGroupBy {
 		// Aggregate query without GROUP BY - all non-aggregate columns must be constants

+ 112 - 0
pkg/executor/dashboard_repro_test.go

@@ -0,0 +1,112 @@
+package executor
+
+import "testing"
+
+// TestSelectColumnTypesGroupBy verifies that GROUP BY projections keep their
+// schema types, so protocol clients can decode timestamp/date columns even
+// though the grouped projection is not a plain table scan.
+func TestSelectColumnTypesGroupBy(t *testing.T) {
+	_, schema, table := newTestDB(t)
+	e := newExec(schema, table)
+	execMust(t, e, `CREATE TABLE hit_counts (
+		site_id INTEGER NOT NULL,
+		path_id INTEGER NOT NULL,
+		hour TIMESTAMP NOT NULL,
+		total INTEGER NOT NULL
+	)`)
+	execMust(t, e, `CREATE TABLE hit_stats (
+		site_id INTEGER NOT NULL,
+		path_id INTEGER NOT NULL,
+		day DATE NOT NULL,
+		stats TEXT
+	)`)
+	execMust(t, e, `INSERT INTO hit_counts VALUES (1, 1, '2024-05-06 07:00:00', 3)`)
+	execMust(t, e, `INSERT INTO hit_stats VALUES (1, 1, '2024-05-06', '[]')`)
+
+	res := execMust(t, e, `SELECT hour, sum(total) FROM hit_counts GROUP BY hour ORDER BY hour`)
+	want := []string{"TIMESTAMP", "TEXT"}
+	if len(res.ColumnTypes) != len(want) {
+		t.Fatalf("group-by column types = %v, want %v", res.ColumnTypes, want)
+	}
+	for i := range want {
+		if res.ColumnTypes[i] != want[i] {
+			t.Fatalf("group-by column types = %v, want %v", res.ColumnTypes, want)
+		}
+	}
+
+	res = execMust(t, e, `SELECT path_id, day, stats FROM hit_stats ORDER BY day`)
+	want = []string{"INTEGER", "DATE", "TEXT"}
+	if len(res.ColumnTypes) != len(want) {
+		t.Fatalf("date column types = %v, want %v", res.ColumnTypes, want)
+	}
+	for i := range want {
+		if res.ColumnTypes[i] != want[i] {
+			t.Fatalf("date column types = %v, want %v", res.ColumnTypes, want)
+		}
+	}
+}
+
+// TestSubstrSQLiteSemantics pins SQLite's zero/negative start behavior, which
+// GoatCounter relies on to derive a country code from a region code.
+func TestSubstrSQLiteSemantics(t *testing.T) {
+	_, schema, table := newTestDB(t)
+	e := newExec(schema, table)
+
+	cases := []struct {
+		sql  string
+		want string
+	}{
+		{`SELECT substr('US-NY', 0, 3)`, "US"},
+		{`SELECT substr('US-NY', 1, 3)`, "US-"},
+		{`SELECT substr('US-NY', 2, 3)`, "S-N"},
+		{`SELECT substr('US-NY', -2, 2)`, "NY"},
+		{`SELECT substr('US-NY', 4)`, "NY"},
+		{`SELECT substr('US-NY', 0)`, "US-NY"},
+	}
+	for _, tc := range cases {
+		res := execMust(t, e, tc.sql)
+		if len(res.Rows) != 1 || res.Rows[0][0] != tc.want {
+			t.Fatalf("%s = %#v, want %q", tc.sql, res.Rows, tc.want)
+		}
+	}
+}
+
+// TestCTEGroupByAlias reproduces the GoatCounter locations query: a CTE that
+// groups by a SELECT alias, with the outer query joining on that alias column.
+func TestCTEGroupByAlias(t *testing.T) {
+	_, schema, table := newTestDB(t)
+	e := newExec(schema, table)
+	execMust(t, e, `CREATE TABLE location_stats (
+		site_id INTEGER NOT NULL,
+		path_id INTEGER NOT NULL,
+		day DATE NOT NULL,
+		location TEXT NOT NULL,
+		count INTEGER NOT NULL
+	)`)
+	execMust(t, e, `CREATE TABLE locations (
+		iso_3166_2 TEXT NOT NULL,
+		country_name TEXT NOT NULL
+	)`)
+	execMust(t, e, `INSERT INTO location_stats VALUES (1, 1, '2024-05-06', 'US-NY', 4)`)
+	execMust(t, e, `INSERT INTO locations VALUES ('US', 'United States')`)
+
+	res := execMust(t, e, `WITH x AS (
+		SELECT substr(location, 0, 3) AS loc, sum(count) AS count
+		FROM location_stats
+		WHERE site_id = 1 AND day >= '2024-01-01' AND day <= '2024-12-31'
+		GROUP BY loc
+		ORDER BY count DESC, loc
+		LIMIT 5
+	)
+	SELECT locations.iso_3166_2 AS id, locations.country_name AS name, x.count AS count
+	FROM x
+	JOIN locations ON locations.iso_3166_2 = x.loc
+	ORDER BY count DESC, name ASC`)
+
+	if res.RowCount != 1 {
+		t.Fatalf("expected 1 row, got %d (%v)", res.RowCount, res.Rows)
+	}
+	if res.Rows[0][0] != "US" || res.Rows[0][2].(int64) != 4 {
+		t.Fatalf("unexpected row %#v", res.Rows[0])
+	}
+}

+ 81 - 16
pkg/executor/executor.go

@@ -1734,17 +1734,28 @@ func (e *Executor) executeGroupBy(stmt *parser.SelectStmt, rows []storage.Row, s
 		}
 	}
 
+	// Populate column types from schema metadata, as the plain SELECT path
+	// does, so protocol clients can decode grouped timestamp/date columns.
+	result.ColumnTypes = make([]string, len(expandedColumns))
+	for i, col := range expandedColumns {
+		result.ColumnTypes[i] = projectionColumnType(col, schema)
+	}
+
+	// Resolve GROUP BY references that name a SELECT alias to the aliased
+	// expression, matching SQLite.
+	groupByExprs := e.resolveGroupByAliases(stmt.GroupBy, expandedColumns)
+
 	// Fast path: use running accumulators instead of collecting rows per group.
 	// Applicable when there is no HAVING clause and all aggregate SELECT columns
 	// are direct FunctionCalls (COUNT/SUM/AVG/MIN/MAX).
 	if e.canUseGroupAccum(stmt, expandedColumns) {
-		return e.executeGroupByAccum(stmt, rows, result, expandedColumns, columnNames)
+		return e.executeGroupByAccum(stmt, rows, result, expandedColumns, columnNames, groupByExprs)
 	}
 
 	// Slow path: collect full rows per group then evaluate aggregates over them.
 	groups := make(map[string][]storage.Row)
 	for _, row := range rows {
-		key := e.buildGroupKey(stmt.GroupBy, row)
+		key := e.buildGroupKey(groupByExprs, row)
 		groups[key] = append(groups[key], row)
 	}
 
@@ -1824,7 +1835,7 @@ type groupAccumState struct {
 
 // executeGroupByAccum is the fast GROUP BY path: increments per-group counters as rows
 // arrive rather than materialising row slices, keeping O(1) state per group.
-func (e *Executor) executeGroupByAccum(stmt *parser.SelectStmt, rows []storage.Row, result *Result, expandedColumns []parser.SelectColumn, columnNames []string) (*Result, error) {
+func (e *Executor) executeGroupByAccum(stmt *parser.SelectStmt, rows []storage.Row, result *Result, expandedColumns []parser.SelectColumn, columnNames []string, groupBy []parser.Expr) (*Result, error) {
 	var aggCols []aggColInfo
 	for i, col := range expandedColumns {
 		if e.isAggregate(col.Expr) {
@@ -1836,7 +1847,7 @@ func (e *Executor) executeGroupByAccum(stmt *parser.SelectStmt, rows []storage.R
 	var keyOrder []string
 
 	for _, row := range rows {
-		key := e.buildGroupKey(stmt.GroupBy, row)
+		key := e.buildGroupKey(groupBy, row)
 		state, exists := states[key]
 		if !exists {
 			accums := make([]*aggAccum, len(aggCols))
@@ -5150,22 +5161,47 @@ func (e *Executor) evalFunctionCall(fn *parser.FunctionCall, row storage.Row) (i
 		}
 	case "SUBSTR", "SUBSTRING":
 		if len(args) >= 2 {
-			s := toString(args[0])
-			start := int(toFloat(args[1])) - 1 // SQL is 1-indexed
-			if start < 0 {
-				start = 0
-			}
-			if start >= len(s) {
-				return "", nil
+			// SQLite substring semantics. Indices are 1-based and inclusive:
+			// a zero or negative Y shifts the window rather than simply
+			// clamping, e.g. substr('US-NY', 0, 3) is 'US', not 'US-'.
+			runes := []rune(toString(args[0]))
+			n := len(runes)
+			y := int(toFloat(args[1]))
+			var start1, end1 int
+			switch {
+			case y < 0:
+				start1 = n + y + 1
+			case y == 0:
+				start1 = 1
+			default:
+				start1 = y
 			}
 			if len(args) >= 3 {
-				length := int(toFloat(args[2]))
-				if start+length > len(s) {
-					length = len(s) - start
+				z := int(toFloat(args[2]))
+				switch {
+				case z < 0:
+					end1 = start1 - 1
+					start1 = end1 + z + 1
+				case y < 0:
+					end1 = start1 + z - 1
+				case y == 0:
+					end1 = z - 1
+				default:
+					end1 = y + z - 1
 				}
-				return s[start : start+length], nil
+			} else {
+				end1 = n
 			}
-			return s[start:], nil
+			if start1 < 1 {
+				start1 = 1
+			}
+			if end1 > n {
+				end1 = n
+			}
+			if start1 > n || start1 > end1 {
+				return "", nil
+			}
+			return string(runes[start1-1 : end1]), nil
 		}
 	case "TRIM":
 		if len(args) > 0 {
@@ -6347,6 +6383,35 @@ func (e *Executor) isAggregate(expr parser.Expr) bool {
 	return false
 }
 
+// resolveGroupByAliases replaces GROUP BY references that name a SELECT alias
+// with the aliased expression. SQLite allows `GROUP BY alias`, so grouping must
+// use the same expression the projection evaluates.
+func (e *Executor) resolveGroupByAliases(groupBy []parser.Expr, selectCols []parser.SelectColumn) []parser.Expr {
+	if len(groupBy) == 0 {
+		return groupBy
+	}
+	aliases := make(map[string]parser.Expr, len(selectCols))
+	for _, col := range selectCols {
+		if col.Alias != "" && col.Expr != nil {
+			aliases[strings.ToUpper(col.Alias)] = col.Expr
+		}
+	}
+	if len(aliases) == 0 {
+		return groupBy
+	}
+	out := make([]parser.Expr, len(groupBy))
+	for i, expr := range groupBy {
+		if ref, ok := expr.(*parser.ColumnRef); ok {
+			if target, ok := aliases[strings.ToUpper(ref.Column)]; ok {
+				out[i] = target
+				continue
+			}
+		}
+		out[i] = expr
+	}
+	return out
+}
+
 func (e *Executor) buildGroupKey(groupBy []parser.Expr, row storage.Row) string {
 	var parts []string
 	for _, expr := range groupBy {

+ 4 - 0
pkg/pgserver/connection.go

@@ -1143,6 +1143,8 @@ func (c *Connection) getOIDForType(typeName string) int32 {
 		// Datetime values are exchanged as UTC RFC3339 (ISO8601 with a timezone
 		// offset), which timestamptz decodes directly.
 		return 1184 // TIMESTAMPTZOID
+	case "DATE":
+		return 1082 // DATEOID
 	default:
 		return 25 // Default to TEXT
 	}
@@ -1163,6 +1165,8 @@ func (c *Connection) getTypeSizeForType(typeName string) int16 {
 		return 1
 	case "DATETIME", "TIMESTAMP":
 		return 8
+	case "DATE":
+		return 4
 	default:
 		return -1 // Variable length
 	}

+ 1 - 0
pkg/pgserver/oid_test.go

@@ -14,6 +14,7 @@ func TestGetOIDForTypeDatetime(t *testing.T) {
 	cases := map[string]int32{
 		"DATETIME":  1184,
 		"TIMESTAMP": 1184,
+		"DATE":      1082,
 		"BIGINT":    20,
 		"INTEGER":   23,
 		"INT":       23,