Adding the option to convert bidirectionally between map type (Protobuf) to list of structs (Arrow).

protarrow/arrow_to_proto.py

@@ -1,7 +1,7 @@
 import collections.abc
 import dataclasses
 import datetime
-from typing import Any, Callable, Iterable, Iterator, List, Optional, Tuple, Type
+from typing import Any, Callable, Iterable, Iterator, List, Optional, Tuple, Type, Union
 
 import pyarrow as pa
 from google.protobuf.descriptor import Descriptor, EnumDescriptor, FieldDescriptor
@@ -458,10 +458,25 @@ def _extract_struct_field(
 
 
 def _extract_map_field(
-    array: pa.MapArray,
+    array: Union[pa.MapArray, pa.ListArray, pa.LargeListArray],
     field_descriptor: FieldDescriptor,
     messages: Iterable[Message],
 ) -> None:
+
+    if pa.types.is_list(array.type) or pa.types.is_large_list(array.type):
+        assert pa.types.is_struct(array.values.type), array.values.type
+        field_names = [field.name for field in array.values.type.fields]
+        assert len(field_names) == 2
+        # Since there's only 2 elements, the index of the key can only be 0 or 1
+        # Therefore, the index of the value is the other choice.
+        key_field_index = field_names.index("key")
+        value_field_index = 1 - key_field_index
+        array = pa.MapArray.from_arrays(
+            offsets=array.offsets,
+            keys=array.values.field(key_field_index),
+            items=array.values.field(value_field_index),
+        )
+
     assert pa.types.is_map(array.type), array.type
     value_descriptor = field_descriptor.message_type.fields_by_name["value"]
 

protarrow/cast_to_proto.py

@@ -107,26 +107,73 @@ def _cast_array(
     config: ProtarrowConfig,
 ) -> pa.Array:
     if is_map(field_descriptor):
-        assert isinstance(array, pa.MapArray)
         key_field, value_field = get_map_descriptors(field_descriptor)
-        map_array = pa.MapArray.from_arrays(
-            # TODO: remove when https://github.com/apache/arrow/issues/40750 is fixed
-            #  and library is pinned to pyarrow>=17.0.0
-            maybe_copy_offsets(array.offsets),
-            _cast_array(array.keys, key_field, config),
-            _cast_array(array.items, value_field, config),
-        )
-        return map_array.cast(
-            pa.map_(
-                map_array.type.key_type,
-                pa.field(
-                    config.map_value_name,
-                    map_array.type.item_type,
-                    nullable=config.map_value_nullable,
-                    metadata=config.field_metadata(field_descriptor.number),
+
+        if pa.types.is_map(array.type):
+            keys = array.keys
+            values = array.items
+        else:
+            assert pa.types.is_list(array.type) or pa.types.is_large_list(array.type), (
+                array.type
+            )
+            assert pa.types.is_struct(array.values.type), array.values.type
+            keys = array.values.field("key")
+            values = array.values.field(config.map_value_name)
+
+        # TODO: remove when https://github.com/apache/arrow/issues/40750 is fixed
+        #  and library is pinned to pyarrow>=17.0.0
+        offsets = maybe_copy_offsets(array.offsets)
+        keys = _cast_array(keys, key_field, config)
+        values = _cast_array(values, value_field, config)
+
+        if config.map_as_list:
+            return config.list_array_type.from_arrays(
+                offsets=offsets,
+                values=pa.StructArray.from_arrays(
+                    arrays=[keys, values],
+                    fields=[
+                        pa.field(
+                            name="key",
+                            type=keys.type,
+                            nullable=False,
+                        ),
+                        pa.field(
+                            name=config.map_value_name,
+                            type=values.type,
+                            nullable=config.map_value_nullable,
+                        ),
+                    ],
                 ),
+            ).cast(
+                config.list_(
+                    item_type=pa.struct(
+                        fields=[
+                            pa.field(
+                                name="key",
+                                type=keys.type,
+                                nullable=False,
+                            ),
+                            pa.field(
+                                name=config.map_value_name,
+                                type=values.type,
+                                nullable=config.map_value_nullable,
+                            ),
+                        ]
+                    )
+                )
+            )
+        else:
+            return pa.MapArray.from_arrays(offsets, keys, values).cast(
+                pa.map_(
+                    keys.type,
+                    pa.field(
+                        config.map_value_name,
+                        values.type,
+                        nullable=config.map_value_nullable,
+                        metadata=config.field_metadata(field_descriptor.number),
+                    ),
+                )
             )
-        )
 
     elif field_descriptor.is_repeated:
         assert isinstance(array, (pa.ListArray, pa.LargeListArray))

protarrow/common.py

@@ -61,6 +61,7 @@ class ProtarrowConfig:
     binary_type: pa.DataType = pa.binary()
     list_array_type: type = pa.ListArray
     skip_recursive_messages: bool = False
+    map_as_list: bool = False
 
     def __post_init__(self):
         _validate_enum_type(self.enum_type, self.string_type, self.binary_type)

protarrow/message_extractor.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, Generic, List, Type, TypeVar
+from typing import Any, Callable, Dict, Generic, List, Type, TypeVar, Union
 
 import pyarrow as pa
 from google.protobuf.descriptor import Descriptor, FieldDescriptor
@@ -65,6 +65,35 @@ def __call__(self, scalar: pa.MapScalar) -> Dict[Any, Any]:
             return {}
 
 
+class MapAsListConverterAdapter:
+    def __init__(
+        self,
+        list_type: Union[pa.ListType, pa.LargeListType],
+        key_descriptor: FieldDescriptor,
+        value_descriptor: FieldDescriptor,
+    ):
+        struct_type = list_type.value_field.type
+        assert pa.types.is_struct(struct_type)
+        key_field, value_field = struct_type.fields
+
+        self._key_converter = get_flat_field_converter(key_field.type, key_descriptor)
+        self._value_converter = get_flat_field_converter(
+            value_field.type, value_descriptor
+        )
+
+    def __call__(
+        self,
+        scalar: Union[pa.ListScalar, pa.LargeListScalar],
+    ) -> Dict[Any, Any]:
+        if scalar.is_valid:
+            return {
+                self._key_converter(item.get(0)): self._value_converter(item.get(1))
+                for item in scalar.values
+            }
+        else:
+            return {}
+
+
 class NullableConverterAdapter:
     def __init__(
         self, converter: Callable[[pa.Scalar], Any], message_type: Type[Message]
@@ -99,7 +128,10 @@ def get_field_converter(
 ) -> Callable[[pa.Scalar], Any]:
     if is_map(field_descriptor):
         key, value = get_map_descriptors(field_descriptor)
-        return MapConverterAdapter(field.type, key, value)
+        if pa.types.is_map(field.type):
+            return MapConverterAdapter(field.type, key, value)
+        else:
+            return MapAsListConverterAdapter(field.type, key, value)
     else:
         if field_descriptor.is_repeated:
             return RepeatedConverterAdapter(

protarrow/proto_to_arrow.py

@@ -275,12 +275,31 @@ def field_descriptor_to_field(
         value_type = field_descriptor_to_data_type(
             value_field, config, descriptor_trace
         )
-        return pa.field(
-            field_descriptor.name,
-            pa.map_(
+        if config.map_as_list:
+            map_type = config.list_(
+                item_type=pa.struct(
+                    fields=[
+                        pa.field(
+                            name="key",
+                            type=key_type,
+                            nullable=False,
+                        ),
+                        pa.field(
+                            name=config.map_value_name,
+                            type=value_type,
+                            nullable=config.map_value_nullable,
+                        ),
+                    ]
+                )
+            )
+        else:
+            map_type = pa.map_(
                 key_type,
                 pa.field(config.map_value_name, value_type, config.map_value_nullable),
-            ),
+            )
+        return pa.field(
+            field_descriptor.name,
+            map_type,
             nullable=config.map_nullable,
             metadata=config.field_metadata(field_descriptor.number),
         )
@@ -498,14 +517,54 @@ def _proto_map_to_array(
         config=config,
         descriptor_trace=descriptor_trace,
     )
-    return pa.MapArray.from_arrays(offsets, keys, values).cast(
-        pa.map_(
-            keys.type,
-            pa.field(
-                config.map_value_name, values.type, nullable=config.map_value_nullable
+    if config.map_as_list:
+        array = config.list_array_type.from_arrays(
+            offsets=offsets,
+            values=pa.StructArray.from_arrays(
+                arrays=[keys, values],
+                fields=[
+                    pa.field(
+                        name="key",
+                        type=keys.type,
+                        nullable=False,
+                    ),
+                    pa.field(
+                        name=config.map_value_name,
+                        type=values.type,
+                        nullable=config.map_value_nullable,
+                    ),
+                ],
             ),
+        ).cast(
+            config.list_(
+                item_type=pa.struct(
+                    fields=[
+                        pa.field(
+                            name="key",
+                            type=keys.type,
+                            nullable=False,
+                        ),
+                        pa.field(
+                            name=config.map_value_name,
+                            type=values.type,
+                            nullable=config.map_value_nullable,
+                        ),
+                    ]
+                )
+            )
         )
-    )
+    else:
+        array = pa.MapArray.from_arrays(offsets, keys, values).cast(
+            pa.map_(
+                keys.type,
+                pa.field(
+                    config.map_value_name,
+                    values.type,
+                    nullable=config.map_value_nullable,
+                ),
+            )
+        )
+    return array
 
 
 def _proto_field_nullable(

tests/test_conversion.py

@@ -51,7 +51,12 @@
 from tests.random_generator import generate_messages, truncate_messages, truncate_nanos
 
 TEST_MESSAGE_COUNT = 5
-MESSAGES = [ExampleMessage, NestedExampleMessage, SuperNestedExampleMessage]
+MESSAGES = [
+    ExampleMessage,
+    NestedExampleMessage,
+    SuperNestedExampleMessage,
+]
+
 CONFIGS = [
     ProtarrowConfig(),
     ProtarrowConfig(enum_type=pa.binary()),
@@ -86,6 +91,7 @@
     ProtarrowConfig(field_number_key=b"PARQUET:field_id"),
     ProtarrowConfig(string_type=pa.large_string()),
     ProtarrowConfig(binary_type=pa.large_binary()),
+    ProtarrowConfig(map_as_list=True),
     ProtarrowConfig(list_array_type=pa.LargeListArray),
 ]
 

tests/test_coverage.py

@@ -29,6 +29,7 @@
 )
 from protarrow.cast_to_proto import get_arrow_default_value
 from protarrow.message_extractor import (
+    MapAsListConverterAdapter,
     MapConverterAdapter,
     NullableConverterAdapter,
     RepeatedConverterAdapter,
@@ -61,6 +62,19 @@ def test_map_converter_adapter():
     assert map_converter_adapter(pa.scalar(None, map_type)) == {}
 
 
+def test_map_as_list_converter_adapter():
+    list_type = pa.list_(pa.struct([("key", pa.int32()), ("value", pa.float64())]))
+    map_field = ExampleMessage.DESCRIPTOR.fields_by_name["double_int32_map"]
+    map_converter_adapter = MapAsListConverterAdapter(
+        list_type=list_type,
+        key_descriptor=map_field.message_type.fields_by_name["key"],
+        value_descriptor=map_field.message_type.fields_by_name["value"],
+    )
+    assert map_converter_adapter(pa.scalar([(123, 1.0)], list_type)) == {123: 1.0}
+    assert map_converter_adapter(pa.scalar([], list_type)) == {}
+    assert map_converter_adapter(pa.scalar(None, list_type)) == {}
+
+
 def test_nullable_converter_adapter():
     nullable_converter_adapter = NullableConverterAdapter(convert_scalar, DoubleValue)
     assert nullable_converter_adapter(pa.scalar(1.0, pa.float64())) == DoubleValue(