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 db2b2dbe8fca..776621dc0843 100644 --- a/ruby/red-arrow-format/lib/arrow-format/record-batch.rb +++ b/ruby/red-arrow-format/lib/arrow-format/record-batch.rb @@ -30,9 +30,14 @@ class RecordBatch attr_reader :message_metadata def initialize(*args, message_metadata: nil) n_args = args.size - args = build(args[0]) if n_args == 1 + case n_args + when 1 + args = build(args[0]) + when 2 + args = build_with_schema(*args) + end if args.size != 3 - message = "wrong number of arguments (given #{n_args}, expected 1 or 3)" + message = "wrong number of arguments (given #{n_args}, expected 1..3)" raise ArgumentError, message end schema, n_rows, columns = args @@ -143,6 +148,57 @@ def all_buffers_enumerator end private + def build_with_schema(schema, data) + fields = schema.fields + name_to_index = {} + fields.each_with_index do |field, i| + name_to_index[field.name] = i + end + if data.is_a?(Hash) + raw_columns = [] + data.each do |name, values| + field_index = name_to_index[name.to_s] + raw_columns[field_index] = values if field_index + end + columns = fields.zip(raw_columns).collect do |field, values| + field.type.build_array(values || []) + end + 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 + else + raw_columns = fields.collect { [] } + data.each_with_index do |record, nth_record| + case record + when nil + when Hash + record.each do |name, value| + field_index = name_to_index[name.to_s] + raw_columns[field_index] << value if field_index + end + else + if record.size > raw_columns.size + message = "row #{nth_record} has more values than schema fields" + raise ArgumentError, message + end + record.each_with_index do |value, field_index| + raw_columns[field_index] << value + end + end + raw_columns.each do |column| + column << nil if column.size != nth_record + 1 + end + end + columns = fields.zip(raw_columns).collect do |field, values| + field.type.build_array(values) + end + end + n_rows = columns.first&.size || 0 + [schema, n_rows, columns] + end + def build(data) records = nil fields = [] diff --git a/ruby/red-arrow-format/test/test-record-batch.rb b/ruby/red-arrow-format/test/test-record-batch.rb index 30bd79fc35af..9a082bebfa6c 100644 --- a/ruby/red-arrow-format/test/test-record-batch.rb +++ b/ruby/red-arrow-format/test/test-record-batch.rb @@ -26,6 +26,92 @@ def setup end sub_test_case("#initialize") do + sub_test_case("[Schema, values]") do + def setup + @schema = ArrowFormat::Schema.new([ + ArrowFormat::Field.new( + "visible", + ArrowFormat::BooleanType.singleton, + ), + ArrowFormat::Field.new( + "count", + ArrowFormat::UInt32Type.singleton, + ), + ]) + end + + test("records") do + record_batch = ArrowFormat::RecordBatch.new( + @schema, + [ + {visible: true, count: 1}, + nil, + [false, 3], + ], + ) + assert_equal(@schema, record_batch.schema) + assert_equal(ArrowFormat::BooleanArray, + record_batch.find_column("visible").class) + assert_equal(ArrowFormat::UInt32Array, + record_batch.find_column("count").class) + assert_equal([ + {"visible" => true, "count" => 1}, + {"visible" => nil, "count" => nil}, + {"visible" => false, "count" => 3}, + ], + record_batch.records.collect(&:to_h)) + end + + test("columns") do + record_batch = ArrowFormat::RecordBatch.new( + @schema, + { + visible: [true, nil, false], + "count" => [1, 2, nil], + }, + ) + assert_equal([ + {"visible" => true, "count" => 1}, + {"visible" => nil, "count" => 2}, + {"visible" => false, "count" => nil}, + ], + record_batch.records.collect(&:to_h)) + end + + test("inconsistent column lengths") do + error = ArgumentError.new("inconsistent the number of rows: 2, 3") + assert_raise(error) do + ArrowFormat::RecordBatch.new( + @schema, + { + visible: [true, nil], + count: [1, 2, 3], + }, + ) + end + end + + test("unknown column") do + record_batch = ArrowFormat::RecordBatch.new( + @schema, + { + visible: [true], + count: [1], + extra: [2], + }, + ) + assert_equal([{"visible" => true, "count" => 1}], + record_batch.records.collect(&:to_h)) + end + + test("too many row values") do + error = ArgumentError.new("row 0 has more values than schema fields") + assert_raise(error) do + ArrowFormat::RecordBatch.new(@schema, [[true, 1, 2]]) + end + end + end + test("{}") do error = ArgumentError.new("no data") assert_raise(error) do