Added test including root names.

substrait-io · Jun 28, 2023 · 1294815 · 1294815
1 parent 6756799
commit 1294815
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 8 deletions.
diff --git a/src/substrait/textplan/SymbolTablePrinter.cpp b/src/substrait/textplan/SymbolTablePrinter.cpp
@@ -715,14 +715,17 @@ ::substrait::proto::Plan SymbolTablePrinter::outputToBinaryPlan(
  if (relationData->newPipelines.empty()) {
  *relation->mutable_root()->mutable_input() = relationData->relation;
  } else {
- // This is a root node, copy the first node in before iterating.
- auto inputRelationData = ANY_CAST(
- std::shared_ptr<RelationData>, relationData->newPipelines[0]->blob);
- *relation->mutable_root()->mutable_input() = inputRelationData->relation;
-
- addInputsToRelation(
- *relationData->newPipelines[0],
- relation->mutable_root()->mutable_input());
+ if (relationData->newPipelines[0]->type != SymbolType::kRoot) {
+ // This is a root node, copy the first node in before iterating.
+ auto inputRelationData = ANY_CAST(
+ std::shared_ptr<RelationData>, relationData->newPipelines[0]->blob);
+ *relation->mutable_root()->mutable_input() =
+ inputRelationData->relation;
+
+ addInputsToRelation(
+ *relationData->newPipelines[0],
+ relation->mutable_root()->mutable_input());
+ }
 
  const auto& rootSymbol =
  symbolTable.nthSymbolByType(0, SymbolType::kRoot);

diff --git a/src/substrait/textplan/parser/ParseText.cpp b/src/substrait/textplan/parser/ParseText.cpp
@@ -71,6 +71,15 @@ ParseResult parseStream(antlr4::ANTLRInputStream stream) {
  *visitor->getSymbolTable(), visitor->getErrorListener());
  try {
  pipelineVisitor->visitPlan(tree);
+ } catch (std::invalid_argument ex) {
+ // Catches the any_cast exception and logs a useful error message.
+ errorListener.syntaxError(
+ &parser,
+ nullptr,
+ /*line=*/1,
+ /*charPositionInLine=*/1,
+ ex.what(),
+ std::current_exception());
  } catch (...) {
  errorListener.syntaxError(
  &parser,

diff --git a/src/substrait/textplan/parser/SubstraitPlanPipelineVisitor.cpp b/src/substrait/textplan/parser/SubstraitPlanPipelineVisitor.cpp
@@ -77,6 +77,9 @@ std::any SubstraitPlanPipelineVisitor::visitPipeline(
 
  // Refetch our symbol table entry to make sure we have the latest version.
  auto* symbol = symbolTable_->lookupSymbolByName(relationName);
+ if (symbol->blob.type() != typeid(std::shared_ptr<RelationData>)) {
+ return defaultResult();
+ }
  auto relationData = ANY_CAST(std::shared_ptr<RelationData>, symbol->blob);
 
  // Check for accidental cross-pipeline use.
@@ -103,6 +106,9 @@ std::any SubstraitPlanPipelineVisitor::visitPipeline(
  }
  const SymbolInfo* rightmostSymbol = rightSymbol;
  if (*rightSymbol != SymbolInfo::kUnknown) {
+ if (rightSymbol->blob.type() != typeid(std::shared_ptr<RelationData>)) {
+errorListener_->addError(ctx->getStart(), "blah");
+ }
  auto rightRelationData =
  ANY_CAST(std::shared_ptr<RelationData>, rightSymbol->blob);
  if (rightRelationData->pipelineStart != nullptr) {

diff --git a/src/substrait/textplan/parser/tests/TextPlanParserTest.cpp b/src/substrait/textplan/parser/tests/TextPlanParserTest.cpp
@@ -1030,6 +1030,28 @@ std::vector<TestCase> getTestCases() {
  HasSymbolsWithTypes(
  {"read", "project", "root"}, {SymbolType::kRelation}),
  ParsesOk()),
+
+ },
+ {
+ "test18-root-and-read",
+ R"(pipelines {
+ root -> read;
+ }
+
+ read relation read {
+ base_schema schemaone;
+ source mynamedtable;
+ }
+
+ root {
+ names = [
+ apple,
+ ]
+ })",
+ AsBinaryPlan(EqualsProto<::substrait::proto::Plan>(
+ R"(relations: {
+ root { names: "apple" }
+ })")),
  },
  };
  return cases;