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

Code before:
from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import TYPE_CHECKING

import pyarrow
import pyarrow.compute as pc

from ray.data.datatype import DataType
from ray.data.expressions import pyarrow_udf

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>>
    # PyArrow represents maps as List<Struct> where the struct's

Review comment (refactor):
Let's create 3 helper functions to make the intent clearer.
1. `_get_child_array` which gets keys and values
2. `_make_empty_list_array`
3. `_rebuild_list_array` with normalized offsets

For each one add an example of what it's doing in the comments.

```py
def _extract_map_component(
    arr: pyarrow.Array, component: MapComponent
) -> pyarrow.Array:
    """Extract keys or values from a MapArray or ListArray<Struct>."""
    
    if isinstance(arr, pyarrow.ChunkedArray):
        return pyarrow.chunked_array(
            [_extract_map_component(chunk, component) for chunk in arr.chunks]
        )

    child_array = _get_child_array(arr, component)
    
    if child_array is None:
        return _make_empty_list_array(arr, component)
    
    return _rebuild_list_array(arr, child_array)
```


Code after:
from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import TYPE_CHECKING, Optional

import numpy as np
import pyarrow
import pyarrow.compute as pc

from ray.data.datatype import DataType
from ray.data.expressions import pyarrow_udf

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