This is an automated email from the ASF dual-hosted git repository. ColinLeeo pushed a commit to branch feature/trusted-index-mode in repository https://gitbox.apache.org/repos/asf/tsfile.git
commit 4159c635de8f29c973f56d06a11771c8c0121ff9 Author: ColinLee <[email protected]> AuthorDate: Fri Aug 28 15:53:51 2026 +0800 feat(python): add trusted dataset index mode --- python/tests/test_dataset_index.py | 83 ++++++++++++++++++++++++++++++++++++++ python/tsfile/dataset/dataframe.py | 50 +++++++++++++++++++++-- python/tsfile/dataset/index.py | 50 ++++++++++++++++++++--- python/tsfile/dataset/runtime.py | 49 +++++++++++++++++----- 4 files changed, 213 insertions(+), 19 deletions(-) diff --git a/python/tests/test_dataset_index.py b/python/tests/test_dataset_index.py index 91761a11a..09a52e3c0 100644 --- a/python/tests/test_dataset_index.py +++ b/python/tests/test_dataset_index.py @@ -311,6 +311,89 @@ def test_persistent_index_path_is_scoped_to_expanded_file_set(tmp_path, monkeypa np.testing.assert_array_equal(dataframe[0][:], np.array([0.0, 1.0, 10.0, 11.0])) +def test_trusted_index_skips_index_and_file_set_validation(tmp_path, monkeypatch): + source = tmp_path / "part.tsfile" + _write_runtime_file(source, 0) + + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as first: + assert len(first) == 1 + + def fail_validation(*_args, **_kwargs): + raise AssertionError("trusted construction must skip validation") + + monkeypatch.setattr(index_module, "index_matches_paths", fail_validation) + monkeypatch.setattr(index_module.MappedDatasetIndex, "_validate", fail_validation) + + with TsFileDataFrame( + str(source), show_progress=False, trust_index=True + ) as dataframe: + assert dataframe._use_index is True + assert dataframe._trust_index is True + assert len(dataframe) == 1 + + +def test_trusted_index_environment_override_enables_index(tmp_path, monkeypatch): + source = tmp_path / "part.tsfile" + _write_runtime_file(source, 0) + + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as first: + assert len(first) == 1 + + def fail_validation(*_args, **_kwargs): + raise AssertionError("trusted construction must skip validation") + + monkeypatch.setenv("TSFILE_DATAFRAME_TRUST_INDEX", "yes") + monkeypatch.setattr(index_module, "index_matches_paths", fail_validation) + monkeypatch.setattr(index_module.MappedDatasetIndex, "_validate", fail_validation) + + # The environment-level option also turns on the persistent-index path, so + # callers do not need to add both use_index=True and trust_index=True. + with TsFileDataFrame(str(source), show_progress=False) as dataframe: + assert dataframe._use_index is True + assert dataframe._trust_index is True + assert len(dataframe) == 1 + + +def test_explicit_trust_index_false_overrides_environment(tmp_path, monkeypatch): + source = tmp_path / "part.tsfile" + _write_runtime_file(source, 0) + monkeypatch.setenv("TSFILE_DATAFRAME_TRUST_INDEX", "1") + + with TsFileDataFrame( + str(source), show_progress=False, trust_index=False + ) as dataframe: + assert dataframe._trust_index is False + assert dataframe._runtime is None + assert len(dataframe) == 1 + + +def test_trusted_index_requires_an_existing_index(tmp_path): + source = tmp_path / "part.tsfile" + _write_runtime_file(source, 0) + + with pytest.raises(FileNotFoundError, match="Trusted Dataset Index not found"): + TsFileDataFrame(str(source), show_progress=False, trust_index=True) + + +def test_trusted_index_skips_reader_generation_revalidation(tmp_path): + source = tmp_path / "part.tsfile" + _write_runtime_file(source, 0) + + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as first: + assert len(first) == 1 + + with TsFileDataFrame( + str(source), show_progress=False, trust_index=True + ) as dataframe: + pool = dataframe._runtime.readers + with pool.acquire(0): + pass + stat = os.stat(source) + os.utime(source, ns=(stat.st_atime_ns, stat.st_mtime_ns + 1_000_000)) + with pool.acquire(0): + pass + + def test_named_selection_reuses_bounded_runtime_descriptor(tmp_path, monkeypatch): source = tmp_path / "part.tsfile" _write_runtime_file(source, 0) diff --git a/python/tsfile/dataset/dataframe.py b/python/tsfile/dataset/dataframe.py index 9e706b0f2..06f7109d2 100644 --- a/python/tsfile/dataset/dataframe.py +++ b/python/tsfile/dataset/dataframe.py @@ -50,6 +50,7 @@ SeriesRef = Tuple[object, int, int] _QUERY_START = np.iinfo(np.int64).min _QUERY_END = np.iinfo(np.int64).max _DATACLASS_SLOTS = {"slots": True} if sys.version_info >= (3, 10) else {} +_TRUST_INDEX_ENV = "TSFILE_DATAFRAME_TRUST_INDEX" # Overlap position reads use chunked k-way merge. Keep the default chunk small # enough to avoid large read amplification for `series[i]` / short slices, but # large enough to avoid excessive query_by_row round-trips when overlap spans @@ -57,6 +58,26 @@ _DATACLASS_SLOTS = {"slots": True} if sys.version_info >= (3, 10) else {} _OVERLAP_ROW_CHUNK_SIZE = 256 +def _resolve_trust_index(value: Optional[bool]) -> bool: + """Resolve the explicit trusted-index option or its process default.""" + if value is not None: + if not isinstance(value, bool): + raise TypeError("trust_index must be a bool or None") + return value + + raw = os.environ.get(_TRUST_INDEX_ENV) + if raw is None: + return False + normalized = raw.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"", "0", "false", "no", "off"}: + return False + raise ValueError( + f"{_TRUST_INDEX_ENV} must be one of 1/0, true/false, yes/no, or on/off" + ) + + @dataclass(**_DATACLASS_SLOTS) class _DataFrameCatalog: """TsFileDataFrame's cross-file unified catalog: merges each tsfile's @@ -772,19 +793,34 @@ class _LocIndexer: class TsFileDataFrame: - """Lazy-loaded unified numeric dataset view over multiple TsFile shards.""" + """Lazy-loaded unified numeric dataset view over multiple TsFile shards. + + ``trust_index=True`` enables the persistent Dataset Index fast path and + skips Dataset Index integrity, path/fingerprint, and reader-session + generation validation. This is intended only for a sealed dataset whose + index and source files are trusted. The underlying TsFile reader still + parses file metadata when a query first opens a series. When + ``trust_index`` is ``None`` (the default), + :envvar:`TSFILE_DATAFRAME_TRUST_INDEX` is consulted; the safe default is + disabled. + """ def __init__( self, paths: Union[str, List[str]], show_progress: bool = True, use_index: bool = False, + trust_index: Optional[bool] = None, ): if not isinstance(use_index, bool): raise TypeError("use_index must be a bool") + self._trust_index = _resolve_trust_index(trust_index) self._paths = _expand_paths(paths) self._show_progress = show_progress - self._use_index = use_index + # A trusted index is meaningful only when the persistent index path is + # used. Make the safe, explicit fast-path convenient for callers by + # enabling it automatically instead of requiring two flags. + self._use_index = use_index or self._trust_index self._readers: Dict[str, object] = {} self._index = _DataFrameCatalog() self._is_view = False @@ -805,6 +841,7 @@ class TsFileDataFrame: obj._paths = parent._paths obj._show_progress = parent._show_progress obj._use_index = parent._use_index + obj._trust_index = parent._trust_index obj._readers = parent._readers subset_refs = list(series_refs) obj._index = SimpleNamespace( @@ -866,7 +903,12 @@ class TsFileDataFrame: from .runtime import DatasetRuntime index_path = index_path_for(self._paths) - if not index_matches_paths(index_path, self._paths): + if self._trust_index: + if not os.path.isfile(index_path): + raise FileNotFoundError( + f"Trusted Dataset Index not found: {index_path}" + ) + elif not index_matches_paths(index_path, self._paths): lock_path = index_path + ".lock" os.makedirs(os.path.dirname(lock_path) or ".", exist_ok=True) with open(lock_path, "a+b") as lock_file: @@ -885,7 +927,7 @@ class TsFileDataFrame: reader.close() self._readers.clear() - self._runtime = DatasetRuntime(index_path) + self._runtime = DatasetRuntime(index_path, trust_index=self._trust_index) self._runtime_lease = self._runtime.lease() self._index = self._runtime.catalog if len(self._index.series) == 0: diff --git a/python/tsfile/dataset/index.py b/python/tsfile/dataset/index.py index 4bd9c3767..790fcd077 100644 --- a/python/tsfile/dataset/index.py +++ b/python/tsfile/dataset/index.py @@ -225,10 +225,23 @@ def write_index_atomic(path: str, section_payloads: Mapping[int, bytes]) -> None class MappedDatasetIndex: - """Validated read-only mmap and typed zero-copy record access.""" + """Read-only mmap and typed zero-copy record access. - def __init__(self, path: str, verify_sections: bool = False): + By default the complete v1 structural validation is performed once while + opening the mapping. ``trust_index=True`` is an explicit unsafe fast path + for callers that already validated the Dataset Index out of band; it only + reads the header and directory needed to locate sections and skips all + integrity and cross-reference checks. + """ + + def __init__( + self, + path: str, + verify_sections: bool = False, + trust_index: bool = False, + ): self.path = path + self._trusted = bool(trust_index) self._file = open(path, "rb") try: stat = os.fstat(self._file.fileno()) @@ -240,7 +253,11 @@ class MappedDatasetIndex: ) self._mmap = mmap.mmap(self._file.fileno(), 0, access=mmap.ACCESS_READ) self._view = memoryview(self._mmap) - self._entries = self._validate(verify_sections) + self._entries = ( + self._map_entries_without_validation() + if self._trusted + else self._validate(verify_sections) + ) except Exception: if getattr(self, "_view", None) is not None: self._view.release() @@ -251,6 +268,27 @@ class MappedDatasetIndex: self._file.close() raise + def _map_entries_without_validation(self): + """Read v1 section descriptors without checking their contents. + + Trusted mode still needs the section descriptors to translate a + logical record into an mmap address. Every other header, range, + checksum, string-pool, and cross-section consistency check is omitted + deliberately; malformed input is outside the trusted-mode contract. + """ + header = HEADER.unpack_from(self._view) + directory_offset = header[4] + section_count = header[5] + directory_entry_size = header[6] + entries = {} + for index in range(section_count): + entry = DIRECTORY.unpack_from( + self._view, + directory_offset + index * directory_entry_size, + ) + entries[entry[0]] = entry + return entries + def _validate(self, verify_sections: bool): if len(self._view) < HEADER_SIZE: raise ValueError("Dataset Index is shorter than its header") @@ -357,7 +395,7 @@ class MappedDatasetIndex: def record(self, section_type: int, record_id: int) -> tuple: entry = self._entries[section_type] - if record_id < 0 or record_id >= entry[4]: + if not self._trusted and (record_id < 0 or record_id >= entry[4]): raise IndexError(record_id) return RECORDS[section_type].unpack_from( self._view, entry[2] + record_id * entry[1] @@ -366,7 +404,7 @@ class MappedDatasetIndex: def records(self, section_type: int, first: int = 0, count: Optional[int] = None): total = self.count(section_type) end = total if count is None else first + count - if first < 0 or end < first or end > total: + if not self._trusted and (first < 0 or end < first or end > total): raise IndexError((first, count)) for record_id in range(first, end): yield self.record(section_type, record_id) @@ -374,7 +412,7 @@ class MappedDatasetIndex: def string_bytes(self, sid: int) -> bytes: offsets = self._entries[STRING_OFFSETS] strings = self._entries[STRING_BYTES] - if sid < 0 or sid + 1 >= offsets[4]: + if not self._trusted and (sid < 0 or sid + 1 >= offsets[4]): raise IndexError(sid) start = RECORDS[STRING_OFFSETS].unpack_from(self._view, offsets[2] + sid * 4)[0] end = RECORDS[STRING_OFFSETS].unpack_from( diff --git a/python/tsfile/dataset/runtime.py b/python/tsfile/dataset/runtime.py index a44bbf49f..1202123e4 100644 --- a/python/tsfile/dataset/runtime.py +++ b/python/tsfile/dataset/runtime.py @@ -148,12 +148,21 @@ class _QueryLease: class _ReaderSession: - def __init__(self, file_id: int, path: str, expected_size: int, fingerprint: int): + def __init__( + self, + file_id: int, + path: str, + expected_size: int, + fingerprint: int, + validate_generation: bool = True, + ): self.file_id = file_id self.path = path self.expected_size = expected_size self.fingerprint = fingerprint - self._validate_generation() + self._validate_generation_enabled = validate_generation + if validate_generation: + self._validate_generation() self.reader = TsFileReaderPy(path) self.active_uses = 0 @@ -175,9 +184,15 @@ class _ReaderSession: class ReaderSessionPool: """Per-Runtime LRU pool with a hard cap on simultaneously open Readers.""" - def __init__(self, index: MappedDatasetIndex, max_open_files: int): + def __init__( + self, + index: MappedDatasetIndex, + max_open_files: int, + validate_generation: bool = True, + ): self._index = index self.max_open_files = max(1, int(max_open_files)) + self._validate_generation = bool(validate_generation) self._sessions: "OrderedDict[int, _ReaderSession]" = OrderedDict() self._condition = threading.Condition() self._closed = False @@ -189,6 +204,7 @@ class ReaderSessionPool: self._index.string(record[0]), record[2], record[3], + validate_generation=self._validate_generation, ) @contextlib.contextmanager @@ -199,7 +215,8 @@ class ReaderSessionPool: raise RuntimeError("ReaderSessionPool is closed") session = self._sessions.get(file_id) if session is not None: - session._validate_generation() + if self._validate_generation: + session._validate_generation() self._sessions.move_to_end(file_id) session.active_uses += 1 break @@ -251,8 +268,13 @@ class ReaderSessionPool: class PreparedSeriesCache: """Runtime-wide single-flight cache of native exact-locator metadata.""" - def __init__(self, index: MappedDatasetIndex): + def __init__( + self, + index: MappedDatasetIndex, + validate_references: bool = True, + ): self._index = index + self._validate_references = bool(validate_references) self._condition = threading.Condition() self._entries = {} self._loading = set() @@ -262,7 +284,7 @@ class PreparedSeriesCache: locator = self._index.record(SERIES_LOCATOR, locator_id) device_span = self._index.record(DEVICE_FILE_SPAN, locator[0]) file_record = self._index.record(TSFILE_RECORD, file_id) - if device_span[1] != file_id: + if self._validate_references and device_span[1] != file_id: raise ValueError("series locator points at another TsFile") return ( id(self._index), @@ -334,8 +356,10 @@ class DatasetRuntime: max_open_files: Optional[int] = None, query_workers: Optional[int] = None, query_parallel_min_rows: Optional[int] = None, + trust_index: bool = False, ): - self.index = MappedDatasetIndex(path) + self.trust_index = bool(trust_index) + self.index = MappedDatasetIndex(path, trust_index=self.trust_index) maximum = ( int(os.environ.get("TSFILE_DATAFRAME_MAX_OPEN_FILES", "16")) if max_open_files is None @@ -366,8 +390,15 @@ class DatasetRuntime: if self.query_workers > 1 else None ) - self.readers = ReaderSessionPool(self.index, maximum) - self.prepared = PreparedSeriesCache(self.index) + self.readers = ReaderSessionPool( + self.index, + maximum, + validate_generation=not self.trust_index, + ) + self.prepared = PreparedSeriesCache( + self.index, + validate_references=not self.trust_index, + ) self._condition = threading.Condition() self._object_leases = 0 self._query_leases = 0
