Function ExecPlan_run_substrait

r/src/compute-exec.cpp:543–576

[[substrait::export]]

// [[substrait::export]]
std::shared_ptr<arrow::Table> ExecPlan_run_substrait(
    const std::shared_ptr<acero::ExecPlan>& plan,
    const std::shared_ptr<arrow::Buffer>& serialized_plan) {
  std::vector<std::shared_ptr<AccumulatingConsumer>> consumers;

  std::function<std::shared_ptr<acero::SinkNodeConsumer>()> consumer_factory = [&] {
    consumers.emplace_back(new AccumulatingConsumer());
    return consumers.back();
  };

  arrow::Result<std::vector<acero::Declaration>> maybe_decls =
      ValueOrStop(arrow::engine::DeserializePlans(*serialized_plan, consumer_factory));
  std::vector<acero::Declaration> decls = std::move(ValueOrStop(maybe_decls));

  // For now, the Substrait plan must include a 'read' that points to
  // a Parquet file (instead of using a source node create in Arrow)
  for (const acero::Declaration& decl : decls) {
    auto node = decl.AddToPlan(plan.get());
    StopIfNotOk(node.status());
  }

  StopIfNotOk(plan->Validate());
  plan->StartProducing();
  StopIfNotOk(plan->finished().status());

  std::vector<std::shared_ptr<arrow::RecordBatch>> all_batches;
  for (const auto& consumer : consumers) {
    for (const auto& batch : consumer->batches()) {
      all_batches.push_back(batch);
    }
  }

  return ValueOrStop(arrow::Table::FromRecordBatches(std::move(all_batches)));
}

#endif