daft-sources 0.1.0__py3-none-any.whl

daft_sources/tar/schema.py

@@ -0,0 +1,16 @@
+from daft import DataType, Schema
+
+TAR_SCHEMA = Schema._from_field_name_and_types(
+    [
+        ("archive_path", DataType.string()),
+        ("member_path", DataType.string()),
+        ("size", DataType.int64()),
+        ("mtime", DataType.float64()),
+        ("mode", DataType.int64()),
+        ("uid", DataType.int64()),
+        ("gid", DataType.int64()),
+        ("uname", DataType.string()),
+        ("gname", DataType.string()),
+        ("bytes", DataType.binary()),
+    ]
+)

daft_sources/tar/source.py

@@ -0,0 +1,46 @@
+from collections.abc import AsyncIterator
+from pathlib import Path
+
+from daft.io import DataSource
+from daft.io.pushdowns import Pushdowns
+from daft.schema import Schema
+
+from daft_sources.tar.schema import TAR_SCHEMA
+from daft_sources.tar.task import TarDataSourceTask
+
+
+class TarDataSource(DataSource):
+    def __init__(
+        self,
+        archive_paths: list[str],
+        *,
+        batch_size: int = 256,
+        max_member_bytes: int = 100 * 1024 * 1024,
+    ):
+        assert batch_size > 0, "Batch_size must be greater than zero"
+        assert max_member_bytes is not None and max_member_bytes > 0, (
+            "max_member_bytes must be non-negative"
+        )
+
+        self.archive_paths = [Path(path) for path in archive_paths]
+        self.batch_size = batch_size
+        self.max_member_bytes = max_member_bytes
+
+    @property
+    def name(self) -> str:
+        return "TAR Data Source"
+
+    @property
+    def schema(self) -> Schema:
+        return TAR_SCHEMA
+
+    async def get_tasks(
+        self,
+        pushdowns: Pushdowns,
+    ) -> AsyncIterator["TarDataSourceTask"]:
+        for archive_path in self.archive_paths:
+            yield TarDataSourceTask(
+                archive_path,
+                batch_size=self.batch_size,
+                max_member_bytes=self.max_member_bytes,
+            )

daft_sources/tar/task.py

@@ -0,0 +1,81 @@
+import tarfile
+from collections.abc import AsyncIterator
+from pathlib import Path
+
+from daft.io import DataSourceTask
+from daft.recordbatch import RecordBatch
+from daft.schema import Schema
+
+from daft_sources.tar.schema import TAR_SCHEMA
+
+
+class TarDataSourceTask(DataSourceTask):
+    """Read one TAR archive."""
+
+    def __init__(
+        self,
+        archive_path: Path,
+        *,
+        batch_size: int,
+        max_member_bytes: int | None,
+    ) -> None:
+        self.archive_path = archive_path
+        self.batch_size = batch_size
+        self.max_member_bytes = max_member_bytes
+
+    @property
+    def schema(self) -> Schema:
+        return TAR_SCHEMA
+
+    async def read(self) -> AsyncIterator[RecordBatch]:
+        columns: dict[str, list[object]] = {
+            "archive_path": [],
+            "member_path": [],
+            "size": [],
+            "mtime": [],
+            "mode": [],
+            "uid": [],
+            "gid": [],
+            "uname": [],
+            "gname": [],
+            "bytes": [],
+        }
+
+        with tarfile.open(self.archive_path, mode="r:*") as archive:
+            for member in archive:
+                if not member.isfile():
+                    continue
+
+                if (
+                    self.max_member_bytes is not None
+                    and member.size > self.max_member_bytes
+                ):
+                    continue
+
+                data: bytes | None = None
+
+                extracted = archive.extractfile(member)
+
+                if extracted is None:
+                    continue
+
+                with extracted:
+                    data = extracted.read()
+
+                columns["archive_path"].append(str(self.archive_path))
+                columns["member_path"].append(member.name)
+                columns["size"].append(member.size)
+                columns["mtime"].append(float(member.mtime))
+                columns["mode"].append(member.mode)
+                columns["uid"].append(member.uid)
+                columns["gid"].append(member.gid)
+                columns["uname"].append(member.uname or None)
+                columns["gname"].append(member.gname or None)
+                columns["bytes"].append(data)
+
+                if len(columns["member_path"]) >= self.batch_size:
+                    yield RecordBatch.from_pydict(columns)
+                    columns = {name: [] for name in columns}
+
+        if columns["member_path"]:
+            yield RecordBatch.from_pydict(columns)

daft_sources/wds/schema.py

@@ -0,0 +1,11 @@
+from daft.datatype import DataType
+from daft.schema import Schema
+
+WDS_SCHEMA = Schema._from_field_name_and_types(
+    [
+        ("__key__", DataType.string()),
+        ("__url__", DataType.string()),
+        ("jpg", DataType.binary()),
+        ("txt", DataType.string()),
+    ]
+)

daft_sources/wds/source.py

@@ -0,0 +1,42 @@
+from collections.abc import AsyncIterator
+from pathlib import Path
+
+from daft.io import DataSource
+from daft.io.pushdowns import Pushdowns
+from daft.schema import Schema
+
+from daft_sources.wds.schema import WDS_SCHEMA
+from daft_sources.wds.task import WebDatasetDataSourceTask
+
+
+class WebDatasetDataSource(DataSource):
+    def __init__(
+        self,
+        shard_paths: list[str],
+        *,
+        batch_size: int = 256,
+    ):
+        assert shard_paths, "At least one WebDataset shard is required"
+        assert batch_size > 0, "batch_size must be greater than zero"
+
+        self.shard_paths = [Path(path) for path in shard_paths]
+        self.batch_size = batch_size
+
+    @property
+    def name(self) -> str:
+        return "WebDataset Data Source"
+
+    @property
+    def schema(self) -> Schema:
+        return WDS_SCHEMA
+
+    async def get_tasks(
+        self,
+        pushdowns: Pushdowns,
+    ) -> AsyncIterator[WebDatasetDataSourceTask]:
+        for shard_path in self.shard_paths:
+            yield WebDatasetDataSourceTask(
+                shard_path=shard_path,
+                batch_size=self.batch_size,
+                schema=self.schema,
+            )

daft_sources/wds/task.py

@@ -0,0 +1,60 @@
+from pathlib import Path
+
+import webdataset as wds
+from daft.io import DataSourceTask
+from daft.recordbatch import RecordBatch
+from daft.schema import Schema
+
+
+class WebDatasetDataSourceTask(DataSourceTask):
+    def __init__(
+        self,
+        *,
+        shard_path: Path,
+        batch_size: int,
+        schema: Schema,
+    ):
+        self.shard_path = shard_path
+        self.batch_size = batch_size
+        self._schema = schema
+
+    @property
+    def schema(self) -> Schema:
+        return self._schema
+
+    async def read(self):
+        dataset = wds.WebDataset(
+            str(self.shard_path),
+            shardshuffle=False,
+            nodesplitter=None,
+            workersplitter=None,
+        )
+
+        columns = self._empty_columns()
+
+        for sample in dataset:
+            txt = sample.get("txt")
+
+            if isinstance(txt, bytes):
+                txt = txt.decode("utf-8")
+
+            columns["__key__"].append(sample["__key__"])
+            columns["__url__"].append(sample.get("__url__", str(self.shard_path)))
+            columns["jpg"].append(sample.get("jpg"))
+            columns["txt"].append(txt)
+
+            if len(columns["__key__"]) >= self.batch_size:
+                yield RecordBatch.from_pydict(columns)
+                columns = self._empty_columns()
+
+        if columns["__key__"]:
+            yield RecordBatch.from_pydict(columns)
+
+    @staticmethod
+    def _empty_columns() -> dict[str, list]:
+        return {
+            "__key__": [],
+            "__url__": [],
+            "jpg": [],
+            "txt": [],
+        }

daft_sources-0.1.0.dist-info/METADATA

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: daft-sources
+Version: 0.1.0
+Summary: Add your description here
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: daft>=0.7.21
+Requires-Dist: google-re2>=1.1.20251105
+Requires-Dist: webdataset>=1.0.2
+
+
+```bash
+uvx ruff format
+uvx ruff check --select I --fix .
+```

daft_sources-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+daft_sources/tar/schema.py,sha256=gAzNjLPuf550wlQHSSfBCFRq40X4OR923jhGdN8MYRM,481
+daft_sources/tar/source.py,sha256=eYNOLSmKZ7kDWAR48NG6M-UzBqEamL6UszChJVQhC4g,1322
+daft_sources/tar/task.py,sha256=CB8h8tCqlMLICzVliSyd7JDWvlKnmFu-C_yFSNvkD1o,2476
+daft_sources/wds/schema.py,sha256=lEjflgZ-5NO_Uw62YlcW5ABunImsdjRZXyX2Ckmivcs,281
+daft_sources/wds/source.py,sha256=Bp0NzKHemnlXjlHaLNjXMf2nD7_BiS9fkwa2Crtd0aU,1172
+daft_sources/wds/task.py,sha256=uq20LswXVZASqEVCLK2BOPno6xt1m5W33AznmWnBjHs,1574
+daft_sources-0.1.0.dist-info/METADATA,sha256=tcsR1CKiuEv5CiSzEcjKay6jUMOBG43m0nX6y-6cg7U,320
+daft_sources-0.1.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+daft_sources-0.1.0.dist-info/top_level.txt,sha256=Jk7UhqTe3ZAHu1s0oAZL78J82zSa5TbDkBpLiRAaBr0,13
+daft_sources-0.1.0.dist-info/RECORD,,

daft_sources-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (83.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

daft_sources-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+daft_sources