[Data] Compute Expressions map by ryankert01 · Pull Request #58743 · ray-project/ray

ryankert01 · 2025-11-18T16:08:28Z

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

TODO

add doc string examples

Additional information

test w/

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

gemini-code-assist

Code Review

This pull request introduces support for map/dict operations on expression columns by adding a .map property to Expr and implementing the _MapNamespace with keys() and values() methods. The changes in expressions.py are straightforward. My review focuses on the new map_namespace.py file, where I've identified significant performance improvement opportunities. I'm suggesting replacing the current loop-based UDFs with vectorized PyArrow compute functions (pc.map_keys and pc.map_values), which is crucial for efficient data processing. I've also included suggestions for more precise return type inference in these methods.

cursor

Bug: Structural equality overlooks critical flag.

The AliasExpr.structurally_equals method contains a comparison bug on the last condition: self._is_rename == self._is_rename always evaluates to True. It should compare self._is_rename == other._is_rename to properly check if two AliasExpr instances have the same rename flag, making the structural equality check incomplete.

python/ray/data/expressions.py#L1042-L1047

https://github.com/ray-project/ray/blob/bcb228a2b9308fc7972b38caeb44a16254f6febe/python/ray/data/expressions.py#L1042-L1047

cursor · 2025-11-22T16:52:07Z

+    # offsets and null-mask are preserved from the parent
+    return pyarrow.ListArray.from_arrays(
+        offsets=arr.offsets, values=child_array, mask=arr.is_null()
+    )


Bug: MapArray lacks offsets attribute in older PyArrow

The code unconditionally accesses arr.offsets on line 85, but MapArray doesn't have an offsets attribute in PyArrow versions before 7.0.0. When Case 1 executes (line 49) for a MapArray in older PyArrow, the code extracts arr.keys or arr.items but then tries to access arr.offsets which doesn't exist. In PyArrow < 7.0.0, MapArray is not a subclass of ListArray and offsets must be accessed through the buffers API instead.

ryankert01 · 2025-11-23T05:36:54Z

cc @goutamvenkat-anyscale
Want to receive some feedbacks!

Signed-off-by: ryankert01 <ryan980053@gmail.com>

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Ryan Huang <ryankert01@gmail.com>

…r PyArrow compatibility Signed-off-by: ryankert01 <ryan980053@gmail.com>

Signed-off-by: ryankert01 <ryan980053@gmail.com>

iamjustinhsu

Thank you for contributing! Left a few comments

iamjustinhsu · 2025-11-24T23:54:37Z


+    @property
+    def map(self) -> "_MapNamespace":
+        """Access map/dict operations for this expression."""


Can we add a docstring example like the _StructNamespace?

Will do! List on todos

iamjustinhsu · 2025-11-25T00:02:54Z

+        if self._expr.data_type.is_arrow_type():
+            arrow_type = self._expr.data_type.to_arrow_dtype()


I'm a little bit confused by this statement. In english, this sounds like: "Convert to arrow type if already arrow type". Did you mean to use not self._expr.dataset_.is_arrow_type()?

It's right to be confused! Maybe I should do some comments. the first is_arrow_type() means check if this is arrow mem. The second to_arrow_dtype() unwrap the ray wrapper and use arrow wrapper.

iamjustinhsu · 2025-11-25T00:05:15Z

+    This handles both formal ``MapArray`` types and ``ListArray<Struct>``
+    types (which represent maps physically) by accessing underlying
+    child arrays directly.


Can you elaborate what this function does?

I want to implement a no copy, robust covert from ray wrapper -> pyarrow wrapper.

Signed-off-by: Ryan Huang <ryan980053@gmail.com>

Signed-off-by: ryankert01 <ryan980053@gmail.com>

cursor · 2025-11-30T02:22:45Z

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


Bug: Sliced array child_array not sliced with normalized offsets

When handling sliced Arrow arrays with non-zero starting offsets, the code normalizes offsets by subtracting start_offset to make them zero-based, but it does not correspondingly slice child_array. Since arr.keys, arr.items, and arr.values return the full underlying flat arrays (due to Arrow's zero-copy slice semantics), using normalized offsets with the full child_array causes each row to reference the wrong elements. For example, a slice of rows 7-9 would incorrectly return data from rows 0-2 of the original array. The child_array needs to be sliced from start_offset to the original offsets[-1] value to match the normalized offsets.

Can we verify if this is True with a simple test?

@ryankert01

have test_map_sliced_offsets in the new pr.

Signed-off-by: ryankert01 <ryan980053@gmail.com>

iamjustinhsu

lgtm! cc: @goutamvenkat-anyscale

iamjustinhsu · 2025-12-03T22:33:02Z

+    offsets = arr.offsets
+    if len(offsets) > 0:  # Handle offsets changes
+        start_offset = offsets[0]
+        if start_offset.as_py() != 0:


to clarify, the start_offset can be > 0 when we recursively go through a chunked array?

Yeah, because as a chunked array, each chunk is a slice of the array which has a non-zero offset.

Signed-off-by: ryankert01 <ryan980053@gmail.com>

goutamvenkat-anyscale · 2025-12-18T17:34:05Z

+    if child_array is None:
+        try:
+            inner_type = arr.type.value_type
+            map_type = pyarrow.map_(inner_type[0].type, inner_type[1].type)
+            return _extract_map_component(arr.cast(map_type), component)
+        except (AttributeError, IndexError, pyarrow.ArrowInvalid):
+            raise ValueError(
+                f"Cannot extract {component} from array of type {arr.type}"
+            )


Can you help me understand why we need to recursively cast to a map in non listarray and maparray cases?

Perhaps a concrete example would help.

@goutamvenkat-anyscale It's an attempt to support unknown mapping, an example is

fixed_list_type = pa.list_( pa.struct([('key', pa.string()), ('value', pa.int64())]), list_size=2 ) arr = pa.array([ [{'key': 'a', 'value': 1}, {'key': 'b', 'value': 2}], [{'key': 'x', 'value': 10}, {'key': 'y', 'value': 20}] ], type=fixed_list_type)

feel free to let me know if this handling is necessary.

I think we can simplify and leave this out for now. If folks have specific workloads that run into this case, we can re-evaluate.

ryankert01 · 2025-12-30T07:26:21Z

Sorry, I re-fork my repo and cause an accidentally close, will open a new one soon.

## 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 ``` <sup><a href="https://cursor.com/dashboard?tab=bugbot">Cursor Bugbot</a> found 1 potential issue for commit <u>7a11478</u></sup> --------- Signed-off-by: Hsien-Cheng Huang <ryankert01@gmail.com> Signed-off-by: Ryan Huang <ryankert01@gmail.com> Signed-off-by: Hsien-Cheng Huang <hcr@apache.org> Co-authored-by: You-Cheng Lin <106612301+owenowenisme@users.noreply.github.com> Co-authored-by: Goutam <goutam@anyscale.com>

ryankert01 requested a review from a team as a code owner November 18, 2025 16:08

gemini-code-assist Bot reviewed Nov 18, 2025

View reviewed changes

Comment thread python/ray/data/namespace_expressions/map_namespace.py

Comment thread python/ray/data/namespace_expressions/map_namespace.py Outdated

Comment thread python/ray/data/namespace_expressions/map_namespace.py Outdated

cursor Bot reviewed Nov 18, 2025

View reviewed changes

Comment thread python/ray/data/namespace_expressions/map_namespace.py Outdated

ryankert01 changed the title ~~[Data]Compute Expressions map~~ Nov 18, 2025

ray-gardener Bot added data Ray Data-related issues community-contribution Contributed by the community labels Nov 18, 2025

ryankert01 marked this pull request as draft November 19, 2025 04:46

ryankert01 marked this pull request as ready for review November 22, 2025 04:02

ryankert01 marked this pull request as draft November 22, 2025 04:02

ryankert01 marked this pull request as ready for review November 22, 2025 16:48

ryankert01 changed the title ~~[WIP] [Data] Compute Expressions map~~ Nov 22, 2025

cursor Bot reviewed Nov 22, 2025

View reviewed changes

ryankert01 marked this pull request as draft November 22, 2025 17:13

ryankert01 marked this pull request as ready for review November 23, 2025 05:33

ryankert01 and others added 8 commits November 24, 2025 09:19

feat:add map expression

dc8d3ef

Signed-off-by: ryankert01 <ryan980053@gmail.com>

Update python/ray/data/namespace_expressions/map_namespace.py

b8fd2f1

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Ryan Huang <ryankert01@gmail.com>

Update python/ray/data/namespace_expressions/map_namespace.py

f1942bc

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Ryan Huang <ryankert01@gmail.com>

Update python/ray/data/namespace_expressions/map_namespace.py

60b6d45

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Ryan Huang <ryankert01@gmail.com>

[namespace] Implement map keys and values extraction with fallback fo…

2396665

…r PyArrow compatibility Signed-off-by: ryankert01 <ryan980053@gmail.com>

refactor

435b0bf

Signed-off-by: ryankert01 <ryan980053@gmail.com>

try to fix ci

2ceccf8

Signed-off-by: ryankert01 <ryan980053@gmail.com>

try to fix ci

81f0678

Signed-off-by: ryankert01 <ryan980053@gmail.com>

cursor Bot reviewed Nov 24, 2025

View reviewed changes

Comment thread python/ray/data/namespace_expressions/map_namespace.py

iamjustinhsu reviewed Nov 25, 2025

View reviewed changes

ryankert01 added 3 commits November 27, 2025 23:42

Merge branch 'master' into feature/expression-map

43d8a12

Signed-off-by: Ryan Huang <ryan980053@gmail.com>

Merge branch 'master' into feature/expression-map

e17647a

refactor, can cleanup code

5058537

Signed-off-by: ryankert01 <ryan980053@gmail.com>

cursor Bot reviewed Nov 30, 2025

View reviewed changes

fix lint

49bb570

Signed-off-by: ryankert01 <ryan980053@gmail.com>

iamjustinhsu approved these changes Dec 3, 2025

View reviewed changes

ryankert01 added 3 commits December 11, 2025 06:44

Merge branch 'master' into feature/expression-map

55a6fbd

add doc string

33691ef

Signed-off-by: ryankert01 <ryan980053@gmail.com>

Merge branch 'master' into feature/expression-map

74a1e32

owenowenisme added the go add ONLY when ready to merge, run all tests label Dec 11, 2025

Merge branch 'master' into feature/expression-map

8a3f215

goutamvenkat-anyscale reviewed Dec 18, 2025

View reviewed changes

ryankert01 added 2 commits December 19, 2025 15:34

Merge branch 'master' into feature/expression-map

e598731

Merge branch 'master' into feature/expression-map

289ccbc

ryankert01 closed this by deleting the head repository Dec 30, 2025

ryankert01 mentioned this pull request Jan 6, 2026

[Data] Add map namespace support for expression operations #59879

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Data] Compute Expressions map#58743

[Data] Compute Expressions map#58743
ryankert01 wants to merge 18 commits into
ray-project:masterfrom
ryankert01:feature/expression-map

ryankert01 commented Nov 18, 2025 •

edited

Loading

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

cursor Bot left a comment

Uh oh!

cursor Bot Nov 22, 2025

ryankert01 commented Nov 23, 2025

Uh oh!

iamjustinhsu left a comment

iamjustinhsu Nov 24, 2025

ryankert01 Nov 30, 2025 •

edited

Loading

Uh oh!

iamjustinhsu Nov 25, 2025

ryankert01 Nov 30, 2025

Uh oh!

iamjustinhsu Nov 25, 2025

ryankert01 Nov 30, 2025

cursor Bot Nov 30, 2025

goutamvenkat-anyscale Dec 18, 2025

goutamvenkat-anyscale Dec 27, 2025

ryankert01 Jan 6, 2026

iamjustinhsu left a comment

iamjustinhsu Dec 3, 2025

ryankert01 Dec 11, 2025

goutamvenkat-anyscale Dec 18, 2025

ryankert01 Dec 19, 2025 •

edited

Loading

goutamvenkat-anyscale Dec 27, 2025

ryankert01 commented Dec 30, 2025

Labels

5 participants

		if self._expr.data_type.is_arrow_type():
		arrow_type = self._expr.data_type.to_arrow_dtype()

Uh oh!

Conversation

ryankert01 commented Nov 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

MapNamespace impl.

Testing

Related issues

TODO

Additional information

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Bug: Structural equality overlooks critical flag.

Uh oh!

cursor Bot Nov 22, 2025

Choose a reason for hiding this comment

Bug: MapArray lacks offsets attribute in older PyArrow

ryankert01 commented Nov 23, 2025

Uh oh!

iamjustinhsu left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ryankert01 Nov 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Choose a reason for hiding this comment

cursor Bot Nov 30, 2025

Choose a reason for hiding this comment

Bug: Sliced array child_array not sliced with normalized offsets

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

iamjustinhsu left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ryankert01 Dec 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ryankert01 commented Dec 30, 2025

Labels

5 participants

ryankert01 commented Nov 18, 2025 •

edited

Loading

`MapNamespace` impl.

ryankert01 Nov 30, 2025 •

edited

Loading

ryankert01 Dec 19, 2025 •

edited

Loading