Added more tests, implemented text plans.

substrait-io · Jul 26, 2023 · 88a90b1 · 88a90b1
1 parent 624dc15
commit 88a90b1
Show file tree

Hide file tree

Showing 12 changed files with 228 additions and 47 deletions.
diff --git a/src/substrait/common/CMakeLists.txt b/src/substrait/common/CMakeLists.txt
@@ -5,8 +5,9 @@ target_link_libraries(substrait_common fmt::fmt-header-only)
 
 add_library(substrait_io Io.cpp)
 add_dependencies(substrait_io substrait_proto substrait_textplan_converter
-                 fmt::fmt-header-only absl::status)
-target_link_libraries(substrait_io substrait_proto substrait_textplan_converter)
+                 substrait_textplan_loader fmt::fmt-header-only absl::status)
+target_link_libraries(substrait_io substrait_proto substrait_textplan_converter
+                      substrait_textplan_loader)
 
 if(${SUBSTRAIT_CPP_BUILD_TESTING})
   add_subdirectory(tests)

diff --git a/src/substrait/common/Io.cpp b/src/substrait/common/Io.cpp
@@ -1,10 +1,10 @@
 #include "substrait/common/Io.h"
 
-#include <string>
 #include <string_view>
 
 #include "substrait/textplan/converter/LoadBinary.h"
 #include "substrait/textplan/converter/SaveBinary.h"
+#include "substrait/textplan/parser/LoadText.h"
 
 namespace io::substrait {
 
@@ -21,6 +21,11 @@ absl::StatusOr<::substrait::proto::Plan> loadPlanWithUnknownEncoding(
     return *planOrError;
   }
 
+  planOrError = textplan::loadFromProtoText(*contentOrError);
+  if (planOrError.ok()) {
+    return *planOrError;
+  }
+
   planOrError = textplan::loadFromText(*contentOrError);
   if (planOrError.ok()) {
     return *planOrError;
@@ -31,8 +36,6 @@ absl::StatusOr<::substrait::proto::Plan> loadPlanWithUnknownEncoding(
     return *planOrError;
   }
 
-  // TODO -- Add support for reading proto text plan files.
-
   return planOrError.status();
 }
 
@@ -42,13 +45,13 @@ absl::Status savePlan(
     PlanFileEncoding encoding) {
   switch (encoding) {
     case kBinary:
-      return textplan::savePlanToBinaryFile(plan, output_filename);
+      return textplan::savePlanToBinary(plan, output_filename);
     case kJson:
-      return textplan::savePlanToJsonFile(plan, output_filename);
+      return textplan::savePlanToJson(plan, output_filename);
     case kProtoText:
-      return textplan::savePlanToProtoTextPlanFile(plan, output_filename);
+      return textplan::savePlanToProtoText(plan, output_filename);
     case kText:
-      return textplan::savePlanToTextPlanFile(plan, output_filename);
+      return textplan::savePlanToText(plan, output_filename);
   }
 }
 

diff --git a/src/substrait/common/tests/CMakeLists.txt b/src/substrait/common/tests/CMakeLists.txt
@@ -16,5 +16,6 @@ add_test_case(
   IoTest.cpp
   EXTRA_LINK_LIBS
   substrait_io
+  protobuf-matchers
   gtest
   gtest_main)
diff --git a/src/substrait/common/tests/IoTest.cpp b/src/substrait/common/tests/IoTest.cpp
@@ -4,10 +4,14 @@
 
 #include <gmock/gmock-matchers.h>
 #include <gtest/gtest.h>
+#include <protobuf-matchers/protocol-buffer-matchers.h>
+
+using ::protobuf_matchers::EqualsProto;
+using ::protobuf_matchers::Partially;
 
 class IoTest : public ::testing::Test {};
 
-TEST_F(IoTest, NonExistantFile) {
+TEST_F(IoTest, LoadMissingFile) {
   auto result =
       ::io::substrait::loadPlanWithUnknownEncoding("non-existent-file");
   ASSERT_FALSE(result.ok());
@@ -16,23 +20,126 @@ TEST_F(IoTest, NonExistantFile) {
       ::testing::ContainsRegex("Failed to open file non-existent-file"));
 }
 
+TEST_F(IoTest, SaveAndLoadBinary) {
+  ::substrait::proto::Plan plan;
+  auto root = plan.add_relations()->mutable_root();
+  auto read = root->mutable_input()->mutable_read();
+  read->mutable_named_table()->add_names("table_name");
+  auto status =
+      ::io::substrait::savePlan(plan, "rwtest.plan", io::substrait::kBinary);
+  ASSERT_TRUE(status.ok()) << status;
+
+  auto result = ::io::substrait::loadPlanWithUnknownEncoding("rwtest.plan");
+  ASSERT_TRUE(result.ok()) << result.status();
+  ASSERT_THAT(
+      *result,
+      Partially(EqualsProto<::substrait::proto::Plan>(
+          R"(relations {
+            root {
+              input {
+                read {
+                  common {
+                    direct {
+                    }
+                  }
+                  named_table {
+                    names: "table_name"
+                  }
+                }
+              }
+            }
+          })")));
+}
+
+TEST_F(IoTest, SaveAndLoadJson) {
+  ::substrait::proto::Plan plan;
+  auto root = plan.add_relations()->mutable_root();
+  auto read = root->mutable_input()->mutable_read();
+  read->mutable_named_table()->add_names("table_name");
+  auto status =
+      ::io::substrait::savePlan(plan, "rwtest.json", io::substrait::kJson);
+  ASSERT_TRUE(status.ok()) << status;
+
+  auto result = ::io::substrait::loadPlanWithUnknownEncoding("rwtest.json");
+  ASSERT_TRUE(result.ok()) << result.status();
+  ASSERT_THAT(
+      *result,
+      Partially(EqualsProto<::substrait::proto::Plan>(
+          R"(relations {
+            root {
+              input {
+                read {
+                  common {
+                    direct {
+                    }
+                  }
+                  named_table {
+                    names: "table_name"
+                  }
+                }
+              }
+            }
+          })")));
+}
+
+TEST_F(IoTest, SaveAndLoadProtoText) {
+  ::substrait::proto::Plan plan;
+  auto root = plan.add_relations()->mutable_root();
+  auto read = root->mutable_input()->mutable_read();
+  read->mutable_named_table()->add_names("table_name");
+  auto status = ::io::substrait::savePlan(
+      plan, "rwtest.protobuf", io::substrait::kProtoText);
+  ASSERT_TRUE(status.ok()) << status;
+
+  auto result = ::io::substrait::loadPlanWithUnknownEncoding("rwtest.protobuf");
+  ASSERT_TRUE(result.ok()) << result.status();
+  ASSERT_THAT(
+      *result,
+      Partially(EqualsProto<::substrait::proto::Plan>(
+          R"(relations {
+            root {
+              input {
+                read {
+                  common {
+                    direct {
+                    }
+                  }
+                  named_table {
+                    names: "table_name"
+                  }
+                }
+              }
+            }
+          })")));
+}
+
 TEST_F(IoTest, SaveAndLoadText) {
   ::substrait::proto::Plan plan;
-  plan.add_relations()->mutable_root()->add_names("hiya");
+  auto root = plan.add_relations()->mutable_root();
+  auto read = root->mutable_input()->mutable_read();
+  read->mutable_named_table()->add_names("table_name");
   auto status =
       ::io::substrait::savePlan(plan, "rwtest.splan", io::substrait::kText);
-  ASSERT_TRUE(status.ok());
+  ASSERT_TRUE(status.ok()) << status;
 
   auto result = ::io::substrait::loadPlanWithUnknownEncoding("rwtest.splan");
-  ASSERT_TRUE(result.ok());
-  ASSERT_THAT(result->relations(0).root().names(0), "hiya");
-}
-
-TEST_F(IoTest, LoadMissingFile) {
-  auto result =
-      ::io::substrait::loadPlanWithUnknownEncoding("non-existent-file");
-  ASSERT_FALSE(result.ok());
+  ASSERT_TRUE(result.ok()) << result.status();
   ASSERT_THAT(
-      result.status().message(),
-      ::testing::ContainsRegex("Failed to open file non-existent-file"));
+      *result,
+      Partially(EqualsProto<::substrait::proto::Plan>(
+          R"(relations {
+            root {
+              input {
+                read {
+                  common {
+                    direct {
+                    }
+                  }
+                  named_table {
+                    names: "table_name"
+                  }
+                }
+              }
+            }
+          })")));
 }
diff --git a/src/substrait/textplan/converter/LoadBinary.cpp b/src/substrait/textplan/converter/LoadBinary.cpp
@@ -74,7 +74,8 @@ absl::StatusOr<::substrait::proto::Plan> loadFromJson(std::string_view json) {
   return plan;
 }
 
-absl::StatusOr<::substrait::proto::Plan> loadFromText(std::string_view text) {
+absl::StatusOr<::substrait::proto::Plan> loadFromProtoText(
+    std::string_view text) {
   ::substrait::proto::Plan plan;
   ::google::protobuf::TextFormat::Parser parser;
   StringErrorCollector collector;
@@ -87,9 +88,10 @@ absl::StatusOr<::substrait::proto::Plan> loadFromText(std::string_view text) {
   return plan;
 }
 
-absl::StatusOr<::substrait::proto::Plan> loadFromBinary(std::string_view text) {
+absl::StatusOr<::substrait::proto::Plan> loadFromBinary(
+    std::string_view bytes) {
   ::substrait::proto::Plan plan;
-  if (!plan.ParseFromString(text)) {
+  if (!plan.ParseFromString(bytes)) {
     return absl::InternalError("Failed to parse as a binary Substrait plan.");
   }
   return plan;

diff --git a/src/substrait/textplan/converter/LoadBinary.h b/src/substrait/textplan/converter/LoadBinary.h
@@ -21,10 +21,11 @@ absl::StatusOr<::substrait::proto::Plan> loadFromJson(std::string_view json);
 
 // Reads a plan encoded as a text protobuf.
 // Returns a list of errors if the file cannot be parsed.
-absl::StatusOr<::substrait::proto::Plan> loadFromText(std::string_view text);
+absl::StatusOr<::substrait::proto::Plan> loadFromProtoText(
+    std::string_view text);
 
 // Reads a plan serialized as a binary protobuf.
 // Returns a list of errors if the file cannot be parsed.
-absl::StatusOr<::substrait::proto::Plan> loadFromBinary(std::string_view text);
+absl::StatusOr<::substrait::proto::Plan> loadFromBinary(std::string_view bytes);
 
 } // namespace io::substrait::textplan
diff --git a/src/substrait/textplan/converter/SaveBinary.cpp b/src/substrait/textplan/converter/SaveBinary.cpp
@@ -2,29 +2,38 @@
 
 #include "substrait/textplan/converter/SaveBinary.h"
 
-#include <google/protobuf/text_format.h>
-#include <google/protobuf/util/json_util.h>
 #include <libc.h>
+#include <fstream>
 #include <limits>
 
+#include <google/protobuf/text_format.h>
+#include <google/protobuf/util/json_util.h>
+#include "substrait/textplan/SymbolTablePrinter.h"
+#include "substrait/textplan/converter/ParseBinary.h"
+
 namespace io::substrait::textplan {
 
-absl::Status savePlanToBinaryFile(
+absl::Status savePlanToBinary(
     const ::substrait::proto::Plan& plan,
     const std::string_view& output_filename) {
-  int outfd = open(std::string{output_filename}.c_str(), O_WRONLY);
-  auto stream = new google::protobuf::io::FileOutputStream(outfd);
+  int outputFileDescriptor =
+      open(std::string{output_filename}.c_str(), O_WRONLY);
+  if (outputFileDescriptor == -1) {
+    return absl::ErrnoToStatus(errno, "Failed to open file for writing");
+  }
+  auto stream =
+      new google::protobuf::io::FileOutputStream(outputFileDescriptor);
 
   if (!plan.SerializeToZeroCopyStream(stream)) {
     return ::absl::UnknownError("Failed to write plan to stream.");
   }
 
   delete stream;
-  close(outfd);
+  close(outputFileDescriptor);
   return absl::OkStatus();
 }
 
-absl::Status savePlanToJsonFile(
+absl::Status savePlanToJson(
     const ::substrait::proto::Plan& plan,
     const std::string_view& output_filename) {
   int outputFileDescriptor =
@@ -47,26 +56,36 @@ absl::Status savePlanToJsonFile(
   return absl::OkStatus();
 }
 
-absl::Status savePlanToTextPlanFile(
+absl::Status savePlanToText(
     const ::substrait::proto::Plan& plan,
     const std::string_view& output_filename) {
-  // MEGAHACK -- Implement.
-  return absl::UnimplementedError(
-      "Writing to text plan files are not yet supported.");
+  std::ofstream stream(output_filename);
+
+  auto result = parseBinaryPlan(plan);
+  auto errors = result.getAllErrors();
+  if (!errors.empty()) {
+    return absl::UnknownError(
+        std::accumulate(errors.begin(), errors.end(), std::string("\n")));
+  }
+  stream << SymbolTablePrinter::outputToText(result.getSymbolTable());
+  stream.close();
+  return absl::OkStatus();
 }
 
-absl::Status savePlanToProtoTextPlanFile(
+absl::Status savePlanToProtoText(
     const ::substrait::proto::Plan& plan,
     const std::string_view& output_filename) {
-  int outfd = open(std::string{output_filename}.c_str(), O_WRONLY);
-  auto stream = new google::protobuf::io::FileOutputStream(outfd);
+  int outputFileDescriptor =
+      open(std::string{output_filename}.c_str(), O_WRONLY);
+  auto stream =
+      new google::protobuf::io::FileOutputStream(outputFileDescriptor);
 
   if (!::google::protobuf::TextFormat::Print(plan, stream)) {
     return absl::UnknownError("Failed to save plan as a text protobuf.");
   }
 
   delete stream;
-  close(outfd);
+  close(outputFileDescriptor);
   return absl::OkStatus();
 }
 

diff --git a/src/substrait/textplan/converter/SaveBinary.h b/src/substrait/textplan/converter/SaveBinary.h
@@ -7,19 +7,23 @@
 
 namespace io::substrait::textplan {
 
-absl::Status savePlanToBinaryFile(
+// Serializes a plan to disk as a binary protobuf.
+absl::Status savePlanToBinary(
     const ::substrait::proto::Plan& plan,
     const std::string_view& output_filename);
 
-absl::Status savePlanToJsonFile(
+// Serializes a plan to disk as a JSON-encoded protobuf.
+absl::Status savePlanToJson(
     const ::substrait::proto::Plan& plan,
     const std::string_view& output_filename);
 
-absl::Status savePlanToTextPlanFile(
+// Calls the converter to store a plan on disk as a text-based substrait plan.
+absl::Status savePlanToText(
     const ::substrait::proto::Plan& plan,
     const std::string_view& output_filename);
 
-absl::Status savePlanToProtoTextPlanFile(
+// Serializes a plan to disk as a text-encoded protobuf.
+absl::Status savePlanToProtoText(
     const ::substrait::proto::Plan& plan,
     const std::string_view& output_filename);
 

diff --git a/src/substrait/textplan/parser/CMakeLists.txt b/src/substrait/textplan/parser/CMakeLists.txt
@@ -12,6 +12,8 @@ add_library(
   SubstraitPlanRelationVisitor.h
   SubstraitPlanTypeVisitor.cpp
   SubstraitPlanTypeVisitor.h
+  LoadText.cpp
+  LoadText.h
   ParseText.cpp
   ParseText.h
   SubstraitParserErrorListener.cpp)