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

) This addresses issue #92 .
substrait-io · Feb 13, 2024 · 9a7ffa7 · 9a7ffa7
1 parent 504ce9e
commit 9a7ffa7
Show file tree

Hide file tree

Showing 6 changed files with 159 additions and 4 deletions.
diff --git a/src/substrait/textplan/converter/PlanPrinterVisitor.cpp b/src/substrait/textplan/converter/PlanPrinterVisitor.cpp
@@ -285,7 +285,7 @@ std::any PlanPrinterVisitor::visitSubquerySetComparison(
     const ::substrait::proto::Expression_Subquery_SetComparison& query) {
   std::stringstream result;
   if (query.has_left()) {
-    result << ANY_CAST(std::string, visitExpression(query.left()));
+    result << ANY_CAST(std::string, visitExpression(query.left())) << " ";
   } else {
     errorListener_->addError(
         "No expression defined for set comparison operation.");
@@ -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/data/set-comparision-any.json b/src/substrait/textplan/data/set-comparision-any.json
@@ -0,0 +1,126 @@
+# "CREATE TABLE NATION (N_REGIONKEY BIGINT NOT NULL)" "SELECT * FROM NATION WHERE N_REGIONKEY < ANY (SELECT N_REGIONKEY FROM NATION)"   
+{
+  "extensionUris": [],
+  "extensions": [],
+  "relations": [{
+    "root": {
+      "input": {
+        "project": {
+          "common": {
+            "emit": {
+              "outputMapping": [1]
+            }
+          },
+          "input": {
+            "filter": {
+              "common": {
+                "direct": {
+                }
+              },
+              "input": {
+                "read": {
+                  "common": {
+                    "direct": {
+                    }
+                  },
+                  "baseSchema": {
+                    "names": ["N_REGIONKEY"],
+                    "struct": {
+                      "types": [{
+                        "i64": {
+                          "typeVariationReference": 0,
+                          "nullability": "NULLABILITY_REQUIRED"
+                        }
+                      }],
+                      "typeVariationReference": 0,
+                      "nullability": "NULLABILITY_REQUIRED"
+                    }
+                  },
+                  "namedTable": {
+                    "names": ["NATION"]
+                  }
+                }
+              },
+              "condition": {
+                "subquery": {
+                  "setComparison": {
+                    "comparisonOp": 3,
+                    "reductionOp": 1,
+                    "left": {
+                      "selection": {
+                        "directReference": {
+                          "structField": {
+                            "field": 0
+                          }
+                        },
+                        "rootReference": {
+                        }
+                      }
+                    },
+                    "right": {
+                      "project": {
+                        "common": {
+                          "emit": {
+                            "outputMapping": [1]
+                          }
+                        },
+                        "input": {
+                          "read": {
+                            "common": {
+                              "direct": {
+                              }
+                            },
+                            "baseSchema": {
+                              "names": ["N_REGIONKEY"],
+                              "struct": {
+                                "types": [{
+                                  "i64": {
+                                    "typeVariationReference": 0,
+                                    "nullability": "NULLABILITY_REQUIRED"
+                                  }
+                                }],
+                                "typeVariationReference": 0,
+                                "nullability": "NULLABILITY_REQUIRED"
+                              }
+                            },
+                            "namedTable": {
+                              "names": ["NATION"]
+                            }
+                          }
+                        },
+                        "expressions": [{
+                          "selection": {
+                            "directReference": {
+                              "structField": {
+                                "field": 0
+                              }
+                            },
+                            "rootReference": {
+                            }
+                          }
+                        }]
+                      }
+                    }
+                  }
+                }
+              }
+            }
+          },
+          "expressions": [{
+            "selection": {
+              "directReference": {
+                "structField": {
+                  "field": 0
+                }
+              },
+              "rootReference": {
+              }
+            }
+          }]
+        }
+      },
+      "names": ["N_REGIONKEY"]
+    }
+  }],
+  "expectedTypeUrls": []
+}
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 "