Infer types from lance blobs (#7966)

infer types from lance blobs

Author: lhoestq

src/datasets/features/audio.py

@@ -252,6 +252,12 @@ def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray]) -> pa.Str
         if pa.types.is_string(storage.type):
             bytes_array = pa.array([None] * len(storage), type=pa.binary())
             storage = pa.StructArray.from_arrays([bytes_array, storage], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_large_binary(storage.type):
+            storage = array_cast(
+                storage, pa.binary()
+            )  # this can fail in case of big audios, paths should be used instead
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())
         elif pa.types.is_binary(storage.type):
             path_array = pa.array([None] * len(storage), type=pa.string())
             storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())

src/datasets/features/image.py

@@ -242,6 +242,12 @@ def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray, pa.ListArr
         if pa.types.is_string(storage.type):
             bytes_array = pa.array([None] * len(storage), type=pa.binary())
             storage = pa.StructArray.from_arrays([bytes_array, storage], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_large_binary(storage.type):
+            storage = array_cast(
+                storage, pa.binary()
+            )  # this can fail in case of big images, paths should be used instead
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())
         elif pa.types.is_binary(storage.type):
             path_array = pa.array([None] * len(storage), type=pa.string())
             storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())

src/datasets/features/video.py

@@ -258,6 +258,12 @@ def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray, pa.ListArr
         if pa.types.is_string(storage.type):
             bytes_array = pa.array([None] * len(storage), type=pa.binary())
             storage = pa.StructArray.from_arrays([bytes_array, storage], ["bytes", "path"], mask=storage.is_null())
+        elif pa.types.is_large_binary(storage.type):
+            storage = array_cast(
+                storage, pa.binary()
+            )  # this can fail in case of big videos, paths should be used instead
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())
         elif pa.types.is_binary(storage.type):
             path_array = pa.array([None] * len(storage), type=pa.string())
             storage = pa.StructArray.from_arrays([storage, path_array], ["bytes", "path"], mask=storage.is_null())

src/datasets/packaged_modules/lance/lance.py

@@ -7,6 +7,7 @@
 from huggingface_hub import HfApi
 
 import datasets
+from datasets import Audio, Image, Video
 from datasets.builder import Key
 from datasets.table import table_cast
 from datasets.utils.file_utils import is_local_path
@@ -18,6 +19,23 @@
 
 logger = datasets.utils.logging.get_logger(__name__)
 
+MAGIC_BYTES_EXTENSION_AND_FEATURE_TYPES = [
+    ("1A 45 DF A3", ".mkv", Video()),
+    ("66 74 79 70 69 73 6F 6D", ".mp4", Video()),
+    ("66 74 79 70 4D 53 4E 56", ".mp4", Video()),
+    ("52 49 46 46", ".avi", Video()),
+    ("00 00 01 BA", ".mpeg", Video()),
+    ("00 00 01 BA", ".mpeg", Video()),
+    ("00 00 01 B3", ".mov", Video()),
+    ("89 50 4E 47", ".png", Image()),
+    ("FF D8", ".jpg", Image()),
+    ("49 49", ".tif", Image()),
+    ("47 49 46 38", ".gif", Image()),
+    ("52 49 46 46", ".wav", Audio()),
+    ("49 44 33", ".mp3", Audio()),
+    ("66 4C 61 43", ".flac", Audio()),
+]
+
 
 @dataclass
 class LanceConfig(datasets.BuilderConfig):
@@ -104,13 +122,23 @@ def _split_generators(self, dl_manager):
 
             lance_dataset_uris = resolve_dataset_uris(files)
             if lance_dataset_uris:
-                fragments = [
-                    frag
-                    for uri in lance_dataset_uris
-                    for frag in lance.dataset(uri, storage_options=storage_options).get_fragments()
-                ]
+                lance_datasets = [lance.dataset(uri, storage_options=storage_options) for uri in lance_dataset_uris]
+                fragments = [frag for lance_dataset in lance_datasets for frag in lance_dataset.get_fragments()]
                 if self.info.features is None:
                     pa_schema = fragments[0]._ds.schema
+                    first_row_first_bytes = {}
+                    for field in pa_schema:
+                        if self.config.columns is not None and field.name not in self.config.columns:
+                            continue
+                        if pa.types.is_binary(field.type) or pa.types.is_large_binary(field.type):
+                            try:
+                                first_row_first_bytes[field.name] = (
+                                    lance_datasets[0].take_blobs(field.name, [0])[0].read(16)
+                                )
+                            except ValueError:
+                                first_row_first_bytes[field.name] = (
+                                    lance_datasets[0].take([0], [field.name]).to_pylist()[0][field.name][:16]
+                                )
                 splits.append(
                     datasets.SplitGenerator(
                         name=split_name,
@@ -124,6 +152,11 @@ def _split_generators(self, dl_manager):
                 ]
                 if self.info.features is None:
                     pa_schema = lance_files[0].metadata().schema
+                    first_row_first_bytes = {
+                        field_name: value[:16]
+                        for field_name, value in lance_files[0].take_rows([0]).to_table().to_pylist()[0].items()
+                        if isinstance(value, bytes)
+                    }
                 splits.append(
                     datasets.SplitGenerator(
                         name=split_name,
@@ -136,7 +169,14 @@ def _split_generators(self, dl_manager):
                         pa_schema.field(name) for name in self.config.columns if pa_schema.get_field_index(name) != -1
                     ]
                     pa_schema = pa.schema(fields)
-                self.info.features = datasets.Features.from_arrow_schema(pa_schema)
+                features = datasets.Features.from_arrow_schema(pa_schema)
+                for field_name, first_bytes in first_row_first_bytes.items():
+                    for magic_bytes_hex, _, feature_type in MAGIC_BYTES_EXTENSION_AND_FEATURE_TYPES:
+                        magic_bytes = bytes.fromhex(magic_bytes_hex)
+                        if magic_bytes in first_bytes[: len(magic_bytes) * 2]:  # allow some padding
+                            features[field_name] = feature_type
+                            break
+                self.info.features = features
 
         return splits