ruby/red-arrow/test/test-record-batch.rb

# 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.

class RecordBatchTest < Test::Unit::TestCase
  sub_test_case(".new") do
    def setup
      @schema = Arrow::Schema.new(visible: :boolean,
                                  count: :uint32)
    end

    test("[raw_table]") do
      raw_table = {
        visible: [true, nil, false],
        count: [1, nil, 3],
      }
      record_batch = Arrow::RecordBatch.new(raw_table)
      assert_equal([
                     {"visible" => true,  "count" => 1},
                     {"visible" => nil,   "count" => nil},
                     {"visible" => false, "count" => 3},
                   ],
                   record_batch.each_record.collect(&:to_h))
    end

    test("[Schema, records]") do
      records = [
        {visible: true, count: 1},
        nil,
        [false, 3],
      ]
      record_batch = Arrow::RecordBatch.new(@schema, records)
      assert_equal([
                     {"visible" => true,  "count" => 1},
                     {"visible" => nil,   "count" => nil},
                     {"visible" => false, "count" => 3},
                   ],
                   record_batch.each_record.collect(&:to_h))
    end

    test("[Schema, columns]") do
      columns = {
        visible: [true, nil, false],
        count: [1, 2, nil],
      }
      record_batch = Arrow::RecordBatch.new(@schema, columns)
      assert_equal([
                     {"visible" => true,  "count" => 1},
                     {"visible" => nil,   "count" => 2},
                     {"visible" => false, "count" => nil},
                   ],
                   record_batch.each_record.collect(&:to_h))
    end

    test("[Schema, n_rows, columns]") do
      columns = [
        Arrow::BooleanArray.new([true, nil, false]),
        Arrow::UInt32Array.new([1, 2, nil]),
      ]
      n_rows = columns[0].length
      record_batch = Arrow::RecordBatch.new(@schema, n_rows, columns)
      assert_equal([
                     {"visible" => true,  "count" => 1},
                     {"visible" => nil,   "count" => 2},
                     {"visible" => false, "count" => nil},
                   ],
                   record_batch.each_record.collect(&:to_h))
    end
  end

  sub_test_case("instance methods") do
    def setup
      @schema = Arrow::Schema.new(count: :uint32)
      @counts = Arrow::UInt32Array.new([1, 2, 4, 8])
      @record_batch = Arrow::RecordBatch.new(@schema, @counts.length, [@counts])
    end

    sub_test_case("#each") do
      test("default") do
        records = []
        @record_batch.each do |record|
          records << [record, record.index]
        end
        assert_equal([
                       [0, 0],
                       [1, 1],
                       [2, 2],
                       [3, 3],
                     ],
                     records.collect {|record, i| [record.index, i]})
      end

      test("reuse_record: true") do
        records = []
        @record_batch.each(reuse_record: true) do |record|
          records << [record, record.index]
        end
        assert_equal([
                       [3, 0],
                       [3, 1],
                       [3, 2],
                       [3, 3],
                     ],
                     records.collect {|record, i| [record.index, i]})
      end
    end

    test("#to_table") do
      assert_equal(Arrow::Table.new(@schema, [@counts]),
                   @record_batch.to_table)
    end

    sub_test_case("#==") do
      test("Arrow::RecordBatch") do
        assert do
          @record_batch == @record_batch
        end
      end

      test("not Arrow::RecordBatch") do
        assert do
          not (@record_batch == 29)
        end
      end
    end

    sub_test_case("#[]") do
      def setup
        @record_batch = Arrow::RecordBatch.new(a: [true],
                                               b: [true],
                                               c: [true],
                                               d: [true],
                                               e: [true],
                                               f: [true],
                                               g: [true])
      end

      test("[String]") do
        assert_equal(Arrow::Column.new(@record_batch, 0),
                     @record_batch["a"])
      end

      test("[Symbol]") do
        assert_equal(Arrow::Column.new(@record_batch, 1),
                     @record_batch[:b])
      end

      test("[Integer]") do
        assert_equal(Arrow::Column.new(@record_batch, 6),
                     @record_batch[-1])
      end

      test("[Range]") do
        assert_equal(Arrow::RecordBatch.new(d: [true],
                                            e: [true]),
                     @record_batch[3..4])
      end

      test("[[Symbol, String, Integer, Range]]") do
        assert_equal(Arrow::RecordBatch.new(c: [true],
                                            a: [true],
                                            g: [true],
                                            d: [true],
                                            e: [true]),
                     @record_batch[[:c, "a", -1, 3..4]])
      end
    end

    sub_test_case("#column") do
      test("#size") do
        assert_equal(@counts.size, @record_batch[:count].size)
      end
    end

    sub_test_case("#merge") do
      def setup
        @count_array = Arrow::UInt8Array.new([1, 2, 4, 8])
        @visible_array = Arrow::BooleanArray.new([true, false, nil, true])

        @record_batch = Arrow::RecordBatch.new(
          count: @count_array,
          visible: @visible_array,
        )
      end

      test("add column") do
        name_array = Arrow::StringArray.new(["a", "b", "c", "d"])

        expected = Arrow::RecordBatch.new(
          count: @count_array,
          visible: @visible_array,
          name: name_array,
        )

        assert_equal(expected,
                     @record_batch.merge(name: name_array))
      end
      test("remove column") do
        expected = Arrow::RecordBatch.new(
          count: @count_array,
        )

        assert_equal(expected,
                     @record_batch.merge(visible: nil))
      end

      test("replace column") do
        visible_array = Arrow::Int32Array.new([1, 1, 1, 1])

        expected = Arrow::RecordBatch.new(
          count: @count_array,
          visible: visible_array,
        )

        assert_equal(expected,
                     @record_batch.merge(visible: visible_array))
      end

      test("merge record batch add") do
        name_array = Arrow::StringArray.new(["a", "b", "c", "d"])

        other = Arrow::RecordBatch.new(
          name: name_array,
        )

        expected = Arrow::RecordBatch.new(
          count: @count_array,
          visible: @visible_array,
          name: name_array,
        )

        assert_equal(expected,
                     @record_batch.merge(other))
      end

      test("merge record batch replace") do
        visible_array = Arrow::Int32Array.new([1, 1, 1, 1])

        other = Arrow::RecordBatch.new(
          visible: visible_array,
        )

        expected = Arrow::RecordBatch.new(
          count: @count_array,
          visible: visible_array,
        )

        assert_equal(expected,
                     @record_batch.merge(other))
      end

      test("invalid target") do
        assert_raise(ArgumentError) do
          @record_batch.merge(29)
        end
      end
    end
  end
end