to_tf_dataset fails with datetime UTC columns even if not included in columns argument
State: closed
Labels: name: bug, name: good first issue

### Describe the bug

There appears to be some eager behavior in `to_tf_dataset` that runs against every column in a dataset even if they aren't included in the columns argument. This is problematic with datetime UTC columns due to them not working with zero copy. If I don't have UTC information in my datetime column, then everything works as expected.

### Steps to reproduce the bug

```python
import numpy as np
import pandas as pd
from datasets import Dataset


df = pd.DataFrame(np.random.rand(2, 1), columns=["x"])
# df["dt"] = pd.to_datetime(["2023-01-01", "2023-01-01"])  # works fine
df["dt"] = pd.to_datetime(["2023-01-01 00:00:00.00000+00:00", "2023-01-01 00:00:00.00000+00:00"])
df.to_parquet("test.pq")


ds = Dataset.from_parquet("test.pq")
tf_ds = ds.to_tf_dataset(columns=["x"], batch_size=2, shuffle=True)
```

```
ArrowInvalid                              Traceback (most recent call last)
Cell In[1], line 12
      8 df.to_parquet("test.pq")
     11 ds = Dataset.from_parquet("test.pq")
---> 12 tf_ds = ds.to_tf_dataset(columns=["r"], batch_size=2, shuffle=True)

File ~/venv/lib/python3.8/site-packages/datasets/arrow_dataset.py:411, in TensorflowDatasetMixin.to_tf_dataset(self, batch_size, columns, shuffle, collate_fn, drop_remainder, collate_fn_args, label_cols, prefetch, num_workers)
    407     dataset = self
    409 # TODO(Matt, QL): deprecate the retention of label_ids and label
--> 411 output_signature, columns_to_np_types = dataset._get_output_signature(
    412     dataset,
    413     collate_fn=collate_fn,
    414     collate_fn_args=collate_fn_args,
    415     cols_to_retain=cols_to_retain,
    416     batch_size=batch_size if drop_remainder else None,
    417 )
    419 if "labels" in output_signature:
    420     if ("label_ids" in columns or "label" in columns) and "labels" not in columns:

File ~/venv/lib/python3.8/site-packages/datasets/arrow_dataset.py:254, in TensorflowDatasetMixin._get_output_signature(dataset, collate_fn, collate_fn_args, cols_to_retain, batch_size, num_test_batches)
    252 for _ in range(num_test_batches):
    253     indices = sample(range(len(dataset)), test_batch_size)
--> 254     test_batch = dataset[indices]
    255     if cols_to_retain is not None:
    256         test_batch = {key: value for key, value in test_batch.items() if key in cols_to_retain}

File ~/venv/lib/python3.8/site-packages/datasets/arrow_dataset.py:2590, in Dataset.__getitem__(self, key)
   2588 def __getitem__(self, key):  # noqa: F811
   2589     """Can be used to index columns (by string names) or rows (by integer index or iterable of indices or bools)."""
-> 2590     return self._getitem(
   2591         key,
   2592     )

File ~/venv/lib/python3.8/site-packages/datasets/arrow_dataset.py:2575, in Dataset._getitem(self, key, **kwargs)
   2573 formatter = get_formatter(format_type, features=self.features, **format_kwargs)
   2574 pa_subtable = query_table(self._data, key, indices=self._indices if self._indices is not None else None)
-> 2575 formatted_output = format_table(
   2576     pa_subtable, key, formatter=formatter, format_columns=format_columns, output_all_columns=output_all_columns
   2577 )
   2578 return formatted_output

File ~/venv/lib/python3.8/site-packages/datasets/formatting/formatting.py:634, in format_table(table, key, formatter, format_columns, output_all_columns)
    632 python_formatter = PythonFormatter(features=None)
    633 if format_columns is None:
--> 634     return formatter(pa_table, query_type=query_type)
    635 elif query_type == "column":
    636     if key in format_columns:

File ~/venv/lib/python3.8/site-packages/datasets/formatting/formatting.py:410, in Formatter.__call__(self, pa_table, query_type)
    408     return self.format_column(pa_table)
    409 elif query_type == "batch":
--> 410     return self.format_batch(pa_table)

File ~/venv/lib/python3.8/site-packages/datasets/formatting/np_formatter.py:78, in NumpyFormatter.format_batch(self, pa_table)
     77 def format_batch(self, pa_table: pa.Table) -> Mapping:
---> 78     batch = self.numpy_arrow_extractor().extract_batch(pa_table)
     79     batch = self.python_features_decoder.decode_batch(batch)
     80     batch = self.recursive_tensorize(batch)

File ~/venv/lib/python3.8/site-packages/datasets/formatting/formatting.py:164, in NumpyArrowExtractor.extract_batch(self, pa_table)
    163 def extract_batch(self, pa_table: pa.Table) -> dict:
--> 164     return {col: self._arrow_array_to_numpy(pa_table[col]) for col in pa_table.column_names}

File ~/venv/lib/python3.8/site-packages/datasets/formatting/formatting.py:164, in <dictcomp>(.0)
    163 def extract_batch(self, pa_table: pa.Table) -> dict:
--> 164     return {col: self._arrow_array_to_numpy(pa_table[col]) for col in pa_table.column_names}

File ~/venv/lib/python3.8/site-packages/datasets/formatting/formatting.py:185, in NumpyArrowExtractor._arrow_array_to_numpy(self, pa_array)
    181     else:
    182         zero_copy_only = _is_zero_copy_only(pa_array.type) and all(
    183             not _is_array_with_nulls(chunk) for chunk in pa_array.chunks
    184         )
--> 185         array: List = [
    186             row for chunk in pa_array.chunks for row in chunk.to_numpy(zero_copy_only=zero_copy_only)
    187         ]
    188 else:
    189     if isinstance(pa_array.type, _ArrayXDExtensionType):
    190         # don't call to_pylist() to preserve dtype of the fixed-size array

File ~/venv/lib/python3.8/site-packages/datasets/formatting/formatting.py:186, in <listcomp>(.0)
    181     else:
    182         zero_copy_only = _is_zero_copy_only(pa_array.type) and all(
    183             not _is_array_with_nulls(chunk) for chunk in pa_array.chunks
    184         )
    185         array: List = [
--> 186             row for chunk in pa_array.chunks for row in chunk.to_numpy(zero_copy_only=zero_copy_only)
    187         ]
    188 else:
    189     if isinstance(pa_array.type, _ArrayXDExtensionType):
    190         # don't call to_pylist() to preserve dtype of the fixed-size array

File ~/venv/lib/python3.8/site-packages/pyarrow/array.pxi:1475, in pyarrow.lib.Array.to_numpy()

File ~/venv/lib/python3.8/site-packages/pyarrow/error.pxi:100, in pyarrow.lib.check_status()

ArrowInvalid: Needed to copy 1 chunks with 0 nulls, but zero_copy_only was True
```

### Expected behavior

I think there are two potential issues/fixes
1. Proper handling of datetime UTC columns (perhaps there is something incorrect with zero copy handling here)
2. Not eagerly running against every column in a dataset when the columns argument of `to_tf_dataset` specifies a subset of columns (although I'm not sure if this is unavoidable)


### Environment info

- `datasets` version: 2.9.0
- Platform: macOS-13.2-x86_64-i386-64bit
- Python version: 3.8.12
- PyArrow version: 11.0.0
- Pandas version: 1.5.3