[Ray Data] Enhance the expression for getting a field or slice

[codingl2k1]

Description

I want to access some fields of a struct, for example, the batch dtype is

file: string                                                                                                              
audio: struct<bytes: binary, path: string>
  child 0, bytes: binary
  child 1, path: string
text: string
speaker_id: int64
chapter_id: int64
id: string

I want to get the bytes from the audio struct, however, the Ray Data expression can't handle this. This is a Daft example:

• Getting a single value:

import daft
df = daft.from_pydict({"struct": [{"x": 1, "y": 2}, {"x": 3, "y": 4}], "list": [[10, 20], [30, 40]]})
df = df.select(df["struct"]["x"], df["list"][0].alias("first"))
df.show()
╭───────┬───────╮
│ x     ┆ first │
│ ---   ┆ ---   │
│ Int64 ┆ Int64 │
╞═══════╪═══════╡
│ 1     ┆ 10    │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 3     ┆ 30    │
╰───────┴───────╯

(Showing first 2 of 2 rows)

• Getting a slice:

df = daft.from_pydict({"x": [[1, 2, 3], [4, 5, 6, 7], [8]]})
df = df.select(df["x"][1:-1])
df.show()
╭─────────────╮
│ x           │
│ ---         │
│ List[Int64] │
╞═════════════╡
│ [2]         │
├╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [5, 6]      │
├╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ []          │
╰─────────────╯

(Showing first 3 of 3 rows)

Use case

from datasets import load_dataset
import ray.data
from ray.data.expressions import col

dataset = load_dataset("openslr/librispeech_asr", "clean", split="validation")
print(dataset)

ds = ray.data.from_huggingface(dataset)
print(ds.take_batch(batch_size=1, batch_format="pyarrow").schema)
ds = ds.with_column("bytes", col("audio")["bytes"])    # This raises an error in Ray.
ds.show(1)

[gvspraveen]

Thanks @codingl2k1 for raising the feature request. cc @goutamvenkat-anyscale who is working on enhanching expression support. We will circle back here when we pick this up.

[goutamvenkat-anyscale]

Hi @codingl2k1, I thought a lit bit more about this feature. Instead of supporting __getitem__ at the Expression base class level, it would make more sense to add it per namespace (struct, list etc.).

So instead of col('audio')['bytes'] it will become col('audio').struct['bytes'].
The point of this tweak is to limit the amount of dynamic dispatching since we'll also add this support for lists, maps etc.

Let me know what you think

[codingl2k1]

Hi @codingl2k1, I thought a lit bit more about this feature. Instead of supporting __getitem__ at the Expression base class level, it would make more sense to add it per namespace (struct, list etc.).

So instead of col('audio')['bytes'] it will become col('audio').struct['bytes']. The point of this tweak is to limit the amount of dynamic dispatching since we'll also add this support for lists, maps etc.

Let me know what you think

It's clearer, but the user needs to know which namespace to use. The daft's __getitem__ method accepts names or slices and automatically handles the dispatch underneath. Both options work well for me :)