[Datasets] [Tensor Story - 1/2] Automatically provide tensor views to UDFs and infer tensor blocks for pure-tensor datasets.

[clarkzinzow]

This PR makes several improvements to the Datasets tensor story. See the issues for each item for more details.

• Automatically infer tensor blocks (single-column tables representing a single tensor) when returning NumPy ndarrays from map_batches(), map(), and flat_map(): [Datasets] Automatically infer tensor blocks when returning NumPy ndarrays from .map_batches() #24208
• Automatically infer tensor columns when building tabular blocks in general: [Datasets] Automatically infer tensor columns when building tabular blocks #24207
• Fixes shuffling and sorting for tensor columns

This should improve the UX/efficiency of the following:

• Working with pure-tensor datasets in general.
• Mapping tensor UDFs over pure-tensor, a better foundation for tensor-native preprocessing for end-users and AIR.

Example of the tensor-native UDF UX

Below is a motivating example for the automatic NumPy conversion for the default "native" format, showing the single-tensor-column case.

Before

ds = ray.data.range_tensor(16, shape=(4, 4), parallelism=4)

def mapper(df: pd.DataFrame) -> pd.DataFrame:
    # I have to know about this `"value"` column and the
    # internal tensor representation as this single-column table.
    df["value"] *= 2
    return df

# I have to convert it to Pandas even though I'm working with tensor data,
# which is not the API that I want to work with AND incurs expensive
# copies during the Arrow --> Pandas conversion.
ds = ds.map_batches(mapper, batch_size=2, batch_format="pandas")

for batch in ds.iter_batches(batch_size=2, batch_format="pandas"):
    # I again have to convert the batches to Pandas and then access this
    # dummy column.
    actual_batch = batch["value"]

After

ds = ray.data.range_tensor(16, shape=(4, 4), parallelism=4)

# UDF now takes and returns an ndarray, which will be zero-copy access
# to the underlying Arrow table, providing far better UX and less
# memory/perf overhead. The returned ndarray will be repacked into an
# Arrow table (again, zero-copy).
ds = ds.map_batches(lambda arr: 2*arr, batch_size=2)

for batch in ds.iter_batches(batch_size=2):
    # The batch is the underlying ndarray batch, as I would expect; I don't
    # have to know or care about the dummy column or the single-column
    # table representation.
    assert isinstance(batch, np.ndarray)

Changed Semantics

This PR changes the semantics on how the default/"native" format is interpreted for pure-tensor datasets; i.e. datasets consisting of a single tensor column in an Arrow/Pandas table, with a dummy internal column name, such as those created by ray.data.range_tensor() or ray.data.from_numpy(). These changed semantics are described below for per-row UDFs, batch UDFs, and per-row and batch iterators.

The following are changed semantics for how UDF returns are interpreted.

• A per-row UDF returns a NumPy ndarray:
  • Before: Creates a simple block (list) of ndarrays.
  • After: Creates a single-tensor-column Arrow table.
• A batch UDF returns a NumPy ndarray:
  • Before: Raises an error.
  • After: Creates a single-tensor-column Arrow table.
• A batch UDF returns a list of NumPy ndarrays:
  • Before: Creates a simple block (list) of ndarrays.
  • After: Creates a single-tensor-column Arrow table.

The following are changed semantics for the format in which UDF data argument is supplied.

• A per-row UDF is called on a pure-tensor dataset.
  • Before: UDF receives a single-column PandasRow containing the tensor column element.
  • After: UDF receives the tensor column element as a NumPy ndarray.
• A batch UDF is called on a pure-tensor dataset:
  • Before: UDF receives a single-column Pandas table containing the tensor column.
  • After: UDF receives the tensor column as a NumPy ndarray.

The following are changed semantics for the format of consumed batches:

• A row iterator is consumed on a pure-tensor dataset:
  • Before: A single column tabular row (either PandasRow or ArrowRow) is returned containing the tensor column element.
  • After: The tensor column element is returned as a NumPy ndarray.
• A batch iterator is consumed on a pure-tensor dataset:
  • Before: A single-column Pandas table is returned containing the tensor column.
  • After: The tensor column is returned as a NumPy ndarray.

Related issue number

Closes #24208, #24207

Checks

• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[ericl]

Main comment is to rejigger this slightly to be more generalizable to handling simple block values in the same way in the future.

[clarkzinzow]

@ericl PR is updated with:

• TENSOR_COL_NAME = "__RAY_TC__"" --> VALUE_COL_NAME = "__value__"
• "numpy" batch format is removed (for follow-up PR)

I'll update the PR description with this downscoping once I open up the stacked follow-up PR (I'm going to want to copy some of this PR description over).

[ericl]

Looks good at high level to me.

[clarkzinzow]

@ericl Feedback addressed! I've also opened a follow-up PR adding the "numpy" format, stacked on this one. #24870

[clarkzinzow]

@ericl @jianoaix Ready for review, PTAL

[clarkzinzow]

@ericl @jianoaix I think those last few commits fixed the remaining bugs after the delegating block builder refactor, should know when the Datasets CI job finishes. My laptop is dead and I won't be able to charge it for another ~hours, so can y'all get this across the line?

[jianoaix]

LG from my side.

[jianoaix]

Use literalinclude?

[clarkzinzow]

I matched the existing inline code blocks style to keep the PR focused and keep the diff small, hoping to do a port all of the code examples in this feature guide in a Working with Tensors feature guide overhaul.

[clarkzinzow]

Datasets failure is unrelated, it's a flaky test that was skipped, but that skip was recently reverted here: 2be45fe. Other failures appear to be unrelated, so I'm going to merge this.

I'll open a separate PR to skip that test.

[clarkzinzow]

PR to skip test here: #25009