support None in cast_storage

Author: lhoestq

src/datasets/features/features.py

@@ -1024,7 +1024,9 @@ def cast_storage(self, storage: Union[pa.StringArray, pa.IntegerArray]) -> pa.In
                     f"Class label {min_max['max']} greater than configured num_classes {self.num_classes}"
                 )
         elif isinstance(storage, pa.StringArray):
-            storage = pa.array(self.str2int(storage.to_pylist()))
+            storage = pa.array(
+                [self._strval2int(label) if label is not None else None for label in storage.to_pylist()]
+            )
         return array_cast(storage, self.pa_type)
 
     @staticmethod

tests/features/test_features.py

@@ -328,6 +328,11 @@ def test_classlabel_cast_storage():
     arr = pa.array(["__label_that_doesnt_exist__"])
     with pytest.raises(ValueError):
         classlabel.cast_storage(arr)
+    # from nulls
+    arr = pa.array([None])
+    result = classlabel.cast_storage(arr)
+    assert result.type == pa.int64()
+    assert result.to_pylist() == [None]
     # from empty
     arr = pa.array([])
     result = classlabel.cast_storage(arr)