Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 3 commits into from
Feb 13, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions src/substrait/textplan/converter/PlanPrinterVisitor.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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.");
Expand Down Expand Up @@ -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,
Expand Down
126 changes: 126 additions & 0 deletions src/substrait/textplan/data/set-comparision-any.json
Original file line number Diff line number Diff line change
@@ -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": []
}
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -235,5 +235,6 @@ simple_id
| EMIT
| NAMED
| ALL
| ANY
| COMPARISON
;
Original file line number Diff line number Diff line change
Expand Up @@ -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 "
Expand Down
Loading