[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

[ryankert01]

row_same operates on pandas that can't handle the mixed None/list column when converting. The to_pandas() path triggers TensorArray casting which fails on the mixed types. Let's keep it!

Although there's workaround, but is too complex for the context of this test:

    ctx = ray.data.context.DataContext.get_current()
    ctx.enable_tensor_extension_casting = False
    try:
        result = (
            ds.with_column("keys", col("m").map.keys())
            .with_column("values", col("m").map.values())
            .to_pandas()
        )
        expected = pd.DataFrame(
            {
                "keys": [["a"], [], None],
                "values": [[1], [], None],
            }
        )
        _assert_result(result, expected, drop_cols=["m"])
    finally:
        ctx.enable_tensor_extension_casting = True

[goutamvenkat-anyscale]

Don't believe you need to call as_py here

[cursor]

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

[goutamvenkat-anyscale]

Please look at open comments. Thanks