Skip to content

Commit 76342c2

Browse files
ornerydorneryd
committed
fix(fabric): apply ORDER BY/SKIP/LIMIT in in-memory APPLY projection
Ensure the Fabric in-memory APPLY projection fast-path preserves trailing result modifiers (ORDER BY, SKIP, LIMIT). Previously the shortcut omitted these modifiers, which could produce non-deterministic row ordering for correlated subqueries and caused an E2E flake (CorrelatedSubqueryJoin_WithThenUse_CollectSemantics).
1 parent 571d86c commit 76342c2

7 files changed

Lines changed: 331 additions & 6 deletions

File tree

‎pkg/cypher/vector_procedures_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -451,7 +451,7 @@ func (m *failingOnLongQueryEmbedder) Embed(ctx context.Context, text string) ([]
451451

452452
func loadLargeDocQuery(t *testing.T) string {
453453
t.Helper()
454-
path := filepath.Join("..", "..", "docs", "plans", "sharding-base-plan.md")
454+
path := filepath.Join("..", "..", "docs", "features", "gpu-acceleration.md")
455455
data, err := os.ReadFile(path)
456456
require.NoError(t, err)
457457
query := string(data)

‎pkg/fabric/executor.go‎

Lines changed: 297 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -950,7 +950,8 @@ func executeApplyInMemoryProjection(inputResult *ResultStream, query string) (*R
950950

951951
if startsWithFold(trimmed, "RETURN ") {
952952
clause := strings.TrimSpace(trimmed[len("RETURN "):])
953-
items := splitTopLevelCSV(clause)
953+
projectionClause, modifierClause := splitTopLevelResultModifiers(clause)
954+
items := splitTopLevelCSV(projectionClause)
954955
if len(items) > 0 {
955956
type retItem struct {
956957
src string
@@ -986,6 +987,13 @@ func executeApplyInMemoryProjection(inputResult *ResultStream, query string) (*R
986987
for i, p := range parsed {
987988
out.Columns[i] = p.alias
988989
}
990+
aliasToCol := make(map[string]string, len(parsed))
991+
for _, p := range parsed {
992+
aliasToCol[p.alias] = p.alias
993+
if _, exists := aliasToCol[p.src]; !exists {
994+
aliasToCol[p.src] = p.alias
995+
}
996+
}
989997
for _, in := range inputResult.Rows {
990998
row := make([]interface{}, len(parsed))
991999
for i, p := range parsed {
@@ -995,6 +1003,7 @@ func executeApplyInMemoryProjection(inputResult *ResultStream, query string) (*R
9951003
}
9961004
out.Rows = append(out.Rows, row)
9971005
}
1006+
applySimpleResultModifiers(out, modifierClause, aliasToCol)
9981007
return out, true
9991008
}
10001009
}
@@ -1067,6 +1076,293 @@ func executeApplyInMemoryProjection(inputResult *ResultStream, query string) (*R
10671076
return nil, false
10681077
}
10691078

1079+
type simpleOrderSpec struct {
1080+
column string
1081+
desc bool
1082+
}
1083+
1084+
func splitTopLevelResultModifiers(clause string) (projection string, modifiers string) {
1085+
projection = strings.TrimSpace(clause)
1086+
modifiers = ""
1087+
paren, bracket, brace := 0, 0, 0
1088+
inSingle, inDouble, inBacktick := false, false, false
1089+
for i := 0; i < len(clause); i++ {
1090+
ch := clause[i]
1091+
switch {
1092+
case inSingle:
1093+
if ch == '\'' {
1094+
inSingle = false
1095+
}
1096+
continue
1097+
case inDouble:
1098+
if ch == '"' {
1099+
inDouble = false
1100+
}
1101+
continue
1102+
case inBacktick:
1103+
if ch == '`' {
1104+
inBacktick = false
1105+
}
1106+
continue
1107+
}
1108+
switch ch {
1109+
case '\'':
1110+
inSingle = true
1111+
case '"':
1112+
inDouble = true
1113+
case '`':
1114+
inBacktick = true
1115+
case '(':
1116+
paren++
1117+
case ')':
1118+
if paren > 0 {
1119+
paren--
1120+
}
1121+
case '[':
1122+
bracket++
1123+
case ']':
1124+
if bracket > 0 {
1125+
bracket--
1126+
}
1127+
case '{':
1128+
brace++
1129+
case '}':
1130+
if brace > 0 {
1131+
brace--
1132+
}
1133+
}
1134+
if paren != 0 || bracket != 0 || brace != 0 {
1135+
continue
1136+
}
1137+
if hasKeywordAt(clause, i, "ORDER BY") || hasKeywordAt(clause, i, "SKIP") || hasKeywordAt(clause, i, "LIMIT") {
1138+
projection = strings.TrimSpace(clause[:i])
1139+
modifiers = strings.TrimSpace(clause[i:])
1140+
return projection, modifiers
1141+
}
1142+
}
1143+
return projection, modifiers
1144+
}
1145+
1146+
func applySimpleResultModifiers(result *ResultStream, modifiers string, aliasToCol map[string]string) {
1147+
if result == nil || len(result.Rows) == 0 {
1148+
return
1149+
}
1150+
orderSpecs, skip, limit := parseSimpleResultModifiers(modifiers, aliasToCol)
1151+
if len(orderSpecs) > 0 {
1152+
colIdx := buildColumnIndex(result.Columns)
1153+
sort.SliceStable(result.Rows, func(i, j int) bool {
1154+
left := result.Rows[i]
1155+
right := result.Rows[j]
1156+
for _, spec := range orderSpecs {
1157+
idx, ok := colIdx[spec.column]
1158+
if !ok {
1159+
continue
1160+
}
1161+
cmp := compareSimpleOrderValues(valueAtRowIndex(left, idx), valueAtRowIndex(right, idx))
1162+
if cmp == 0 {
1163+
continue
1164+
}
1165+
if spec.desc {
1166+
return cmp > 0
1167+
}
1168+
return cmp < 0
1169+
}
1170+
return false
1171+
})
1172+
}
1173+
if skip > 0 {
1174+
if skip >= len(result.Rows) {
1175+
result.Rows = result.Rows[:0]
1176+
return
1177+
}
1178+
result.Rows = result.Rows[skip:]
1179+
}
1180+
if limit >= 0 && limit < len(result.Rows) {
1181+
result.Rows = result.Rows[:limit]
1182+
}
1183+
}
1184+
1185+
func parseSimpleResultModifiers(modifiers string, aliasToCol map[string]string) ([]simpleOrderSpec, int, int) {
1186+
if strings.TrimSpace(modifiers) == "" {
1187+
return nil, 0, -1
1188+
}
1189+
orderSpecs := []simpleOrderSpec{}
1190+
skip := 0
1191+
limit := -1
1192+
remaining := strings.TrimSpace(modifiers)
1193+
for remaining != "" {
1194+
switch {
1195+
case startsWithFold(remaining, "ORDER BY"):
1196+
orderClause, rest := splitLeadingModifierClause(remaining, "ORDER BY")
1197+
orderSpecs = parseSimpleOrderByClause(strings.TrimSpace(orderClause[len("ORDER BY"):]), aliasToCol)
1198+
remaining = strings.TrimSpace(rest)
1199+
case startsWithFold(remaining, "SKIP"):
1200+
skipClause, rest := splitLeadingModifierClause(remaining, "SKIP")
1201+
if n, ok := parseSimplePositiveInt(strings.TrimSpace(skipClause[len("SKIP"):])); ok {
1202+
skip = n
1203+
}
1204+
remaining = strings.TrimSpace(rest)
1205+
case startsWithFold(remaining, "LIMIT"):
1206+
limitClause, rest := splitLeadingModifierClause(remaining, "LIMIT")
1207+
if n, ok := parseSimplePositiveInt(strings.TrimSpace(limitClause[len("LIMIT"):])); ok {
1208+
limit = n
1209+
}
1210+
remaining = strings.TrimSpace(rest)
1211+
default:
1212+
remaining = ""
1213+
}
1214+
}
1215+
return orderSpecs, skip, limit
1216+
}
1217+
1218+
func splitLeadingModifierClause(s string, keyword string) (clause string, rest string) {
1219+
trimmed := strings.TrimSpace(s)
1220+
if trimmed == "" {
1221+
return "", ""
1222+
}
1223+
end := len(trimmed)
1224+
if strings.EqualFold(keyword, "ORDER BY") {
1225+
for i := len(keyword); i < len(trimmed); i++ {
1226+
if hasKeywordAt(trimmed, i, "SKIP") || hasKeywordAt(trimmed, i, "LIMIT") {
1227+
end = i
1228+
break
1229+
}
1230+
}
1231+
} else {
1232+
for i := len(keyword); i < len(trimmed); i++ {
1233+
if hasKeywordAt(trimmed, i, "ORDER BY") || hasKeywordAt(trimmed, i, "SKIP") || hasKeywordAt(trimmed, i, "LIMIT") {
1234+
end = i
1235+
break
1236+
}
1237+
}
1238+
}
1239+
return strings.TrimSpace(trimmed[:end]), strings.TrimSpace(trimmed[end:])
1240+
}
1241+
1242+
func parseSimpleOrderByClause(clause string, aliasToCol map[string]string) []simpleOrderSpec {
1243+
parts := splitTopLevelCSV(clause)
1244+
specs := make([]simpleOrderSpec, 0, len(parts))
1245+
for _, part := range parts {
1246+
item := strings.TrimSpace(part)
1247+
if item == "" {
1248+
continue
1249+
}
1250+
desc := false
1251+
if len(item) > 5 && strings.EqualFold(strings.TrimSpace(item[len(item)-5:]), " DESC") {
1252+
desc = true
1253+
item = strings.TrimSpace(item[:len(item)-5])
1254+
} else if len(item) > 4 && strings.EqualFold(strings.TrimSpace(item[len(item)-4:]), " ASC") {
1255+
item = strings.TrimSpace(item[:len(item)-4])
1256+
}
1257+
if mapped, ok := aliasToCol[item]; ok {
1258+
item = mapped
1259+
}
1260+
if !isSimpleIdentifier(item) {
1261+
continue
1262+
}
1263+
specs = append(specs, simpleOrderSpec{column: item, desc: desc})
1264+
}
1265+
return specs
1266+
}
1267+
1268+
func parseSimplePositiveInt(s string) (int, bool) {
1269+
if s == "" {
1270+
return 0, false
1271+
}
1272+
value := 0
1273+
for i := 0; i < len(s); i++ {
1274+
if s[i] < '0' || s[i] > '9' {
1275+
return 0, false
1276+
}
1277+
value = value*10 + int(s[i]-'0')
1278+
}
1279+
return value, true
1280+
}
1281+
1282+
func valueAtRowIndex(row []interface{}, idx int) interface{} {
1283+
if idx < 0 || idx >= len(row) {
1284+
return nil
1285+
}
1286+
return row[idx]
1287+
}
1288+
1289+
func compareSimpleOrderValues(left interface{}, right interface{}) int {
1290+
if left == nil && right == nil {
1291+
return 0
1292+
}
1293+
if left == nil {
1294+
return -1
1295+
}
1296+
if right == nil {
1297+
return 1
1298+
}
1299+
if lf, ok := asComparableFloat(left); ok {
1300+
if rf, ok := asComparableFloat(right); ok {
1301+
switch {
1302+
case lf < rf:
1303+
return -1
1304+
case lf > rf:
1305+
return 1
1306+
default:
1307+
return 0
1308+
}
1309+
}
1310+
}
1311+
if lb, ok := left.(bool); ok {
1312+
if rb, ok := right.(bool); ok {
1313+
switch {
1314+
case lb == rb:
1315+
return 0
1316+
case !lb && rb:
1317+
return -1
1318+
default:
1319+
return 1
1320+
}
1321+
}
1322+
}
1323+
ls := fmt.Sprint(left)
1324+
rs := fmt.Sprint(right)
1325+
switch {
1326+
case ls < rs:
1327+
return -1
1328+
case ls > rs:
1329+
return 1
1330+
default:
1331+
return 0
1332+
}
1333+
}
1334+
1335+
func asComparableFloat(v interface{}) (float64, bool) {
1336+
switch x := v.(type) {
1337+
case int:
1338+
return float64(x), true
1339+
case int8:
1340+
return float64(x), true
1341+
case int16:
1342+
return float64(x), true
1343+
case int32:
1344+
return float64(x), true
1345+
case int64:
1346+
return float64(x), true
1347+
case uint:
1348+
return float64(x), true
1349+
case uint8:
1350+
return float64(x), true
1351+
case uint16:
1352+
return float64(x), true
1353+
case uint32:
1354+
return float64(x), true
1355+
case uint64:
1356+
return float64(x), true
1357+
case float32:
1358+
return float64(x), true
1359+
case float64:
1360+
return x, true
1361+
default:
1362+
return 0, false
1363+
}
1364+
}
1365+
10701366
func projectInputRowsAsMaps(input *ResultStream, mapSpec string) ([]map[string]interface{}, bool) {
10711367
colIdx := buildColumnIndex(input.Columns)
10721368
entries := splitTopLevelCSV(mapSpec)

‎pkg/fabric/executor_test.go‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -658,6 +658,35 @@ func TestExecuteApplyInMemoryProjection_WithCollectMapReturnAlias(t *testing.T)
658658
}
659659
}
660660

661+
func TestExecuteApplyInMemoryProjection_ReturnOrderBySkipLimit(t *testing.T) {
662+
input := &ResultStream{
663+
Columns: []string{"textKey128", "texts"},
664+
Rows: [][]interface{}{
665+
{"a2", []interface{}{}},
666+
{"a1", []interface{}{"ORD-001"}},
667+
{"a3", []interface{}{"ORD-003"}},
668+
},
669+
}
670+
query := "RETURN textKey128, texts ORDER BY textKey128 SKIP 1 LIMIT 1"
671+
res, handled := executeApplyInMemoryProjection(input, query)
672+
if !handled {
673+
t.Fatal("expected ordered projection to be handled")
674+
}
675+
if res == nil {
676+
t.Fatal("expected non-nil projection result")
677+
}
678+
if len(res.Rows) != 1 {
679+
t.Fatalf("expected one row after SKIP/LIMIT, got %d", len(res.Rows))
680+
}
681+
if got := res.Rows[0][0]; got != "a2" {
682+
t.Fatalf("expected middle ordered row a2, got %#v", got)
683+
}
684+
texts, ok := res.Rows[0][1].([]interface{})
685+
if !ok || len(texts) != 0 {
686+
t.Fatalf("expected empty texts slice for a2, got %#v", res.Rows[0][1])
687+
}
688+
}
689+
661690
func TestMergeBindings(t *testing.T) {
662691
parent := map[string]interface{}{"rows": []interface{}{1, 2}, "x": "parent"}
663692
row := map[string]interface{}{"x": "row", "k": "v"}

‎pkg/heimdall/metrics_discover_chunk_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ func (s *testSearcher) GetNode(ctx context.Context, nodeID string) (*NodeData, e
9696

9797
func loadLargeDocQuery(t *testing.T) string {
9898
t.Helper()
99-
path := filepath.Join("..", "..", "docs", "plans", "sharding-base-plan.md")
99+
path := filepath.Join("..", "..", "docs", "features", "gpu-acceleration.md")
100100
data, err := os.ReadFile(path)
101101
require.NoError(t, err)
102102
query := string(data)

‎pkg/nornicdb/embed_query_chunk_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,7 @@ func (e *scriptedBatchEmbedder) Model() string { return "scripted-batch" }
115115

116116
func loadLargeDocQuery(t *testing.T) string {
117117
t.Helper()
118-
path := filepath.Join("..", "..", "docs", "plans", "sharding-base-plan.md")
118+
path := filepath.Join("..", "..", "docs", "features", "gpu-acceleration.md")
119119
data, err := os.ReadFile(path)
120120
require.NoError(t, err)
121121
query := string(data)

‎pkg/nornicgrpc/search_service_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -320,7 +320,7 @@ func TestService_SearchText_ChunksLongQueryAndFusesAcrossChunks(t *testing.T) {
320320

321321
// Load a real large document as the base query text to validate chunking
322322
// behavior on natural content instead of synthetic repeated characters.
323-
path := filepath.Join("..", "..", "docs", "plans", "sharding-base-plan.md")
323+
path := filepath.Join("..", "..", "docs", "features", "gpu-acceleration.md")
324324
data, readErr := os.ReadFile(path)
325325
require.NoError(t, readErr)
326326
base := string(data)

0 commit comments

Comments
 (0)