Addressed some of the review notes.

substrait-io · Jul 31, 2023 · 6fa011e · 6fa011e
1 parent 2dfee3b
commit 6fa011e
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 21 deletions.
diff --git a/include/substrait/common/Io.h b/include/substrait/common/Io.h
@@ -9,15 +9,15 @@
 
 namespace io::substrait {
 
-enum PlanFileEncoding {
+enum class PlanFileEncoding {
  kBinary = 0,
  kJson = 1,
  kProtoText = 2,
  kText = 3,
 };
 
 // Loads a Substrait plan consisting of any encoding type from the given file.
-absl::StatusOr<::substrait::proto::Plan> loadPlanWithUnknownEncoding(
+absl::StatusOr<::substrait::proto::Plan> loadPlan(
  std::string_view input_filename);
 
 // Writes the provided plan to the specified location with the specified

diff --git a/src/substrait/common/Io.cpp b/src/substrait/common/Io.cpp
@@ -14,28 +14,29 @@ namespace io::substrait {
 
 namespace {
 
-const std::regex kIsJson(R"(("extensionUris"|"extensions"|"relations"))");
+const std::regex kIsJson(
+ R"(("extensionUris"|"extension_uris"|"extensions"|"relations"))");
 const std::regex kIsProtoText(
  R"((^|\n)((relations|extensions|extension_uris|expected_type_urls) \{))");
 const std::regex kIsText(
  R"((^|\n) *(pipelines|[a-z]+ *relation|schema|source|extension_space) *)");
 
 PlanFileEncoding detectEncoding(std::string_view content) {
  if (std::regex_search(content.begin(), content.end(), kIsJson)) {
- return kJson;
+ return PlanFileEncoding::kJson;
  }
  if (std::regex_search(content.begin(), content.end(), kIsProtoText)) {
- return kProtoText;
+ return PlanFileEncoding::kProtoText;
  }
  if (std::regex_search(content.begin(), content.end(), kIsText)) {
- return kText;
+ return PlanFileEncoding::kText;
  }
- return kBinary;
+ return PlanFileEncoding::kBinary;
 }
 
 } // namespace
 
-absl::StatusOr<::substrait::proto::Plan> loadPlanWithUnknownEncoding(
+absl::StatusOr<::substrait::proto::Plan> loadPlan(
  std::string_view input_filename) {
  auto contentOrError = textplan::readFromFile(input_filename.data());
  if (!contentOrError.ok()) {
@@ -45,13 +46,13 @@ absl::StatusOr<::substrait::proto::Plan> loadPlanWithUnknownEncoding(
  auto encoding = detectEncoding(*contentOrError);
  absl::StatusOr<::substrait::proto::Plan> planOrError;
  switch (encoding) {
- case kBinary:
+ case PlanFileEncoding::kBinary:
  return textplan::loadFromBinary(*contentOrError);
- case kJson:
+ case PlanFileEncoding::kJson:
  return textplan::loadFromJson(*contentOrError);
- case kProtoText:
+ case PlanFileEncoding::kProtoText:
  return textplan::loadFromProtoText(*contentOrError);
- case kText:
+ case PlanFileEncoding::kText:
  return textplan::loadFromText(*contentOrError);
  }
  return absl::UnimplementedError("Unexpected encoding requested.");
@@ -62,13 +63,13 @@ absl::Status savePlan(
  std::string_view output_filename,
  PlanFileEncoding encoding) {
  switch (encoding) {
- case kBinary:
+ case PlanFileEncoding::kBinary:
  return textplan::savePlanToBinary(plan, output_filename);
- case kJson:
+ case PlanFileEncoding::kJson:
  return textplan::savePlanToJson(plan, output_filename);
- case kProtoText:
+ case PlanFileEncoding::kProtoText:
  return textplan::savePlanToProtoText(plan, output_filename);
- case kText:
+ case PlanFileEncoding::kText:
  return textplan::savePlanToText(plan, output_filename);
  }
  return absl::UnimplementedError("Unexpected encoding requested.");

diff --git a/src/substrait/common/tests/IoTest.cpp b/src/substrait/common/tests/IoTest.cpp
@@ -33,8 +33,7 @@ constexpr const char* planFileEncodingToString(PlanFileEncoding e) noexcept {
 class IoTest : public ::testing::Test {};
 
 TEST_F(IoTest, LoadMissingFile) {
- auto result =
- ::io::substrait::loadPlanWithUnknownEncoding("non-existent-file");
+ auto result = ::io::substrait::loadPlan("non-existent-file");
  ASSERT_FALSE(result.ok());
  ASSERT_THAT(
  result.status().message(),
@@ -71,7 +70,7 @@ TEST_P(SaveAndLoadTestFixture, SaveAndLoad) {
  auto status = ::io::substrait::savePlan(plan, tempFilename, encoding);
  ASSERT_TRUE(status.ok()) << "Save failed.\n" << status;
 
- auto result = ::io::substrait::loadPlanWithUnknownEncoding(tempFilename);
+ auto result = ::io::substrait::loadPlan(tempFilename);
  ASSERT_TRUE(result.ok()) << "Load failed.\n" << result.status();
  ASSERT_THAT(
  *result,
@@ -96,7 +95,11 @@ TEST_P(SaveAndLoadTestFixture, SaveAndLoad) {
 INSTANTIATE_TEST_SUITE_P(
  SaveAndLoadTests,
  SaveAndLoadTestFixture,
- testing::Values(kBinary, kJson, kProtoText, kText),
+ testing::Values(
+ PlanFileEncoding::kBinary,
+ PlanFileEncoding::kJson,
+ PlanFileEncoding::kProtoText,
+ PlanFileEncoding::kText),
  [](const testing::TestParamInfo<SaveAndLoadTestFixture::ParamType>& info) {
  return planFileEncodingToString(info.param);
  });

diff --git a/src/substrait/textplan/converter/Tool.cpp b/src/substrait/textplan/converter/Tool.cpp
@@ -15,7 +15,7 @@ namespace io::substrait::textplan {
 namespace {
 
 void convertPlanToText(const char* filename) {
- auto planOrError = loadPlanWithUnknownEncoding(filename);
+ auto planOrError = loadPlan(filename);
  if (!planOrError.ok()) {
  std::cerr << planOrError.status() << std::endl;
  return;