text: // Licensed to the Apache Software Foundation (ASF) under one
  // or more contributor license agreements.  See the NOTICE file
  // distributed with this work for additional information
  // regarding copyright ownership.  The ASF licenses this file
  // to you under the Apache License, Version 2.0 (the
  // "License"); you may not use this file except in compliance
  // with the License.  You may obtain a copy of the License at
  //
  //   http://www.apache.org/licenses/LICENSE-2.0
  //
  // Unless required by applicable law or agreed to in writing,
  // software distributed under the License is distributed on an
  // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  // KIND, either express or implied.  See the License for the
  // specific language governing permissions and limitations
  // under the License.
  
  #include "arrow/compute/exec/exec_plan.h"
  
  #include <sstream>
  #include <unordered_map>
  #include <unordered_set>
  
  #include "arrow/compute/exec.h"
  #include "arrow/compute/exec/expression.h"
  #include "arrow/compute/exec_internal.h"
  #include "arrow/compute/registry.h"
  #include "arrow/datum.h"
  #include "arrow/record_batch.h"
  #include "arrow/result.h"
  #include "arrow/util/async_generator.h"
  #include "arrow/util/checked_cast.h"
  #include "arrow/util/logging.h"
  #include "arrow/util/optional.h"
  
  namespace arrow {
  
  using internal::checked_cast;
  
  namespace compute {
  
  namespace {
  
  struct ExecPlanImpl : public ExecPlan {
    explicit ExecPlanImpl(ExecContext* exec_context) : ExecPlan(exec_context) {}
  
    ~ExecPlanImpl() override {
      if (started_ && !finished_.is_finished()) {
        ARROW_LOG(WARNING) << "Plan was destroyed before finishing";
        StopProducing();
        finished().Wait();
      }
    }
  
    ExecNode* AddNode(std::unique_ptr<ExecNode> node) {
      if (node->label().empty()) {
        node->SetLabel(std::to_string(auto_label_counter_++));
      }
      if (node->num_inputs() == 0) {
        sources_.push_back(node.get());
      }
      if (node->num_outputs() == 0) {
        sinks_.push_back(node.get());
      }
      nodes_.push_back(std::move(node));
      return nodes_.back().get();
    }
  
    Status Validate() const {
      if (nodes_.empty()) {
        return Status::Invalid("ExecPlan has no node");
      }
      for (const auto& node : nodes_) {
        RETURN_NOT_OK(node->Validate());
      }
      return Status::OK();
    }
  
    Status StartProducing() {
      if (started_) {
        return Status::Invalid("restarted ExecPlan");
      }
      started_ = true;
  
      // producers precede consumers
      sorted_nodes_ = TopoSort();
  
      std::vector<Future<>> futures;
  
      Status st = Status::OK();
  
      using rev_it = std::reverse_iterator<NodeVector::iterator>;
      for (rev_it it(sorted_nodes_.end()), end(sorted_nodes_.begin()); it != end; ++it) {
        auto node = *it;
  
        st = node->StartProducing();
        if (!st.ok()) {
          // Stop nodes that successfully started, in reverse order
          stopped_ = true;
          StopProducingImpl(it.base(), sorted_nodes_.end());
          break;
        }
  
        futures.push_back(node->finished());
      }
  
      finished_ = AllFinished(futures);
      return st;
    }
  
    void StopProducing() {
      DCHECK(started_) << "stopped an ExecPlan which never started";
      stopped_ = true;
  
      StopProducingImpl(sorted_nodes_.begin(), sorted_nodes_.end());
    }
  
    template <typename It>
    void StopProducingImpl(It begin, It end) {
      for (auto it = begin; it != end; ++it) {
        auto node = *it;
        node->StopProducing();
      }
    }
  
    NodeVector TopoSort() const {
      struct Impl {
        const std::vector<std::unique_ptr<ExecNode>>& nodes;
        std::unordered_set<ExecNode*> visited;
        NodeVector sorted;
  
        explicit Impl(const std::vector<std::unique_ptr<ExecNode>>& nodes) : nodes(nodes) {
          visited.reserve(nodes.size());
          sorted.resize(nodes.size());
  
          for (const auto& node : nodes) {
            Visit(node.get());
          }
  
          DCHECK_EQ(visited.size(), nodes.size());
        }
  
        void Visit(ExecNode* node) {
          if (visited.count(node) != 0) return;
  
          for (auto input : node->inputs()) {
            // Ensure that producers are inserted before this consumer
            Visit(input);
          }
  
          sorted[visited.size()] = node;
          visited.insert(node);
        }
      };
  
      return std::move(Impl{nodes_}.sorted);
    }
  
    std::string ToString() const {
      std::stringstream ss;
      ss << "ExecPlan with " << nodes_.size() << " nodes:" << std::endl;
      for (const auto& node : TopoSort()) {
        ss << node->ToString() << std::endl;
      }
      return ss.str();
    }
  
    Future<> finished_ = Future<>::MakeFinished();
    bool started_ = false, stopped_ = false;
    std::vector<std::unique_ptr<ExecNode>> nodes_;
    NodeVector sources_, sinks_;
    NodeVector sorted_nodes_;
    uint32_t auto_label_counter_ = 0;
  };
  
  ExecPlanImpl* ToDerived(ExecPlan* ptr) { return checked_cast<ExecPlanImpl*>(ptr); }
  
  const ExecPlanImpl* ToDerived(const ExecPlan* ptr) {
    return checked_cast<const ExecPlanImpl*>(ptr);
  }
  
  util::optional<int> GetNodeIndex(const std::vector<ExecNode*>& nodes,
                                   const ExecNode* node) {
    for (int i = 0; i < static_cast<int>(nodes.size()); ++i) {
      if (nodes[i] == node) return i;
    }
    return util::nullopt;
  }
  
  }  // namespace
  
  Result<std::shared_ptr<ExecPlan>> ExecPlan::Make(ExecContext* ctx) {
    return std::shared_ptr<ExecPlan>(new ExecPlanImpl{ctx});
  }
  
  ExecNode* ExecPlan::AddNode(std::unique_ptr<ExecNode> node) {
    return ToDerived(this)->AddNode(std::move(node));
  }
  
  const ExecPlan::NodeVector& ExecPlan::sources() const {
    return ToDerived(this)->sources_;
  }
  
  const ExecPlan::NodeVector& ExecPlan::sinks() const { return ToDerived(this)->sinks_; }
  
  Status ExecPlan::Validate() { return ToDerived(this)->Validate(); }
  
  Status ExecPlan::StartProducing() { return ToDerived(this)->StartProducing(); }
  
  void ExecPlan::StopProducing() { ToDerived(this)->StopProducing(); }
  
  Future<> ExecPlan::finished() { return ToDerived(this)->finished_; }
  
  std::string ExecPlan::ToString() const { return ToDerived(this)->ToString(); }
  
  ExecNode::ExecNode(ExecPlan* plan, NodeVector inputs,
                     std::vector<std::string> input_labels,
                     std::shared_ptr<Schema> output_schema, int num_outputs)
      : plan_(plan),
        inputs_(std::move(inputs)),
        input_labels_(std::move(input_labels)),
        output_schema_(std::move(output_schema)),
        num_outputs_(num_outputs) {
    for (auto input : inputs_) {
      input->outputs_.push_back(this);
    }
  }
  
  Status ExecNode::Validate() const {
    if (inputs_.size() != input_labels_.size()) {
      return Status::Invalid("Invalid number of inputs for '", label(), "' (expected ",
                             num_inputs(), ", actual ", input_labels_.size(), ")");
    }
  
    if (static_cast<int>(outputs_.size()) != num_outputs_) {
      return Status::Invalid("Invalid number of outputs for '", label(), "' (expected ",
                             num_outputs(), ", actual ", outputs_.size(), ")");
    }
  
    for (auto out : outputs_) {
      auto input_index = GetNodeIndex(out->inputs(), this);
      if (!input_index) {
        return Status::Invalid("Node '", label(), "' outputs to node '", out->label(),
                               "' but is not listed as an input.");
      }
    }
  
    return Status::OK();
  }
  
  std::string ExecNode::ToString() const {
    std::stringstream ss;
    ss << kind_name() << "{\"" << label_ << '"';
    if (!inputs_.empty()) {
      ss << ", inputs=[";
      for (size_t i = 0; i < inputs_.size(); i++) {
        if (i > 0) ss << ", ";
        ss << input_labels_[i] << ": \"" << inputs_[i]->label() << '"';
      }
      ss << ']';
    }
  
    if (!outputs_.empty()) {
      ss << ", outputs=[";
      for (size_t i = 0; i < outputs_.size(); i++) {
        if (i > 0) ss << ", ";
        ss << "\"" << outputs_[i]->label() << "\"";
      }
      ss << ']';
    }
  
    const std::string extra = ToStringExtra();
    if (!extra.empty()) ss << ", " << extra;
  
    ss << '}';
    return ss.str();
  }
  
  std::string ExecNode::ToStringExtra() const { return ""; }
  
  bool ExecNode::ErrorIfNotOk(Status status) {
    if (status.ok()) return false;
  
    for (auto out : outputs_) {
      out->ErrorReceived(this, out == outputs_.back() ? std::move(status) : status);
    }
    return true;
  }
  
  MapNode::MapNode(ExecPlan* plan, std::vector<ExecNode*> inputs,
                   std::shared_ptr<Schema> output_schema, bool async_mode)
      : ExecNode(plan, std::move(inputs), /*input_labels=*/{"target"},
                 std::move(output_schema),
                 /*num_outputs=*/1) {
    if (async_mode) {
      executor_ = plan_->exec_context()->executor();
    } else {
      executor_ = nullptr;
    }
  }
  
  void MapNode::ErrorReceived(ExecNode* input, Status error) {
    DCHECK_EQ(input, inputs_[0]);
    outputs_[0]->ErrorReceived(this, std::move(error));
  }
  
  void MapNode::InputFinished(ExecNode* input, int total_batches) {
    DCHECK_EQ(input, inputs_[0]);
    outputs_[0]->InputFinished(this, total_batches);
    if (input_counter_.SetTotal(total_batches)) {
      this->Finish();
    }
  }
  
  Status MapNode::StartProducing() { return Status::OK(); }
  
  void MapNode::PauseProducing(ExecNode* output) {}
  
  void MapNode::ResumeProducing(ExecNode* output) {}
  
  void MapNode::StopProducing(ExecNode* output) {
    DCHECK_EQ(output, outputs_[0]);
    StopProducing();
  }
  
  void MapNode::StopProducing() {
    if (executor_) {
      this->stop_source_.RequestStop();
    }
    if (input_counter_.Cancel()) {
      this->Finish();
    }
    inputs_[0]->StopProducing(this);
  }
  
  Future<> MapNode::finished() { return finished_; }
  
  void MapNode::SubmitTask(std::function<Result<ExecBatch>(ExecBatch)> map_fn,
                           ExecBatch batch) {
    Status status;
    // This will be true if the node is stopped early due to an error or manual
    // cancellation
    if (input_counter_.Completed()) {
      return;
    }
    auto task = [this, map_fn, batch]() {
      auto guarantee = batch.guarantee;
      auto output_batch = map_fn(std::move(batch));
      if (ErrorIfNotOk(output_batch.status())) {
        return output_batch.status();
      }
      output_batch->guarantee = guarantee;
      outputs_[0]->InputReceived(this, output_batch.MoveValueUnsafe());
      return Status::OK();
    };
  
    if (executor_) {
      status = task_group_.AddTask([this, task]() -> Result<Future<>> {
        return this->executor_->Submit(this->stop_source_.token(), [this, task]() {
          auto status = task();
          if (this->input_counter_.Increment()) {
            this->Finish(status);
          }
          return status;
        });
      });
    } else {
      status = task();
      if (input_counter_.Increment()) {
        this->Finish(status);
      }
    }
    // If we get a cancelled status from AddTask it means this node was stopped
    // or errored out already so we can just drop the task.
    if (!status.ok() && !status.IsCancelled()) {
      if (input_counter_.Cancel()) {
        this->Finish(status);
      }
      inputs_[0]->StopProducing(this);
      return;
    }
  }
  
  void MapNode::Finish(Status finish_st /*= Status::OK()*/) {
    if (executor_) {
      task_group_.End().AddCallback([this, finish_st](const Status& st) {
        Status final_status = finish_st & st;
        this->finished_.MarkFinished(final_status);
      });
    } else {
      this->finished_.MarkFinished(finish_st);
    }
  }
  
  std::shared_ptr<RecordBatchReader> MakeGeneratorReader(
      std::shared_ptr<Schema> schema,
      std::function<Future<util::optional<ExecBatch>>()> gen, MemoryPool* pool) {
    struct Impl : RecordBatchReader {
      std::shared_ptr<Schema> schema() const override { return schema_; }
  
      Status ReadNext(std::shared_ptr<RecordBatch>* record_batch) override {
        ARROW_ASSIGN_OR_RAISE(auto batch, iterator_.Next());
        if (batch) {
          ARROW_ASSIGN_OR_RAISE(*record_batch, batch->ToRecordBatch(schema_, pool_));
        } else {
          *record_batch = IterationEnd<std::shared_ptr<RecordBatch>>();
        }
        return Status::OK();
      }
  
      MemoryPool* pool_;
      std::shared_ptr<Schema> schema_;
      Iterator<util::optional<ExecBatch>> iterator_;
    };
  
    auto out = std::make_shared<Impl>();
    out->pool_ = pool;
    out->schema_ = std::move(schema);
    out->iterator_ = MakeGeneratorIterator(std::move(gen));
    return out;
  }
  
  Result<ExecNode*> Declaration::AddToPlan(ExecPlan* plan,
                                           ExecFactoryRegistry* registry) const {
    std::vector<ExecNode*> inputs(this->inputs.size());
  
    size_t i = 0;
    for (const Input& input : this->inputs) {
      if (auto node = util::get_if<ExecNode*>(&input)) {
        inputs[i++] = *node;
        continue;
      }
      ARROW_ASSIGN_OR_RAISE(inputs[i++],
                            util::get<Declaration>(input).AddToPlan(plan, registry));
    }
  
    ARROW_ASSIGN_OR_RAISE(
        auto node, MakeExecNode(this->factory_name, plan, std::move(inputs), *this->options,
                                registry));
    node->SetLabel(this->label);
    return node;
  }
  
  Declaration Declaration::Sequence(std::vector<Declaration> decls) {
    DCHECK(!decls.empty());
  
    Declaration out = std::move(decls.back());
    decls.pop_back();
    auto receiver = &out;
    while (!decls.empty()) {
      Declaration input = std::move(decls.back());
      decls.pop_back();
  
      receiver->inputs.emplace_back(std::move(input));
      receiver = &util::get<Declaration>(receiver->inputs.front());
    }
    return out;
  }
  
  namespace internal {
  
  void RegisterSourceNode(ExecFactoryRegistry*);
  void RegisterFilterNode(ExecFactoryRegistry*);
  void RegisterProjectNode(ExecFactoryRegistry*);
  void RegisterUnionNode(ExecFactoryRegistry*);
  void RegisterAggregateNode(ExecFactoryRegistry*);
  void RegisterSinkNode(ExecFactoryRegistry*);
  void RegisterHashJoinNode(ExecFactoryRegistry*);
  
  }  // namespace internal
  
  ExecFactoryRegistry* default_exec_factory_registry() {
    class DefaultRegistry : public ExecFactoryRegistry {
     public:
      DefaultRegistry() {
        internal::RegisterSourceNode(this);
        internal::RegisterFilterNode(this);
        internal::RegisterProjectNode(this);
        internal::RegisterUnionNode(this);
        internal::RegisterAggregateNode(this);
        internal::RegisterSinkNode(this);
        internal::RegisterHashJoinNode(this);
      }
  
      Result<Factory> GetFactory(const std::string& factory_name) override {
        auto it = factories_.find(factory_name);
        if (it == factories_.end()) {
          return Status::KeyError("ExecNode factory named ", factory_name,
                                  " not present in registry.");
        }
        return it->second;
      }
  
      Status AddFactory(std::string factory_name, Factory factory) override {
        auto it_success = factories_.emplace(std::move(factory_name), std::move(factory));
  
        if (!it_success.second) {
          const auto& factory_name = it_success.first->first;
          return Status::KeyError("ExecNode factory named ", factory_name,
                                  " already registered.");
        }
  
        return Status::OK();
      }
  
     private:
      std::unordered_map<std::string, Factory> factories_;
    };
  
    static DefaultRegistry instance;
    return &instance;
  }
  
  Result<std::function<Future<util::optional<ExecBatch>>()>> MakeReaderGenerator(
      std::shared_ptr<RecordBatchReader> reader, ::arrow::internal::Executor* io_executor,
      int max_q, int q_restart) {
    auto batch_it = MakeMapIterator(
        [](std::shared_ptr<RecordBatch> batch) {
          return util::make_optional(ExecBatch(*batch));
        },
        MakeIteratorFromReader(reader));
  
    return MakeBackgroundGenerator(std::move(batch_it), io_executor, max_q, q_restart);
  }
  }  // namespace compute
  
  }  // namespace arrow
  