nullable list

Author: gabriel.g.robin

dataframely/columns/list.py

@@ -147,4 +147,7 @@ def _sample_unchecked(self, generator: Generator, n: int) -> pl.Series:
                 chain([0], element_lengths.cum_sum()), element_lengths
             )
         ]
-        return pl.Series(list_elements)
+        # Finally, apply a null mask
+        return generator._apply_null_mask(
+            pl.Series(list_elements), null_probability=self._null_probability
+        )

dataframely/random.py

@@ -392,8 +392,10 @@ def sample_duration(
     def _apply_null_mask(self, series: pl.Series, null_probability: float) -> pl.Series:
         if null_probability == 0:
             return series
-        null_mask = self.numpy_generator.random(series.len()) < null_probability
-        return series.scatter(np.where(null_mask)[0], None)
+        null_mask = (
+            pl.Series(self.numpy_generator.random(series.len())) > null_probability
+        )
+        return pl.select(pl.when(null_mask).then(series)).to_series()
 
 
 # --------------------------------------- UTILS -------------------------------------- #

tests/columns/test_sample.py

@@ -178,7 +178,7 @@ def test_sample_enum(generator: Generator) -> None:
 def test_sample_list(generator: Generator) -> None:
     column = dy.List(dy.String(regex="[abc]"), min_length=5, max_length=10)
     samples = sample_and_validate(column, generator, n=10_000)
-    assert set(samples.list.len()) == set(range(5, 11))
+    assert set(samples.list.len()) == set(range(5, 11)) | {None}
 
 
 def test_sample_struct(generator: Generator) -> None: