diff --git a/ruby/red-arrow-format/lib/arrow-format/record-batch.rb b/ruby/red-arrow-format/lib/arrow-format/record-batch.rb index 23de1b205857..46d39a98e839 100644 --- a/ruby/red-arrow-format/lib/arrow-format/record-batch.rb +++ b/ruby/red-arrow-format/lib/arrow-format/record-batch.rb @@ -26,7 +26,14 @@ class RecordBatch alias_method :length, :n_rows attr_reader :columns attr_reader :message_metadata - def initialize(schema, n_rows, columns, message_metadata: nil) + def initialize(*args, message_metadata: nil) + n_args = args.size + args = build(args[0]) if n_args == 1 + if args.size != 3 + message = "wrong number of arguments (given #{n_args}, expected 1 or 3)" + raise ArgumentError, message + end + schema, n_rows, columns = args @schema = schema @n_rows = n_rows @columns = columns @@ -101,5 +108,57 @@ def all_buffers_enumerator end end end + + private + def build(data) + records = nil + fields = [] + columns = [] + mode = nil + i = 0 + data.each do |name, column| + if i.zero? + if column.nil? + mode = :record + records = {} + else + mode = :column + end + end + + if mode == :record + record = name + record.each do |n, value| + values = (records[n] ||= []) + while values.size < i + values << nil + end + values << value + end + else + fields << Field.new(name, column.type) + columns << column + end + i += 1 + end + + if mode == :record + records.each do |name, values| + column = Array.build(values) + fields << Field.new(name, column.type) + columns << column + end + end + + raise ArgumentError, "no data" if columns.empty? + all_n_rows = columns.collect(&:size) + if all_n_rows.uniq.size != 1 + message = + "inconsistent the number of rows: #{all_n_rows.join(", ")}" + raise ArgumentError, message + end + + return Schema.new(fields), n_rows, columns + end end end diff --git a/ruby/red-arrow-format/test/test-record-batch.rb b/ruby/red-arrow-format/test/test-record-batch.rb new file mode 100644 index 000000000000..61e907e3a5c1 --- /dev/null +++ b/ruby/red-arrow-format/test/test-record-batch.rb @@ -0,0 +1,75 @@ +# 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 TestRecordBatch < Test::Unit::TestCase + sub_test_case("#initialize") do + def setup + @boolean_array = ArrowFormat::BooleanArray.new([true, nil, false]) + @int32_array = ArrowFormat::Int32Array.new([-(2 ** 31), 0, (2 ** 31) - 1]) + end + + test("{}") do + error = ArgumentError.new("no data") + assert_raise(error) do + ArrowFormat::RecordBatch.new({}) + end + end + + test("{Symbol => Array}") do + raw_records = { + boolean: @boolean_array, + int32: @int32_array, + } + record_batch = ArrowFormat::RecordBatch.new(raw_records) + assert_equal(raw_records, record_batch.to_h) + end + + test("{String => Array}") do + raw_records = { + "boolean" => @boolean_array, + "int32" => @int32_array, + } + record_batch = ArrowFormat::RecordBatch.new(raw_records) + assert_equal(raw_records, record_batch.to_h) + end + + test("[{}]") do + raw_records = [ + {boolean: true, int32: -(2 ** 31)}, + { int32: 0}, + {boolean: false, int32: (2 ** 31) - 1}, + ] + record_batch = ArrowFormat::RecordBatch.new(raw_records) + assert_equal({ + boolean: @boolean_array, + int32: @int32_array, + }, + record_batch.to_h) + end + + test("inconsistent n_rows") do + raw_records = { + boolean: ArrowFormat::BooleanArray.new([true, nil]), + int32: ArrowFormat::Int32Array.new([-(2 ** 31), 0, (2 ** 31) - 1]), + } + error = ArgumentError.new("inconsistent the number of rows: 2, 3") + assert_raise(error) do + ArrowFormat::RecordBatch.new(raw_records) + end + end + end +end