Skip to content

Commit 4b53cb0

Browse files
committed
fixing an issue with aliases in aggregation queries
1 parent 81f4b3c commit 4b53cb0

File tree

4 files changed

+259
-10
lines changed

4 files changed

+259
-10
lines changed

pkg/cypher/match.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1315,10 +1315,10 @@ func (e *StorageExecutor) executeMatchRelationshipsWithClause(ctx context.Contex
13151315
// Check for WHERE between WITH and RETURN (post-aggregation filter, like SQL HAVING)
13161316
var withClause string
13171317
var postWithWhere string
1318-
postWhereIdx := findKeywordNotInBrackets(strings.ToUpper(withSection), " WHERE ")
1318+
postWhereIdx := findKeywordIndex(withSection, "WHERE")
13191319
if postWhereIdx > 0 {
13201320
withClause = strings.TrimSpace(withSection[:postWhereIdx])
1321-
postWithWhere = strings.TrimSpace(withSection[postWhereIdx+7:]) // Skip " WHERE "
1321+
postWithWhere = strings.TrimSpace(withSection[postWhereIdx+5:]) // Skip "WHERE"
13221322
} else {
13231323
withClause = withSection
13241324
}

pkg/cypher/optimized_executors.go

Lines changed: 43 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -65,8 +65,12 @@ func (e *StorageExecutor) ExecuteOptimized(ctx context.Context, query string, pa
6565
// =============================================================================
6666

6767
func (e *StorageExecutor) executeMutualRelationshipOptimized(ctx context.Context, query string, info PatternInfo) (*ExecuteResult, error) {
68+
// Parse RETURN clause to get actual column names/aliases
69+
returnItems := e.extractReturnItemsFromQuery(query)
70+
columns := e.buildColumnsFromReturnItems(returnItems)
71+
6872
result := &ExecuteResult{
69-
Columns: []string{info.StartVar + ".name", info.EndVar + ".name"},
73+
Columns: columns,
7074
Rows: [][]interface{}{},
7175
Stats: &QueryStats{},
7276
}
@@ -132,8 +136,12 @@ func (e *StorageExecutor) executeMutualRelationshipOptimized(ctx context.Context
132136
// =============================================================================
133137

134138
func (e *StorageExecutor) executeIncomingCountOptimized(ctx context.Context, query string, info PatternInfo) (*ExecuteResult, error) {
139+
// Parse RETURN clause to get actual column names/aliases
140+
returnItems := e.extractReturnItemsFromQuery(query)
141+
columns := e.buildColumnsFromReturnItems(returnItems)
142+
135143
result := &ExecuteResult{
136-
Columns: []string{info.StartVar + ".name", "count"},
144+
Columns: columns,
137145
Rows: [][]interface{}{},
138146
Stats: &QueryStats{},
139147
}
@@ -196,8 +204,12 @@ func (e *StorageExecutor) executeIncomingCountOptimized(ctx context.Context, que
196204
// =============================================================================
197205

198206
func (e *StorageExecutor) executeOutgoingCountOptimized(ctx context.Context, query string, info PatternInfo) (*ExecuteResult, error) {
207+
// Parse RETURN clause to get actual column names/aliases
208+
returnItems := e.extractReturnItemsFromQuery(query)
209+
columns := e.buildColumnsFromReturnItems(returnItems)
210+
199211
result := &ExecuteResult{
200-
Columns: []string{info.StartVar + ".name", "count"},
212+
Columns: columns,
201213
Rows: [][]interface{}{},
202214
Stats: &QueryStats{},
203215
}
@@ -326,11 +338,9 @@ func (e *StorageExecutor) executeEdgePropertyAggOptimized(ctx context.Context, q
326338
}
327339
}
328340

329-
// Build columns based on aggregation functions
330-
result.Columns = append(result.Columns, "name")
331-
for _, fn := range info.AggFunctions {
332-
result.Columns = append(result.Columns, fn)
333-
}
341+
// Build columns from RETURN clause (respects aliases)
342+
returnItems := e.extractReturnItemsFromQuery(query)
343+
result.Columns = e.buildColumnsFromReturnItems(returnItems)
334344

335345
// Convert to result rows
336346
type aggRow struct {
@@ -391,6 +401,31 @@ func (e *StorageExecutor) executeEdgePropertyAggOptimized(ctx context.Context, q
391401
// Helper Functions
392402
// =============================================================================
393403

404+
// extractReturnItemsFromQuery extracts RETURN items from a Cypher query
405+
func (e *StorageExecutor) extractReturnItemsFromQuery(query string) []returnItem {
406+
upperQuery := strings.ToUpper(query)
407+
returnIdx := strings.Index(upperQuery, "RETURN")
408+
if returnIdx == -1 {
409+
return nil
410+
}
411+
412+
returnPart := strings.TrimSpace(query[returnIdx+6:])
413+
return e.parseReturnItems(returnPart)
414+
}
415+
416+
// buildColumnsFromReturnItems builds column names from parsed return items
417+
func (e *StorageExecutor) buildColumnsFromReturnItems(items []returnItem) []string {
418+
columns := make([]string, len(items))
419+
for i, item := range items {
420+
if item.alias != "" {
421+
columns[i] = item.alias
422+
} else {
423+
columns[i] = item.expr
424+
}
425+
}
426+
return columns
427+
}
428+
394429
// getNodeCached retrieves a node, using cache to avoid repeated lookups
395430
func (e *StorageExecutor) getNodeCached(id storage.NodeID, cache map[storage.NodeID]*storage.Node) *storage.Node {
396431
if node, exists := cache[id]; exists {

pkg/cypher/query_patterns.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,13 @@ func DetectQueryPattern(query string) PatternInfo {
122122

123123
upperQuery := strings.ToUpper(query)
124124

125+
// Don't optimize queries with WITH clause - they have complex aggregation
126+
// semantics that the optimized executors don't handle (aliases, collect, etc.)
127+
// Use word boundary check to avoid matching "STARTS WITH" or "ENDS WITH"
128+
if containsKeywordOutsideStrings(query, "WITH") {
129+
return info
130+
}
131+
125132
// Extract LIMIT first (affects multiple patterns)
126133
if matches := patternLimitRegex.FindStringSubmatch(query); matches != nil {
127134
limit, _ := strconv.Atoi(matches[1])

pkg/cypher/return_alias_test.go

Lines changed: 207 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,207 @@
1+
package cypher
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"github.com/orneryd/nornicdb/pkg/storage"
8+
"github.com/stretchr/testify/assert"
9+
"github.com/stretchr/testify/require"
10+
)
11+
12+
func TestReturnAliases_WithAggregation(t *testing.T) {
13+
store := storage.NewMemoryEngine()
14+
exec := NewStorageExecutor(store)
15+
ctx := context.Background()
16+
17+
// Setup test data
18+
exec.Execute(ctx, `CREATE (p1:Person {name: 'Alice'})`, nil)
19+
exec.Execute(ctx, `CREATE (p2:Person {name: 'Bob'})`, nil)
20+
exec.Execute(ctx, `CREATE (a1:Area {name: 'Engineering'})`, nil)
21+
exec.Execute(ctx, `CREATE (a2:Area {name: 'Sales'})`, nil)
22+
exec.Execute(ctx, `CREATE (a3:Area {name: 'Marketing'})`, nil)
23+
exec.Execute(ctx, `MATCH (p:Person {name: 'Alice'}), (a:Area {name: 'Engineering'}) CREATE (p)-[:WORKS_IN]->(a)`, nil)
24+
exec.Execute(ctx, `MATCH (p:Person {name: 'Alice'}), (a:Area {name: 'Sales'}) CREATE (p)-[:WORKS_IN]->(a)`, nil)
25+
exec.Execute(ctx, `MATCH (p:Person {name: 'Bob'}), (a:Area {name: 'Marketing'}) CREATE (p)-[:WORKS_IN]->(a)`, nil)
26+
27+
t.Run("WITH count and collect with aliases", func(t *testing.T) {
28+
query := `
29+
MATCH (p:Person)-[:WORKS_IN]->(a:Area)
30+
WITH p, count(a) as area_count, collect(a.name) as areas
31+
RETURN p.name as name, area_count, areas
32+
`
33+
result, err := exec.Execute(ctx, query, nil)
34+
require.NoError(t, err)
35+
36+
// Check column names respect aliases
37+
assert.Equal(t, []string{"name", "area_count", "areas"}, result.Columns)
38+
39+
// Check we have results
40+
assert.Len(t, result.Rows, 2)
41+
})
42+
43+
t.Run("WITH count alias only", func(t *testing.T) {
44+
query := `
45+
MATCH (p:Person)-[:WORKS_IN]->(a:Area)
46+
WITH p, count(a) as area_count
47+
RETURN p.name as name, area_count
48+
`
49+
result, err := exec.Execute(ctx, query, nil)
50+
require.NoError(t, err)
51+
52+
assert.Equal(t, []string{"name", "area_count"}, result.Columns)
53+
assert.Len(t, result.Rows, 2)
54+
})
55+
56+
t.Run("simple RETURN with alias", func(t *testing.T) {
57+
query := `MATCH (p:Person) RETURN p.name as personName`
58+
result, err := exec.Execute(ctx, query, nil)
59+
require.NoError(t, err)
60+
61+
assert.Equal(t, []string{"personName"}, result.Columns)
62+
assert.Len(t, result.Rows, 2)
63+
})
64+
65+
t.Run("multiple aliases in RETURN", func(t *testing.T) {
66+
query := `MATCH (p:Person) RETURN p.name as name, 'constant' as type`
67+
result, err := exec.Execute(ctx, query, nil)
68+
require.NoError(t, err)
69+
70+
assert.Equal(t, []string{"name", "type"}, result.Columns)
71+
})
72+
}
73+
74+
func TestReturnAliases_OptimizedExecutors(t *testing.T) {
75+
store := storage.NewMemoryEngine()
76+
exec := NewStorageExecutor(store)
77+
ctx := context.Background()
78+
79+
// Setup test data for optimized executor tests
80+
exec.Execute(ctx, `CREATE (p1:Person {name: 'Alice'})`, nil)
81+
exec.Execute(ctx, `CREATE (p2:Person {name: 'Bob'})`, nil)
82+
exec.Execute(ctx, `CREATE (p3:Person {name: 'Charlie'})`, nil)
83+
exec.Execute(ctx, `MATCH (a:Person {name: 'Alice'}), (b:Person {name: 'Bob'}) CREATE (a)-[:FOLLOWS]->(b)`, nil)
84+
exec.Execute(ctx, `MATCH (a:Person {name: 'Bob'}), (b:Person {name: 'Alice'}) CREATE (a)-[:FOLLOWS]->(b)`, nil)
85+
exec.Execute(ctx, `MATCH (a:Person {name: 'Charlie'}), (b:Person {name: 'Alice'}) CREATE (a)-[:FOLLOWS]->(b)`, nil)
86+
87+
t.Run("incoming count with alias", func(t *testing.T) {
88+
query := `MATCH (p:Person)<-[:FOLLOWS]-(follower:Person) RETURN p.name as person, count(follower) as followers`
89+
result, err := exec.Execute(ctx, query, nil)
90+
require.NoError(t, err)
91+
92+
assert.Equal(t, []string{"person", "followers"}, result.Columns)
93+
})
94+
95+
t.Run("outgoing count with alias", func(t *testing.T) {
96+
query := `MATCH (p:Person)-[:FOLLOWS]->(following:Person) RETURN p.name as person, count(following) as following_count`
97+
result, err := exec.Execute(ctx, query, nil)
98+
require.NoError(t, err)
99+
100+
assert.Equal(t, []string{"person", "following_count"}, result.Columns)
101+
})
102+
103+
t.Run("mutual relationship with aliases", func(t *testing.T) {
104+
query := `MATCH (a:Person)-[:FOLLOWS]->(b:Person)-[:FOLLOWS]->(a) RETURN a.name as person1, b.name as person2`
105+
result, err := exec.Execute(ctx, query, nil)
106+
require.NoError(t, err)
107+
108+
assert.Equal(t, []string{"person1", "person2"}, result.Columns)
109+
})
110+
}
111+
112+
func TestReturnAliases_WithWhereFilter(t *testing.T) {
113+
store := storage.NewMemoryEngine()
114+
exec := NewStorageExecutor(store)
115+
ctx := context.Background()
116+
117+
// Setup test data
118+
exec.Execute(ctx, `CREATE (p1:Person {name: 'Alice'})`, nil)
119+
exec.Execute(ctx, `CREATE (p2:Person {name: 'Bob'})`, nil)
120+
exec.Execute(ctx, `CREATE (a1:Area {name: 'Engineering'})`, nil)
121+
exec.Execute(ctx, `CREATE (a2:Area {name: 'Sales'})`, nil)
122+
exec.Execute(ctx, `CREATE (a3:Area {name: 'Marketing'})`, nil)
123+
exec.Execute(ctx, `MATCH (p:Person {name: 'Alice'}), (a:Area {name: 'Engineering'}) CREATE (p)-[:WORKS_IN]->(a)`, nil)
124+
exec.Execute(ctx, `MATCH (p:Person {name: 'Alice'}), (a:Area {name: 'Sales'}) CREATE (p)-[:WORKS_IN]->(a)`, nil)
125+
exec.Execute(ctx, `MATCH (p:Person {name: 'Bob'}), (a:Area {name: 'Marketing'}) CREATE (p)-[:WORKS_IN]->(a)`, nil)
126+
127+
t.Run("WITH WHERE filter on aggregated value", func(t *testing.T) {
128+
query := `
129+
MATCH (p:Person)-[:WORKS_IN]->(a:Area)
130+
WITH p, count(a) as area_count, collect(a.name) as areas
131+
WHERE area_count > 1
132+
RETURN p.name as name, area_count, areas
133+
LIMIT 5
134+
`
135+
result, err := exec.Execute(ctx, query, nil)
136+
require.NoError(t, err)
137+
138+
// Check column names
139+
assert.Equal(t, []string{"name", "area_count", "areas"}, result.Columns)
140+
141+
// Should only return Alice (2 areas), not Bob (1 area)
142+
assert.Len(t, result.Rows, 1)
143+
if len(result.Rows) > 0 {
144+
assert.Equal(t, "Alice", result.Rows[0][0])
145+
assert.EqualValues(t, 2, result.Rows[0][1])
146+
}
147+
})
148+
149+
t.Run("WITH WHERE filter area_count >= 1", func(t *testing.T) {
150+
query := `
151+
MATCH (p:Person)-[:WORKS_IN]->(a:Area)
152+
WITH p, count(a) as area_count
153+
WHERE area_count >= 1
154+
RETURN p.name as name, area_count
155+
`
156+
result, err := exec.Execute(ctx, query, nil)
157+
require.NoError(t, err)
158+
159+
assert.Equal(t, []string{"name", "area_count"}, result.Columns)
160+
assert.Len(t, result.Rows, 2) // Both Alice and Bob
161+
})
162+
}
163+
164+
func TestExtractReturnItemsFromQuery(t *testing.T) {
165+
store := storage.NewMemoryEngine()
166+
exec := NewStorageExecutor(store)
167+
168+
tests := []struct {
169+
name string
170+
query string
171+
expectedColumns []string
172+
}{
173+
{
174+
name: "simple property with alias",
175+
query: "MATCH (n) RETURN n.name as name",
176+
expectedColumns: []string{"name"},
177+
},
178+
{
179+
name: "multiple items with aliases",
180+
query: "MATCH (n) RETURN n.name as name, n.age as age",
181+
expectedColumns: []string{"name", "age"},
182+
},
183+
{
184+
name: "count with alias",
185+
query: "MATCH (n)-[:REL]->(m) RETURN n.name, count(m) as total",
186+
expectedColumns: []string{"n.name", "total"},
187+
},
188+
{
189+
name: "mixed aliases and no aliases",
190+
query: "MATCH (n) RETURN n.id, n.name as name, n.age",
191+
expectedColumns: []string{"n.id", "name", "n.age"},
192+
},
193+
{
194+
name: "with ORDER BY and LIMIT",
195+
query: "MATCH (n) RETURN n.name as name ORDER BY name LIMIT 10",
196+
expectedColumns: []string{"name"},
197+
},
198+
}
199+
200+
for _, tt := range tests {
201+
t.Run(tt.name, func(t *testing.T) {
202+
items := exec.extractReturnItemsFromQuery(tt.query)
203+
columns := exec.buildColumnsFromReturnItems(items)
204+
assert.Equal(t, tt.expectedColumns, columns)
205+
})
206+
}
207+
}

0 commit comments

Comments
 (0)