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

Code before:
    # 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}."
            )
        return pyarrow.ListArray.from_arrays(
            offsets=[0] * (len(arr) + 1),
            values=pyarrow.array([], type=pyarrow.null()),
            mask=pyarrow.array([True] * len(arr)),
        )

    # Reconstruct ListArray & Normalize Offsets
    offsets = arr.offsets
    if len(offsets) > 0:  # Handle offsets changes
        start_offset = offsets[0]
        if start_offset.as_py() != 0:
            # Slice child_array to match normalized offsets
            end_offset = offsets[-1]
            child_array = child_array.slice(
                offset=start_offset.as_py(), length=(end_offset - start_offset).as_py()
            )
            offsets = pc.subtract(offsets, start_offset)

    return pyarrow.ListArray.from_arrays(
        offsets=offsets, values=child_array, mask=arr.is_null()
    )


@dataclass
class _MapNamespace:

Review comment (suggestion):
```suggestion
			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}."
			)
```

Code after:
            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()),
        mask=pyarrow.array(np.repeat(True, len(arr))),
    )


def _rebuild_list_array(
    arr: pyarrow.Array, child_array: pyarrow.Array
) -> pyarrow.Array:
    """Rebuild a ListArray from parent offsets and child values, normalizing sliced offsets.

    Example: offsets [5, 7, 10] -> slice child to [5:10], normalize offsets to [0, 2, 5]
    """
    offsets = arr.offsets
    if len(offsets) > 0:
        start_offset = offsets[0]
        if start_offset != 0:
            end_offset = offsets[-1]
            child_array = child_array.slice(
                offset=int(start_offset), length=int(end_offset) - int(start_offset)
            )
            offsets = pc.subtract(offsets, start_offset)

    factory = (
        pyarrow.LargeListArray.from_arrays
        if isinstance(arr, pyarrow.LargeListArray)
        else pyarrow.ListArray.from_arrays
    )