partitioned-aggregation-node.h

Go to the documentation of this file.

// Copyright 2012 Cloudera Inc.
//
// 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.


#ifndef IMPALA_EXEC_PARTITIONED_AGGREGATION_NODE_H
#define IMPALA_EXEC_PARTITIONED_AGGREGATION_NODE_H

#include <functional>
#include <boost/scoped_ptr.hpp>

#include "exec/exec-node.h"
#include "exec/hash-table.h"
#include "runtime/buffered-block-mgr.h"
#include "runtime/buffered-tuple-stream.h"
#include "runtime/descriptors.h"  // for TupleId
#include "runtime/mem-pool.h"
#include "runtime/string-value.h"

namespace llvm {
  class Function;
}

namespace impala {

class AggFnEvaluator;
class LlvmCodeGen;
class RowBatch;
class RuntimeState;
struct StringValue;
class Tuple;
class TupleDescriptor;
class SlotDescriptor;

//
//
//
class PartitionedAggregationNode : public ExecNode {
 public:
  PartitionedAggregationNode(ObjectPool* pool,
      const TPlanNode& tnode, const DescriptorTbl& descs);

  virtual Status Init(const TPlanNode& tnode);
  virtual Status Prepare(RuntimeState* state);
  virtual Status Open(RuntimeState* state);
  virtual Status GetNext(RuntimeState* state, RowBatch* row_batch, bool* eos);
  virtual Status Reset(RuntimeState* state);
  virtual void Close(RuntimeState* state);

  static const char* LLVM_CLASS_NAME;

 protected:
  virtual Status QueryMaintenance(RuntimeState* state);

  virtual void DebugString(int indentation_level, std::stringstream* out) const;

 private:
  struct Partition;

  static const int PARTITION_FANOUT = 16;

  static const int NUM_PARTITIONING_BITS = 4;

  static const int MAX_PARTITION_DEPTH = 16;

  TupleId intermediate_tuple_id_;
  TupleDescriptor* intermediate_tuple_desc_;

  boost::scoped_ptr<RowDescriptor> intermediate_row_desc_;

  TupleId output_tuple_id_;
  TupleDescriptor* output_tuple_desc_;

  const bool needs_finalize_;

  bool needs_serialize_;

  std::vector<AggFnEvaluator*> aggregate_evaluators_;

  std::vector<impala_udf::FunctionContext*> agg_fn_ctxs_;
  boost::scoped_ptr<MemPool> agg_fn_pool_;

  std::vector<ExprContext*> probe_expr_ctxs_;

  std::vector<ExprContext*> build_expr_ctxs_;

  bool contains_var_len_grouping_exprs_;

  RuntimeState* state_;
  BufferedBlockMgr::Client* block_mgr_client_;

  bool using_small_buffers_;

  Tuple* singleton_output_tuple_;
  bool singleton_output_tuple_returned_;

  boost::scoped_ptr<MemPool> mem_pool_;

  boost::scoped_ptr<HashTableCtx> ht_ctx_;

  Partition* output_partition_;
  HashTable::Iterator output_iterator_;

  typedef Status (*ProcessRowBatchFn)(
      PartitionedAggregationNode*, RowBatch*, HashTableCtx*);
  ProcessRowBatchFn process_row_batch_fn_;

  RuntimeProfile::Counter* build_timer_;

  RuntimeProfile::Counter* ht_resize_timer_;

  RuntimeProfile::Counter* get_results_timer_;

  RuntimeProfile::Counter* num_hash_buckets_;

  RuntimeProfile::Counter* partitions_created_;

  RuntimeProfile::HighWaterMarkCounter* max_partition_level_;

  RuntimeProfile::Counter* num_row_repartitioned_;

  RuntimeProfile::Counter* num_repartitions_;

  RuntimeProfile::Counter* num_spilled_partitions_;

  RuntimeProfile::HighWaterMarkCounter* largest_partition_percent_;

  struct Partition {
    Partition(PartitionedAggregationNode* parent, int level)
      : parent(parent), is_closed(false), level(level) {}

    Status InitStreams();

    bool InitHashTable();

    void Close(bool finalize_rows);

    Status Spill(Tuple* tuple = NULL);

    bool is_spilled() const { return hash_tbl.get() == NULL; }

    PartitionedAggregationNode* parent;

    bool is_closed;

    const int level;

    boost::scoped_ptr<HashTable> hash_tbl;

    std::vector<impala_udf::FunctionContext*> agg_fn_ctxs;
    boost::scoped_ptr<MemPool> agg_fn_pool;

    boost::scoped_ptr<BufferedTupleStream> aggregated_row_stream;

    boost::scoped_ptr<BufferedTupleStream> unaggregated_row_stream;
  };

  std::vector<Partition*> hash_partitions_;

  std::list<Partition*> spilled_partitions_;

  std::list<Partition*> aggregated_partitions_;

  boost::scoped_ptr<BufferedTupleStream> serialize_stream_;

  Tuple* ConstructIntermediateTuple(
      const std::vector<impala_udf::FunctionContext*>& agg_fn_ctxs,
      MemPool* pool, BufferedTupleStream* stream);

  void UpdateTuple(impala_udf::FunctionContext** agg_fn_ctxs, Tuple* tuple, TupleRow* row,
                   bool is_merge = false);

  Tuple* GetOutputTuple(const std::vector<impala_udf::FunctionContext*>& agg_fn_ctxs,
                        Tuple* tuple, MemPool* pool);

  Status ProcessBatchNoGrouping(RowBatch* batch, HashTableCtx* ht_ctx = NULL);

  //
  template<bool AGGREGATED_ROWS>
  Status IR_ALWAYS_INLINE ProcessBatch(RowBatch* batch, HashTableCtx* ht_ctx);

  template<bool AGGREGATED_ROWS>
  Status ProcessStream(BufferedTupleStream* input_stream);

  Status CreateHashPartitions(int level);

  int64_t LargestSpilledPartition() const;

  Status NextPartition();

  Status SpillPartition(Partition* curr_partition = NULL,
      Tuple* curr_intermediate_tuple = NULL);

  Status MoveHashPartitions(int64_t input_rows);

  void CleanupHashTbl(const std::vector<impala_udf::FunctionContext*>& fn_ctxs,
      HashTable::Iterator it);

  llvm::Function* CodegenUpdateSlot(AggFnEvaluator* evaluator, SlotDescriptor* slot_desc);

  llvm::Function* CodegenUpdateTuple();

  llvm::Function* CodegenProcessBatch();

  Status ProcessBatch_false(RowBatch* batch, HashTableCtx* ht_ctx);
  Status ProcessBatch_true(RowBatch* batch, HashTableCtx* ht_ctx);

  int MinRequiredBuffers() const {
    return 2 * PARTITION_FANOUT + 1 + (needs_serialize_ ? 1 : 0);
  }
};

}

#endif