[Data] Add map namespace support for expression operations

[ryankert01]

Description

MapNamespace impl.

• Implemented _extract_map_component as a robust, vectorized fallback since native pc.map_keys kernels are not standard in PyArrow yet.
• Support: Handles both Logical Maps (MapArray) and Physical Maps (List<Struct>).

Testing

• test_map_keys / test_map_values: Standard extraction.
• test_physical_map_extraction: Verifies support for List<Struct>.
• test_map_sliced_offsets: Verifies the critical fix for sliced data.
• test_map_nulls_and_empty: Verifies handling of None and empty maps {}.
• test_map_chaining: Verifies composition with List namespace (e.g., .map.keys().list.len()).

Related issues

Related to #58674
Continues #58743

Additional information

test w/

python -m pytest -v -s python/ray/data/tests/test_namespace_expressions.py::TestMapNamespace

^{Cursor Bugbot found 1 potential issue for commit 7a11478}

[gemini-code-assist]

Code Review

This pull request introduces support for map/dict operations on expression columns by adding a map namespace. The implementation is well-structured, adding a _MapNamespace with keys() and values() methods that work on both logical MapArray and physical List<Struct> representations. The handling of sliced arrays with non-zero offsets is a great detail that ensures correctness. The accompanying tests are thorough, covering various representations, edge cases like nulls and empty maps, and integration with other namespaces.

I've added a couple of suggestions to map_namespace.py to further improve the robustness of the implementation by handling LargeListArray and providing clearer errors for unsupported types. Overall, this is a solid contribution that enhances Ray Data's expression capabilities.

[ryankert01]

PTAL @goutamvenkat-anyscale @owenowenisme

[owenowenisme]

Minor fixes, overall LGTM

[goutamvenkat-anyscale]

let's use https://numpy.org/devdocs/reference/generated/numpy.repeat.html.

[goutamvenkat-anyscale]

same here

[goutamvenkat-anyscale]

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.

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)

[goutamvenkat-anyscale]

Let's use rows_same

[goutamvenkat-anyscale]

Don't believe you need to call as_py here

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

[cursor]

Duplicated type inference logic between functions

Medium Severity

_create_projection_udf (lines 179-201) duplicates the map type detection and inner type extraction logic that already exists in _get_result_type (lines 90-104). Both check for map types, physical maps (List), and extract key/value types using identical patterns. The method could reuse _get_result_type and convert the result: return_dtype = DataType.from_arrow(_get_result_type(arrow_type, component)).

 

[goutamvenkat-anyscale]

+1

[goutamvenkat-anyscale]

Please look at open comments. Thanks