src/dctv/query_execution.h

// Copyright (C) 2020 The Android Open Source Project
//
// Licensed 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.
#pragma once
#include "dctv.h"

#include <boost/intrusive/set.hpp>
#include <boost/intrusive/slist.hpp>
#include <utility>

#include "npy.h"
#include "operator_context.h"
#include "optional.h"
#include "perf.h"
#include "pyutil.h"
#include "query.h"
#include "vector.h"

namespace dctv {

// Define our queue data structures.  The nested namespace saves
// typing without polluting the dctv namespace with boost crud.

struct OperatorContextCompare {
  inline bool operator()(const OperatorContext& left,
                         const OperatorContext& right) const noexcept;
};

namespace oplink {
using namespace boost::intrusive;
using ReScoreQueue = slist<
  OperatorContext,
  constant_time_size<false>,
  member_hook<OperatorContext,
              oplink::ReScoreQueueLink,
              &OperatorContext::link_re_score_queue>
  >;
using RunQueue = set<
  OperatorContext,
  constant_time_size<false>,
  compare<OperatorContextCompare>,
  member_hook<OperatorContext,
              oplink::RunQueueLink,
              &OperatorContext::link_run_queue>
  >;
}  // namespace oplink

struct StringTable;

struct QueryExecution final : BasePyObject,
                              SupportsGcClear,
                              SupportsWeakRefs
{
  using LinkState = OperatorContext::LinkState;

  struct Config {
    Vector<unique_pyref> plan;
    unique_obj_pyref<StringTable> st;
    unique_obj_pyref<QueryCache> qc;
    unique_pyref env;
    unique_pyref progress_callback;
    unique_pyref perf_callback;
  };

  // Named constructor: see pyutil.h comment on object exceptions.
  static unique_obj_pyref<QueryExecution> make(Config config);

  ~QueryExecution() noexcept;

  // Extract the next iterator value (since QueryExecution is just one
  // big iterator)
  unique_pyref next();

  // Add an operator to the queue for the first time.  The new
  // operator must be in the orphaned state; it ends up on the
  // score-recompute queue.
  void add_operator(unique_op_ref op);
  void remove_operator_on_close(OperatorContext* oc) noexcept;
  void add_to_re_score_queue(OperatorContext* oc) noexcept;

  inline OperatorContext* get_current_operator() const;

  int py_traverse(visitproc visit, void* arg) const noexcept;
  int py_clear() noexcept;

  inline QueryCache* get_qc() const noexcept;

  inline const perf::Sampler* get_perf_sampler() const noexcept;
  void accumulate_perf(pyref query_action,
                       const perf::Sample& accumulated_perf);

  // Size of all non-final, non-broadcast blocks used in this query.
  const npy_intp block_size;

  template<typename Functor>
  void enumerate_operators(Functor&& functor);

  static PyTypeObject pytype;

 private:
  friend struct unique_obj_pyref<QueryExecution>;  // For constructor

  struct PerfInfo final {
    unique_pyref callback;
    perf::Sampler sampler;
  };

  explicit QueryExecution(Config config);
  void flush_pending_score_recomputations();
  void remove_from_re_score_queue(OperatorContext* oc) noexcept;

  // Operator queues.  The run queue is for operators whose score is
  // up-to-date; the rescore queue is for operators whose scores we
  // need to recompute.  We do it this way to batch score
  // recomputations, since a score can change many times during a
  // single "turn of the crank" in query executing.
  oplink::ReScoreQueue re_score_queue;
  oplink::RunQueue run_queue;

  // Currently-executing operator.
  OperatorContext* current_operator = nullptr;

  // Query cache backing this execution.  Set during initialization.
  unique_obj_pyref<QueryCache> qc;

  // Miscellaneous dict holding user-specified associations.
  unique_pyref env;

  // Progress callback.
  unique_pyref progress_callback;

  optional<PerfInfo> perf_info;

  static PyGetSetDef  pygetset[];
  static PyMethodDef pymethods[];
  static PyMemberDef pymembers[];
};

void init_query_execution(pyref m);

}  // namespace dctv

#include "query_execution-inl.h"