ceph/src/arrow/c_glib/test/gandiva/test-selectable-projector.rb

   1 # Licensed to the Apache Software Foundation (ASF) under one
   2 # or more contributor license agreements.  See the NOTICE file
   3 # distributed with this work for additional information
   4 # regarding copyright ownership.  The ASF licenses this file
   5 # to you under the Apache License, Version 2.0 (the
   6 # "License"); you may not use this file except in compliance
   7 # with the License.  You may obtain a copy of the License at
   8 #
   9 #   http://www.apache.org/licenses/LICENSE-2.0
  10 #
  11 # Unless required by applicable law or agreed to in writing,
  12 # software distributed under the License is distributed on an
  13 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  14 # KIND, either express or implied.  See the License for the
  15 # specific language governing permissions and limitations
  16 # under the License.
  17
  18 class TestGandivaSelectableProjector < Test::Unit::TestCase
  19   include Helper::Buildable
  20
  21   def setup
  22     omit("Gandiva is required") unless defined?(::Gandiva)
  23
  24     field1 = Arrow::Field.new("field1", Arrow::Int32DataType.new)
  25     field2 = Arrow::Field.new("field2", Arrow::Int32DataType.new)
  26     @schema = Arrow::Schema.new([field1, field2])
  27
  28     input_arrays = [
  29       build_int32_array([1, 2, 3, 4]),
  30       build_int32_array([11, 13, 15, 17]),
  31     ]
  32     @record_batch = Arrow::RecordBatch.new(@schema,
  33                                            input_arrays[0].length,
  34                                            input_arrays)
  35
  36     @field_node1 = Gandiva::FieldNode.new(field1)
  37     @field_node2 = Gandiva::FieldNode.new(field2)
  38     add_function_node =
  39       Gandiva::FunctionNode.new("add",
  40                                 [@field_node1, @field_node2],
  41                                 Arrow::Int32DataType.new)
  42     subtract_function_node =
  43       Gandiva::FunctionNode.new("subtract",
  44                                 [@field_node1, @field_node2],
  45                                 Arrow::Int32DataType.new)
  46     add_result = Arrow::Field.new("add_result", Arrow::Int32DataType.new)
  47     add_expression = Gandiva::Expression.new(add_function_node, add_result)
  48     subtract_result = Arrow::Field.new("subtract_result",
  49                                        Arrow::Int32DataType.new)
  50     subtract_expression = Gandiva::Expression.new(subtract_function_node,
  51                                                   subtract_result)
  52     @selection_vector = Gandiva::UInt16SelectionVector.new(@record_batch.n_rows)
  53     @projector =
  54       Gandiva::SelectableProjector.new(@schema,
  55                                        [add_expression, subtract_expression],
  56                                        @selection_vector.mode)
  57   end
  58
  59   def test_evaluate
  60     two_node = Gandiva::Int32LiteralNode.new(2)
  61     condition_node = Gandiva::FunctionNode.new("greater_than",
  62                                                [@field_node1, two_node],
  63                                                Arrow::BooleanDataType.new)
  64     condition = Gandiva::Condition.new(condition_node)
  65     filter = Gandiva::Filter.new(@schema, condition)
  66     filter.evaluate(@record_batch, @selection_vector)
  67     outputs = @projector.evaluate(@record_batch, @selection_vector)
  68     assert_equal([
  69                    [18, 21],
  70                    [-12, -13],
  71                  ],
  72                  outputs.collect(&:values))
  73   end
  74 end