uda-test-harness.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_UDA_TEST_HARNESS_H
#define IMPALA_UDA_TEST_HARNESS_H

#include <string>
#include <sstream>
#include <vector>

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

#include "udf/udf.h"
#include "udf/udf-debug.h"
#include "udf/udf-test-harness.h"

namespace impala_udf {

enum UdaExecutionMode {
  ALL = 0,
  SINGLE_NODE = 1,
  ONE_LEVEL = 2,
  TWO_LEVEL = 3,
};

template<typename RESULT, typename INTERMEDIATE>
class UdaTestHarnessBase {
 public:
  typedef void (*InitFn)(FunctionContext* context, INTERMEDIATE* result);

  typedef void (*MergeFn)(FunctionContext* context, const INTERMEDIATE& src,
      INTERMEDIATE* dst);

  typedef const INTERMEDIATE (*SerializeFn)(FunctionContext* context,
      const INTERMEDIATE& type);

  typedef RESULT (*FinalizeFn)(FunctionContext* context, const INTERMEDIATE& value);

  typedef bool (*ResultComparator)(const RESULT& x, const RESULT& y);

  void SetResultComparator(ResultComparator fn) {
    result_comparator_fn_ = fn;
  }

  void SetIntermediateSize(int byte_size) {
    fixed_buffer_byte_size_ = byte_size;
  }

  const std::string& GetErrorMsg() const { return error_msg_; }

 protected:
  UdaTestHarnessBase(InitFn init_fn, MergeFn merge_fn,
      SerializeFn serialize_fn, FinalizeFn finalize_fn)
    : init_fn_(init_fn),
      merge_fn_(merge_fn),
      serialize_fn_(serialize_fn),
      finalize_fn_(finalize_fn),
      result_comparator_fn_(NULL),
      num_input_values_(0) {
  }

  struct ScopedFunctionContext {
    ScopedFunctionContext(FunctionContext* context, UdaTestHarnessBase* harness)
        : context_(context), harness_(harness) { }

    ~ScopedFunctionContext() {
      UdfTestHarness::CloseContext(context_);
      harness_->CheckContext(context_);
      delete context_;
    }

    FunctionContext* get() { return context_; }

   private:
    FunctionContext* context_;
    UdaTestHarnessBase* harness_;
  };

  bool Execute(const RESULT& expected, UdaExecutionMode mode);

  bool CheckContext(FunctionContext* context);

  bool CheckResult(const RESULT& x, const RESULT& y);

  RESULT ExecuteSingleNode(ScopedFunctionContext* result_context);

  RESULT ExecuteOneLevel(int num_nodes, ScopedFunctionContext* result_context);

  RESULT ExecuteTwoLevel(int num1, int num2, ScopedFunctionContext* result_context);

  virtual void Update(int idx, FunctionContext* context, INTERMEDIATE* dst) = 0;

  InitFn init_fn_;
  MergeFn merge_fn_;
  SerializeFn serialize_fn_;
  FinalizeFn finalize_fn_;

  ResultComparator result_comparator_fn_;

  int num_input_values_;

  int fixed_buffer_byte_size_;

  std::string error_msg_;
};

template<typename RESULT, typename INTERMEDIATE, typename INPUT>
class UdaTestHarness : public UdaTestHarnessBase<RESULT, INTERMEDIATE> {
 public:
  typedef void (*UpdateFn)(FunctionContext* context, const INPUT& input,
      INTERMEDIATE* result);

  typedef UdaTestHarnessBase<RESULT, INTERMEDIATE> BaseClass;

  UdaTestHarness(
      typename BaseClass::InitFn init_fn,
      UpdateFn update_fn,
      typename BaseClass::MergeFn merge_fn,
      typename BaseClass::SerializeFn serialize_fn,
      typename BaseClass::FinalizeFn finalize_fn)
    : BaseClass(init_fn, merge_fn, serialize_fn, finalize_fn),
      update_fn_(update_fn) {
  }

  bool Execute(const std::vector<INPUT>& values, const RESULT& expected,
      UdaExecutionMode mode = ALL);

  template<typename T>
  bool Execute(const std::vector<T>& values, const RESULT& expected,
      UdaExecutionMode mode = ALL) {
    input_.resize(values.size());
    BaseClass::num_input_values_ = input_.size();
    for (int i = 0; i < values.size(); ++i) {
      input_[i] = &values[i];
    }
    return BaseClass::Execute(expected, mode);
  }

 protected:
  virtual void Update(int idx, FunctionContext* context, INTERMEDIATE* dst);

 private:
  UpdateFn update_fn_;
  std::vector<const INPUT*> input_;
};

template<typename RESULT, typename INTERMEDIATE, typename INPUT1, typename INPUT2>
class UdaTestHarness2 : public UdaTestHarnessBase<RESULT, INTERMEDIATE> {
 public:
  typedef void (*UpdateFn)(FunctionContext* context, const INPUT1& input1,
      const INPUT2& input2, INTERMEDIATE* result);

  typedef UdaTestHarnessBase<RESULT, INTERMEDIATE> BaseClass;

  UdaTestHarness2(
      typename BaseClass::InitFn init_fn,
      UpdateFn update_fn,
      typename BaseClass::MergeFn merge_fn,
      typename BaseClass::SerializeFn serialize_fn,
      typename BaseClass::FinalizeFn finalize_fn)
    : BaseClass(init_fn, merge_fn, serialize_fn, finalize_fn),
      update_fn_(update_fn) {
  }

  bool Execute(const std::vector<INPUT1>& values1, const std::vector<INPUT2>& values2,
      const RESULT& expected, UdaExecutionMode mode = ALL);

 protected:
  virtual void Update(int idx, FunctionContext* context, INTERMEDIATE* dst);

 private:
  UpdateFn update_fn_;
  const std::vector<INPUT1>* input1_;
  const std::vector<INPUT2>* input2_;
};

template<typename RESULT, typename INTERMEDIATE, typename INPUT1, typename INPUT2,
    typename INPUT3>
class UdaTestHarness3 : public UdaTestHarnessBase<RESULT, INTERMEDIATE> {
 public:
  typedef void (*UpdateFn)(FunctionContext* context, const INPUT1& input1,
      const INPUT2& input2, const INPUT3& input3, INTERMEDIATE* result);

  typedef UdaTestHarnessBase<RESULT, INTERMEDIATE> BaseClass;

  UdaTestHarness3(
      typename BaseClass::InitFn init_fn,
      UpdateFn update_fn,
      typename BaseClass::MergeFn merge_fn,
      typename BaseClass::SerializeFn serialize_fn,
      typename BaseClass::FinalizeFn finalize_fn)
    : BaseClass(init_fn, merge_fn, serialize_fn, finalize_fn),
      update_fn_(update_fn) {
  }

  bool Execute(const std::vector<INPUT1>& values1, const std::vector<INPUT2>& values2,
      const std::vector<INPUT3>& values3,
      const RESULT& expected, UdaExecutionMode mode = ALL);

 protected:
  virtual void Update(int idx, FunctionContext* context, INTERMEDIATE* dst);

 private:
  UpdateFn update_fn_;
  const std::vector<INPUT1>* input1_;
  const std::vector<INPUT2>* input2_;
  const std::vector<INPUT3>* input3_;
};

template<typename RESULT, typename INTERMEDIATE, typename INPUT1, typename INPUT2,
    typename INPUT3, typename INPUT4>
class UdaTestHarness4 : public UdaTestHarnessBase<RESULT, INTERMEDIATE> {
 public:
  typedef void (*UpdateFn)(FunctionContext* context, const INPUT1& input1,
      const INPUT2& input2, const INPUT3& input3, const INPUT4& input4,
      INTERMEDIATE* result);

  typedef UdaTestHarnessBase<RESULT, INTERMEDIATE> BaseClass;

  UdaTestHarness4(
      typename BaseClass::InitFn init_fn,
      UpdateFn update_fn,
      typename BaseClass::MergeFn merge_fn,
      typename BaseClass::SerializeFn serialize_fn,
      typename BaseClass::FinalizeFn finalize_fn)
    : BaseClass(init_fn, merge_fn, serialize_fn, finalize_fn),
      update_fn_(update_fn) {
  }

  bool Execute(const std::vector<INPUT1>& values1, const std::vector<INPUT2>& values2,
      const std::vector<INPUT3>& values3, const std::vector<INPUT4>& values4,
      const RESULT& expected, UdaExecutionMode mode = ALL);

 protected:
  virtual void Update(int idx, FunctionContext* context, INTERMEDIATE* dst);

 private:
  UpdateFn update_fn_;
  const std::vector<INPUT1>* input1_;
  const std::vector<INPUT2>* input2_;
  const std::vector<INPUT3>* input3_;
  const std::vector<INPUT4>* input4_;
};

}

#include "udf/uda-test-harness-impl.h"

#endif