feat: Add support for reduction operation in set comparison queries (…

…as mentioned in #92)
substrait-io · Feb 8, 2024 · 2b70788 · 2b70788
1 parent 4f6ed2f
commit 2b70788
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 3 deletions.
diff --git a/src/substrait/textplan/converter/PlanPrinterVisitor.cpp b/src/substrait/textplan/converter/PlanPrinterVisitor.cpp
@@ -324,7 +324,25 @@ std::any PlanPrinterVisitor::visitSubquerySetComparison(
  errorListener_->addError("Did not recognize the subquery comparison.");
  return std::string("UNSUPPORTED SUBQUERY");
  }
- result << "ALL SUBQUERY ";
+ switch (query.reduction_op()) {
+ case ::substrait::proto::
+ Expression_Subquery_SetComparison_ReductionOp_REDUCTION_OP_ANY:
+ result << "ANY ";
+ break;
+ case ::substrait::proto::
+ Expression_Subquery_SetComparison_ReductionOp_REDUCTION_OP_ALL:
+ result << "ALL ";
+ break;
+ case ::substrait::proto::
+ Expression_Subquery_SetComparison_ReductionOp_REDUCTION_OP_UNSPECIFIED:
+ case ::substrait::proto::
+ Expression_Subquery_SetComparison_ReductionOp_Expression_Subquery_SetComparison_ReductionOp_INT_MIN_SENTINEL_DO_NOT_USE_:
+ case ::substrait::proto::
+ Expression_Subquery_SetComparison_ReductionOp_Expression_Subquery_SetComparison_ReductionOp_INT_MAX_SENTINEL_DO_NOT_USE_:
+ errorListener_->addError("Did not recognize the subquery reduction op.");
+ return std::string("UNSUPPORTED SUBQUERY");
+ }
+ result << "SUBQUERY ";
  if (query.has_right()) {
  const SymbolInfo* symbol = symbolTable_->lookupSymbolByParentQueryAndType(
  currentScope_->sourceLocation,

diff --git a/src/substrait/textplan/parser/SubstraitPlanSubqueryRelationVisitor.cpp b/src/substrait/textplan/parser/SubstraitPlanSubqueryRelationVisitor.cpp
@@ -1362,6 +1362,15 @@ SubstraitPlanSubqueryRelationVisitor::visitExpressionSetComparisonSubquery(
  ::substrait::proto::Expression, visitExpression(ctx->expression()));
  expr.mutable_subquery()->mutable_set_comparison()->set_comparison_op(
  comparisonToProto(ctx->COMPARISON()->getText()));
+ if (ctx->ANY() != nullptr) {
+ expr.mutable_subquery()->mutable_set_comparison()->set_reduction_op(
+ ::substrait::proto::
+ Expression_Subquery_SetComparison_ReductionOp_REDUCTION_OP_ANY);
+ } else {
+ expr.mutable_subquery()->mutable_set_comparison()->set_reduction_op(
+ ::substrait::proto::
+ Expression_Subquery_SetComparison_ReductionOp_REDUCTION_OP_ALL);
+ }
  // Next find the relation created in a previous step.
  auto symbol =
  symbolTable_->lookupSymbolByName(ctx->relation_ref()->getText());

diff --git a/src/substrait/textplan/parser/grammar/SubstraitPlanLexer.g4 b/src/substrait/textplan/parser/grammar/SubstraitPlanLexer.g4
@@ -51,6 +51,7 @@ EXISTS: 'EXISTS';
 UNIQUE: 'UNIQUE';
 IN: 'IN';
 ALL: 'ALL';
+ANY: 'ANY';
 COMPARISON: 'EQ'|'NE'|'LT'|'GT'|'LE'|'GE';
 
 VIRTUAL_TABLE: 'VIRTUAL_TABLE';

diff --git a/src/substrait/textplan/parser/grammar/SubstraitPlanParser.g4 b/src/substrait/textplan/parser/grammar/SubstraitPlanParser.g4
@@ -93,7 +93,7 @@ expression
  | SUBQUERY relation_ref # expressionScalarSubquery
  | expression_list IN SUBQUERY relation_ref # expressionInPredicateSubquery
  | (UNIQUE|EXISTS) IN SUBQUERY relation_ref # expressionSetPredicateSubquery
- | expression COMPARISON ALL SUBQUERY relation_ref  # expressionSetComparisonSubquery
+ | expression COMPARISON (ALL|ANY) SUBQUERY relation_ref # expressionSetComparisonSubquery
  ;
 
 expression_list
@@ -235,5 +235,6 @@ simple_id
  | EMIT
  | NAMED
  | ALL
+ | ANY
  | COMPARISON
  ;
diff --git a/src/substrait/textplan/parser/tests/TextPlanParserTest.cpp b/src/substrait/textplan/parser/tests/TextPlanParserTest.cpp
@@ -1061,7 +1061,7 @@ std::vector<TestCase> getTestCases() {
  "1:0 → extraneous input 'relation' expecting {<EOF>, "
  "'EXTENSION_SPACE', 'NAMED', 'SCHEMA', 'PIPELINES', 'FILTER', "
  "'GROUPING', 'MEASURE', 'SORT', 'COUNT', 'TYPE', 'EMIT', "
- "'ALL', COMPARISON, 'SOURCE', 'ROOT', 'NULL', IDENTIFIER}",
+ "'ALL', 'ANY', COMPARISON, 'SOURCE', 'ROOT', 'NULL', IDENTIFIER}",
  "1:24 → mismatched input '{' expecting 'RELATION'",
  "1:9 → Unrecognized relation type: notyperelation",
  "1:9 → Internal error: Previously encountered symbol "