Skip to content

Commit 023d31f

Browse files
committed
[Feature] Add table function relation to SQL grammar for vectorSearch()
Add table function relation support to the SQL parser: - New `tableFunctionRelation` alternative in `relation` grammar rule - Named argument syntax: `key=value` (e.g., table='index', field='vec') - Alias is required by grammar (FROM func(...) AS alias) - AstBuilder emits existing TableFunction + SubqueryAlias AST nodes - 3 parser unit tests: basic parse, with WHERE/ORDER BY/LIMIT, alias required This is a pure grammar change — no execution support yet. Queries will parse successfully but fail at the Analyzer with "unsupported function". Signed-off-by: Eric Wei <mengwei.eric@gmail.com>
1 parent c2c97db commit 023d31f

3 files changed

Lines changed: 86 additions & 0 deletions

File tree

sql/src/main/antlr/OpenSearchSQLParser.g4

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,15 @@ fromClause
111111
relation
112112
: tableName (AS? alias)? # tableAsRelation
113113
| LR_BRACKET subquery = querySpecification RR_BRACKET AS? alias # subqueryAsRelation
114+
| qualifiedName LR_BRACKET tableFunctionArgs RR_BRACKET AS alias # tableFunctionRelation
115+
;
116+
117+
tableFunctionArgs
118+
: tableFunctionArg (COMMA tableFunctionArg)*
119+
;
120+
121+
tableFunctionArg
122+
: ident EQUAL_SYMBOL functionArg
114123
;
115124

116125
whereClause

sql/src/main/java/org/opensearch/sql/sql/parser/AstBuilder.java

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.SelectElementContext;
1414
import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.SubqueryAsRelationContext;
1515
import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.TableAsRelationContext;
16+
import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.TableFunctionRelationContext;
1617
import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.WhereClauseContext;
1718
import static org.opensearch.sql.sql.parser.ParserUtils.getTextInQuery;
1819
import static org.opensearch.sql.utils.SystemIndexUtils.TABLE_INFO;
@@ -26,6 +27,7 @@
2627
import org.opensearch.sql.ast.expression.Alias;
2728
import org.opensearch.sql.ast.expression.AllFields;
2829
import org.opensearch.sql.ast.expression.Function;
30+
import org.opensearch.sql.ast.expression.UnresolvedArgument;
2931
import org.opensearch.sql.ast.expression.UnresolvedExpression;
3032
import org.opensearch.sql.ast.tree.DescribeRelation;
3133
import org.opensearch.sql.ast.tree.Filter;
@@ -34,6 +36,7 @@
3436
import org.opensearch.sql.ast.tree.Relation;
3537
import org.opensearch.sql.ast.tree.RelationSubquery;
3638
import org.opensearch.sql.ast.tree.SubqueryAlias;
39+
import org.opensearch.sql.ast.tree.TableFunction;
3740
import org.opensearch.sql.ast.tree.UnresolvedPlan;
3841
import org.opensearch.sql.ast.tree.Values;
3942
import org.opensearch.sql.common.antlr.SyntaxCheckException;
@@ -189,6 +192,23 @@ public UnresolvedPlan visitSubqueryAsRelation(SubqueryAsRelationContext ctx) {
189192
return new RelationSubquery(visit(ctx.subquery), subqueryAlias);
190193
}
191194

195+
@Override
196+
public UnresolvedPlan visitTableFunctionRelation(TableFunctionRelationContext ctx) {
197+
ImmutableList.Builder<UnresolvedExpression> args = ImmutableList.builder();
198+
ctx.tableFunctionArgs()
199+
.tableFunctionArg()
200+
.forEach(
201+
arg -> {
202+
String argName = arg.ident().getText();
203+
UnresolvedExpression argValue = visitAstExpression(arg.functionArg());
204+
args.add(new UnresolvedArgument(argName, argValue));
205+
});
206+
TableFunction tableFunction =
207+
new TableFunction(visitAstExpression(ctx.qualifiedName()), args.build());
208+
String alias = StringUtils.unquoteIdentifier(ctx.alias().getText());
209+
return new SubqueryAlias(alias, tableFunction);
210+
}
211+
192212
@Override
193213
public UnresolvedPlan visitWhereClause(WhereClauseContext ctx) {
194214
return new Filter(visitAstExpression(ctx.expression()));

sql/src/test/java/org/opensearch/sql/sql/parser/AstBuilderTest.java

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,9 @@
77

88
import static java.util.Collections.emptyList;
99
import static org.junit.jupiter.api.Assertions.assertEquals;
10+
import static org.junit.jupiter.api.Assertions.assertNotNull;
1011
import static org.junit.jupiter.api.Assertions.assertThrows;
12+
import static org.junit.jupiter.api.Assertions.assertTrue;
1113
import static org.opensearch.sql.ast.dsl.AstDSL.agg;
1214
import static org.opensearch.sql.ast.dsl.AstDSL.aggregate;
1315
import static org.opensearch.sql.ast.dsl.AstDSL.alias;
@@ -40,6 +42,10 @@
4042
import org.opensearch.sql.ast.expression.DataType;
4143
import org.opensearch.sql.ast.expression.Literal;
4244
import org.opensearch.sql.ast.expression.NestedAllTupleFields;
45+
import org.opensearch.sql.ast.expression.UnresolvedArgument;
46+
import org.opensearch.sql.ast.tree.SubqueryAlias;
47+
import org.opensearch.sql.ast.tree.TableFunction;
48+
import org.opensearch.sql.ast.tree.UnresolvedPlan;
4349
import org.opensearch.sql.common.antlr.SyntaxCheckException;
4450

4551
class AstBuilderTest extends AstBuilderTestBase {
@@ -131,6 +137,57 @@ public void can_build_from_index_with_alias_quoted() {
131137
buildAST("SELECT `t`.name FROM test `t` WHERE `t`.age = 30"));
132138
}
133139

140+
@Test
141+
public void can_build_from_table_function() {
142+
assertEquals(
143+
project(
144+
new SubqueryAlias(
145+
"v",
146+
new TableFunction(
147+
qualifiedName("vectorSearch"),
148+
ImmutableList.of(
149+
new UnresolvedArgument("table", stringLiteral("products")),
150+
new UnresolvedArgument("field", stringLiteral("embedding")),
151+
new UnresolvedArgument("vector", stringLiteral("[0.1,0.2]")),
152+
new UnresolvedArgument("option", stringLiteral("k=10"))))),
153+
AllFields.of()),
154+
buildAST(
155+
"SELECT * FROM vectorSearch("
156+
+ "table='products', field='embedding', "
157+
+ "vector='[0.1,0.2]', option='k=10') AS v"));
158+
}
159+
160+
@Test
161+
public void can_build_from_table_function_with_where_and_order() {
162+
// Verify parsing succeeds for table function with WHERE, ORDER BY, and LIMIT
163+
UnresolvedPlan plan =
164+
buildAST(
165+
"SELECT s.title, s._score FROM vectorSearch("
166+
+ "table='products', field='embedding', "
167+
+ "vector='[0.1,0.2]', option='k=10') AS s "
168+
+ "WHERE s.category = 'shoes' "
169+
+ "ORDER BY s._score DESC "
170+
+ "LIMIT 5");
171+
assertNotNull(plan);
172+
// Verify the plan contains the expected structure
173+
String planStr = plan.toString();
174+
assertTrue(planStr.contains("SubqueryAlias(alias=s"));
175+
assertTrue(planStr.contains("TableFunction(functionName=vectorSearch"));
176+
assertTrue(planStr.contains("UnresolvedArgument(argName=table, value=products)"));
177+
assertTrue(planStr.contains("Filter(condition==(s.category, shoes)"));
178+
}
179+
180+
@Test
181+
public void table_function_relation_requires_alias() {
182+
assertThrows(
183+
SyntaxCheckException.class,
184+
() ->
185+
buildAST(
186+
"SELECT * FROM vectorSearch("
187+
+ "table='products', field='embedding', "
188+
+ "vector='[0.1,0.2]', option='k=10')"));
189+
}
190+
134191
@Test
135192
public void can_build_where_clause() {
136193
assertEquals(

0 commit comments

Comments
 (0)