import quincy 17.2.0

[ceph.git] / ceph / src / arrow / cpp / src / gandiva / tests / filter_project_test.cc

// 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 <gtest/gtest.h>
#include "arrow/memory_pool.h"
#include "gandiva/filter.h"
#include "gandiva/projector.h"
#include "gandiva/selection_vector.h"
#include "gandiva/tests/test_util.h"
#include "gandiva/tree_expr_builder.h"

namespace gandiva {

using arrow::boolean;
using arrow::float32;
using arrow::int32;

class TestFilterProject : public ::testing::Test {
 public:
  void SetUp() { pool_ = arrow::default_memory_pool(); }

 protected:
  arrow::MemoryPool* pool_;
};

TEST_F(TestFilterProject, TestSimple16) {
  // schema for input fields
  auto field0 = field("f0", int32());
  auto field1 = field("f1", int32());
  auto field2 = field("f2", int32());
  auto resultField = field("result", int32());
  auto schema = arrow::schema({field0, field1, field2});

  // Build condition f0 < f1
  auto node_f0 = TreeExprBuilder::MakeField(field0);
  auto node_f1 = TreeExprBuilder::MakeField(field1);
  auto node_f2 = TreeExprBuilder::MakeField(field2);
  auto less_than_function =
      TreeExprBuilder::MakeFunction("less_than", {node_f0, node_f1}, arrow::boolean());
  auto condition = TreeExprBuilder::MakeCondition(less_than_function);
  auto sum_expr = TreeExprBuilder::MakeExpression("add", {field1, field2}, resultField);

  auto configuration = TestConfiguration();

  std::shared_ptr<Filter> filter;
  std::shared_ptr<Projector> projector;

  auto status = Filter::Make(schema, condition, configuration, &filter);
  EXPECT_TRUE(status.ok());

  status = Projector::Make(schema, {sum_expr}, SelectionVector::MODE_UINT16,
                           configuration, &projector);
  EXPECT_TRUE(status.ok());

  // Create a row-batch with some sample data
  int num_records = 5;
  auto array0 = MakeArrowArrayInt32({1, 2, 6, 40, 3}, {true, true, true, true, true});
  auto array1 = MakeArrowArrayInt32({5, 9, 3, 17, 6}, {true, true, true, true, true});
  auto array2 = MakeArrowArrayInt32({1, 2, 6, 40, 3}, {true, true, true, true, false});
  // expected output
  auto result = MakeArrowArrayInt32({6, 11, 0}, {true, true, false});
  // prepare input record batch
  auto in_batch = arrow::RecordBatch::Make(schema, num_records, {array0, array1, array2});

  std::shared_ptr<SelectionVector> selection_vector;
  status = SelectionVector::MakeInt16(num_records, pool_, &selection_vector);
  EXPECT_TRUE(status.ok());
  // Evaluate expression
  status = filter->Evaluate(*in_batch, selection_vector);
  EXPECT_TRUE(status.ok());

  // Evaluate expression
  arrow::ArrayVector outputs;

  status = projector->Evaluate(*in_batch, selection_vector.get(), pool_, &outputs);
  EXPECT_TRUE(status.ok());

  // Validate results
  EXPECT_ARROW_ARRAY_EQUALS(result, outputs.at(0));
}

TEST_F(TestFilterProject, TestSimple32) {
  // schema for input fields
  auto field0 = field("f0", int32());
  auto field1 = field("f1", int32());
  auto field2 = field("f2", int32());
  auto resultField = field("result", int32());
  auto schema = arrow::schema({field0, field1, field2});

  // Build condition f0 < f1
  auto node_f0 = TreeExprBuilder::MakeField(field0);
  auto node_f1 = TreeExprBuilder::MakeField(field1);
  auto node_f2 = TreeExprBuilder::MakeField(field2);
  auto less_than_function =
      TreeExprBuilder::MakeFunction("less_than", {node_f0, node_f1}, arrow::boolean());
  auto condition = TreeExprBuilder::MakeCondition(less_than_function);
  auto sum_expr = TreeExprBuilder::MakeExpression("add", {field1, field2}, resultField);

  auto configuration = TestConfiguration();

  std::shared_ptr<Filter> filter;
  std::shared_ptr<Projector> projector;

  auto status = Filter::Make(schema, condition, configuration, &filter);
  EXPECT_TRUE(status.ok());

  status = Projector::Make(schema, {sum_expr}, SelectionVector::MODE_UINT32,
                           configuration, &projector);
  EXPECT_TRUE(status.ok());

  // Create a row-batch with some sample data
  int num_records = 5;
  auto array0 = MakeArrowArrayInt32({1, 2, 6, 40, 3}, {true, true, true, true, true});
  auto array1 = MakeArrowArrayInt32({5, 9, 3, 17, 6}, {true, true, true, true, true});
  auto array2 = MakeArrowArrayInt32({1, 2, 6, 40, 3}, {true, true, true, true, false});
  // expected output
  auto result = MakeArrowArrayInt32({6, 11, 0}, {true, true, false});
  // prepare input record batch
  auto in_batch = arrow::RecordBatch::Make(schema, num_records, {array0, array1, array2});

  std::shared_ptr<SelectionVector> selection_vector;
  status = SelectionVector::MakeInt32(num_records, pool_, &selection_vector);
  EXPECT_TRUE(status.ok());
  // Evaluate expression
  status = filter->Evaluate(*in_batch, selection_vector);
  EXPECT_TRUE(status.ok());

  // Evaluate expression
  arrow::ArrayVector outputs;

  status = projector->Evaluate(*in_batch, selection_vector.get(), pool_, &outputs);
  ASSERT_OK(status);

  // Validate results
  EXPECT_ARROW_ARRAY_EQUALS(result, outputs.at(0));
}

TEST_F(TestFilterProject, TestSimple64) {
  // schema for input fields
  auto field0 = field("f0", int32());
  auto field1 = field("f1", int32());
  auto field2 = field("f2", int32());
  auto resultField = field("result", int32());
  auto schema = arrow::schema({field0, field1, field2});

  // Build condition f0 < f1
  auto node_f0 = TreeExprBuilder::MakeField(field0);
  auto node_f1 = TreeExprBuilder::MakeField(field1);
  auto node_f2 = TreeExprBuilder::MakeField(field2);
  auto less_than_function =
      TreeExprBuilder::MakeFunction("less_than", {node_f0, node_f1}, arrow::boolean());
  auto condition = TreeExprBuilder::MakeCondition(less_than_function);
  auto sum_expr = TreeExprBuilder::MakeExpression("add", {field1, field2}, resultField);

  auto configuration = TestConfiguration();

  std::shared_ptr<Filter> filter;
  std::shared_ptr<Projector> projector;

  auto status = Filter::Make(schema, condition, configuration, &filter);
  EXPECT_TRUE(status.ok());

  status = Projector::Make(schema, {sum_expr}, SelectionVector::MODE_UINT64,
                           configuration, &projector);
  ASSERT_OK(status);

  // Create a row-batch with some sample data
  int num_records = 5;
  auto array0 = MakeArrowArrayInt32({1, 2, 6, 40, 3}, {true, true, true, true, true});
  auto array1 = MakeArrowArrayInt32({5, 9, 3, 17, 6}, {true, true, true, true, true});
  auto array2 = MakeArrowArrayInt32({1, 2, 6, 40, 3}, {true, true, true, true, false});
  // expected output
  auto result = MakeArrowArrayInt32({6, 11, 0}, {true, true, false});
  // prepare input record batch
  auto in_batch = arrow::RecordBatch::Make(schema, num_records, {array0, array1, array2});

  std::shared_ptr<SelectionVector> selection_vector;
  status = SelectionVector::MakeInt64(num_records, pool_, &selection_vector);
  EXPECT_TRUE(status.ok());
  // Evaluate expression
  status = filter->Evaluate(*in_batch, selection_vector);
  EXPECT_TRUE(status.ok());

  // Evaluate expression
  arrow::ArrayVector outputs;

  status = projector->Evaluate(*in_batch, selection_vector.get(), pool_, &outputs);
  EXPECT_TRUE(status.ok());

  // Validate results
  EXPECT_ARROW_ARRAY_EQUALS(result, outputs.at(0));
}

TEST_F(TestFilterProject, TestSimpleIf) {
  // schema for input fields
  auto fielda = field("a", int32());
  auto fieldb = field("b", int32());
  auto fieldc = field("c", int32());
  auto schema = arrow::schema({fielda, fieldb, fieldc});

  // output fields
  auto field_result = field("res", int32());

  auto node_a = TreeExprBuilder::MakeField(fielda);
  auto node_b = TreeExprBuilder::MakeField(fieldb);
  auto node_c = TreeExprBuilder::MakeField(fieldc);

  auto greater_than_function =
      TreeExprBuilder::MakeFunction("greater_than", {node_a, node_b}, boolean());
  auto filter_condition = TreeExprBuilder::MakeCondition(greater_than_function);

  auto project_condition =
      TreeExprBuilder::MakeFunction("less_than", {node_b, node_c}, boolean());
  auto if_node = TreeExprBuilder::MakeIf(project_condition, node_b, node_c, int32());

  auto expr = TreeExprBuilder::MakeExpression(if_node, field_result);
  auto configuration = TestConfiguration();

  // Build a filter for the expressions.
  std::shared_ptr<Filter> filter;
  auto status = Filter::Make(schema, filter_condition, configuration, &filter);
  EXPECT_TRUE(status.ok());

  // Build a projector for the expressions.
  std::shared_ptr<Projector> projector;
  status = Projector::Make(schema, {expr}, SelectionVector::MODE_UINT32, configuration,
                           &projector);
  ASSERT_OK(status);

  // Create a row-batch with some sample data
  int num_records = 6;
  auto array0 =
      MakeArrowArrayInt32({10, 12, -20, 5, 21, 29}, {true, true, true, true, true, true});
  auto array1 =
      MakeArrowArrayInt32({5, 15, 15, 17, 12, 3}, {true, true, true, true, true, true});
  auto array2 = MakeArrowArrayInt32({1, 25, 11, 30, -21, 30},
                                    {true, true, true, true, true, false});

  // Create a selection vector
  std::shared_ptr<SelectionVector> selection_vector;
  status = SelectionVector::MakeInt32(num_records, pool_, &selection_vector);
  EXPECT_TRUE(status.ok());

  // expected output
  auto exp = MakeArrowArrayInt32({1, -21, 0}, {true, true, false});

  // prepare input record batch
  auto in_batch = arrow::RecordBatch::Make(schema, num_records, {array0, array1, array2});

  // Evaluate filter
  status = filter->Evaluate(*in_batch, selection_vector);
  EXPECT_TRUE(status.ok());

  // Evaluate project
  arrow::ArrayVector outputs;
  status = projector->Evaluate(*in_batch, selection_vector.get(), pool_, &outputs);
  EXPECT_TRUE(status.ok());

  // Validate results
  EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));
}
}  // namespace gandiva