[Datasets] Simplify to_tf interface

[bveeramani]

Signed-off-by: Balaji Veeramani [email protected]

Why are these changes needed?

The TensorFlow UX is bad because to_tf is confusing and iter_tf_batches is boilerplate-y.

1. to_tf automatically concatenates and unsqueezes columns for the user. The original motivation was to let users convert tabular datasets to TensorFlow datasets. But, the semantics are complicated, and it often doesn't work the way you want it to. More importantly, the concatenation and unsqueezing functionality is obsolete with the introduction of preprocessors.Concatenator.

2. The currently recommended iter_tf_batches requires lots of boilerplate. You need to specify the output signature, call prepare_dataset_shard, and yield tensors from iter_tf_batchse. We can improve the UX by eliminating calling prepare_dataset_shard for the user and inferring the output_signature from the dataset schema.

Before

def to_tf_dataset(dataset, batch_size):
    def to_tensor_iterator():
        data_iterator = dataset.iter_tf_batches(
            batch_size=batch_size, dtypes=tf.float32
        )
        for d in data_iterator:
            # "concat_out" is the output column of the Concatenator.
            yield d["concat_out"], d["target"]

    output_signature = (
        tf.TensorSpec(shape=(None, num_features), dtype=tf.float32),
        tf.TensorSpec(shape=(None), dtype=tf.float32),
    )
    tf_dataset = tf.data.Dataset.from_generator(
        to_tensor_iterator, output_signature=output_signature
    )
    return prepare_dataset_shard(tf_dataset)

to_tf_dataset(dataset, batch_size=64)

After

dataset.to_tf("concat_out", "target", batch_size=64)

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• 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 :(

[matthewdeng]

Add this to BUILD file.

[bveeramani]

Looks like test_dataset_tf is already running in CI. Maybe it's captured by this?

ray/python/ray/data/BUILD

Lines 31 to 42 in 1bd3f94

	py_test_module_list(
	files = glob(
	include=["tests/test_*.py"],
	exclude=[
	"tests/test_preprocessors.py",
	"tests/test_dataset_formats.py",
	],
	),
	size = "large",
	tags = ["team:core", "exclusive"],
	deps = ["//:ray_lib", ":conftest"],
	)

[clarkzinzow]

LGTM overall, mostly small nits on error handling and test coverage.

[clarkzinzow]

Now that the variable-shaped tensor column (simple ragged tensors) PR is unreverted, could we add test coverage of variable-shaped tensors back? #29071

[bveeramani]

Discussed offline. We'll do this in follow-up PRs:

1. [Datasets] Make ragged tensor shape more descriptive #29135
2. [Datasets] Make to_tf work with ragged tensors #29136

[clarkzinzow]

LGTM with the two changes given!

stale

[clarkzinzow]

All relevant tests are passing, merging!