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;