Address second review round: clearer docs, eager check in SqliteITarEntryReader

- data_prep.md: replace "consumed sequentially" with the precise constraint
  ("not used as auxiliary or mounted").
- SqliteIndexReader: expose has_sample_tables as a constructor-time attribute
  (mirrors db_has_sample_parts); drop the per-method _check_samples_table guard.
- SqliteITarEntryReader: raise MissingSamplesTableError at __init__ when the
  samples table is missing — fail fast at the boundary that actually requires it.
- Test updated to assert at SqliteITarEntryReader construction.

Signed-off-by: Pei Li <pei.li@kaiko.ai>

Author: pei-li-hedgehog

docs/source/basic/data_prep.md

@@ -688,13 +688,13 @@ The `media_metadata` table is used to store the media metadata for the selected
 
 #### Skipping the samples tables for very large datasets
 
-For very large datasets (100M+ samples), populating the `samples` and `sample_parts` tables can dominate `energon prepare` runtime. If the dataset will only be consumed sequentially (not as an auxiliary dataset), you can skip these tables:
+For very large datasets (100M+ samples), populating the `samples` and `sample_parts` tables can dominate `energon prepare` runtime. If the dataset will not be used as an auxiliary dataset or mounted, you can skip these tables:
 
 ```sh
 > energon prepare --no-sample-tables /path/to/dataset
 ```
 
-A dataset prepared this way cannot be used as an [auxiliary dataset](aux-data).
+A dataset prepared this way cannot be used as an [auxiliary dataset](aux-data) or with `energon mount`.
 
 (data-on-disk-jsonl)=
 ## Dataset Format on Disk for JSONL Datasets

src/megatron/energon/flavors/webdataset/indexing.py

@@ -397,7 +397,6 @@ class SqliteIndexReader:
 
     sqlite_path: EPath
     db: ThreadLocalSqlite
-    _samples_table_checked: bool
 
     def __init__(self, sqlite_path: EPath):
         """Initialize the SQLite database reader.
@@ -413,25 +412,20 @@ def __init__(self, sqlite_path: EPath):
         path = f"file:{path}?mode=ro&immutable=1"
 
         self.db = ThreadLocalSqlite(path, is_uri=True)
-        self._samples_table_checked = False
 
-    def _check_samples_table(self) -> None:
-        """Verify the ``samples`` table is present, raising a clear error if not.
+    def db_has_sample_tables(self) -> bool:
+        """Check if the database has a samples table.
 
-        Called by every method that queries the ``samples`` / ``sample_parts`` tables, so callers
-        accessing a dataset prepared with ``--no-sample-tables`` get a descriptive error instead
-        of a raw ``sqlite3.OperationalError: no such table: samples``. The check runs once per
-        reader instance.
+        Returns:
+            True if samples table exists, False otherwise.
         """
-        if self._samples_table_checked:
-            return
         assert self.db is not None, "Database is closed"
-        row = self.db.select_one(
+
+        db_exists = self.db.select_one(
             "SELECT name FROM sqlite_master WHERE type='table' AND name='samples'"
         )
-        if row is None:
-            raise MissingSamplesTableError(self.sqlite_path)
-        self._samples_table_checked = True
+        self.db.thread_close()
+        return db_exists is not None
 
     def db_has_sample_parts(self) -> bool:
         """Check if the database has a sample_parts table.
@@ -466,7 +460,6 @@ def list_all_samples(self) -> Generator[Tuple[str, int, int], None, None]:
         """
 
         assert self.db is not None, "Database is closed"
-        self._check_samples_table()
 
         for row in self.db.select_all("SELECT sample_key, byte_size, tar_file_id FROM samples"):
             yield row[0], row[1], row[2]
@@ -479,7 +472,6 @@ def list_all_sample_parts(self) -> Generator[Tuple[str, int, int], None, None]:
         """
 
         assert self.db is not None, "Database is closed"
-        self._check_samples_table()
 
         # Select all parts (sorted by tar_file_id, sample_index) but joined with the sample_key names
         for row in self.db.select_all(
@@ -505,7 +497,6 @@ def list_sample_parts(self, sample_key: str) -> Generator[Tuple[str, int, int],
         """
 
         assert self.db is not None, "Database is closed"
-        self._check_samples_table()
 
         # Select all parts (sorted by tar_file_id, sample_index) but joined with the sample_key names
         for row in self.db.select_all(
@@ -525,15 +516,13 @@ def list_sample_parts(self, sample_key: str) -> Generator[Tuple[str, int, int],
     def get_total_size(self) -> int:
         """Get the total size of all samples in the database."""
         assert self.db is not None, "Database is closed"
-        self._check_samples_table()
 
         count = self.db.select_one("SELECT SUM(byte_size) FROM samples")
         return count[0] if count else 0
 
     def get_sample_count(self) -> int:
         """Get the total number of samples in the database."""
         assert self.db is not None, "Database is closed"
-        self._check_samples_table()
 
         count = self.db.select_one("SELECT COUNT(*) FROM samples")
         return count[0] if count else 0
@@ -549,7 +538,6 @@ def get_sample_part(self, key: str, part_name: str) -> ITarRawSamplePartPointer:
             Pointer to the sample part raw data.
         """
         assert self.db is not None, "Database is closed"
-        self._check_samples_table()
 
         row = self.db.select_one(
             "SELECT sp.tar_file_id, sp.content_byte_offset, sp.content_byte_size "
@@ -579,7 +567,6 @@ def get_sample_pointer_by_key(self, key: str) -> ITarSamplePointer:
             Tuple of (tar_file_id, sample_key, sample_index, byte_offset, byte_size)
         """
         assert self.db is not None, "Database is closed"
-        self._check_samples_table()
 
         sample = self.db.select_one(
             "SELECT tar_file_id, sample_key, sample_index, byte_offset, byte_size "

src/megatron/energon/flavors/webdataset/itar_reader.py

@@ -484,7 +484,8 @@ class SqliteITarEntryReader(ITarReader[str]):
     """
 
     sqlite_reader: SqliteIndexReader
-    db_has_sample_parts: int
+    db_has_sample_parts: bool
+    db_has_sample_tables: bool
 
     def __init__(
         self,
@@ -496,7 +497,10 @@ def __init__(
         disable_cache: bool = False,
     ):
         from megatron.energon.flavors.webdataset.config import MAIN_FOLDER_NAME
-        from megatron.energon.flavors.webdataset.indexing import SqliteIndexReader
+        from megatron.energon.flavors.webdataset.indexing import (
+            MissingSamplesTableError,
+            SqliteIndexReader,
+        )
 
         # shard_name_to_info_idx = {name: i for i, name in enumerate(wds_meta.info_shard_files)}
         tar_filenames = get_info_shard_files(base_path)
@@ -507,6 +511,10 @@ def __init__(
         self.sqlite_reader = SqliteIndexReader(sqlite_path)
 
         self.db_has_sample_parts = self.sqlite_reader.db_has_sample_parts()
+        self.db_has_sample_tables = self.sqlite_reader.db_has_sample_tables()
+
+        if not self.db_has_sample_tables:
+            raise MissingSamplesTableError(sqlite_path)
 
         self.key_is_full_entryname = key_is_full_entryname
 

tests/test_dataset.py

@@ -50,6 +50,7 @@
 from megatron.energon.edataclass import edataclass
 from megatron.energon.flavors import BaseWebdatasetFactory
 from megatron.energon.flavors.webdataset.config import MAIN_FOLDER_NAME
+from megatron.energon.flavors.webdataset.itar_reader import SqliteITarEntryReader
 from megatron.energon.task_encoder.base import stateless
 from megatron.energon.tools.analyze_debug import command as analyze_debug_command
 from megatron.energon.tools.info import command as info_command
@@ -1862,25 +1863,19 @@ def test_prepare_dataset_no_sample_tables(self):
             result.stdout + (result.stderr or "")
         ) or isinstance(result.exception, click.UsageError)
 
-        # Sample-key operations against the SQLite reader must raise a descriptive error,
-        # not the raw `sqlite3.OperationalError: no such table: samples`.
+        # SqliteIndexReader exposes db_has_sample_tables() as part of its public surface.
         from megatron.energon.epathlib import EPath
         from megatron.energon.flavors.webdataset.indexing import (
             MissingSamplesTableError,
             SqliteIndexReader,
         )
 
-        reader = SqliteIndexReader(EPath(str(sqlite_path)))
-        try:
-            with self.assertRaises(MissingSamplesTableError) as ctx:
-                reader.get_sample_pointer_by_key("any-key")
-            assert "no-sample-tables" in str(ctx.exception)
-            with self.assertRaises(MissingSamplesTableError):
-                reader.get_sample_count()
-            with self.assertRaises(MissingSamplesTableError):
-                reader.get_sample_part("any-key", "txt")
-        finally:
-            reader.close()
+        index_reader = SqliteIndexReader(EPath(str(sqlite_path)))
+        assert index_reader.db_has_sample_tables() is False
+
+        with self.assertRaises(MissingSamplesTableError) as ctx:
+            SqliteITarEntryReader(EPath(str(self.dataset_path)))
+        assert "no-sample-tables" in str(ctx.exception)
 
     def test_prepare_dataset_no_sample_tables_save_restore(self):
         """Resume after a mid-iteration checkpoint must reach the same samples in the
@@ -1894,8 +1889,6 @@ def test_prepare_dataset_no_sample_tables_save_restore(self):
         a save/restore round-trip against a reference run.
         """
 
-        from megatron.energon import get_savable_loader, get_train_dataset
-
         runner = CliRunner()
         result = runner.invoke(
             prepare_command,