This is an automated email from the ASF dual-hosted git repository. ColinLeeo pushed a commit to branch fix/dataframe-index-opt-in in repository https://gitbox.apache.org/repos/asf/tsfile.git
commit 97ba5d367ca709693ea12d366a4911647ac03930 Author: ColinLee <[email protected]> AuthorDate: Tue Aug 25 18:21:06 2026 +0800 fix(python): make dataframe index opt-in --- python/tests/test_dataset_index.py | 77 +++++++++++++++++++++++++++++------ python/tests/test_tsfile_dataset.py | 14 +++---- python/tsfile/dataset/dataframe.py | 80 +++++++++++++++++++++++++++---------- python/tsfile/dataset/index.py | 16 ++++++-- 4 files changed, 144 insertions(+), 43 deletions(-) diff --git a/python/tests/test_dataset_index.py b/python/tests/test_dataset_index.py index 417b54d58..91761a11a 100644 --- a/python/tests/test_dataset_index.py +++ b/python/tests/test_dataset_index.py @@ -244,14 +244,14 @@ def _write_runtime_devices_file(path): def test_hot_construction_maps_index_without_opening_readers(tmp_path, monkeypatch): source = tmp_path / "part.tsfile" _write_runtime_file(source, 0) - with TsFileDataFrame(str(source), show_progress=False) as first: + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as first: assert len(first) == 1 def fail_legacy_scan(*_args, **_kwargs): raise AssertionError("hot construction must not build a legacy catalog") monkeypatch.setattr("tsfile.dataset.reader.TsFileSeriesReader", fail_legacy_scan) - with TsFileDataFrame(str(source), show_progress=False) as second: + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as second: assert len(second) == 1 assert second._runtime.readers.open_count == 0 series = second[0] @@ -264,11 +264,58 @@ def test_hot_construction_maps_index_without_opening_readers(tmp_path, monkeypat series.close() -def test_named_selection_reuses_bounded_runtime_descriptor(tmp_path, monkeypatch): +def test_dataframe_does_not_use_or_create_index_by_default(tmp_path): source = tmp_path / "part.tsfile" _write_runtime_file(source, 0) + index_path = index_module.index_path_for([str(source)]) with TsFileDataFrame(str(source), show_progress=False) as dataframe: + assert dataframe._runtime is None + assert len(dataframe) == 1 + np.testing.assert_array_equal(dataframe[0][:], np.array([0.0, 1.0])) + aligned = dataframe.loc[0:1, [0]] + np.testing.assert_array_equal(aligned.timestamps, np.array([0, 1])) + np.testing.assert_array_equal(aligned.values, np.array([[0.0], [1.0]])) + + assert not os.path.exists(index_path) + + +def test_persistent_index_path_is_scoped_to_expanded_file_set(tmp_path, monkeypatch): + first = tmp_path / "part1.tsfile" + second = tmp_path / "part2.tsfile" + third = tmp_path / "part3.tsfile" + _write_runtime_file(first, 0) + _write_runtime_file(second, 2) + _write_runtime_file(third, 10) + + first_set = [str(first), str(second)] + second_set = [str(first), str(third)] + first_index = index_module.index_path_for(first_set) + second_index = index_module.index_path_for(second_set) + assert first_index != second_index + + with TsFileDataFrame(first_set, show_progress=False, use_index=True) as dataframe: + assert len(dataframe) == 1 + with TsFileDataFrame(second_set, show_progress=False, use_index=True) as dataframe: + assert len(dataframe) == 1 + assert os.path.exists(first_index) + assert os.path.exists(second_index) + + def fail_legacy_scan(*_args, **_kwargs): + raise AssertionError("matching file-set index should be reused") + + monkeypatch.setattr("tsfile.dataset.reader.TsFileSeriesReader", fail_legacy_scan) + with TsFileDataFrame(first_set, show_progress=False, use_index=True) as dataframe: + np.testing.assert_array_equal(dataframe[0][:], np.array([0.0, 1.0, 2.0, 3.0])) + with TsFileDataFrame(second_set, show_progress=False, use_index=True) as dataframe: + np.testing.assert_array_equal(dataframe[0][:], np.array([0.0, 1.0, 10.0, 11.0])) + + +def test_named_selection_reuses_bounded_runtime_descriptor(tmp_path, monkeypatch): + source = tmp_path / "part.tsfile" + _write_runtime_file(source, 0) + + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: name = str(dataframe.list_timeseries()[0]) find_device_calls = 0 original_find_device = dataframe._runtime.index.find_device_id @@ -307,7 +354,7 @@ def test_listed_series_path_resolves_directly_by_snapshot_series_id( source = tmp_path / "part.tsfile" _write_runtime_file(source, 0) - with TsFileDataFrame(str(source), show_progress=False) as dataframe: + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: path = dataframe.list_timeseries()[0] assert isinstance(path, str) assert path.series_id == 0 @@ -346,9 +393,13 @@ def test_series_path_from_another_index_falls_back_to_its_name(tmp_path, monkeyp _write_runtime_file(first_source, 0) _write_runtime_file(second_source, 10) - with TsFileDataFrame(str(first_source), show_progress=False) as first: + with TsFileDataFrame( + str(first_source), show_progress=False, use_index=True + ) as first: foreign_path = first.list_timeseries()[0] - with TsFileDataFrame(str(second_source), show_progress=False) as second: + with TsFileDataFrame( + str(second_source), show_progress=False, use_index=True + ) as second: assert foreign_path._index_identity != second._runtime.index.identity find_device_calls = 0 original_find_device = second._runtime.index.find_device_id @@ -372,7 +423,7 @@ def test_runtime_descriptor_cache_evicts_least_recent_name(tmp_path, monkeypatch _write_runtime_devices_file(source) monkeypatch.setattr(runtime_module, "_SERIES_DESCRIPTOR_CACHE_SIZE", 2) - with TsFileDataFrame(str(source), show_progress=False) as dataframe: + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: names = [str(name) for name in dataframe.list_timeseries()] find_device_calls = 0 original_find_device = dataframe._runtime.index.find_device_id @@ -403,7 +454,9 @@ def test_reader_pool_enforces_open_file_cap(tmp_path, monkeypatch): _write_runtime_file(first, 0) _write_runtime_file(second, 2) monkeypatch.setenv("TSFILE_DATAFRAME_MAX_OPEN_FILES", "1") - with TsFileDataFrame([str(first), str(second)], show_progress=False) as dataframe: + with TsFileDataFrame( + [str(first), str(second)], show_progress=False, use_index=True + ) as dataframe: series = dataframe[0] assert list(series[:]) == [0.0, 1.0, 2.0, 3.0] assert dataframe._runtime.readers.open_count == 1 @@ -481,7 +534,7 @@ def test_prepared_query_reads_nullable_offset_window_in_arrow_batches(tmp_path): ) ) - with TsFileDataFrame(str(source), show_progress=False) as dataframe: + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: runtime = dataframe._runtime series = runtime.index.record(LOGICAL_SERIES, 0) span = runtime.index.record(SERIES_FILE_SPAN, series[2]) @@ -514,7 +567,7 @@ def test_prepared_query_reads_nullable_offset_window_in_arrow_batches(tmp_path): def test_prepared_locator_rejects_stale_generation_and_bad_range(tmp_path): source = tmp_path / "part.tsfile" _write_runtime_file(source, 0) - with TsFileDataFrame(str(source), show_progress=False) as dataframe: + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: runtime = dataframe._runtime series = runtime.index.record(LOGICAL_SERIES, 0) span = runtime.index.record(SERIES_FILE_SPAN, series[2]) @@ -534,7 +587,7 @@ def test_prepared_locator_rejects_stale_generation_and_bad_range(tmp_path): def test_reader_session_revalidates_generation_when_reused(tmp_path): source = tmp_path / "part.tsfile" _write_runtime_file(source, 0) - with TsFileDataFrame(str(source), show_progress=False) as dataframe: + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: pool = dataframe._runtime.readers with pool.acquire(0): pass @@ -548,7 +601,7 @@ def test_reader_session_revalidates_generation_when_reused(tmp_path): def test_runtime_lease_close_waits_for_query_lease(tmp_path): source = tmp_path / "part.tsfile" _write_runtime_file(source, 0) - with TsFileDataFrame(str(source), show_progress=False) as dataframe: + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: runtime = DatasetRuntime(str(dataframe._runtime.index.path), query_workers=1) lease = runtime.lease() entered = threading.Event() diff --git a/python/tests/test_tsfile_dataset.py b/python/tests/test_tsfile_dataset.py index 55115987e..55dceaee5 100644 --- a/python/tests/test_tsfile_dataset.py +++ b/python/tests/test_tsfile_dataset.py @@ -315,7 +315,7 @@ def test_dataset_loc_aligns_timestamp_union_and_preserves_requested_order(tmp_pa }, ) - with TsFileDataFrame(str(path), show_progress=False) as tsdf: + with TsFileDataFrame(str(path), show_progress=False, use_index=True) as tsdf: aligned = tsdf.loc[ 0:2, [ @@ -358,7 +358,7 @@ def test_dataset_loc_batches_aligned_fields_per_device_then_unions_devices(tmp_p "weather.device_a.humidity", "weather.device_b.humidity", ] - with TsFileDataFrame(str(path), show_progress=False) as tsdf: + with TsFileDataFrame(str(path), show_progress=False, use_index=True) as tsdf: aligned = tsdf.loc[0:2, requested] assert aligned.series_names == requested @@ -407,7 +407,7 @@ def test_dataset_loc_runs_independent_device_groups_concurrently(tmp_path, monke observed_read, ) - with TsFileDataFrame(str(path), show_progress=False) as tsdf: + with TsFileDataFrame(str(path), show_progress=False, use_index=True) as tsdf: aligned = tsdf.loc[ 0:1, [ @@ -452,7 +452,7 @@ def test_dataset_loc_keeps_small_device_groups_inline(tmp_path, monkeypatch): observed_read, ) - with TsFileDataFrame(str(path), show_progress=False) as tsdf: + with TsFileDataFrame(str(path), show_progress=False, use_index=True) as tsdf: tsdf.loc[ 0:1, [ @@ -1009,7 +1009,7 @@ def test_dataset_close_only_releases_current_handle(tmp_path): path = tmp_path / "weather.tsfile" _write_weather_file(path, 0) - tsdf = TsFileDataFrame(str(path), show_progress=False) + tsdf = TsFileDataFrame(str(path), show_progress=False, use_index=True) series = tsdf[0] tsdf.close() @@ -1026,7 +1026,7 @@ def test_subset_close_releases_only_subset_lease(tmp_path): path = tmp_path / "weather.tsfile" _write_weather_file(path, 0) - with TsFileDataFrame(str(path), show_progress=False) as tsdf: + with TsFileDataFrame(str(path), show_progress=False, use_index=True) as tsdf: subset = tsdf[:1] subset.close() @@ -1103,7 +1103,7 @@ def test_dataset_close_waits_for_an_active_public_query(tmp_path, monkeypatch): blocked_read, ) - dataframe = TsFileDataFrame(str(path), show_progress=False) + dataframe = TsFileDataFrame(str(path), show_progress=False, use_index=True) query_done = threading.Event() def run_query(): diff --git a/python/tsfile/dataset/dataframe.py b/python/tsfile/dataset/dataframe.py index e3ce78372..9e706b0f2 100644 --- a/python/tsfile/dataset/dataframe.py +++ b/python/tsfile/dataset/dataframe.py @@ -663,8 +663,28 @@ class _LocIndexer: _, table_entry, _ = self._df._get_series_components(series_ref) field_name = table_entry.field_columns[field_idx] - descriptor = self._df._index.series_shards.describe(series_ref) - for shard in descriptor.shards: + describe = getattr(self._df._index.series_shards, "describe", None) + if describe is None: + shard_entries = [] + for ( + reader, + device_id, + reader_field_idx, + ) in self._df._index.series_shards[series_ref]: + info = reader.get_series_info_by_ref(device_id, reader_field_idx) + shard_entries.append( + SimpleNamespace( + reader=reader, + device_id=device_id, + column_id=reader_field_idx, + min_time=info["timeline_min_time"], + max_time=info["timeline_max_time"], + timeline_length=info["timeline_length"], + ) + ) + else: + shard_entries = describe(series_ref).shards + for shard in shard_entries: overlap_start = max(start_time, shard.min_time) overlap_end = min(end_time, shard.max_time) if shard.timeline_length <= 0 or overlap_start > overlap_end: @@ -710,13 +730,16 @@ class _LocIndexer: return entries, ts_arr, field_vals group_entries = list(groups.values()) - group_results = self._df._runtime.map_query_groups( - query_group, - group_entries, - estimated_rows=[ - max(entry[6] for entry in entries) for entries in group_entries - ], - ) + if self._df._runtime is None: + group_results = [query_group(entries) for entries in group_entries] + else: + group_results = self._df._runtime.map_query_groups( + query_group, + group_entries, + estimated_rows=[ + max(entry[6] for entry in entries) for entries in group_entries + ], + ) series_time_parts = defaultdict(list) series_value_parts = defaultdict(list) @@ -751,9 +774,17 @@ class _LocIndexer: class TsFileDataFrame: """Lazy-loaded unified numeric dataset view over multiple TsFile shards.""" - def __init__(self, paths: Union[str, List[str]], show_progress: bool = True): + def __init__( + self, + paths: Union[str, List[str]], + show_progress: bool = True, + use_index: bool = False, + ): + if not isinstance(use_index, bool): + raise TypeError("use_index must be a bool") self._paths = _expand_paths(paths) self._show_progress = show_progress + self._use_index = use_index self._readers: Dict[str, object] = {} self._index = _DataFrameCatalog() self._is_view = False @@ -773,6 +804,7 @@ class TsFileDataFrame: obj._is_view = True obj._paths = parent._paths obj._show_progress = parent._show_progress + obj._use_index = parent._use_index obj._readers = parent._readers subset_refs = list(series_refs) obj._index = SimpleNamespace( @@ -808,9 +840,24 @@ class TsFileDataFrame: with self._runtime_lease.query_lease(): yield + def _load_metadata_without_index(self, reader_class): + if len(self._paths) >= 2: + self._load_metadata_parallel(reader_class) + else: + self._load_metadata_serial(reader_class) + + if not self._index.series: + raise ValueError("No valid time series found in the provided TsFile files") + _validate_unique_shard_timestamps(self._index) + def _load_metadata(self): - """Map a valid persistent index, or build it once under a file lock.""" + """Load metadata, optionally through a persistent mmap-backed index.""" from .reader import TsFileSeriesReader + + if not self._use_index: + self._load_metadata_without_index(TsFileSeriesReader) + return + from .index import ( build_index_from_dataframe, index_matches_paths, @@ -830,17 +877,8 @@ class TsFileDataFrame: except ImportError: pass if not index_matches_paths(index_path, self._paths): - if len(self._paths) >= 2: - self._load_metadata_parallel(TsFileSeriesReader) - else: - self._load_metadata_serial(TsFileSeriesReader) - - if not self._index.series: - raise ValueError( - "No valid time series found in the provided TsFile files" - ) + self._load_metadata_without_index(TsFileSeriesReader) try: - _validate_unique_shard_timestamps(self._index) build_index_from_dataframe(self, index_path) finally: for reader in self._readers.values(): diff --git a/python/tsfile/dataset/index.py b/python/tsfile/dataset/index.py index 394860d54..4bd9c3767 100644 --- a/python/tsfile/dataset/index.py +++ b/python/tsfile/dataset/index.py @@ -28,6 +28,7 @@ from __future__ import annotations from collections import defaultdict import contextlib +import hashlib import mmap import os import struct @@ -43,7 +44,8 @@ HEADER_SIZE = 64 DIRECTORY_ENTRY_SIZE = 32 SECTION_COUNT = 13 ALIGNMENT = 64 -INDEX_FILE_NAME = ".tsfile_dataframe_index.tsidx" +INDEX_FILE_PREFIX = ".tsfile_dataframe_index" +INDEX_FILE_SUFFIX = ".tsidx" STRING_OFFSETS = 1 STRING_BYTES = 2 @@ -460,10 +462,18 @@ class MappedDatasetIndex: def index_path_for(paths: Sequence[str]) -> str: - common = os.path.commonpath([os.path.abspath(path) for path in paths]) + canonical_paths = sorted(os.path.abspath(path) for path in paths) + common = os.path.commonpath(canonical_paths) if not os.path.isdir(common): common = os.path.dirname(common) - return os.path.join(common, INDEX_FILE_NAME) + digest = hashlib.sha256() + for path in canonical_paths: + digest.update(path.encode("utf-8")) + digest.update(b"\0") + return os.path.join( + common, + f"{INDEX_FILE_PREFIX}.{digest.hexdigest()[:32]}{INDEX_FILE_SUFFIX}", + ) def index_matches_paths(path: str, paths: Sequence[str]) -> bool:
