ceph/src/arrow/ruby/red-gandiva/test/test-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 TestProjector < Test::Unit::TestCase
  19   def test_evaluate
  20     table = Arrow::Table.new(:field1 => Arrow::Int32Array.new([1, 13, 3, 17]),
  21                              :field2 => Arrow::Int32Array.new([11, 2, 15, 17]),
  22                              :field3 => Arrow::Int32Array.new([1, 10, 2, 2]))
  23     schema = table.schema
  24
  25     expression1 = schema.build_expression do |record|
  26       record.field1 + record.field2
  27     end
  28
  29     expression2 = schema.build_expression do |record, context|
  30       context.if(record.field1 > record.field2)
  31         .then(record.field1 + record.field2 * record.field3)
  32         .elsif(record.field1 == record.field2)
  33         .then(record.field1 - record.field2 / record.field3)
  34         .else(record.field2)
  35     end
  36
  37     projector = Gandiva::Projector.new(schema,
  38                                        [expression1, expression2])
  39
  40     table.each_record_batch do |record_batch|
  41       outputs = projector.evaluate(record_batch)
  42       assert_equal([
  43                      [12, 15, 18, 34],
  44                      [11, 33, 15, 9]
  45                    ],
  46                    outputs.collect(&:values))
  47     end
  48   end
  49 end