GH-48388: [Ruby] Add support for reading map array

Author: kou

ruby/red-arrow-format/lib/arrow-format/array.rb

@@ -162,7 +162,7 @@ def encoding
     end
   end
 
-  class ListArray < Array
+  class VariableSizeListArray < Array
     def initialize(type, size, validity_buffer, offsets_buffer, child)
       super(type, size, validity_buffer)
       @offsets_buffer = offsets_buffer
@@ -181,6 +181,9 @@ def to_a
     end
   end
 
+  class ListArray < VariableSizeListArray
+  end
+
   class StructArray < Array
     def initialize(type, size, validity_buffer, children)
       super(type, size, validity_buffer)
@@ -197,4 +200,20 @@ def to_a
       apply_validity(values)
     end
   end
+
+  class MapArray < VariableSizeListArray
+    def to_a
+      super.collect do |entries|
+        if entries.nil?
+          entries
+        else
+          hash = {}
+          entries.each do |key, value|
+            hash[key] = value
+          end
+          hash
+        end
+      end
+    end
+  end
 end

ruby/red-arrow-format/lib/arrow-format/error.rb

@@ -18,11 +18,14 @@ module ArrowFormat
   class Error < StandardError
   end
 
-  class ReadError < StandardError
+  class ReadError < Error
     attr_reader :buffer
     def initialize(buffer, message)
       @buffer = buffer
       super("#{message}: #{@buffer}")
     end
   end
+
+  class TypeError < Error
+  end
 end

ruby/red-arrow-format/lib/arrow-format/field.rb

@@ -18,9 +18,14 @@ module ArrowFormat
   class Field
     attr_reader :name
     attr_reader :type
-    def initialize(name, type)
+    def initialize(name, type, nullable)
       @name = name
       @type = type
+      @nullable = nullable
+    end
+
+    def nullable?
+      @nullable
     end
   end
 end

ruby/red-arrow-format/lib/arrow-format/file-reader.rb

@@ -29,6 +29,7 @@
 require_relative "org/apache/arrow/flatbuf/int"
 require_relative "org/apache/arrow/flatbuf/large_binary"
 require_relative "org/apache/arrow/flatbuf/list"
+require_relative "org/apache/arrow/flatbuf/map"
 require_relative "org/apache/arrow/flatbuf/message"
 require_relative "org/apache/arrow/flatbuf/null"
 require_relative "org/apache/arrow/flatbuf/precision"
@@ -163,14 +164,16 @@ def read_field(fb_field)
       when Org::Apache::Arrow::Flatbuf::Struct
         children = fb_field.children.collect {|child| read_field(child)}
         type = StructType.new(children)
+      when Org::Apache::Arrow::Flatbuf::Map
+        type = MapType.new(read_field(fb_field.children[0]))
       when Org::Apache::Arrow::Flatbuf::Binary
         type = BinaryType.singleton
       when Org::Apache::Arrow::Flatbuf::LargeBinary
         type = LargeBinaryType.singleton
       when Org::Apache::Arrow::Flatbuf::Utf8
         type = UTF8Type.singleton
       end
-      Field.new(fb_field.name, type)
+      Field.new(fb_field.name, type, fb_field.nullable?)
     end
 
     def read_schema(fb_schema)
@@ -199,7 +202,7 @@ def read_column(field, nodes, buffers, body)
         values_buffer = buffers.shift
         values = body.slice(values_buffer.offset, values_buffer.length)
         field.type.build_array(length, validity, values)
-      when ListType
+      when VariableSizeListType
         offsets_buffer = buffers.shift
         offsets = body.slice(offsets_buffer.offset, offsets_buffer.length)
         child = read_column(field.type.child, nodes, buffers, body)

ruby/red-arrow-format/lib/arrow-format/type.rb

@@ -195,13 +195,20 @@ def build_array(size, validity_buffer, offsets_buffer, values_buffer)
     end
   end
 
-  class ListType < Type
+  class VariableSizeListType < Type
     attr_reader :child
-    def initialize(child)
-      super("List")
+    def initialize(name, child)
+      super(name)
       @child = child
     end
 
+  end
+
+  class ListType < VariableSizeListType
+    def initialize(child)
+      super("List", child)
+    end
+
     def build_array(size, validity_buffer, offsets_buffer, child)
       ListArray.new(self, size, validity_buffer, offsets_buffer, child)
     end
@@ -218,4 +225,30 @@ def build_array(size, validity_buffer, children)
       StructArray.new(self, size, validity_buffer, children)
     end
   end
+
+  class MapType < VariableSizeListType
+    def initialize(child)
+      if child.nullable?
+        raise TypeError.new("Map entry field must not be nullable: " +
+                            child.inspect)
+      end
+      type = child.type
+      unless type.is_a?(StructType)
+        raise TypeError.new("Map entry type must be struct: #{type.inspect}")
+      end
+      unless type.children.size == 2
+        raise TypeError.new("Map entry struct type must have 2 children: " +
+                            type.inspect)
+      end
+      if type.children[0].nullable?
+        raise TypeError.new("Map key field must not be nullable: " +
+                            type.children[0].inspect)
+      end
+      super("Map", child)
+    end
+
+    def build_array(size, validity_buffer, offsets_buffer, child)
+      MapArray.new(self, size, validity_buffer, offsets_buffer, child)
+    end
+  end
 end

ruby/red-arrow-format/test/test-file-reader.rb

@@ -160,4 +160,29 @@ def test_read
                    read)
     end
   end
+
+  sub_test_case("Map") do
+    def build_array
+      data_type = Arrow::MapDataType.new(:string, :int8)
+      Arrow::MapArray.new(data_type,
+                          [
+                            {"a" => -128, "b" => 127},
+                            nil,
+                            {"c" => nil},
+                          ])
+    end
+
+    def test_read
+      assert_equal([
+                     {
+                       "value" => [
+                         {"a" => -128, "b" => 127},
+                         nil,
+                         {"c" => nil},
+                       ],
+                     },
+                   ],
+                   read)
+    end
+  end
 end