Add table_cast to loaders

Author: mariosasko

src/datasets/packaged_modules/csv/csv.py

@@ -7,6 +7,8 @@
 
 import datasets
 import datasets.config
+from datasets.features.features import require_cast_storage
+from datasets.table import table_cast
 
 
 logger = datasets.utils.logging.get_logger(__name__)
@@ -66,6 +68,10 @@ def __post_init__(self):
         if self.column_names is not None:
             self.names = self.column_names
 
+    @property
+    def schema(self):
+        return self.features.arrow_schema if self.features is not None else None
+
     @property
     def read_csv_kwargs(self):
         read_csv_kwargs = dict(
@@ -146,19 +152,37 @@ def _split_generators(self, dl_manager):
             splits.append(datasets.SplitGenerator(name=split_name, gen_kwargs={"files": dl_manager.iter_files(files)}))
         return splits
 
+    def _cast_table(self, pa_table: pa.Table) -> pa.Table:
+        if self.config.features is not None:
+            schema = self.config.schema
+            if all(not require_cast_storage(feature) for feature in self.config.features.values()):
+                # cheaper cast
+                pa_table = pa.Table.from_arrays([pa_table[field.name] for field in schema], schema=schema)
+            else:
+                # more expensive cast; allows str <-> int/float or str to Audio for example
+                pa_table = table_cast(pa_table, schema)
+        return pa_table
+
     def _generate_tables(self, files):
-        schema = pa.schema(self.config.features.type) if self.config.features is not None else None
+        schema = self.config.schema
         # dtype allows reading an int column as str
-        dtype = {name: dtype.to_pandas_dtype() for name, dtype in zip(schema.names, schema.types)} if schema else None
+        dtype = (
+            {
+                name: dtype.to_pandas_dtype() if not require_cast_storage(feature) else object
+                for name, dtype, feature in zip(schema.names, schema.types, self.config.features.values())
+            }
+            if schema
+            else None
+        )
         for file_idx, file in enumerate(files):
             csv_file_reader = pd.read_csv(file, iterator=True, dtype=dtype, **self.config.read_csv_kwargs)
             try:
                 for batch_idx, df in enumerate(csv_file_reader):
-                    pa_table = pa.Table.from_pandas(df, schema=schema)
+                    pa_table = pa.Table.from_pandas(df)
                     # Uncomment for debugging (will print the Arrow table size and elements)
                     # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                     # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
-                    yield (file_idx, batch_idx), pa_table
+                    yield (file_idx, batch_idx), self._cast_table(pa_table)
             except ValueError as e:
                 logger.error(f"Failed to read file '{file}' with error {type(e)}: {e}")
                 raise

src/datasets/packaged_modules/json/json.py

@@ -7,6 +7,7 @@
 import pyarrow.json as paj
 
 import datasets
+from datasets.features.features import require_cast_storage
 from datasets.table import table_cast
 from datasets.utils.file_utils import readline
 
@@ -64,21 +65,15 @@ def _split_generators(self, dl_manager):
             splits.append(datasets.SplitGenerator(name=split_name, gen_kwargs={"files": dl_manager.iter_files(files)}))
         return splits
 
-    def _cast_classlabels(self, pa_table: pa.Table) -> pa.Table:
-        if self.config.features:
-            # Encode column if ClassLabel
-            for i, col in enumerate(self.config.features.keys()):
-                if isinstance(self.config.features[col], datasets.ClassLabel):
-                    if pa_table[col].type == pa.string():
-                        pa_table = pa_table.set_column(
-                            i, self.config.schema.field(col), [self.config.features[col].str2int(pa_table[col])]
-                        )
-                    elif pa_table[col].type != self.config.schema.field(col).type:
-                        raise ValueError(
-                            f"Field '{col}' from the JSON data of type {pa_table[col].type} is not compatible with ClassLabel. Compatible types are int64 and string."
-                        )
-            # Cast allows str <-> int/float or str to Audio for example
-            pa_table = table_cast(pa_table, self.config.schema)
+    def _cast_table(self, pa_table: pa.Table) -> pa.Table:
+        if self.config.features is not None:
+            schema = self.config.schema
+            if all(not require_cast_storage(feature) for feature in self.config.features.values()):
+                # cheaper cast
+                pa_table = pa.Table.from_arrays([pa_table[field.name] for field in schema], schema=schema)
+            else:
+                # more expensive cast; allows str <-> int/float or str to Audio for example
+                pa_table = table_cast(pa_table, schema)
         return pa_table
 
     def _generate_tables(self, files):
@@ -98,7 +93,7 @@ def _generate_tables(self, files):
                 else:
                     mapping = dataset
                 pa_table = pa.Table.from_pydict(mapping=mapping)
-                yield file_idx, self._cast_classlabels(pa_table)
+                yield file_idx, self._cast_table(pa_table)
 
             # If the file has one json object per line
             else:
@@ -153,5 +148,5 @@ def _generate_tables(self, files):
                         # Uncomment for debugging (will print the Arrow table size and elements)
                         # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                         # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
-                        yield (file_idx, batch_idx), self._cast_classlabels(pa_table)
+                        yield (file_idx, batch_idx), self._cast_table(pa_table)
                         batch_idx += 1

src/datasets/packaged_modules/pandas/pandas.py

@@ -1,12 +1,30 @@
+from dataclasses import dataclass
+from typing import Optional
+
 import pandas as pd
 import pyarrow as pa
 
 import datasets
+from datasets.features.features import require_cast_storage
+from datasets.table import table_cast
+
+
+@dataclass
+class PandasConfig(datasets.BuilderConfig):
+    """BuilderConfig for Pandas."""
+
+    features: Optional[datasets.Features] = None
+
+    @property
+    def schema(self):
+        return self.features.arrow_schema if self.features is not None else None
 
 
 class Pandas(datasets.ArrowBasedBuilder):
+    BUILDER_CONFIG_CLASS = PandasConfig
+
     def _info(self):
-        return datasets.DatasetInfo()
+        return datasets.DatasetInfo(features=self.config.features)
 
     def _split_generators(self, dl_manager):
         """We handle string, list and dicts in datafiles"""
@@ -25,8 +43,19 @@ def _split_generators(self, dl_manager):
             splits.append(datasets.SplitGenerator(name=split_name, gen_kwargs={"files": files}))
         return splits
 
+    def _cast_table(self, pa_table: pa.Table) -> pa.Table:
+        if self.config.features is not None:
+            schema = self.config.schema
+            if all(not require_cast_storage(feature) for feature in self.config.features.values()):
+                # cheaper cast
+                pa_table = pa.Table.from_arrays([pa_table[field.name] for field in schema], schema=schema)
+            else:
+                # more expensive cast; allows str <-> int/float or str to Audio for example
+                pa_table = table_cast(pa_table, schema)
+        return pa_table
+
     def _generate_tables(self, files):
         for i, file in enumerate(files):
             with open(file, "rb") as f:
                 pa_table = pa.Table.from_pandas(pd.read_pickle(f))
-                yield i, pa_table
+                yield i, self._cast_table(pa_table)

src/datasets/packaged_modules/parquet/parquet.py

@@ -5,6 +5,8 @@
 import pyarrow.parquet as pq
 
 import datasets
+from datasets.features.features import require_cast_storage
+from datasets.table import table_cast
 
 
 logger = datasets.utils.logging.get_logger(__name__)
@@ -18,6 +20,10 @@ class ParquetConfig(datasets.BuilderConfig):
     columns: Optional[List[str]] = None
     features: Optional[datasets.Features] = None
 
+    @property
+    def schema(self):
+        return self.features.arrow_schema if self.features is not None else None
+
 
 class Parquet(datasets.ArrowBasedBuilder):
     BUILDER_CONFIG_CLASS = ParquetConfig
@@ -42,6 +48,17 @@ def _split_generators(self, dl_manager):
             splits.append(datasets.SplitGenerator(name=split_name, gen_kwargs={"files": files}))
         return splits
 
+    def _cast_table(self, pa_table: pa.Table) -> pa.Table:
+        if self.config.features is not None:
+            schema = self.config.schema
+            if all(not require_cast_storage(feature) for feature in self.config.features.values()):
+                # cheaper cast
+                pa_table = pa.Table.from_arrays([pa_table[field.name] for field in schema], schema=schema)
+            else:
+                # more expensive cast; allows str <-> int/float or str to Audio for example
+                pa_table = table_cast(pa_table, schema)
+        return pa_table
+
     def _generate_tables(self, files):
         schema = pa.schema(self.config.features.type) if self.config.features is not None else None
         if self.config.features is not None and self.config.columns is not None:
@@ -57,12 +74,10 @@ def _generate_tables(self, files):
                         parquet_file.iter_batches(batch_size=self.config.batch_size, columns=self.config.columns)
                     ):
                         pa_table = pa.Table.from_batches([record_batch])
-                        if self.config.features is not None:
-                            pa_table = pa.Table.from_arrays([pa_table[field.name] for field in schema], schema=schema)
                         # Uncomment for debugging (will print the Arrow table size and elements)
                         # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                         # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
-                        yield f"{file_idx}_{batch_idx}", pa_table
+                        yield f"{file_idx}_{batch_idx}", self._cast_table(pa_table)
                 except ValueError as e:
                     logger.error(f"Failed to read file '{file}' with error {type(e)}: {e}")
                     raise

src/datasets/packaged_modules/text/text.py

@@ -5,6 +5,8 @@
 import pyarrow as pa
 
 import datasets
+from datasets.features.features import require_cast_storage
+from datasets.table import table_cast
 
 
 logger = datasets.utils.logging.get_logger(__name__)
@@ -20,6 +22,10 @@ class TextConfig(datasets.BuilderConfig):
     keep_linebreaks: bool = False
     sample_by: str = "line"
 
+    @property
+    def schema(self):
+        return self.features.arrow_schema if self.features is not None else None
+
 
 class Text(datasets.ArrowBasedBuilder):
     BUILDER_CONFIG_CLASS = TextConfig
@@ -50,8 +56,20 @@ def _split_generators(self, dl_manager):
             splits.append(datasets.SplitGenerator(name=split_name, gen_kwargs={"files": dl_manager.iter_files(files)}))
         return splits
 
+    def _cast_table(self, pa_table: pa.Table) -> pa.Table:
+        if self.config.features is not None:
+            schema = self.config.schema
+            if all(not require_cast_storage(feature) for feature in self.config.features.values()):
+                # cheaper cast
+                pa_table = pa_table.cast(schema)
+            else:
+                # more expensive cast; allows str <-> int/float or str to Audio for example
+                pa_table = table_cast(pa_table, schema)
+            return pa_table
+        else:
+            return pa_table.cast(pa.schema({"text": pa.string()}))
+
     def _generate_tables(self, files):
-        schema = pa.schema(self.config.features.type if self.config.features is not None else {"text": pa.string()})
         for file_idx, file in enumerate(files):
             # open in text mode, by default translates universal newlines ("\n", "\r\n" and "\r") into "\n"
             with open(file, encoding=self.config.encoding) as f:
@@ -66,11 +84,11 @@ def _generate_tables(self, files):
                         batch = StringIO(batch).readlines()
                         if not self.config.keep_linebreaks:
                             batch = [line.rstrip("\n") for line in batch]
-                        pa_table = pa.Table.from_arrays([pa.array(batch)], schema=schema)
+                        pa_table = pa.Table.from_arrays([pa.array(batch)], names=["text"])
                         # Uncomment for debugging (will print the Arrow table size and elements)
                         # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                         # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
-                        yield (file_idx, batch_idx), pa_table
+                        yield (file_idx, batch_idx), self._cast_table(pa_table)
                         batch_idx += 1
                 elif self.config.sample_by == "paragraph":
                     batch_idx = 0
@@ -82,15 +100,15 @@ def _generate_tables(self, files):
                         batch += f.readline()  # finish current line
                         batch = batch.split("\n\n")
                         pa_table = pa.Table.from_arrays(
-                            [pa.array([example for example in batch[:-1] if example])], schema=schema
+                            [pa.array([example for example in batch[:-1] if example])], names=["text"]
                         )
                         # Uncomment for debugging (will print the Arrow table size and elements)
                         # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                         # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
-                        yield (file_idx, batch_idx), pa_table
+                        yield (file_idx, batch_idx), self._cast_table(pa_table)
                         batch_idx += 1
                         batch = batch[-1]
                 elif self.config.sample_by == "document":
                     text = f.read()
-                    pa_table = pa.Table.from_arrays([pa.array([text])], schema=schema)
-                    yield file_idx, pa_table
+                    pa_table = pa.Table.from_arrays([pa.array([text])], names=["text"])
+                    yield file_idx, self._cast_table(pa_table)