agg-fn-evaluator.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_EXPRS_AGG_FN_EVALUATOR_H
#define IMPALA_EXPRS_AGG_FN_EVALUATOR_H

#include <string>

#include <boost/scoped_ptr.hpp>
#include "common/status.h"
#include "runtime/descriptors.h"
#include "runtime/lib-cache.h"
#include "runtime/tuple-row.h"
#include "runtime/types.h"
#include "udf/udf.h"
#include "udf/udf-internal.h"

#include "gen-cpp/Exprs_types.h"
#include "gen-cpp/PlanNodes_types.h"
#include "gen-cpp/Types_types.h"

using namespace impala_udf;

namespace impala {

class AggregationNode;
class Expr;
class ExprContext;
class MemPool;
class MemTracker;
class ObjectPool;
class RowDescriptor;
class RuntimeState;
class SlotDescriptor;
class Tuple;
class TupleRow;
class TExprNode;

//
//
class AggFnEvaluator {
 public:
  enum AggregationOp {
    COUNT,
    MIN,
    MAX,
    SUM,
    AVG,
    NDV,
    OTHER,
  };

  static Status Create(ObjectPool* pool, const TExpr& desc, AggFnEvaluator** result);

  static Status Create(ObjectPool* pool, const TExpr& desc, bool is_analytic_fn,
      AggFnEvaluator** result);

  Status Prepare(RuntimeState* state, const RowDescriptor& desc,
      const SlotDescriptor* intermediate_slot_desc,
      const SlotDescriptor* output_slot_desc,
      MemPool* agg_fn_pool, FunctionContext** agg_fn_ctx);

  ~AggFnEvaluator();

  Status Open(RuntimeState* state, FunctionContext* agg_fn_ctx);

  void Close(RuntimeState* state);

  const ColumnType& intermediate_type() const { return intermediate_slot_desc_->type(); }
  bool is_merge() const { return is_merge_; }
  AggregationOp agg_op() const { return agg_op_; }
  const std::vector<ExprContext*>& input_expr_ctxs() const { return input_expr_ctxs_; }
  bool is_count_star() const { return agg_op_ == COUNT && input_expr_ctxs_.empty(); }
  bool is_builtin() const { return fn_.binary_type == TFunctionBinaryType::BUILTIN; }
  bool SupportsRemove() const { return remove_fn_ != NULL; }
  bool SupportsSerialize() const { return serialize_fn_ != NULL; }
  const std::string& fn_name() const { return fn_.name.function_name; }
  const std::string& update_symbol() const { return fn_.aggregate_fn.update_fn_symbol; }
  const std::string& merge_symbol() const { return fn_.aggregate_fn.merge_fn_symbol; }

  static std::string DebugString(const std::vector<AggFnEvaluator*>& exprs);
  std::string DebugString() const;

  void Init(FunctionContext* agg_fn_ctx, Tuple* dst);

  void Add(FunctionContext* agg_fn_ctx, TupleRow* src, Tuple* dst);

  void Remove(FunctionContext* agg_fn_ctx, TupleRow* src, Tuple* dst);

  void Merge(FunctionContext* agg_fn_ctx, Tuple* src, Tuple* dst);

  void Serialize(FunctionContext* agg_fn_ctx, Tuple* dst);
  void Finalize(FunctionContext* agg_fn_ctx, Tuple* src, Tuple* dst);

  void GetValue(FunctionContext* agg_fn_ctx, Tuple* src, Tuple* dst);

  static void Init(const std::vector<AggFnEvaluator*>& evaluators,
      const std::vector<FunctionContext*>& fn_ctxs, Tuple* dst);
  static void Add(const std::vector<AggFnEvaluator*>& evaluators,
      const std::vector<FunctionContext*>& fn_ctxs, TupleRow* src, Tuple* dst);
  static void Remove(const std::vector<AggFnEvaluator*>& evaluators,
      const std::vector<FunctionContext*>& fn_ctxs, TupleRow* src, Tuple* dst);
  static void Serialize(const std::vector<AggFnEvaluator*>& evaluators,
      const std::vector<FunctionContext*>& fn_ctxs, Tuple* dst);
  static void GetValue(const std::vector<AggFnEvaluator*>& evaluators,
      const std::vector<FunctionContext*>& fn_ctxs, Tuple* src, Tuple* dst);
  static void Finalize(const std::vector<AggFnEvaluator*>& evaluators,
      const std::vector<FunctionContext*>& fn_ctxs, Tuple* src, Tuple* dst);


 private:
  const TFunction fn_;
  const bool is_merge_;
  const bool is_analytic_fn_;

  const SlotDescriptor* intermediate_slot_desc_;

  const SlotDescriptor* output_slot_desc_;

  std::vector<ExprContext*> input_expr_ctxs_;

  AggregationOp agg_op_;

  std::vector<impala_udf::AnyVal*> staging_input_vals_;
  impala_udf::AnyVal* staging_intermediate_val_;
  impala_udf::AnyVal* staging_merge_input_val_;

  LibCache::LibCacheEntry* cache_entry_;

  void* init_fn_;
  void* update_fn_;
  void* remove_fn_;
  void* merge_fn_;
  void* serialize_fn_;
  void* get_value_fn_;
  void* finalize_fn_;

  AggFnEvaluator(const TExprNode& desc, bool is_analytic_fn);


  void Update(FunctionContext* agg_fn_ctx, TupleRow* row, Tuple* dst, void* fn);

  void SerializeOrFinalize(FunctionContext* agg_fn_ctx, Tuple* src,
      const SlotDescriptor* dst_slot_desc, Tuple* dst, void* fn);

  void SetDstSlot(FunctionContext* ctx, const impala_udf::AnyVal* src,
      const SlotDescriptor* dst_slot_desc, Tuple* dst);
};

inline void AggFnEvaluator::Add(
    FunctionContext* agg_fn_ctx, TupleRow* row, Tuple* dst) {
  agg_fn_ctx->impl()->IncrementNumUpdates();
  Update(agg_fn_ctx, row, dst, is_merge() ? merge_fn_ : update_fn_);
}
inline void AggFnEvaluator::Remove(
    FunctionContext* agg_fn_ctx, TupleRow* row, Tuple* dst) {
  agg_fn_ctx->impl()->IncrementNumRemoves();
  Update(agg_fn_ctx, row, dst, remove_fn_);
}
inline void AggFnEvaluator::Serialize(
    FunctionContext* agg_fn_ctx, Tuple* tuple) {
  SerializeOrFinalize(agg_fn_ctx, tuple, intermediate_slot_desc_, tuple, serialize_fn_);
}
inline void AggFnEvaluator::Finalize(
    FunctionContext* agg_fn_ctx, Tuple* src, Tuple* dst) {
  SerializeOrFinalize(agg_fn_ctx, src, output_slot_desc_, dst, finalize_fn_);
}
inline void AggFnEvaluator::GetValue(
    FunctionContext* agg_fn_ctx, Tuple* src, Tuple* dst) {
  SerializeOrFinalize(agg_fn_ctx, src, output_slot_desc_, dst, get_value_fn_);
}

inline void AggFnEvaluator::Init(const std::vector<AggFnEvaluator*>& evaluators,
      const std::vector<FunctionContext*>& fn_ctxs, Tuple* dst) {
  DCHECK_EQ(evaluators.size(), fn_ctxs.size());
  for (int i = 0; i < evaluators.size(); ++i) {
    evaluators[i]->Init(fn_ctxs[i], dst);
  }
}
inline void AggFnEvaluator::Add(const std::vector<AggFnEvaluator*>& evaluators,
      const std::vector<FunctionContext*>& fn_ctxs, TupleRow* src, Tuple* dst) {
  DCHECK_EQ(evaluators.size(), fn_ctxs.size());
  for (int i = 0; i < evaluators.size(); ++i) {
    evaluators[i]->Add(fn_ctxs[i], src, dst);
  }
}
inline void AggFnEvaluator::Remove(const std::vector<AggFnEvaluator*>& evaluators,
      const std::vector<FunctionContext*>& fn_ctxs, TupleRow* src, Tuple* dst) {
  DCHECK_EQ(evaluators.size(), fn_ctxs.size());
  for (int i = 0; i < evaluators.size(); ++i) {
    evaluators[i]->Remove(fn_ctxs[i], src, dst);
  }
}
inline void AggFnEvaluator::Serialize(const std::vector<AggFnEvaluator*>& evaluators,
      const std::vector<FunctionContext*>& fn_ctxs, Tuple* dst) {
  DCHECK_EQ(evaluators.size(), fn_ctxs.size());
  for (int i = 0; i < evaluators.size(); ++i) {
    evaluators[i]->Serialize(fn_ctxs[i], dst);
  }
}
inline void AggFnEvaluator::GetValue(const std::vector<AggFnEvaluator*>& evaluators,
      const std::vector<FunctionContext*>& fn_ctxs, Tuple* src, Tuple* dst) {
  DCHECK_EQ(evaluators.size(), fn_ctxs.size());
  for (int i = 0; i < evaluators.size(); ++i) {
    evaluators[i]->GetValue(fn_ctxs[i], src, dst);
  }
}
inline void AggFnEvaluator::Finalize(const std::vector<AggFnEvaluator*>& evaluators,
      const std::vector<FunctionContext*>& fn_ctxs, Tuple* src, Tuple* dst) {
  DCHECK_EQ(evaluators.size(), fn_ctxs.size());
  for (int i = 0; i < evaluators.size(); ++i) {
    evaluators[i]->Finalize(fn_ctxs[i], src, dst);
  }
}

}

#endif