mldataforge 0.1.5__tar.gz → 0.1.7__tar.gz

{mldataforge-0.1.5 → mldataforge-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mldataforge
-Version: 0.1.5
+Version: 0.1.7
 Summary: swiss army knife of scripts for transforming and processing datasets for machine learning.
 Project-URL: Homepage, https://github.com/schneiderkamplab/mldataforge
 Project-URL: Bug Tracker, https://github.com/schneiderkamplab/mldataforge/issues

{mldataforge-0.1.5 → mldataforge-0.1.7}/mldataforge/commands/join.py

@@ -53,7 +53,7 @@ def join_mds(output_dir, mds_directories, compression, processes, overwrite, yes
         compression=compression,
         buf_size=buf_size,
         shard_size=shard_size,
-        pigz=use_pigz(compression, no_pigz)
+        pigz=use_pigz(compression, no_pigz),
     )
 
 @join.command()

{mldataforge-0.1.5 → mldataforge-0.1.7}/mldataforge/compression.py

@@ -30,7 +30,7 @@ JSONL_COMPRESSIONS = dict(
 )
 MDS_COMPRESSIONS = dict(
     default=None,
-    choices=["none", "brotli", "bz2", "gzip", "pigz", "snappy", "zstd"],
+    choices=["none", "brotli", "bz2", "gzip", "pigz", "snappy", "zstd", "sample::brotli", "sample::bz2", "sample::gzip", "sample::snappy", "sample::zstd"],
 )
 PARQUET_COMPRESSIONS = dict(
     default="snappy",
@@ -55,6 +55,10 @@ def determine_compression(fmt, file_path, compression="infer", no_pigz=False):
             return "gz"
         if compression == "brotli":
             return "br"
+        if compression == "sample::gzip":
+            return "gz"
+        if compression == "sample::brotli":
+            return "br"
         return compression
     if fmt == "parquet":
         return compression

mldataforge-0.1.7/mldataforge/mds.py

@@ -0,0 +1,448 @@
+from copy import deepcopy
+import json
+import numpy as np
+import os
+import shutil
+from streaming.base.compression import compress, decompress, get_compression_extension, is_compression
+from streaming.base.format import _readers
+from streaming.base.format.base.reader import FileInfo, JointReader
+from streaming.base.format.index import get_index_basename
+from streaming.base.format.mds.encodings import mds_decode, mds_encode, is_mds_encoding, is_mds_encoding_safe, get_mds_encodings, get_mds_encoded_size
+from streaming.base.hashing import get_hash, is_hash
+from streaming.base.util import bytes_to_int
+from typing import Any, Optional, Generator, Self, Union
+
+from .utils import open_compression
+
+__all__ = [
+    "MDSBulkReader",
+    "MDSBulkShardReader",
+    "MDSReader",
+    "MDSWriter",
+]
+
+class MDSBulkReader:
+    def __init__(
+        self,
+        dirnames: list[str],
+        split: Optional[str],
+    ) -> None:
+        self.shards = []
+        self.samples = 0
+        for dirname in dirnames:
+            if split is not None:
+                dirname = os.path.join(dirname, split)
+            index = json.load(open(os.path.join(dirname, "index.json"), 'rt'))
+            for shard in index["shards"]:
+                basename = shard['raw_data']['basename'] if shard['zip_data'] is None else shard['zip_data']['basename']
+                filename = os.path.join(dirname, basename)
+                self.shards.append({
+                    "filename": filename,
+                    "compression": shard['compression'],
+                })
+                self.samples += shard['samples']
+
+    def __len__(self) -> int:
+        return self.samples
+
+    def __iter__(self) -> Generator[dict[str, Any], None, None]:
+        for shard in self.shards:
+            with MDSBulkShardReader(**shard) as reader:
+                for sample in reader:
+                    yield sample
+
+class MDSBulkShardReader:
+    def __init__(
+        self,
+        filename: str,
+        compression: Optional[str],
+    ) -> None:
+        self.sample_compression = None
+        if compression is not None and compression.startswith("sample::"):
+            compression, self.sample_compression = None, compression.removeprefix("sample::")
+        self.fp = open_compression(filename, "rb", compression=compression)
+        self.samples = np.frombuffer(self.fp.read(4), np.uint32)[0]
+        self.index = np.frombuffer(self.fp.read((1+self.samples)*4), np.uint32)
+        info = json.loads(self.fp.read(self.index[0]-self.fp.tell()))
+        self.column_encodings = info["column_encodings"]
+        self.column_names = info["column_names"]
+        self.column_sizes = info["column_sizes"]
+        assert self.fp.tell() == self.index[0]
+
+    def decode_sample(self, data: bytes) -> dict[str, Any]:
+        sizes = []
+        idx = 0
+        for key, size in zip(self.column_names, self.column_sizes):
+            if size:
+                sizes.append(size)
+            else:
+                size, = np.frombuffer(data[idx:idx + 4], np.uint32)
+                sizes.append(size)
+                idx += 4
+        sample = {}
+        for key, encoding, size in zip(self.column_names, self.column_encodings, sizes):
+            value = data[idx:idx + size]
+            sample[key] = mds_decode(encoding, value)
+            idx += size
+        return sample
+
+    def get_sample_data(self, idx: int) -> bytes:
+        begin, end = self.index[idx:idx+2]
+        assert self.fp.tell() == begin
+        data = self.fp.read(end - begin)
+        assert self.fp.tell() == end
+        assert data
+        return data
+
+    def get_item(self, idx: int) -> dict[str, Any]:
+        data = self.get_sample_data(idx)
+        if self.sample_compression is not None:
+            data = decompress(self.sample_compression, data)
+        return self.decode_sample(data)
+
+    def __iter__(self) -> Generator[dict[str, Any], None, None]:
+        for i in range(self.samples):
+            yield self.get_item(i)
+
+    def __enter__(self) -> "MDSBulkShardReader":
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        self.fp.close()
+
+
+class MDSWriter:
+
+    format = 'mds'
+    extra_bytes_per_sample = 4
+
+    def __init__(self,
+                 *,
+                 columns: dict[str, str],
+                 out: Union[str, tuple[str, str]],
+                 compression: Optional[str] = None,
+                 hashes: Optional[list[str]] = None,
+                 size_limit: Optional[Union[int, str]] = 1 << 26,
+                 **kwargs: Any) -> None:
+        compression = compression or None
+        sample_compression = None
+        if compression is not None and compression.startswith("sample::"):
+            compression, sample_compression = None, compression.removeprefix("sample::")
+        if compression:
+            if not is_compression(compression):
+                raise ValueError(f'Invalid compression: {compression}.')
+        if sample_compression:
+            if not is_compression(sample_compression):
+                raise ValueError(f'Invalid sample compression: {sample_compression}.')
+        hashes = hashes or []
+        if list(hashes) != sorted(hashes):
+            raise ValueError('Hashes must be unique and in sorted order.')
+        for algo in hashes:
+            if not is_hash(algo):
+                raise ValueError(f'Invalid hash: {algo}.')
+
+        size_limit_value = None
+        if size_limit:
+            size_limit_value = bytes_to_int(size_limit)
+            if size_limit_value < 0:
+                raise ValueError(f'`size_limit` must be greater than zero, instead, ' +
+                                 f'found as {size_limit_value}.')
+            if size_limit_value >= 2**32:
+                raise ValueError(f'`size_limit` must be less than 2**32, instead, ' +
+                                 f'found as {size_limit_value}. This is because sample ' +
+                                 f'byte offsets are stored with uint32.')
+
+        # Validate keyword arguments
+        invalid_kwargs = [
+            arg for arg in kwargs.keys()
+            if arg not in ('progress_bar', 'exist_ok')
+        ]
+        if invalid_kwargs:
+            raise ValueError(f'Invalid Writer argument(s): {invalid_kwargs} ')
+
+        self.compression = compression
+        self.sample_compression = sample_compression
+        self.hashes = hashes
+        self.size_limit = size_limit_value
+        self.new_samples: list[bytes]
+        self.new_shard_size: int
+
+        self.shards = []
+
+        # Remove local directory if requested prior to creating writer
+        self.local = os.path.expanduser(out)
+        if os.path.exists(self.local) and len(os.listdir(self.local)) != 0:
+            if kwargs.get('exist_ok', False):
+                raise FileExistsError(f'Directory is not empty: {self.local}')
+            shutil.rmtree(self.local)
+        os.makedirs(self.local, exist_ok=True)
+    
+        self.columns = columns
+        self.column_names = []
+        self.column_encodings = []
+        self.column_sizes = []
+        for name in sorted(columns):
+            encoding = columns[name]
+            if not is_mds_encoding(encoding):
+                raise TypeError(f'MDSWriter passed column `{name}` with encoding `{encoding}` ' +
+                                f'is unsupported. Supported encodings are {get_mds_encodings()}')
+            size = get_mds_encoded_size(encoding)
+            self.column_names.append(name)
+            self.column_encodings.append(encoding)
+            self.column_sizes.append(size)
+
+        obj = self.get_config()
+        text = json.dumps(obj, sort_keys=True)
+        self.config_data = text.encode('utf-8')
+        self.extra_bytes_per_shard = 4 + 4 + len(self.config_data)
+        self._reset_cache()
+
+    def encode_sample(self, sample: dict[str, Any]) -> bytes:
+        sizes = []
+        data = []
+        for key, encoding, size in zip(self.column_names, self.column_encodings,
+                                       self.column_sizes):
+            value = sample[key]
+            datum = mds_encode(encoding, value)
+            if size is None:
+                size = len(datum)
+                sizes.append(size)
+            else:
+                if size != len(datum):
+                    raise KeyError(f'Unexpected data size; was this data typed with the correct ' +
+                                   f'encoding ({encoding})?')
+            data.append(datum)
+        head = np.array(sizes, np.uint32).tobytes()
+        body = b''.join(data)
+        sample_data = head + body
+        if self.sample_compression:
+            sample_data = compress(self.sample_compression, sample_data)
+        return sample_data
+
+    def encode_joint_shard(self) -> bytes:
+        num_samples = np.uint32(len(self.new_samples))
+        sizes = list(map(len, self.new_samples))
+        offsets = np.array([0] + sizes).cumsum().astype(np.uint32)
+        offsets += len(num_samples.tobytes()) + len(offsets.tobytes()) + len(self.config_data)
+        sample_data = b''.join(self.new_samples)
+        return num_samples.tobytes() + offsets.tobytes() + self.config_data + sample_data
+
+    def flush_shard(self) -> None:
+        raw_data_basename, zip_data_basename = self._name_next_shard()
+        raw_data = self.encode_joint_shard()
+        raw_data_info, zip_data_info = self._process_file(raw_data, raw_data_basename,
+                                                          zip_data_basename)
+        obj = {
+            'samples': len(self.new_samples),
+            'raw_data': raw_data_info,
+            'zip_data': zip_data_info
+        }
+        obj.update(self.get_config())
+        self.shards.append(obj)
+
+    def _reset_cache(self) -> None:
+        self.new_samples = []
+        self.new_shard_size = self.extra_bytes_per_shard
+
+    def _name_next_shard(self, extension: Optional[str] = None) -> tuple[str, Optional[str]]:
+        shard = len(self.shards)
+        parts = ['shard', f'{shard:05}', self.format]
+        if extension:
+            parts.append(extension)
+        raw_basename = '.'.join(parts)
+        if self.compression:
+            ext = get_compression_extension(self.compression)
+            parts.append(ext)
+            zip_basename = '.'.join(parts)
+        else:
+            zip_basename = None
+        return raw_basename, zip_basename
+
+    def _hash(self, data: bytes, basename: str) -> dict[str, Any]:
+        hashes = {}
+        for algo in self.hashes:
+            hashes[algo] = get_hash(algo, data)
+        return {'basename': basename, 'bytes': len(data), 'hashes': hashes}
+
+    def _process_file(self, raw_data: bytes, raw_basename: str,
+                      zip_basename: Optional[str]) -> tuple[dict, Optional[dict]]:
+        raw_info = self._hash(raw_data, raw_basename)
+        if zip_basename:
+            zip_data = compress(self.compression, raw_data)
+            zip_info = self._hash(zip_data, zip_basename)
+            data = zip_data
+            basename = zip_basename
+        else:
+            zip_info = None
+            data = raw_data
+            basename = raw_basename
+        filename = os.path.join(self.local, basename)
+        with open(filename, 'wb') as out:
+            out.write(data)
+        return raw_info, zip_info
+
+    def get_config(self) -> dict[str, Any]:
+        return {
+            'version': 2,
+            'format': self.format,
+            'compression': self.compression if self.sample_compression is None else f"sample::{self.sample_compression}",
+            'hashes': self.hashes,
+            'size_limit': self.size_limit,
+            'column_names': self.column_names,
+            'column_encodings': self.column_encodings,
+            'column_sizes': self.column_sizes,
+        }
+
+    def write(self, sample: dict[str, Any]) -> None:
+        new_sample = self.encode_sample(sample)
+        new_sample_size = len(new_sample) + self.extra_bytes_per_sample
+        if self.size_limit and self.size_limit < self.new_shard_size + new_sample_size:
+            self.flush_shard()
+            self._reset_cache()
+        self.new_samples.append(new_sample)
+        self.new_shard_size += new_sample_size
+
+    def _write_index(self) -> None:
+        if self.new_samples:
+            raise RuntimeError('Internal error: not all samples have been written.')
+        basename = get_index_basename()
+        filename = os.path.join(self.local, basename)
+        obj = {
+            'version': 2,
+            'shards': self.shards,
+        }
+        with open(filename, 'w') as out:
+            json.dump(obj, out, sort_keys=True)
+
+    def finish(self) -> None:
+        if self.new_samples:
+            self.flush_shard()
+            self._reset_cache()
+        self._write_index()
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, exc_type, exc, traceback):
+        self.finish()
+
+class MDSReader(JointReader):
+
+    def __init__(
+        self,
+        dirname: str,
+        split: Optional[str],
+        column_encodings: list[str],
+        column_names: list[str],
+        column_sizes: list[Optional[int]],
+        compression: Optional[str],
+        hashes: list[str],
+        raw_data: FileInfo,
+        samples: int,
+        size_limit: Optional[Union[int, str]],
+        zip_data: Optional[FileInfo],
+    ) -> None:
+        self.sample_compression = None
+        if compression and compression.startswith("sample::"):
+            compression, self.sample_compression = None, compression.removeprefix("sample::")
+        super().__init__(dirname, split, compression, hashes, raw_data, samples, size_limit,
+                         zip_data)
+        self.column_encodings = column_encodings
+        self.column_names = column_names
+        self.column_sizes = column_sizes
+
+    @classmethod
+    def from_json(cls, dirname: str, split: Optional[str], obj: dict[str, Any]) -> Self:
+        """Initialize from JSON object.
+
+        Args:
+            dirname (str): Local directory containing shards.
+            split (str, optional): Which dataset split to use, if any.
+            obj (Dict[str, Any]): JSON object to load.
+
+        Returns:
+            Self: Loaded MDSReader.
+        """
+        args = deepcopy(obj)
+        args_version = args['version']
+        if args_version != 2:
+            raise ValueError(
+                f'Unsupported streaming data version: {args_version}. Expected version 2.')
+        del args['version']
+        args_format = args['format']
+        if args_format != 'mds':
+            raise ValueError(f'Unsupported data format: {args_format}. Expected to be `mds`.')
+        del args['format']
+        args['dirname'] = dirname
+        args['split'] = split
+        for key in ['raw_data', 'zip_data']:
+            arg = args[key]
+            args[key] = FileInfo(**arg) if arg else None
+        return cls(**args)
+
+    def validate(self, allow_unsafe_types: bool) -> None:
+        """Check whether this shard is acceptable to be part of some Stream.
+
+        Args:
+            allow_unsafe_types (bool): If a shard contains Pickle, which allows arbitrary code
+                execution during deserialization, whether to keep going if ``True`` or raise an
+                error if ``False``.
+        """
+        if not allow_unsafe_types:
+            for column_id, encoding in enumerate(self.column_encodings):
+                if not is_mds_encoding_safe(encoding):
+                    name = self.column_names[column_id]
+                    raise ValueError(f'Column {name} contains an unsafe type: {encoding}. To ' +
+                                     f'proceed anyway, set ``allow_unsafe_types=True``.')
+
+    def decode_sample(self, data: bytes) -> dict[str, Any]:
+        """Decode a sample dict from bytes.
+
+        Args:
+            data (bytes): The sample encoded as bytes.
+
+        Returns:
+            Dict[str, Any]: Sample dict.
+        """
+        sizes = []
+        idx = 0
+        for key, size in zip(self.column_names, self.column_sizes):
+            if size:
+                sizes.append(size)
+            else:
+                size, = np.frombuffer(data[idx:idx + 4], np.uint32)
+                sizes.append(size)
+                idx += 4
+        sample = {}
+        for key, encoding, size in zip(self.column_names, self.column_encodings, sizes):
+            value = data[idx:idx + size]
+            sample[key] = mds_decode(encoding, value)
+            idx += size
+        return sample
+
+    def get_sample_data(self, idx: int) -> bytes:
+        """Get the raw sample data at the index.
+
+        Args:
+            idx (int): Sample index.
+
+        Returns:
+            bytes: Sample data.
+        """
+        filename = os.path.join(self.dirname, self.split, self.raw_data.basename)
+        offset = (1 + idx) * 4
+        with open(filename, 'rb', 0) as fp:
+            fp.seek(offset)
+            pair = fp.read(8)
+            begin, end = np.frombuffer(pair, np.uint32)
+            fp.seek(begin)
+            data = fp.read(end - begin)
+        if not data:
+            raise IndexError(
+                f'Relative sample index {idx} is not present in the {self.raw_data.basename} file.'
+            )
+        if self.sample_compression:
+            data = decompress(self.sample_compression, data)
+        return data
+
+_readers["mds"] = MDSReader

{mldataforge-0.1.5 → mldataforge-0.1.7}/mldataforge/snappy.py

@@ -36,60 +36,6 @@ class _SnappyWriteWrapper(io.RawIOBase):
     def writable(self):
         return True
 
-
-# class _SnappyReadWrapper(io.RawIOBase):
-#     def __init__(self, fileobj):
-#         self.fileobj = fileobj
-#         self.buffer = io.BytesIO()
-#         self.eof = False
-
-#     def _fill_buffer_if_needed(self, min_bytes):
-#         self.buffer.seek(0, io.SEEK_END)
-#         while not self.eof and self.buffer.tell() < min_bytes:
-#             length_bytes = self.fileobj.read(4)
-#             if not length_bytes:
-#                 self.eof = True
-#                 break
-#             if len(length_bytes) < 4:
-#                 self.eof = True  # mark as EOF even if last chunk is malformed
-#                 break
-
-#             try:
-#                 length = struct.unpack(">I", length_bytes)[0]
-#                 compressed = self.fileobj.read(length)
-#                 if len(compressed) < length:
-#                     self.eof = True
-#                     break
-
-#                 decompressed = snappy.decompress(compressed)
-#                 self.buffer.write(decompressed)
-#             except Exception:
-#                 self.eof = True
-#                 break
-
-#         self.buffer.seek(0)
-
-#     def read(self, size=-1):
-#         if size == -1:
-#             while not self.eof:
-#                 self._fill_buffer_if_needed(_CHUNK_SIZE)
-#             result = self.buffer.read()
-#             self.buffer = io.BytesIO()
-#             return result
-
-#         self._fill_buffer_if_needed(size)
-#         data = self.buffer.read(size)
-#         rest = self.buffer.read()
-#         self.buffer = io.BytesIO()
-#         self.buffer.write(rest)
-#         return data
-
-#     def readable(self):
-#         return True
-
-#     def close(self):
-#         self.fileobj.close()
-
 class _SnappyReadWrapper(io.RawIOBase):
     def __init__(self, fileobj):
         self.fileobj = fileobj

{mldataforge-0.1.5 → mldataforge-0.1.7}/mldataforge/utils.py

@@ -6,11 +6,11 @@ import pyarrow as pa
 import pyarrow.parquet as pq
 import os
 import shutil
-from streaming import MDSWriter, StreamingDataset
+from streaming import StreamingDataset
 from tqdm import tqdm
 
 from .compression import determine_compression, open_compression, pigz_compress
-from .mds import MDSBulkReader
+from .mds import MDSBulkReader, MDSWriter
 from .pigz import pigz_open
 
 __all__ = [
@@ -23,6 +23,11 @@ __all__ = [
     "save_parquet",
 ]
 
+_NO_PROGESS = False
+def set_progress(value):
+    global _NO_PROGESS
+    _NO_PROGESS = value
+
 def _batch_iterable(iterable, batch_size):
     batch = []
     for item in iterable:
@@ -73,7 +78,7 @@ def _infer_mds_encoding(value):
     return 'pkl'
 
 def _streaming_jsonl(jsonl_files, compressions):
-    for jsonl_file, compression in tqdm(zip(jsonl_files, compressions), desc="Loading JSONL files", unit="file"):
+    for jsonl_file, compression in tqdm(zip(jsonl_files, compressions), desc="Loading JSONL files", unit="file", disable=_NO_PROGESS):
         for line in open_compression(jsonl_file, mode="rt", compression=compression):
             yield json.loads(line)
 
@@ -109,7 +114,7 @@ def load_mds_directories(mds_directories, split='.', batch_size=2**16, bulk=True
 def save_jsonl(iterable, output_file, compression=None, processes=64, size_hint=None, overwrite=True, yes=True):
     f = None
     part = 0
-    for item in tqdm(iterable, desc="Writing to JSONL", unit="sample"):
+    for item in tqdm(iterable, desc="Writing to JSONL", unit="sample", disable=_NO_PROGESS):
         if f is None:
             part_file = output_file.format(part=part)
             check_arguments(part_file, overwrite, yes)
@@ -127,7 +132,7 @@ def save_mds(it, output_dir, processes=64, compression=None, buf_size=2**24, pig
     writer = None
     part = 0
     files = []
-    for sample in tqdm(it, desc="Writing to MDS", unit="sample"):
+    for sample in tqdm(it, desc="Writing to MDS", unit="sample", disable=_NO_PROGESS):
         if writer is None:
             part_dir = output_dir.format(part=part)
             check_arguments(part_dir, overwrite, yes)
@@ -151,7 +156,7 @@ def save_mds(it, output_dir, processes=64, compression=None, buf_size=2**24, pig
             name2info = {shard["raw_data"]["basename"]: shard for shard in index["shards"]}
             file_names = [file for file in os.listdir(output_dir) if file.endswith(".mds")]
             assert set(file_names) == set(name2info.keys())
-            for file_name in tqdm(file_names, desc="Compressing with pigz", unit="file"):
+            for file_name in tqdm(file_names, desc="Compressing with pigz", unit="file", disable=_NO_PROGESS):
                 compressed_file_name = file_name + ".gz"
                 file_path = os.path.join(output_dir, file_name)
                 compressed_file_path = os.path.join(output_dir, compressed_file_name)
@@ -169,7 +174,7 @@ def save_parquet(it, output_file, compression=None, batch_size=2**16, size_hint=
     compression = determine_compression("parquet", output_file, compression)
     writer = None
     part = 0
-    it = tqdm(it, desc="Writing to Parquet", unit="sample")
+    it = tqdm(it, desc="Writing to Parquet", unit="sample", disable=_NO_PROGESS)
     for batch in _batch_iterable(it, batch_size):
         table = pa.Table.from_pylist(batch)
         if writer is None:

{mldataforge-0.1.5 → mldataforge-0.1.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "mldataforge"
-version = "0.1.5"
+version = "0.1.7"
 authors = [
   { name = "Peter Schneider-Kamp" }
 ]

mldataforge-0.1.5/mldataforge/mds.py

@@ -1,95 +0,0 @@
-import gzip
-import json
-from mltiming import timing
-import numpy as np
-import os
-import snappy
-from streaming.base.format.mds.encodings import mds_decode
-from typing import Any, Optional, Generator
-
-from .options import MDS_COMPRESSIONS
-from .utils import open_compression
-
-class MDSBulkReader:
-    def __init__(
-        self,
-        dirnames: list[str],
-        split: Optional[str],
-    ) -> None:
-        self.shards = []
-        self.samples = 0
-        for dirname in dirnames:
-            if split is not None:
-                dirname = os.path.join(dirname, split)
-            index = json.load(open(os.path.join(dirname, "index.json"), 'rt'))
-            for shard in index["shards"]:
-                basename = shard['raw_data']['basename'] if shard['zip_data'] is None else shard['zip_data']['basename']
-                filename = os.path.join(dirname, basename)
-                self.shards.append({
-                    "filename": filename,
-                    "compression": shard['compression'],
-                })
-                self.samples += shard['samples']
-
-    def __len__(self) -> int:
-        return self.samples
-
-    def __iter__(self) -> Generator[dict[str, Any], None, None]:
-        for shard in self.shards:
-            with MDSShardReader(**shard) as reader:
-                for sample in reader:
-                    yield sample
-
-class MDSShardReader:
-    def __init__(
-        self,
-        filename: str,
-        compression: Optional[str],
-    ) -> None:
-        self.fp = open_compression(filename, "rb", compression=compression)
-        self.samples = np.frombuffer(self.fp.read(4), np.uint32)[0]
-        self.index = np.frombuffer(self.fp.read((1+self.samples)*4), np.uint32)
-        info = json.loads(self.fp.read(self.index[0]-self.fp.tell()))
-        self.column_encodings = info["column_encodings"]
-        self.column_names = info["column_names"]
-        self.column_sizes = info["column_sizes"]
-        assert self.fp.tell() == self.index[0]
-
-    def decode_sample(self, data: bytes) -> dict[str, Any]:
-        sizes = []
-        idx = 0
-        for key, size in zip(self.column_names, self.column_sizes):
-            if size:
-                sizes.append(size)
-            else:
-                size, = np.frombuffer(data[idx:idx + 4], np.uint32)
-                sizes.append(size)
-                idx += 4
-        sample = {}
-        for key, encoding, size in zip(self.column_names, self.column_encodings, sizes):
-            value = data[idx:idx + size]
-            sample[key] = mds_decode(encoding, value)
-            idx += size
-        return sample
-
-    def get_sample_data(self, idx: int) -> bytes:
-        begin, end = self.index[idx:idx+2]
-        assert self.fp.tell() == begin
-        data = self.fp.read(end - begin)
-        assert self.fp.tell() == end
-        assert data
-        return data
-
-    def get_item(self, idx: int) -> dict[str, Any]:
-        data = self.get_sample_data(idx)
-        return self.decode_sample(data)
-
-    def __iter__(self) -> Generator[dict[str, Any], None, None]:
-        for i in range(self.samples):
-            yield self.get_item(i)
-
-    def __enter__(self) -> "MDSShardReader":
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback) -> None:
-        self.fp.close()