GH-50571: [Ruby] Add schema-aware ArrowFormat::RecordBatch construction

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 = []

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