[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 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

[goutamvenkat-anyscale]

Please look at open comments. Thanks

[goutamvenkat-anyscale]

lgtm!

[alexeykudinkin]

Help me understand why we need to do this?

[ryankert01]

It's because when we slice a MapArray or ListArray, the child arrays (keys/values) remain unchanged and offsets still reference positions in the original buffer. (zero-copy slicing by pyArrow)

-> we have to re-build it to 0-based.

[goutamvenkat-anyscale]

@ryankert01 There seems to be some open comments. Please address those

[cursor]

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

[ryankert01]

PTAL @goutamvenkat-anyscale @alexeykudinkin