Make sure to pass features into the embed storage branch as well.

HaukurPall · HaukurPall · commit 4c3231ded086 · 2026-03-05T21:48:09.000Z
diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py
@@ -4450,6 +4450,7 @@ def _push_parquet_shards_to_hub_single(
                     partial(embed_table_storage, token_per_repo_id=self._token_per_repo_id),
                     batched=True,
                     batch_size=get_arrow_writer_batch_size_from_features(shard.features),
+                    features=shard.features,
                 )
             shard_path_in_repo = f"{data_dir}/{split}-{index:05d}-of-{num_shards:05d}.parquet"
             tmp_file = tempfile.NamedTemporaryFile(suffix=".parquet", delete=False)

Original file line number	Diff line number	Diff line change
`@@ -4450,6 +4450,7 @@ def _push_parquet_shards_to_hub_single(`
`4450`	`4450`	`partial(embed_table_storage, token_per_repo_id=self._token_per_repo_id),`
`4451`	`4451`	`batched=True,`
`4452`	`4452`	`batch_size=get_arrow_writer_batch_size_from_features(shard.features),`
	`4453`	`+ features=shard.features,`
`4453`	`4454`	`)`
`4454`	`4455`	`shard_path_in_repo = f"{data_dir}/{split}-{index:05d}-of-{num_shards:05d}.parquet"`
`4455`	`4456`	`tmp_file = tempfile.NamedTemporaryFile(suffix=".parquet", delete=False)`