Skip to content

Commit e0564e0

Browse files
committed
Route pushDownSort count through LIMIT > k validation
pushDownSort() called requestBuilder.pushDownLimit() directly, bypassing the LIMIT > k guard in pushDownLimit(). Extract validateLimitWithinK() helper and call it from both paths so the invariant holds when PPL or future callers set a non-zero sort count. Signed-off-by: Eric Wei <mengwei.eric@gmail.com>
1 parent bea6607 commit e0564e0

2 files changed

Lines changed: 35 additions & 7 deletions

File tree

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/VectorSearchQueryBuilder.java

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -89,13 +89,7 @@ public boolean pushDownFilter(LogicalFilter filter) {
8989

9090
@Override
9191
public boolean pushDownLimit(LogicalLimit limit) {
92-
if (options.containsKey("k")) {
93-
int k = Integer.parseInt(options.get("k"));
94-
if (limit.getLimit() > k) {
95-
throw new ExpressionEvaluationException(
96-
String.format("LIMIT %d exceeds k=%d in top-k vector search", limit.getLimit(), k));
97-
}
98-
}
92+
validateLimitWithinK(limit.getLimit());
9993
limitPushed = true;
10094
return super.pushDownLimit(limit);
10195
}
@@ -123,11 +117,23 @@ public boolean pushDownSort(LogicalSort sort) {
123117
// Preserve the parent's sort.getCount() → limit pushdown contract: SQL always sets count=0,
124118
// but PPL or future callers may set a non-zero count to combine sort+limit in one node.
125119
if (sort.getCount() != 0) {
120+
validateLimitWithinK(sort.getCount());
126121
requestBuilder.pushDownLimit(sort.getCount(), 0);
127122
}
128123
return true;
129124
}
130125

126+
/** Validates that the requested limit does not exceed k in top-k mode. */
127+
private void validateLimitWithinK(int limit) {
128+
if (options.containsKey("k")) {
129+
int k = Integer.parseInt(options.get("k"));
130+
if (limit > k) {
131+
throw new ExpressionEvaluationException(
132+
String.format("LIMIT %d exceeds k=%d in top-k vector search", limit, k));
133+
}
134+
}
135+
}
136+
131137
@Override
132138
public OpenSearchRequestBuilder build() {
133139
if (filterTypeExplicit && !filterPushed) {

opensearch/src/test/java/org/opensearch/sql/opensearch/storage/scan/VectorSearchQueryBuilderTest.java

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,28 @@ void pushDownSortPreservesSortCountAsLimit() {
183183
"sort.getCount() should be pushed down as request size");
184184
}
185185

186+
@Test
187+
void pushDownSortCountExceedingKRejects() {
188+
var requestBuilder = createRequestBuilder();
189+
var knnQuery = new WrapperQueryBuilder("{\"knn\":{}}");
190+
var builder = new VectorSearchQueryBuilder(requestBuilder, knnQuery, Map.of("k", "5"));
191+
192+
var dummyChild = new LogicalValues(Collections.emptyList());
193+
// LogicalSort with count=10 exceeds k=5 — should be rejected
194+
var sort =
195+
new org.opensearch.sql.planner.logical.LogicalSort(
196+
dummyChild,
197+
10,
198+
List.of(
199+
org.apache.commons.lang3.tuple.ImmutablePair.of(
200+
org.opensearch.sql.ast.tree.Sort.SortOption.DEFAULT_DESC,
201+
new ReferenceExpression("_score", ExprCoreType.FLOAT))));
202+
203+
ExpressionEvaluationException ex =
204+
assertThrows(ExpressionEvaluationException.class, () -> builder.pushDownSort(sort));
205+
assertTrue(ex.getMessage().contains("LIMIT 10 exceeds k=5"));
206+
}
207+
186208
@Test
187209
void pushDownSortNonScoreFieldRejected() {
188210
var requestBuilder = createRequestBuilder();

0 commit comments

Comments
 (0)