[Tabular] Easy way to iterate on inputs and target (X, y)

[mbignotti]

Hi!
I'm currently working on a tabular regression task and I'm using sklearn's SGDRegressor, which implements the partial_fit method.
I would like to ask if there exists a simple method to iterate on the dataset (with batches, not single rows) and automatically split the X and the y. All without using the pytorch Dataloader. Something like:

for X, y in ds.iter(batch_size = 100):
  print(X, y)

It would also be nice to automatically format the X and y batches to pandas and/or numpy but, as far as I understand, at the moment this is only possible for map-style datasets.

At the moment, I am solving the issue as follows, but I'm wondering if there's a better way to do it:

from datasets import load_dataset

# Original dataset is a csv file with "a", "b", "c", "d", "e" columns. The target is the "e" column.

features=["a", "b", "c", "d"]
ds = load_dataset("my_dataset", split="train", streaming=False)
X = ds.select_columns(features)
y = ds.select_columns(["e"])

X = X.with_format("pandas")
y = y.with_format("pandas")

ds_iter = zip(X.iter(batch_size=100), y.iter(batch_size=100))

for X, y in ds_iter:
  print(X, y)

For streaming datasets this works in the same way, except for the fact that the batch is not converted to pandas (note that the conversion simply does not work, no errors raised).

Thanks a lot!

[mariosasko]

Hi! Yes, your approach to iteration is the right one. Regarding with_format with IterableDataset, this issue is still a to-do, as we don't operate on Arrow tables there, so mimicking Dataset's formatting is not as easy since we cannot just call to_pandas/to_numpy on the underlying table.

[mbignotti]

Of course, it would be nice to have that feature for IterableDataset. However, to be fair, converting the dictionary to a pandas dataframe is rather trivial. So I guess I can add a step in the previous loop that transforms the batch.
Thanks a lot for the clarification!

[mariosasko]

It's not that trivial if we want to ensure the conversion is the same for the same input map-style and iterable-style datasets.

[mbignotti]

Yes, sorry. I did not mean that it's trivial to include the conversion in the IterableDataset. What I meant is that it's trivial to add the conversion in my loop, like this:

from datasets import load_dataset
import pandas as pd

# Original dataset is a csv file with "a", "b", "c", "d", "e" columns. The target is the "e" column.

features=["a", "b", "c", "d"]
ds = load_dataset("my_dataset", split="train", streaming=False)
X = ds.select_columns(features)
y = ds.select_columns(["e"])

X = X.with_format("pandas")
y = y.with_format("pandas")

ds_iter = zip(X.iter(batch_size=100), y.iter(batch_size=100))

for X, y in ds_iter:
  X = pd.DataFrame(X) # Here we convert from dict to dataframe
  y = pd.Series(y)
  print(X, y)

I didn't express that well ☺️.