Repository: ray-project/ray
File: python/ray/data/namespace_expressions/map_namespace.py (Python)

Code before:
if TYPE_CHECKING:
    from ray.data.expressions import Expr, UDFExpr


class MapComponent(str, Enum):
    KEYS = "keys"
    VALUES = "values"


def _extract_map_component(
    arr: pyarrow.Array, component: MapComponent
) -> pyarrow.Array:
    """
    Extracts keys or values from a MapArray or ListArray<Struct>.

    This serves as the primary implementation since PyArrow does not yet
    expose dedicated compute kernels for map projection in the Python API.
    """
    # 1. Handle Chunked Arrays (Recursion)
    if isinstance(arr, pyarrow.ChunkedArray):
        return pyarrow.chunked_array(
            [_extract_map_component(chunk, component) for chunk in arr.chunks]
        )

    child_array = None

    # Case 1: MapArray
    if isinstance(arr, pyarrow.MapArray):
        if component == MapComponent.KEYS:
            child_array = arr.keys
        else:
            child_array = arr.items

    # Case 2: ListArray<Struct<Key, Value>>
    elif isinstance(arr, (pyarrow.ListArray, pyarrow.LargeListArray)):
        flat_values = arr.values
        if (
            isinstance(flat_values, pyarrow.StructArray)
            and flat_values.type.num_fields >= 2
        ):
            idx = 0 if component == MapComponent.KEYS else 1
            child_array = flat_values.field(idx)

    if child_array is None:
        # This can happen if the input array is not a supported map type.
        # We allow this to proceed only if the array is empty or all-nulls,
        # in which case we'll produce an empty or all-nulls output.
        if len(arr) > 0 and arr.null_count < len(arr):
            raise TypeError(
                f"Expression is not a map type. .map.{component.value}() can only be "
                f"called on MapArray or List<Struct<key, value>> types, but got {arr.type}."

Review comment (suggestion):
I think this is kind of implicit, we should document that we use the idx 0 item as key and idx 1 as value.

Code after:

if TYPE_CHECKING:
    from ray.data.expressions import Expr, UDFExpr


class MapComponent(str, Enum):
    KEYS = "keys"
    VALUES = "values"


def _get_child_array(
    arr: pyarrow.Array, component: MapComponent
) -> Optional[pyarrow.Array]:
    """Extract the flat keys or values array from a map-like array.

    Example: MapArray [{"a": 1}, {"b": 2}] -> keys ["a", "b"] or values [1, 2]
    """
    if isinstance(arr, pyarrow.MapArray):
        if component == MapComponent.KEYS:
            return arr.keys
        else:
            return arr.items

    if isinstance(arr, (pyarrow.ListArray, pyarrow.LargeListArray)):
        flat_values = arr.values
        if (
            isinstance(flat_values, pyarrow.StructArray)
            and flat_values.type.num_fields >= 2
        ):
            idx = 0 if component == MapComponent.KEYS else 1
            return flat_values.field(idx)

    return None


def _make_empty_list_array(
    arr: pyarrow.Array, component: MapComponent
) -> pyarrow.Array:
    """Create an all-null ListArray matching the input length.

    Example: arr of length 3 -> ListArray [null, null, null]
    """
    if len(arr) > 0 and arr.null_count < len(arr):
        raise TypeError(
            f"Expression is not a valid map type. .map.{component.value}() requires "
            f"pyarrow.MapArray or pyarrow.ListArray<Struct> with at least 2 fields "
            f"(key and value), but got: {arr.type}."
        )
    return pyarrow.ListArray.from_arrays(
        offsets=np.repeat(0, len(arr) + 1),
        values=pyarrow.array([], type=pyarrow.null()),