Skip to content

Commit

Permalink
IF function should support complex predicates in PPL (opensearch-proj…
Browse files Browse the repository at this point in the history
…ect#2756)

Signed-off-by: Lantao Jin <ltjin@amazon.com>
  • Loading branch information
LantaoJin authored and manasvinibs committed Aug 14, 2024
1 parent 28f77bf commit 9dc17ba
Show file tree
Hide file tree
Showing 6 changed files with 154 additions and 8 deletions.
11 changes: 11 additions & 0 deletions docs/user/ppl/functions/condition.rst
Original file line number Diff line number Diff line change
Expand Up @@ -181,3 +181,14 @@ Example::
| Bates | Nanette | Bates |
| Adams | Dale | Adams |
+----------+-------------+------------+

os> source=accounts | eval is_vip = if(age > 30 AND isnotnull(employer), true, false) | fields is_vip, firstname, lastname
fetched rows / total rows = 4/4
+----------+-------------+------------+
| is_vip | firstname | lastname |
|----------+-------------+------------|
| True | Amber | Duke |
| True | Hattie | Bond |
| False | Nanette | Bates |
| False | Dale | Adams |
+----------+-------------+------------+
16 changes: 11 additions & 5 deletions ppl/src/main/antlr/OpenSearchPPLParser.g4
Original file line number Diff line number Diff line change
Expand Up @@ -255,6 +255,7 @@ expression
| valueExpression
;

// predicates
logicalExpression
: comparisonExpression # comparsion
| NOT logicalExpression # logicalNot
Expand Down Expand Up @@ -362,7 +363,7 @@ dataTypeFunctionCall

// boolean functions
booleanFunctionCall
: conditionFunctionBase LT_PRTHS functionArgs RT_PRTHS
: conditionFunctionName LT_PRTHS functionArgs RT_PRTHS
;

convertedDataType
Expand All @@ -382,7 +383,8 @@ evalFunctionName
: mathematicalFunctionName
| dateTimeFunctionName
| textFunctionName
| conditionFunctionBase
| conditionFunctionName
| flowControlFunctionName
| systemFunctionName
| positionFunctionName
;
Expand All @@ -392,7 +394,7 @@ functionArgs
;

functionArg
: (ident EQUAL)? valueExpression
: (ident EQUAL)? expression
;

relevanceArg
Expand Down Expand Up @@ -623,11 +625,15 @@ timestampFunctionName
;

// condition function return boolean value
conditionFunctionBase
conditionFunctionName
: LIKE
| IF
| ISNULL
| ISNOTNULL
;

// flow control function return non-boolean value
flowControlFunctionName
: IF
| IFNULL
| NULLIF
;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -332,8 +332,7 @@ public UnresolvedPlan visitTableFunction(TableFunctionContext ctx) {
arg -> {
String argName = (arg.ident() != null) ? arg.ident().getText() : null;
builder.add(
new UnresolvedArgument(
argName, this.internalVisitExpression(arg.valueExpression())));
new UnresolvedArgument(argName, this.internalVisitExpression(arg.expression())));
});
return new TableFunction(this.internalVisitExpression(ctx.qualifiedName()), builder.build());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,7 @@ public UnresolvedExpression visitTakeAggFunctionCall(
/** Eval function. */
@Override
public UnresolvedExpression visitBooleanFunctionCall(BooleanFunctionCallContext ctx) {
final String functionName = ctx.conditionFunctionBase().getText().toLowerCase();
final String functionName = ctx.conditionFunctionName().getText().toLowerCase();
return buildFunction(
FUNCTION_NAME_MAPPING.getOrDefault(functionName, functionName),
ctx.functionArgs().functionArg());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,120 @@ public void testEvalFunctionExprNoArgs() {
assertEqual("source=t | eval f=PI()", eval(relation("t"), let(field("f"), function("PI"))));
}

@Test
public void testEvalIfFunctionExpr() {
assertEqual(
"source=t | eval f=if(true, 1, 0)",
eval(
relation("t"),
let(field("f"), function("if", booleanLiteral(true), intLiteral(1), intLiteral(0)))));
assertEqual(
"source=t | eval f=if(1>2, 1, 0)",
eval(
relation("t"),
let(
field("f"),
function(
"if",
compare(">", intLiteral(1), intLiteral(2)),
intLiteral(1),
intLiteral(0)))));
assertEqual(
"source=t | eval f=if(1<=2, 1, 0)",
eval(
relation("t"),
let(
field("f"),
function(
"if",
compare("<=", intLiteral(1), intLiteral(2)),
intLiteral(1),
intLiteral(0)))));
assertEqual(
"source=t | eval f=if(1=2, 1, 0)",
eval(
relation("t"),
let(
field("f"),
function(
"if",
compare("=", intLiteral(1), intLiteral(2)),
intLiteral(1),
intLiteral(0)))));
assertEqual(
"source=t | eval f=if(1!=2, 1, 0)",
eval(
relation("t"),
let(
field("f"),
function(
"if",
compare("!=", intLiteral(1), intLiteral(2)),
intLiteral(1),
intLiteral(0)))));
assertEqual(
"source=t | eval f=if(isnull(a), 1, 0)",
eval(
relation("t"),
let(
field("f"),
function("if", function("is null", field("a")), intLiteral(1), intLiteral(0)))));
assertEqual(
"source=t | eval f=if(isnotnull(a), 1, 0)",
eval(
relation("t"),
let(
field("f"),
function(
"if", function("is not null", field("a")), intLiteral(1), intLiteral(0)))));
assertEqual(
"source=t | eval f=if(not 1>2, 1, 0)",
eval(
relation("t"),
let(
field("f"),
function(
"if",
not(compare(">", intLiteral(1), intLiteral(2))),
intLiteral(1),
intLiteral(0)))));
assertEqual(
"source=t | eval f=if(not a in (0, 1), 1, 0)",
eval(
relation("t"),
let(
field("f"),
function(
"if",
not(in(field("a"), intLiteral(0), intLiteral(1))),
intLiteral(1),
intLiteral(0)))));
assertEqual(
"source=t | eval f=if(not a in (0, 1) OR isnull(a), 1, 0)",
eval(
relation("t"),
let(
field("f"),
function(
"if",
or(
not(in(field("a"), intLiteral(0), intLiteral(1))),
function("is null", field("a"))),
intLiteral(1),
intLiteral(0)))));
assertEqual(
"source=t | eval f=if(like(a, '_a%b%c_d_'), 1, 0)",
eval(
relation("t"),
let(
field("f"),
function(
"if",
function("like", field("a"), stringLiteral("_a%b%c_d_")),
intLiteral(1),
intLiteral(0)))));
}

@Test
public void testPositionFunctionExpr() {
assertEqual(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -719,6 +719,22 @@ public void canParseMultiMatchAlternateSyntax() {
assertNotNull(parser.parse("SELECT * FROM test WHERE Field = multimatch(\"query\")"));
}

@Test
public void canParseIfFunction() {
assertNotNull(parser.parse("SELECT IF(1 > 2, 1, 0)"));
assertNotNull(parser.parse("SELECT IF(1 < 2, 1, 0)"));
assertNotNull(parser.parse("SELECT IF(1 >= 2, 1, 0)"));
assertNotNull(parser.parse("SELECT IF(1 <= 2, 1, 0)"));
assertNotNull(parser.parse("SELECT IF(1 <> 2, 1, 0)"));
assertNotNull(parser.parse("SELECT IF(1 != 2, 1, 0)"));
assertNotNull(parser.parse("SELECT IF(1 = 2, 1, 0)"));
assertNotNull(parser.parse("SELECT IF(true, 1, 0)"));
assertNotNull(parser.parse("SELECT IF(1 IS NOT NULL, 1, 0)"));
assertNotNull(parser.parse("SELECT IF(NOT 1 > 2, 1, 0)"));
assertNotNull(parser.parse("SELECT IF(NOT 1 IN (0, 1), 1, 0)"));
assertNotNull(parser.parse("SELECT IF(NOT 1 IN (0, 1) OR 1 IS NOT NULL, 1, 0)"));
}

private static Stream<String> matchPhraseQueryComplexQueries() {
return Stream.of(
"SELECT * FROM t WHERE matchphrasequery(c, 3)",
Expand Down

0 comments on commit 9dc17ba

Please sign in to comment.