Webdataset backend by GaganNarula · Pull Request #212 · earthspecies/esp-data

GaganNarula · 2026-01-14T09:01:29Z

Adds tarfile dataset read via webdataset library as a backend
Adds StreamingBackend protocol which differs from DataBackend (should it be StreamingDataBackend?)

mil-ad

few comments while looking at it more carefully

mil-ad · 2026-01-21T13:59:14Z

+            else:
+                return all(sample.get(col) is not None for col in subset)
+
+        self._dropna_fn = dropna_fn


I haven't looked at other backends but doesn't this pattern of storing in self._dropna_fn() means that we won't be able to do any chaining?

I'd expect to be able to do:

foo.backend.dropna_fn(...).dropna_fn()

and right now we only apply the last one.

same thing for filter_isin()

yeah the implementation is a bit incomplete right now (these methods have to applied to the WebDataset obj) because I wanted you to have a quick look. But you're right we should instead return a new WebDatasetBackend(). This is the pattern in the other backends

can we raise an issue for this?

i've changed the implementation in the recent commits, that allows for chaining . wdyt ?

Paulchen-git · 2026-05-12T07:22:31Z

+        # Imported here to avoid circular import: generic_dataset imports Dataset
+        from esp_data.generic_dataset import GenericDataset


I'm not used to keep circular import, I feel like it isn't clean and safe generally, but here I'm still thinking how I could avoid that, wdyt ?

good question, this shouldn't really be done.. one solution would be to not add @register_dataset decorator to GenericDataset (and also not subclass it to Dataset). that breaks the pattern..lets keep it like this for now

GaganNarula · 2026-05-12T07:39:01Z

+            backend=dataset_config.backend,
+            streaming=dataset_config.streaming,
+        )
+        if dataset_config.transformations:


transformations shouldn't be applied here.. we have to think about whether transformations can be applied at all, but if they can they have to be done within from_path using dataset_config.transformations

Ok for now, I've removed it

GaganNarula · 2026-05-12T09:09:12Z

+        if resolved.is_dir():
+            data = WebDatasetBackend.from_path(path, **kwargs)
+        else:
+            data = get_backend(backend).from_path(path, streaming=streaming, **kwargs)


you can just use get_backend instead of explicitly importing and instantiating WebDatasetBackend (no if - else needed),. We shouldnt check for is_dir() .. path is expected to be a dir anyway

actually this whole bit from line 525 to 539 should be moved to init of GenericDataset. then init only takes (path, **kwargs), not (data, info)

GaganNarula · 2026-05-12T09:19:41Z

+
+        info = None
+        config_dir = resolved if resolved.is_dir() else resolved.parent
+        for config_name in ("config.json", "config.yaml"):


lets assume its yaml

GaganNarula · 2026-05-12T09:23:46Z

+            If provided, overrides the class-level generic `info`.
+        """
+        self._is_webdataset = isinstance(data, WebDatasetBackend)
+        streaming = self._is_webdataset


hmm, i think streaming should be an init param, not dependent on a BackendType. We should not mention any strict types in GenericDataset except literals like backend="polars"

GaganNarula · 2026-05-12T09:23:54Z

+        pass
+
+    def __len__(self) -> int:
+        if self._is_webdataset:


same issue here

GaganNarula · 2026-05-12T09:23:59Z

+        return len(self._data)
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        if self._is_webdataset:


GaganNarula · 2026-05-12T09:25:13Z

+        )
+
+    @classmethod
+    def from_config(cls, dataset_config: DatasetConfig) -> tuple["GenericDataset", dict[str, Any]]:


hmm this is getting quite mixed up.. GenericDataset.from_config is calling Dataset.from_path within it.. we need a rethink

re my comment on dataset_from_config: if we introduce GenericDatasetConfig, then we dont need to use Dataset.from_path here. Rather just create cls(path)

GaganNarula

some comments

GaganNarula · 2026-05-12T09:31:13Z

    the type of dataset configuration:
    - 'dataset' for DatasetConfig
    - 'concat' for ConcatConfig
    - 'chain' for ChainedDatasetConfig


maybe we can directly use a 'generic' key and do GenericDataset.from_config in dataset_from_config

GaganNarula · 2026-05-12T09:36:00Z

+    def __str__(self) -> str:
+        return (
+            f"GenericDataset(source={self.info.split_paths}, "
+            f"backend={'webdataset' if self._is_webdataset else 'tabular'})"


again _is_webdataset check.. we should maybe just print the self.info information in str.. reason is that if we overwrite self.info when we load from path, then print(ds) should reflect the actual dataset.. so it could be anuraset or something specific

Paulchen-git · 2026-05-12T14:36:57Z

After playing around with esp_data and webdataset backend, I'm currently confused. The save_to export I've made is working but is useless, you can't do much from a GenericDataset after having it loaded from tar files. maybe I've misunderstood the purpose of the export and load methods ?

GenericDataset yields raw dicts with no per-sample processing. Round-trip AnuraSetStrong → save_to → from_path → GenericDataset loses all the domain logic (audio decoding, label mapping, typed fields).

GaganNarula · 2026-05-18T09:08:18Z

+        """
+        self.path = anypath(path)
+        fs = filesystem_from_path(self.path)
+        info_path = self.path / "info.yaml"


hmm, the info object itself doesnt contain the params used to construct the dataset object.. for e.g. we need:

backend_type

split

streaming

sample_rate

.. other dataset specific params like version in audioset

So i think we should expect one file called config.yaml, which contains these params along with an info key which is the dict representing DatasetInfo

GaganNarula · 2026-05-18T09:08:52Z

+                f"No info.yaml found at {info_path}. Cannot infer backend or streaming mode."
+            )
+
+        super().__init__(backend=self.info.backend, streaming=self.info.streaming)


same issue, DatasetInfo doesnt contain these so self.info wont have them

GaganNarula · 2026-05-18T09:09:46Z

+        yield from self._data
+
+    def __str__(self) -> str:
+        return f"GenericDataset(name={self.info.name}, version={self.info.version})"


we can add split and sample_rate etc to it

GaganNarula · 2026-05-18T09:11:19Z

+    def save_to(self, path: str, format: str = "webdataset", **kwargs: Any) -> int:
+        """Save the dataset to a file.
+
+        Writes sharded tar files and an ``info.yaml`` containing `DatasetInfo`


should not be specific to tar because that's just one format. As mentioned, we probably need config.yaml not info.yaml

GaganNarula · 2026-05-18T09:13:35Z

+            info_dict["sample_rate"] = sample_rate
+
+        # Record backend so from_path can reload without the user specifying it.
+        info_dict["backend"] = backend


oh! i see you're overloading the info object here..

GaganNarula · 2026-05-18T10:54:58Z

+            If `format` is not supported
+        """
+        self._ensure_not_streaming("save_to")
+        if format == "webdataset":


if our exporters are functions like write_to_webdataset (which i prefer) , then we can make an abstraction like export_to(path, format) which is a public function that does this selection .. and then we dont have to import a webdataset utility in a pandas backend! which is kind of a confusing pattern

I took a look to your tar_exporter PR, maybe we could merge it into the webdataset_backend PR, then I would add the export abstraction and I could use the code you wrote for tar files export

wdyt ?

Or I can add the export abstraction directly here

and we can put this export_to into an export.py and make it public api from there.. the stuff like webdataset utils or format specific stuff shouldnt be public

oh i just oh i just saw your message :D .. yes i think that works to merge this and the other branch, but maybe we should stack the export on this one, exposing a dummy export function for now

#291 I open this stacked PR for the export abstraction

…amingBackend Protocol

…or default

Clean #234. Rebased on on #212 with a clean commit stack

GaganNarula marked this pull request as ready for review January 14, 2026 11:02

mil-ad reviewed Jan 21, 2026

View reviewed changes

mil-ad reviewed Jan 23, 2026

View reviewed changes

Comment thread esp_data/backends/__init__.py Outdated

mil-ad reviewed Jan 23, 2026

View reviewed changes

Comment thread esp_data/backends/webdataset_backend.py Outdated

Paulchen-git force-pushed the gagan/webdataset-backend branch from 4929968 to 6b01b23 Compare May 11, 2026 07:51

Paulchen-git requested a review from a team as a code owner May 11, 2026 07:51

Paulchen-git force-pushed the gagan/webdataset-backend branch 3 times, most recently from 4d249bd to ef58864 Compare May 11, 2026 13:51

Paulchen-git reviewed May 12, 2026

View reviewed changes

GaganNarula commented May 12, 2026

View reviewed changes

Comment thread esp_data/dataset.py Outdated

GaganNarula commented May 12, 2026

View reviewed changes

Comment thread esp_data/dataset.py Outdated

GaganNarula commented May 12, 2026

View reviewed changes

Paulchen-git force-pushed the gagan/webdataset-backend branch from ef58864 to e15a3a8 Compare May 12, 2026 12:46

GaganNarula commented May 18, 2026

View reviewed changes

GaganNarula and others added 26 commits May 22, 2026 07:42

Remove bad doctest, needs to much work to make a doc test here

645300c

Allow for chaining filter_isin and drop_na

c0f9b43

Return a new WebDatasetBackend obj after applying funcs, matches Stre…

c471488

…amingBackend Protocol

Refactor: move out utility funcs, remove from module level export

196c1a7

Fix prek

abf93ef

feat: add write_to_webdataset with cloud support

2ef3366

fix: fix docstrings, _columns reset on copy, add save_to

d538bdf

feat: add from_path and save_to to tabular backends and protocol

0e729bc

feat: add GenericDataset for path-based loading

6b06063

feat: add Dataset.from_path, save_to, fix _make_dataset_from_config

39503b0

test: add tests for save_to, GenericDataset, and Dataset.from_path

150c203

feat: add GenericDatasetConfig and wire 'generic' YAML key

8fd8576

fix: add save_to to StreamingBackend protocol; normalize data_process…

64a5b1a

…or default

fix: fix some bugs raised by Claude

5340ad9

feat: auto-detect backend/streaming from info.yaml on load

c6957f6

fix: fix integration test

dfc33c0

extend changes to all backends

9d99379

fix prek

03481b5

small refactoring of GenericDataset

3adc1d7

refactor Dataset.save_to

8277271

fix some tests

4df52c2

fix integration test

8ab19c3

update docstring Dataset.save_to and how parameters and ino are saved

8f642d6

update from_path to read config.yaml

92a3b72

update test_dataset

cb06463

Export abastraction (#291)

818848a

Paulchen-git force-pushed the gagan/webdataset-backend branch from 0e6d73b to 818848a Compare May 22, 2026 07:47

Paulchen-git mentioned this pull request May 22, 2026

Tar exporter rebased #295

Merged

Paulchen-git added 2 commits May 28, 2026 12:33

Tar exporter rebased (#295)

dac5610

Clean #234. Rebased on on #212 with a clean commit stack

remove integration test that need to be improved

076ca1d

		# Imported here to avoid circular import: generic_dataset imports Dataset
		from esp_data.generic_dataset import GenericDataset

Conversation

GaganNarula commented Jan 14, 2026

Uh oh!

mil-ad left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

GaganNarula left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Paulchen-git commented May 12, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects