From d9ae7bd12deda94bead062f50cd2b5c34e87c300 Mon Sep 17 00:00:00 2001 From: Yifan Chen Date: Mon, 7 Sep 2026 06:36:41 -0700 Subject: [PATCH 1/3] GH-50571: [Ruby] Add schema-aware ArrowFormat RecordBatch construction --- .../lib/arrow-format/record-batch.rb | 55 +++++++++++- .../test/test-record-batch.rb | 85 +++++++++++++++++++ 2 files changed, 138 insertions(+), 2 deletions(-) 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..1bccdd9ee12c 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,52 @@ 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 = fields.collect { [] } + data.each do |name, values| + raw_columns[name_to_index[name.to_s]] = values + end + columns = fields.zip(raw_columns).collect do |field, values| + field.type.build_array(values) + 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 + 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 + 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 + [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..ad25679a5c1d 100644 --- a/ruby/red-arrow-format/test/test-record-batch.rb +++ b/ruby/red-arrow-format/test/test-record-batch.rb @@ -26,6 +26,91 @@ 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 + assert_raise(TypeError) do + ArrowFormat::RecordBatch.new( + @schema, + { + visible: [true], + count: [1], + extra: [2], + }, + ) + end + end + + test("too many row values") do + assert_raise(NoMethodError) do + ArrowFormat::RecordBatch.new(@schema, [[true, 1, 2]]) + end + end + end + test("{}") do error = ArgumentError.new("no data") assert_raise(error) do From 3f5fb39542da81fc0fdc1b5921c5fd3a46ed922a Mon Sep 17 00:00:00 2001 From: Yifan Chen Date: Mon, 7 Sep 2026 22:19:34 -0700 Subject: [PATCH 2/3] GH-51198: Address schema construction review --- .../lib/arrow-format/record-batch.rb | 13 +++++++---- .../test/test-record-batch.rb | 23 ++++++++++--------- 2 files changed, 21 insertions(+), 15 deletions(-) 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 1bccdd9ee12c..f3a2cb8b45c2 100644 --- a/ruby/red-arrow-format/lib/arrow-format/record-batch.rb +++ b/ruby/red-arrow-format/lib/arrow-format/record-batch.rb @@ -155,12 +155,13 @@ def build_with_schema(schema, data) name_to_index[field.name] = i end if data.is_a?(Hash) - raw_columns = fields.collect { [] } + raw_columns = [] data.each do |name, values| - raw_columns[name_to_index[name.to_s]] = 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) + field.type.build_array(values || []) end else raw_columns = fields.collect { [] } @@ -174,6 +175,10 @@ def build_with_schema(schema, data) end else record.each_with_index do |value, field_index| + if field_index >= raw_columns.size + message = "row #{nth_record} has more values than schema fields" + raise ArgumentError, message + end raw_columns[field_index] << value end end @@ -187,7 +192,7 @@ def build_with_schema(schema, data) end n_rows = columns.first&.size || 0 all_n_rows = columns.collect(&:size) - if all_n_rows.uniq.size != 1 + if data.is_a?(Hash) && all_n_rows.uniq.size != 1 message = "inconsistent the number of rows: #{all_n_rows.join(", ")}" raise ArgumentError, message end diff --git a/ruby/red-arrow-format/test/test-record-batch.rb b/ruby/red-arrow-format/test/test-record-batch.rb index ad25679a5c1d..8297e34154b2 100644 --- a/ruby/red-arrow-format/test/test-record-batch.rb +++ b/ruby/red-arrow-format/test/test-record-batch.rb @@ -92,20 +92,21 @@ def setup end test("unknown column") do - assert_raise(TypeError) do - ArrowFormat::RecordBatch.new( - @schema, - { - visible: [true], - count: [1], - extra: [2], - }, - ) - end + 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 - assert_raise(NoMethodError) 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 From 37d29caa040913ca55617e3399793ea8a1339feb Mon Sep 17 00:00:00 2001 From: Yifan Chen Date: Tue, 8 Sep 2026 04:09:01 -0700 Subject: [PATCH 3/3] GH-51198: Refine schema input validation --- .../lib/arrow-format/record-batch.rb | 18 +++++++++--------- .../red-arrow-format/test/test-record-batch.rb | 8 ++++---- 2 files changed, 13 insertions(+), 13 deletions(-) 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 f3a2cb8b45c2..776621dc0843 100644 --- a/ruby/red-arrow-format/lib/arrow-format/record-batch.rb +++ b/ruby/red-arrow-format/lib/arrow-format/record-batch.rb @@ -163,6 +163,11 @@ def build_with_schema(schema, data) 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| @@ -174,11 +179,11 @@ def build_with_schema(schema, data) 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| - if field_index >= raw_columns.size - message = "row #{nth_record} has more values than schema fields" - raise ArgumentError, message - end raw_columns[field_index] << value end end @@ -191,11 +196,6 @@ def build_with_schema(schema, data) end end n_rows = columns.first&.size || 0 - all_n_rows = columns.collect(&:size) - if data.is_a?(Hash) && all_n_rows.uniq.size != 1 - message = "inconsistent the number of rows: #{all_n_rows.join(", ")}" - raise ArgumentError, message - end [schema, n_rows, columns] end diff --git a/ruby/red-arrow-format/test/test-record-batch.rb b/ruby/red-arrow-format/test/test-record-batch.rb index 8297e34154b2..9a082bebfa6c 100644 --- a/ruby/red-arrow-format/test/test-record-batch.rb +++ b/ruby/red-arrow-format/test/test-record-batch.rb @@ -55,8 +55,8 @@ def setup assert_equal(ArrowFormat::UInt32Array, record_batch.find_column("count").class) assert_equal([ - {"visible" => true, "count" => 1}, - {"visible" => nil, "count" => nil}, + {"visible" => true, "count" => 1}, + {"visible" => nil, "count" => nil}, {"visible" => false, "count" => 3}, ], record_batch.records.collect(&:to_h)) @@ -71,8 +71,8 @@ def setup }, ) assert_equal([ - {"visible" => true, "count" => 1}, - {"visible" => nil, "count" => 2}, + {"visible" => true, "count" => 1}, + {"visible" => nil, "count" => 2}, {"visible" => false, "count" => nil}, ], record_batch.records.collect(&:to_h))