feat: create an exportable version of plan loading

CMakeLists.txt

@@ -49,3 +49,4 @@ if(${SUBSTRAIT_CPP_BUILD_TESTING})
 endif()
 
 add_subdirectory(src/substrait)
+add_subdirectory(export)

export/CMakeLists.txt

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+
+add_subdirectory(planloader)

export/planloader/CMakeLists.txt

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+
+if(NOT BUILD_SUBDIR_NAME EQUAL "release")
+  message(
+    SEND_ERROR,
+    "The planloader library does not work in Debug mode due to its dependencies."
+  )
+endif()
+
+add_library(planloader SHARED planloader.cpp)
+add_compile_options(-FPIC)
+set_target_properties(planloader PROPERTIES SUBVERSION 1)
+
+add_dependencies(planloader substrait_io)
+target_link_libraries(planloader substrait_io)
+
+install(
+  TARGETS planloader
+  LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+  PRIVATE_HEADER DESTINATION ${CMAKE_INSTALL_INCDIR})
+
+if(${SUBSTRAIT_CPP_BUILD_TESTING})
+  add_subdirectory(tests)
+endif()

export/planloader/planloader.cpp

@@ -0,0 +1,62 @@
+/* SPDX-License-Identifier: Apache-2.0 */
+
+#include "planloader.h"
+
+#include <substrait/common/Io.h>
+
+extern "C" {
+
+// Load a Substrait plan (in any format) from disk.
+// Stores the Substrait plan in planBuffer in serialized form.
+// Returns a SerializedPlan structure containing either the serialized plan or
+// an error message.  errorMessage is nullptr upon success.
+SerializedPlan* load_substrait_plan(const char* filename) {
+  auto newPlan = new SerializedPlan();
+  newPlan->buffer = nullptr;
+  newPlan->size = 0;
+  newPlan->errorMessage = nullptr;
+
+  auto planOrError = io::substrait::loadPlan(filename);
+  if (!planOrError.ok()) {
+    auto errMsg = planOrError.status().message();
+    newPlan->errorMessage = new char[errMsg.length()+1];
+    strncpy(newPlan->errorMessage, errMsg.data(), errMsg.length()+1);
+    return newPlan;
+  }
+  ::substrait::proto::Plan plan = *planOrError;
+  std::string text = plan.SerializeAsString();
+  newPlan->buffer = new char[text.length()+1];
+  memcpy(newPlan->buffer, text.data(), text.length()+1);
+  newPlan->size = text.length();
+  return newPlan;
+}
+
+void free_substrait_plan(SerializedPlan* plan) {
+  if (plan->buffer) {
+    free(plan->buffer);
+  }
+  if (plan->errorMessage) {
+    free(plan->errorMessage);
+  }
+  free(plan);
+}
+
+// Write a serialized Substrait plan to disk in the specified format.
+// On error returns a non-empty error message.
+// On success a nullptr is returned.
+const char* save_substrait_plan(
+    const char* planData,
+    uint32_t planDataLength,
+    const char* filename,
+    io::substrait::PlanFileFormat format) {
+  ::substrait::proto::Plan plan;
+  std::string data(planData, planDataLength);
+  plan.ParseFromString(data);
+  auto result = io::substrait::savePlan(plan, filename, format);
+  if (!result.ok()) {
+    return result.message().data();
+  }
+  return nullptr;
+}
+
+}  // extern "C"

export/planloader/planloader.h

@@ -0,0 +1,35 @@
+/* SPDX-License-Identifier: Apache-2.0 */
+
+#include <substrait/common/Io.h>
+
+extern "C" {
+
+using SerializedPlan = struct {
+  char *buffer;
+  uint32_t size;
+  char *errorMessage;
+};
+
+// Since this is actually C code, stick to C style names for exporting.
+// NOLINTBEGIN(readability-identifier-naming)
+
+// Load a Substrait plan (in any format) from disk.
+// Stores the Substrait plan in planBuffer in serialized form.
+// Returns a SerializedPlan structure containing either the serialized plan or
+// an error message.
+SerializedPlan* load_substrait_plan(const char* filename);
+
+void free_substrait_plan(SerializedPlan* plan);
+
+// Write a serialized Substrait plan to disk in the specified format.
+// On error returns a non-empty error message.
+// On success an empty string is returned.
+const char* save_substrait_plan(
+    const char* planData,
+    uint32_t planDataLength,
+    const char* filename,
+    io::substrait::PlanFileFormat format);
+
+// NOLINTEND(readability-identifier-naming)
+
+}  // extern "C"

export/planloader/tests/CMakeLists.txt

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+
+cmake_path(GET CMAKE_CURRENT_BINARY_DIR PARENT_PATH
+           CMAKE_CURRENT_BINARY_PARENT_DIR)
+cmake_path(GET CMAKE_CURRENT_BINARY_PARENT_DIR PARENT_PATH
+           CMAKE_CURRENT_BINARY_TOPLEVEL_DIR)
+
+set(CMAKE_RUNTIME_OUTPUT_DIRECTORY
+    "${CMAKE_CURRENT_BINARY_TOPLEVEL_DIR}/${BUILD_SUBDIR_NAME}")
+
+add_test_case(
+  planloader_test
+  SOURCES
+  PlanLoaderTest.cpp
+  EXTRA_LINK_LIBS
+  planloader
+  gtest
+  gtest_main)
+
+set(TEXTPLAN_SOURCE_DIR "${CMAKE_SOURCE_DIR}/src/substrait/textplan")
+
+add_custom_command(
+  TARGET planloader_test
+  POST_BUILD
+  COMMAND ${CMAKE_COMMAND} -E echo "Copying unit test data.."
+  COMMAND ${CMAKE_COMMAND} -E make_directory
+          "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/tests/data"
+  COMMAND
+    ${CMAKE_COMMAND} -E copy
+    "${TEXTPLAN_SOURCE_DIR}/converter/data/q6_first_stage.json"
+    "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/tests/data/q6_first_stage.json")
+
+message(
+  STATUS "test data will be here:  ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/tests/data"
+)

export/planloader/tests/PlanLoaderTest.cpp

@@ -0,0 +1,32 @@
+/* SPDX-License-Identifier: Apache-2.0 */
+
+#include <gtest/gtest.h>
+#include <functional>
+
+#include "../planloader.h"
+#include "substrait/proto/plan.pb.h"
+
+namespace io::substrait::textplan {
+namespace {
+
+TEST(PlanLoaderTest, LoadAndSave) {
+  auto serializedPlan = load_substrait_plan("data/q6_first_stage.json");
+  ASSERT_EQ(serializedPlan->errorMessage, nullptr);
+
+  ::substrait::proto::Plan plan;
+  bool parseStatus =
+      plan.ParseFromArray(serializedPlan->buffer, serializedPlan->size);
+  ASSERT_TRUE(parseStatus) << "Failed to parse the plan.";
+
+  const char* saveStatus = save_substrait_plan(
+      serializedPlan->buffer,
+      serializedPlan->size,
+      "outfile.splan",
+      PlanFileFormat::kText);
+  ASSERT_EQ(saveStatus, nullptr);
+
+  free_substrait_plan(serializedPlan);
+}
+
+} // namespace
+} // namespace io::substrait::textplan

src/substrait/textplan/ParseResult.h

@@ -46,6 +46,10 @@ class ParseResult {
     return errors;
   }
 
+  void addErrors(const std::vector<std::string>& errors) {
+    syntaxErrors_.insert(syntaxErrors_.end(), errors.begin(), errors.end());
+  }
+
   // Add the capability for ::testing::PrintToString to print ParseResult.
   friend std::ostream& operator<<(std::ostream& os, const ParseResult& result);
 

src/substrait/textplan/StructuredSymbolData.h

@@ -24,6 +24,8 @@ struct RelationData {
   const SymbolInfo* continuingPipeline{nullptr};
   // The next nodes in the pipelines that this node starts.
   std::vector<const SymbolInfo*> newPipelines;
+  // Expressions in this relation consume subqueries with these symbols.
+  std::vector<const SymbolInfo*> subQueryPipelines;
 
   // The information corresponding to the relation without any references to
   // other relations or inputs.
@@ -46,6 +48,10 @@ struct RelationData {
   // references to the symbol would use the alias.)
   std::map<size_t, std::string> generatedFieldReferenceAlternativeExpression;
 
+  // Temporary storage for global aliases for expressions.  Used during the
+  // construction of a relation.
+  std::map<size_t, std::string> generatedFieldReferenceAliases;
+
   // If populated, supersedes the combination of fieldReferences and
   // generatedFieldReferences for the field symbols exposed by this relation.
   std::vector<const SymbolInfo*> outputFieldReferences;

src/substrait/textplan/SubstraitErrorListener.h

@@ -30,6 +30,10 @@ class SubstraitErrorListener {
     addError(-1, -1, msg);
   };
 
+  void addErrorInstances(const std::vector<ErrorInstance>& errors) {
+    errors_.insert(errors_.end(), errors.begin(), errors.end());
+  }
+
   const std::vector<ErrorInstance>& getErrors() {
     return errors_;
   };

src/substrait/textplan/SymbolTable.cpp

@@ -61,7 +61,8 @@ const SymbolInfo SymbolInfo::kUnknown = {
     std::nullopt};
 
 bool operator==(const SymbolInfo& left, const SymbolInfo& right) {
-  return (left.name == right.name) && (left.location == right.location) &&
+  return (left.name == right.name) &&
+      (left.sourceLocation == right.sourceLocation) &&
       (left.type == right.type);
 }
 
@@ -118,13 +119,31 @@ size_t SymbolTable::findSymbolIndex(const SymbolInfo& symbol) {
   return symbols_.size();
 }
 
-void SymbolTable::updateLocation(
+void SymbolTable::addPermanentLocation(
     const SymbolInfo& symbol,
     const Location& location) {
   auto index = findSymbolIndex(symbol);
+  symbols_[index]->permanentLocation = location;
   symbolsByLocation_.insert(std::make_pair(location, index));
 }
 
+void SymbolTable::setParentQueryLocation(
+    const io::substrait::textplan::SymbolInfo& symbol,
+    const io::substrait::textplan::Location& location) {
+  auto index = findSymbolIndex(symbol);
+  symbols_[index]->parentQueryLocation = location;
+
+  int highestIndex = -1;
+  for (const auto& sym : symbols_) {
+    if (sym->parentQueryLocation == location) {
+      if (sym->parentQueryIndex > highestIndex) {
+        highestIndex = sym->parentQueryIndex;
+      }
+    }
+  }
+  symbols_[index]->parentQueryIndex = highestIndex + 1;
+}
+
 void SymbolTable::addAlias(const std::string& alias, const SymbolInfo* symbol) {
   auto index = findSymbolIndex(*symbol);
   symbols_[index]->alias = alias;
@@ -169,6 +188,19 @@ const SymbolInfo* SymbolTable::lookupSymbolByLocationAndTypes(
   return nullptr;
 }
 
+const SymbolInfo* SymbolTable::lookupSymbolByParentQueryAndType(
+    const Location& location,
+    int index,
+    SymbolType type) const {
+  for (const auto& symbol : symbols_) {
+    if (symbol->parentQueryLocation == location &&
+        symbol->parentQueryIndex == index && symbol->type == type) {
+      return symbol.get();
+    }
+  }
+  return nullptr;
+}
+
 const SymbolInfo& SymbolTable::nthSymbolByType(uint32_t n, SymbolType type)
     const {
   int count = 0;
@@ -200,7 +232,11 @@ std::string SymbolTable::toDebugString() const {
     }
     auto relationData = ANY_CAST(std::shared_ptr<RelationData>, symbol->blob);
     result << std::left << std::setw(4) << relationCount++;
-    result << std::left << std::setw(20) << symbol->name << std::endl;
+    result << std::left << std::setw(20) << symbol->name;
+    if (!relationData->subQueryPipelines.empty()) {
+      result << " SQC=" << relationData->subQueryPipelines.size();
+    }
+    result << std::endl;
 
     int32_t fieldNum = 0;
     for (const auto& field : relationData->fieldReferences) {

src/substrait/textplan/SymbolTable.h

@@ -33,7 +33,7 @@ enum class SymbolType {
 };
 
 enum class RelationType {
-  // Logical plans
+  // Logical
   kUnknown = 0,
   kRead = 1,
   kProject = 2,
@@ -45,11 +45,11 @@ enum class RelationType {
   kFilter = 8,
   kSet = 9,
 
-  // Physical plans
+  // Physical
   kHashJoin = 31,
   kMergeJoin = 32,
 
-  // Write relations, currently unreachable in Plan protos.
+  // Write
   kExchange = 50,
   kDdl = 51,
   kWrite = 52,
@@ -79,7 +79,10 @@ struct SymbolInfo {
   std::string name;
   std::string alias{}; // If present, use this instead of name.
   const SymbolInfo* schema{nullptr}; // The related schema symbol if present.
-  Location location;
+  Location sourceLocation;
+  Location permanentLocation{Location::kUnknownLocation};
+  Location parentQueryLocation{Location::kUnknownLocation};
+  int parentQueryIndex{-1};
   SymbolType type;
   std::any subtype;
   std::any blob;
@@ -91,7 +94,7 @@ struct SymbolInfo {
       std::any newSubtype,
       std::any newBlob)
       : name(std::move(newName)),
-        location(newLocation),
+        sourceLocation(newLocation),
         type(newType),
         subtype(std::move(newSubtype)),
         blob(std::move(newBlob)){};
@@ -145,8 +148,14 @@ class SymbolTable {
       const std::any& subtype,
       const std::any& blob);
 
-  // Changes the location for a specified existing symbol.
-  void updateLocation(const SymbolInfo& symbol, const Location& location);
+  // Changes the permanent location (the version stored in the symbol table)
+  // for a specified existing symbol.
+  void addPermanentLocation(const SymbolInfo& symbol, const Location& location);
+
+  // Sets the location of the parent query.
+  void setParentQueryLocation(
+      const SymbolInfo& symbol,
+      const Location& location);
 
   // Adds an alias to the given symbol.
   void addAlias(const std::string& alias, const SymbolInfo* symbol);
@@ -165,6 +174,11 @@ class SymbolTable {
       const Location& location,
       std::unordered_set<SymbolType> types) const;
 
+  [[nodiscard]] const SymbolInfo* lookupSymbolByParentQueryAndType(
+      const Location& location,
+      int index,
+      SymbolType type) const;
+
   [[nodiscard]] const SymbolInfo& nthSymbolByType(uint32_t n, SymbolType type)
       const;