From be87d3135b438af1afadde1bc0bef4c8d75e9bc5 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Mon, 14 Sep 2026 11:57:30 +0800 Subject: [PATCH 1/4] perf(python): accelerate Dataset Index metadata lookups --- python/setup.py | 5 + python/tests/test_dataset_index.py | 250 ++++++++++++ python/tsfile/dataset/_index_fast.pyx | 531 ++++++++++++++++++++++++++ python/tsfile/dataset/index.py | 119 ++++++ python/tsfile/dataset/runtime.py | 230 +++++++---- 5 files changed, 1051 insertions(+), 84 deletions(-) create mode 100644 python/tsfile/dataset/_index_fast.pyx diff --git a/python/setup.py b/python/setup.py index c7b4cceaa..7f41a07f1 100644 --- a/python/setup.py +++ b/python/setup.py @@ -273,6 +273,11 @@ def finalize_options(self): exts = [ Extension("tsfile.dataset._merge", ["tsfile/dataset/_merge.pyx"], **merge_common), + Extension( + "tsfile.dataset._index_fast", + ["tsfile/dataset/_index_fast.pyx"], + **merge_common, + ), Extension("tsfile.tsfile_py_cpp", ["tsfile/tsfile_py_cpp.pyx"], **common), Extension( "tsfile.tsfile_reader", diff --git a/python/tests/test_dataset_index.py b/python/tests/test_dataset_index.py index 91761a11a..c521958af 100644 --- a/python/tests/test_dataset_index.py +++ b/python/tests/test_dataset_index.py @@ -124,6 +124,218 @@ def test_build_publish_map_and_lookup(tmp_path): assert index.string(file_record[0]) == str(source) +@pytest.mark.skipif( + index_module.FastIndexLookup is None, + reason="requires the built Cython Dataset Index lookup extension", +) +def test_hot_lookup_does_not_unpack_python_record_tuples(tmp_path, monkeypatch): + source = tmp_path / "source.tsfile" + source.write_bytes(b"T" * 4096) + output = tmp_path / "dataset.tsidx" + dataframe = _synthetic_dataframe(str(source)) + write_index_atomic(str(output), build_sections_from_dataframe(dataframe)) + + with MappedDatasetIndex(str(output), verify_sections=True) as index: + + def fail_record(*_args, **_kwargs): + raise AssertionError("hot lookup should not unpack Python record tuples") + + monkeypatch.setattr(index, "record", fail_record) + assert index.find_device_id(0, "root.") == 0 + assert index.find_column_id(0, "s1") == 0 + assert index.find_series_id(0, 0) == 0 + + +@pytest.mark.skipif( + index_module.FastIndexLookup is None, + reason="requires the built Cython Dataset Index lookup extension", +) +def test_fast_series_description_does_not_unpack_python_record_tuples( + tmp_path, monkeypatch +): + source = tmp_path / "source.tsfile" + source.write_bytes(b"T" * 4096) + output = tmp_path / "dataset.tsidx" + dataframe = _synthetic_dataframe(str(source)) + write_index_atomic(str(output), build_sections_from_dataframe(dataframe)) + + with MappedDatasetIndex(str(output), verify_sections=True) as index: + series = index.record(LOGICAL_SERIES, 0) + span = index.record(SERIES_FILE_SPAN, series[2]) + locator = index.record(SERIES_LOCATOR, span[2]) + device_span = index.record(DEVICE_FILE_SPAN, locator[0]) + expected = (series[0], series[1], series[4], series[5]) + expected_count = device_span[6] if device_span[4] == 1 else span[6] + + def fail_record(*_args, **_kwargs): + raise AssertionError("series description should not unpack records") + + monkeypatch.setattr(index, "record", fail_record) + device_id, column_id, min_time, max_time, count, shards = index.describe_series( + 0 + ) + + assert (device_id, column_id, min_time, max_time) == expected + assert count == expected_count + assert shards == [(span[1], span[2], expected_count, span[4], span[5])] + + +@pytest.mark.skipif( + index_module.FastIndexLookup is None, + reason="requires the built Cython Dataset Index lookup extension", +) +def test_fast_span_and_locator_metadata_do_not_unpack_python_records( + tmp_path, monkeypatch +): + source = tmp_path / "source.tsfile" + source.write_bytes(b"T" * 4096) + output = tmp_path / "dataset.tsidx" + dataframe = _synthetic_dataframe(str(source)) + write_index_atomic(str(output), build_sections_from_dataframe(dataframe)) + + with MappedDatasetIndex(str(output), verify_sections=True) as index: + series = index.record(LOGICAL_SERIES, 0) + span = index.record(SERIES_FILE_SPAN, series[2]) + locator = index.record(SERIES_LOCATOR, span[2]) + device_span = index.record(DEVICE_FILE_SPAN, locator[0]) + + def fail_record(*_args, **_kwargs): + raise AssertionError("span metadata should not unpack records") + + monkeypatch.setattr(index, "record", fail_record) + assert index.find_series_span(0, span[1]) == ( + span[2], + span[4], + span[5], + span[6], + ) + assert index.locator_metadata(span[2]) == ( + locator[0], + locator[1], + device_span[1], + device_span[4], + device_span[6], + ) + + +@pytest.mark.skipif( + index_module.FastIndexLookup is None, + reason="requires the built Cython Dataset Index lookup extension", +) +def test_runtime_reader_span_lookup_uses_fast_metadata(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: + reader = dataframe._runtime.catalog.reader_for(0) + series = dataframe._runtime.index.record(LOGICAL_SERIES, 0) + + def fail_record(*_args, **_kwargs): + raise AssertionError("runtime span lookup should not unpack records") + + monkeypatch.setattr(dataframe._runtime.index, "record", fail_record) + locator_id, min_time, max_time, length = reader._span(series[0], series[1]) + + assert locator_id >= 0 + assert min_time == 0 + assert max_time == 1 + assert length >= 0 + + +@pytest.mark.skipif( + index_module.FastIndexLookup is None, + reason="requires the built Cython Dataset Index lookup extension", +) +def test_fast_prepared_locator_metadata_does_not_unpack_python_records( + tmp_path, monkeypatch +): + source = tmp_path / "source.tsfile" + source.write_bytes(b"T" * 4096) + output = tmp_path / "dataset.tsidx" + dataframe = _synthetic_dataframe(str(source)) + write_index_atomic(str(output), build_sections_from_dataframe(dataframe)) + + with MappedDatasetIndex(str(output), verify_sections=True) as index: + span = index.record(SERIES_FILE_SPAN, 0) + locator = index.record(SERIES_LOCATOR, span[2]) + device_span = index.record(DEVICE_FILE_SPAN, locator[0]) + file_record = index.record(TSFILE_RECORD, span[1]) + + def fail_record(*_args, **_kwargs): + raise AssertionError("prepared locator should not unpack records") + + monkeypatch.setattr(index, "record", fail_record) + assert index.prepared_locator_metadata(span[1], span[2]) == ( + file_record[2], + file_record[3], + locator[1], + locator[2], + locator[3], + locator[4], + device_span[2], + device_span[3], + ) + + +@pytest.mark.skipif( + index_module.FastIndexLookup is None, + reason="requires the built Cython Dataset Index lookup extension", +) +def test_fast_identity_and_device_bounds_do_not_unpack_python_records( + tmp_path, monkeypatch +): + source = tmp_path / "source.tsfile" + source.write_bytes(b"T" * 4096) + output = tmp_path / "dataset.tsidx" + dataframe = _synthetic_dataframe(str(source)) + write_index_atomic(str(output), build_sections_from_dataframe(dataframe)) + + with MappedDatasetIndex(str(output), verify_sections=True) as index: + device = index.record(index_module.DEVICE_RECORD, 0) + table = index.record(index_module.TABLE_RECORD, device[0]) + column = index.record(COLUMN_SCHEMA, 0) + expected_device = (device[0], device[1]) + expected_table_name_id = table[0] + expected_column_name_id = column[1] + expected_bounds = (device[8], device[9]) + expected_name = index.string(expected_column_name_id) + + def fail_record(*_args, **_kwargs): + raise AssertionError("identity metadata should not unpack records") + + def fail_string_bytes(*_args, **_kwargs): + raise AssertionError("fast string should not copy bytes") + + monkeypatch.setattr(index, "record", fail_record) + monkeypatch.setattr(index, "string_bytes", fail_string_bytes) + assert index.device_route(0) == expected_device + assert index.table_name_id(device[0]) == expected_table_name_id + assert index.column_name_id(0) == expected_column_name_id + assert index.device_time_bounds(0) == expected_bounds + assert index.string(expected_column_name_id) == expected_name + + +@pytest.mark.skipif( + index_module.FastIndexLookup is None, + reason="requires the built Cython Dataset Index lookup extension", +) +def test_fast_series_identity_does_not_unpack_python_records(tmp_path, monkeypatch): + source = tmp_path / "source.tsfile" + source.write_bytes(b"T" * 4096) + output = tmp_path / "dataset.tsidx" + dataframe = _synthetic_dataframe(str(source)) + write_index_atomic(str(output), build_sections_from_dataframe(dataframe)) + + with MappedDatasetIndex(str(output), verify_sections=True) as index: + series = index.record(LOGICAL_SERIES, 0) + + def fail_record(*_args, **_kwargs): + raise AssertionError("series identity should not unpack records") + + monkeypatch.setattr(index, "record", fail_record) + assert index.series_identity(0) == (series[0], series[1]) + + def test_child_lookup_checks_full_bytes_for_hash_collisions(monkeypatch): rows = [ (0, 10, 42, 0, 0), @@ -448,6 +660,44 @@ def count_find_device(*args, **kwargs): assert find_device_calls == 4 +def test_runtime_descriptor_objects_use_slots(): + shard = runtime_module.RuntimeSeriesShard(None, 1, 2, 3, 4, 5, 6) + descriptor = runtime_module.RuntimeSeriesDescriptor((1, 2), 3, 4, (shard,), 5, 6, 7) + + assert not hasattr(shard, "__dict__") + assert not hasattr(descriptor, "__dict__") + + +def test_runtime_descriptor_cache_stats_report_hits_and_evictions( + tmp_path, monkeypatch +): + source = tmp_path / "devices.tsfile" + _write_runtime_devices_file(source) + monkeypatch.setattr(runtime_module, "_SERIES_DESCRIPTOR_CACHE_SIZE", 1) + monkeypatch.setenv("TSFILE_DATAFRAME_DESCRIPTOR_CACHE_STATS", "1") + + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: + names = [str(name) for name in dataframe.list_timeseries()] + dataframe[names[0]].close() + dataframe[names[0]].close() + dataframe[names[1]].close() + + stats = dataframe._index.descriptor_cache_stats + assert stats["hits"] >= 1 + assert stats["misses"] >= 2 + assert stats["evictions"] >= 1 + + +def test_runtime_descriptor_cache_size_can_be_configured(tmp_path, monkeypatch): + source = tmp_path / "part.tsfile" + _write_runtime_file(source, 0) + monkeypatch.setenv("TSFILE_DATAFRAME_DESCRIPTOR_CACHE_SIZE", "0") + + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: + assert dataframe._index._descriptor_cache_size == 0 + assert dataframe._index.series_shards._cache_size == 0 + + def test_reader_pool_enforces_open_file_cap(tmp_path, monkeypatch): first = tmp_path / "part1.tsfile" second = tmp_path / "part2.tsfile" diff --git a/python/tsfile/dataset/_index_fast.pyx b/python/tsfile/dataset/_index_fast.pyx new file mode 100644 index 000000000..329d48b15 --- /dev/null +++ b/python/tsfile/dataset/_index_fast.pyx @@ -0,0 +1,531 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# cython: boundscheck=False, wraparound=False, cdivision=True, language_level=3 + +"""Cython hot lookups over the existing read-only Dataset Index mmap. + +This module deliberately owns no index-sized data structures. It receives the +already-mapped byte view and section directory from ``MappedDatasetIndex`` and +only removes Python tuple/bytes allocation from the binary-search loops. +""" + +from libc.stdint cimport int64_t, uint8_t, uint16_t, uint32_t, uint64_t +from cpython.unicode cimport PyUnicode_AsUTF8AndSize, PyUnicode_DecodeUTF8 + + +cdef inline uint32_t _read_u32( + const uint8_t[:] view, uint64_t offset) noexcept nogil: + return ( + view[offset] + | (view[offset + 1] << 8) + | (view[offset + 2] << 16) + | (view[offset + 3] << 24) + ) + + +cdef inline uint16_t _read_u16( + const uint8_t[:] view, uint64_t offset) noexcept nogil: + return ( + view[offset] + | (view[offset + 1] << 8) + ) + + +cdef inline uint64_t _read_u64( + const uint8_t[:] view, uint64_t offset) noexcept nogil: + return ( + view[offset] + | (view[offset + 1] << 8) + | (view[offset + 2] << 16) + | (view[offset + 3] << 24) + | (view[offset + 4] << 32) + | (view[offset + 5] << 40) + | (view[offset + 6] << 48) + | (view[offset + 7] << 56) + ) + + +cdef inline uint64_t _fnv1a( + const char* data, Py_ssize_t length) noexcept nogil: + cdef uint64_t result = 1469598103934665603 + cdef uint64_t prime = 1099511628211 + cdef Py_ssize_t index + for index in range(length): + result ^= data[index] + result = result * prime + return result + + +cdef inline bint _string_equals( + const uint8_t[:] view, + uint64_t offsets_base, + uint64_t strings_base, + uint32_t string_count, + uint32_t string_id, + const char* data, + Py_ssize_t length, +) noexcept nogil: + cdef uint32_t start + cdef uint32_t end + cdef Py_ssize_t index + if string_id + 1 >= string_count: + return False + start = _read_u32(view, offsets_base + string_id * 4) + end = _read_u32(view, offsets_base + (string_id + 1) * 4) + if end - start != length: + return False + for index in range(length): + if view[strings_base + start + index] != data[index]: + return False + return True + + +cdef inline int64_t _find_child( + const uint8_t[:] view, + uint64_t section_base, + uint32_t record_size, + uint32_t first, + uint32_t count, + uint32_t table_id, + uint64_t target_hash, + const char* data, + Py_ssize_t length, + uint64_t string_offsets_base, + uint64_t string_bytes_base, + uint32_t string_count, +) noexcept nogil: + cdef uint32_t low = first + cdef uint32_t high = first + count + cdef uint32_t middle + cdef uint64_t row_base + cdef uint64_t row_hash + cdef uint32_t row_table_id + cdef uint32_t row_id + cdef uint32_t row_string_id + + while low < high: + middle = low + (high - low) // 2 + row_base = section_base + middle * record_size + row_hash = _read_u64(view, row_base + 8) + if row_hash < target_hash: + low = middle + 1 + else: + high = middle + + high = first + count + while low < high: + row_base = section_base + low * record_size + if _read_u64(view, row_base + 8) != target_hash: + break + row_table_id = _read_u32(view, row_base) + row_string_id = _read_u32(view, row_base + 16) + if row_table_id == table_id and _string_equals( + view, + string_offsets_base, + string_bytes_base, + string_count, + row_string_id, + data, + length, + ): + row_id = _read_u32(view, row_base + 4) + return row_id + low += 1 + return -1 + + +cdef const char* _utf8_name(object name, Py_ssize_t* length): + cdef const char* data + if not isinstance(name, str): + raise TypeError("Dataset Index names must be str") + data = PyUnicode_AsUTF8AndSize(name, length) + if data == NULL: + raise UnicodeEncodeError("utf-8", name, 0, len(name), "invalid UTF-8") + return data + + +cdef class FastIndexLookup: + """Typed lookup kernel over a ``MappedDatasetIndex`` memoryview.""" + + cdef const uint8_t[:] _view + cdef uint64_t _section_offsets[14] + cdef uint32_t _section_sizes[14] + cdef uint32_t _section_counts[14] + + def __cinit__(self, object view, object entries): + cdef int section + self._view = view + for section in range(1, 14): + self._section_offsets[section] = entries[section][2] + self._section_sizes[section] = entries[section][1] + self._section_counts[section] = entries[section][4] + + def find_device_id(self, uint32_t table_id, object name): + cdef const char* data + cdef Py_ssize_t length + cdef const uint8_t[:] view = self._view + cdef uint64_t table_base + cdef uint32_t first + cdef uint32_t count + cdef uint64_t target_hash + cdef int64_t result + data = _utf8_name(name, &length) + if table_id >= self._section_counts[4]: + raise IndexError(table_id) + table_base = self._section_offsets[4] + table_id * self._section_sizes[4] + first = _read_u32(view, table_base + 8) + count = _read_u32(view, table_base + 12) + target_hash = _fnv1a(data, length) + with nogil: + result = _find_child( + view, + self._section_offsets[5], + self._section_sizes[5], + first, + count, + table_id, + target_hash, + data, + length, + self._section_offsets[1], + self._section_offsets[2], + self._section_counts[1], + ) + if result < 0: + raise KeyError(name) + return result + + def find_column_id(self, uint32_t table_id, object name): + cdef const char* data + cdef Py_ssize_t length + cdef const uint8_t[:] view = self._view + cdef uint64_t table_base + cdef uint32_t first + cdef uint32_t count + cdef uint64_t target_hash + cdef int64_t result + data = _utf8_name(name, &length) + if table_id >= self._section_counts[4]: + raise IndexError(table_id) + table_base = self._section_offsets[4] + table_id * self._section_sizes[4] + first = _read_u32(view, table_base + 16) + count = _read_u32(view, table_base + 20) + target_hash = _fnv1a(data, length) + with nogil: + result = _find_child( + view, + self._section_offsets[7], + self._section_sizes[7], + first, + count, + table_id, + target_hash, + data, + length, + self._section_offsets[1], + self._section_offsets[2], + self._section_counts[1], + ) + if result < 0: + raise KeyError(name) + return result + + def find_series_id(self, uint32_t device_id, uint32_t column_id): + cdef const uint8_t[:] view = self._view + cdef uint64_t device_base + cdef uint32_t first + cdef uint32_t count + cdef uint32_t low + cdef uint32_t high + cdef uint32_t middle + cdef uint64_t row_base + cdef uint32_t row_column_id + cdef uint32_t row_device_id + cdef int64_t result = -1 + if device_id >= self._section_counts[6]: + raise IndexError(device_id) + device_base = self._section_offsets[6] + device_id * self._section_sizes[6] + first = _read_u32(view, device_base + 16) + count = _read_u32(view, device_base + 20) + low = first + high = first + count + with nogil: + while low < high: + middle = low + (high - low) // 2 + row_base = self._section_offsets[9] + middle * self._section_sizes[9] + row_column_id = _read_u32(view, row_base + 4) + if row_column_id < column_id: + low = middle + 1 + else: + high = middle + if low < first + count: + row_base = self._section_offsets[9] + low * self._section_sizes[9] + row_device_id = _read_u32(view, row_base) + row_column_id = _read_u32(view, row_base + 4) + if row_device_id == device_id and row_column_id == column_id: + result = low + if result < 0: + raise KeyError(column_id) + return result + + def describe_series(self, uint32_t series_id): + """Expand one logical series without creating intermediate records.""" + cdef const uint8_t[:] view = self._view + cdef uint64_t series_base + cdef uint32_t device_id + cdef uint32_t column_id + cdef uint32_t first_span + cdef uint32_t span_count + cdef uint32_t index + cdef uint32_t span_id + cdef uint64_t span_base + cdef uint32_t file_id + cdef uint32_t locator_id + cdef uint32_t device_span_id + cdef uint64_t locator_base + cdef uint64_t device_span_base + cdef uint64_t span_length + cdef uint64_t timeline_length + cdef int64_t min_time + cdef int64_t max_time + cdef int64_t span_min_time + cdef int64_t span_max_time + cdef uint64_t count = 0 + cdef list shards = [] + + if series_id >= self._section_counts[9]: + raise IndexError(series_id) + + series_base = ( + self._section_offsets[9] + + series_id * self._section_sizes[9] + ) + device_id = _read_u32(view, series_base) + column_id = _read_u32(view, series_base + 4) + first_span = _read_u32(view, series_base + 8) + span_count = _read_u32(view, series_base + 12) + min_time = _read_u64(view, series_base + 16) + max_time = _read_u64(view, series_base + 24) + + if first_span + span_count > self._section_counts[12]: + raise IndexError(series_id) + + for index in range(span_count): + span_id = first_span + index + span_base = ( + self._section_offsets[12] + + span_id * self._section_sizes[12] + ) + file_id = _read_u32(view, span_base + 4) + locator_id = _read_u32(view, span_base + 8) + span_min_time = _read_u64(view, span_base + 16) + span_max_time = _read_u64(view, span_base + 24) + span_length = _read_u64(view, span_base + 32) + + if locator_id >= self._section_counts[13]: + raise IndexError(locator_id) + locator_base = ( + self._section_offsets[13] + + locator_id * self._section_sizes[13] + ) + device_span_id = _read_u32(view, locator_base) + if device_span_id >= self._section_counts[11]: + raise IndexError(device_span_id) + device_span_base = ( + self._section_offsets[11] + + device_span_id * self._section_sizes[11] + ) + if _read_u16(view, device_span_base + 20) == 1: + timeline_length = _read_u64(view, device_span_base + 24) + else: + timeline_length = span_length + count += timeline_length + shards.append( + (file_id, locator_id, timeline_length, span_min_time, span_max_time) + ) + + return device_id, column_id, min_time, max_time, count, shards + + def series_identity(self, uint32_t series_id): + """Return device and column ids for one logical series.""" + cdef const uint8_t[:] view = self._view + cdef uint64_t base + if series_id >= self._section_counts[9]: + raise IndexError(series_id) + base = self._section_offsets[9] + series_id * self._section_sizes[9] + return _read_u32(view, base), _read_u32(view, base + 4) + + def find_series_span(self, uint32_t series_id, uint32_t file_id): + """Find one series span and return locator/time/length scalars.""" + cdef const uint8_t[:] view = self._view + cdef uint64_t series_base + cdef uint32_t first_span + cdef uint32_t span_count + cdef uint32_t index + cdef uint32_t span_id + cdef uint64_t span_base + + if series_id >= self._section_counts[9]: + raise IndexError(series_id) + series_base = ( + self._section_offsets[9] + + series_id * self._section_sizes[9] + ) + first_span = _read_u32(view, series_base + 8) + span_count = _read_u32(view, series_base + 12) + if first_span + span_count > self._section_counts[12]: + raise IndexError(series_id) + for index in range(span_count): + span_id = first_span + index + span_base = ( + self._section_offsets[12] + + span_id * self._section_sizes[12] + ) + if _read_u32(view, span_base + 4) == file_id: + return ( + _read_u32(view, span_base + 8), + _read_u64(view, span_base + 16), + _read_u64(view, span_base + 24), + _read_u64(view, span_base + 32), + ) + raise KeyError((series_id, file_id)) + + def locator_metadata(self, uint32_t locator_id): + """Return locator and owning device-span scalars.""" + cdef const uint8_t[:] view = self._view + cdef uint64_t locator_base + cdef uint32_t device_span_id + cdef uint64_t device_span_base + + if locator_id >= self._section_counts[13]: + raise IndexError(locator_id) + locator_base = ( + self._section_offsets[13] + + locator_id * self._section_sizes[13] + ) + device_span_id = _read_u32(view, locator_base) + if device_span_id >= self._section_counts[11]: + raise IndexError(device_span_id) + device_span_base = ( + self._section_offsets[11] + + device_span_id * self._section_sizes[11] + ) + return ( + device_span_id, + _read_u16(view, locator_base + 4), + _read_u32(view, device_span_base + 4), + _read_u16(view, device_span_base + 20), + _read_u64(view, device_span_base + 24), + ) + + def prepared_locator_metadata(self, uint32_t file_id, uint32_t locator_id): + """Return the generation and locator fields used by native prepare.""" + cdef const uint8_t[:] view = self._view + cdef uint64_t locator_base + cdef uint64_t device_span_base + cdef uint64_t file_base + cdef uint32_t device_span_id + cdef uint32_t device_file_id + + if locator_id >= self._section_counts[13]: + raise IndexError(locator_id) + if file_id >= self._section_counts[10]: + raise IndexError(file_id) + locator_base = ( + self._section_offsets[13] + + locator_id * self._section_sizes[13] + ) + device_span_id = _read_u32(view, locator_base) + if device_span_id >= self._section_counts[11]: + raise IndexError(device_span_id) + device_span_base = ( + self._section_offsets[11] + + device_span_id * self._section_sizes[11] + ) + device_file_id = _read_u32(view, device_span_base + 4) + if device_file_id != file_id: + raise ValueError("series locator points at another TsFile") + file_base = ( + self._section_offsets[10] + + file_id * self._section_sizes[10] + ) + return ( + _read_u64(view, file_base + 8), + _read_u64(view, file_base + 16), + _read_u16(view, locator_base + 4), + _read_u16(view, locator_base + 6), + _read_u64(view, locator_base + 8), + _read_u32(view, locator_base + 16), + _read_u64(view, device_span_base + 8), + _read_u32(view, device_span_base + 16), + ) + + def device_route(self, uint32_t device_id): + """Return table id and logical-path string id for one device.""" + cdef const uint8_t[:] view = self._view + cdef uint64_t base + if device_id >= self._section_counts[6]: + raise IndexError(device_id) + base = self._section_offsets[6] + device_id * self._section_sizes[6] + return _read_u32(view, base), _read_u32(view, base + 4) + + def table_name_id(self, uint32_t table_id): + """Return the string-pool id for one table name.""" + cdef const uint8_t[:] view = self._view + cdef uint64_t base + if table_id >= self._section_counts[4]: + raise IndexError(table_id) + base = self._section_offsets[4] + table_id * self._section_sizes[4] + return _read_u32(view, base) + + def column_name_id(self, uint32_t column_id): + """Return the string-pool id for one column name.""" + cdef const uint8_t[:] view = self._view + cdef uint64_t base + if column_id >= self._section_counts[8]: + raise IndexError(column_id) + base = self._section_offsets[8] + column_id * self._section_sizes[8] + return _read_u32(view, base + 4) + + def device_time_bounds(self, uint32_t device_id): + """Return min/max timestamps for one device.""" + cdef const uint8_t[:] view = self._view + cdef uint64_t base + if device_id >= self._section_counts[6]: + raise IndexError(device_id) + base = self._section_offsets[6] + device_id * self._section_sizes[6] + return ( + _read_u64(view, base + 32), + _read_u64(view, base + 40), + ) + + def string(self, uint32_t string_id): + """Decode one string pool entry without a temporary bytes slice.""" + cdef const uint8_t[:] view = self._view + cdef uint64_t offsets_base = self._section_offsets[1] + cdef uint64_t strings_base = self._section_offsets[2] + cdef uint32_t string_count = self._section_counts[1] + cdef uint32_t start + cdef uint32_t end + cdef const char* data + if string_id + 1 >= string_count: + raise IndexError(string_id) + start = _read_u32(view, offsets_base + string_id * 4) + end = _read_u32(view, offsets_base + (string_id + 1) * 4) + data = &view[strings_base + start] + return PyUnicode_DecodeUTF8(data, end - start, "strict") diff --git a/python/tsfile/dataset/index.py b/python/tsfile/dataset/index.py index 4bd9c3767..151c5cacb 100644 --- a/python/tsfile/dataset/index.py +++ b/python/tsfile/dataset/index.py @@ -37,6 +37,11 @@ from ..constants import ColumnCategory from .metadata import MODEL_TREE, _join_series_path +try: + from ._index_fast import FastIndexLookup +except ImportError: # pragma: no cover - source checkouts may lack built extensions + FastIndexLookup = None + MAGIC = b"TSIDX\0\0\0" VERSION_MAJOR = 1 VERSION_MINOR = 0 @@ -229,6 +234,7 @@ class MappedDatasetIndex: def __init__(self, path: str, verify_sections: bool = False): self.path = path + self._fast_lookup = None self._file = open(path, "rb") try: stat = os.fstat(self._file.fileno()) @@ -241,7 +247,10 @@ def __init__(self, path: str, verify_sections: bool = False): self._mmap = mmap.mmap(self._file.fileno(), 0, access=mmap.ACCESS_READ) self._view = memoryview(self._mmap) self._entries = self._validate(verify_sections) + if FastIndexLookup is not None: + self._fast_lookup = FastIndexLookup(self._view, self._entries) except Exception: + self._fast_lookup = None if getattr(self, "_view", None) is not None: self._view.release() self._view = None @@ -336,6 +345,7 @@ def _validate(self, verify_sections: bool): return entries def close(self): + self._fast_lookup = None if getattr(self, "_view", None) is not None: self._view.release() self._view = None @@ -383,6 +393,8 @@ def string_bytes(self, sid: int) -> bytes: return bytes(self._view[strings[2] + start : strings[2] + end]) def string(self, sid: int) -> str: + if self._fast_lookup is not None: + return self._fast_lookup.string(sid) return self.string_bytes(sid).decode("utf-8") def _equal_hash_range(self, section_type: int, hash_index: int, value_hash: int): @@ -437,14 +449,20 @@ def _find_child( raise KeyError(name) def find_device_id(self, table_id: int, name: str) -> int: + if self._fast_lookup is not None: + return self._fast_lookup.find_device_id(table_id, name) table = self.record(TABLE_RECORD, table_id) return self._find_child(DEVICE_NAME_INDEX, table_id, name, table[2], table[3]) def find_column_id(self, table_id: int, name: str) -> int: + if self._fast_lookup is not None: + return self._fast_lookup.find_column_id(table_id, name) table = self.record(TABLE_RECORD, table_id) return self._find_child(COLUMN_NAME_INDEX, table_id, name, table[4], table[5]) def find_series_id(self, device_id: int, column_id: int) -> int: + if self._fast_lookup is not None: + return self._fast_lookup.find_series_id(device_id, column_id) device = self.record(DEVICE_RECORD, device_id) low, high = device[4], device[4] + device[5] while low < high: @@ -460,6 +478,107 @@ def find_series_id(self, device_id: int, column_id: int) -> int: return low raise KeyError(column_id) + def describe_series(self, series_id: int): + """Return scalar route metadata for one logical series. + + The optional Cython backend reads the packed mmap directly and avoids + creating intermediate ``struct.unpack`` tuples for every span. The + Python fallback preserves the same shape for source checkouts without + a compiled extension. + """ + if self._fast_lookup is not None: + return self._fast_lookup.describe_series(series_id) + + series = self.record(LOGICAL_SERIES, series_id) + shards = [] + count = 0 + for span_id in range(series[2], series[2] + series[3]): + span = self.record(SERIES_FILE_SPAN, span_id) + locator = self.record(SERIES_LOCATOR, span[2]) + device_span = self.record(DEVICE_FILE_SPAN, locator[0]) + timeline_length = device_span[6] if device_span[4] == 1 else span[6] + count += timeline_length + shards.append((span[1], span[2], timeline_length, span[4], span[5])) + return series[0], series[1], series[4], series[5], count, shards + + def series_identity(self, series_id: int): + """Return device and column ids for one logical series.""" + if self._fast_lookup is not None: + return self._fast_lookup.series_identity(series_id) + series = self.record(LOGICAL_SERIES, series_id) + return series[0], series[1] + + def find_series_span(self, series_id: int, file_id: int): + """Return one series span without exposing its full record tuple.""" + if self._fast_lookup is not None: + return self._fast_lookup.find_series_span(series_id, file_id) + series = self.record(LOGICAL_SERIES, series_id) + for span_id in range(series[2], series[2] + series[3]): + span = self.record(SERIES_FILE_SPAN, span_id) + if span[1] == file_id: + return span[2], span[4], span[5], span[6] + raise KeyError((series_id, file_id)) + + def locator_metadata(self, locator_id: int): + """Return locator/device-span fields needed by the runtime reader.""" + if self._fast_lookup is not None: + return self._fast_lookup.locator_metadata(locator_id) + locator = self.record(SERIES_LOCATOR, locator_id) + device_span = self.record(DEVICE_FILE_SPAN, locator[0]) + return ( + locator[0], + locator[1], + device_span[1], + device_span[4], + device_span[6], + ) + + def prepared_locator_metadata(self, file_id: int, locator_id: int): + """Return generation and locator fields used by native prepare.""" + if self._fast_lookup is not None: + return self._fast_lookup.prepared_locator_metadata(file_id, locator_id) + locator = self.record(SERIES_LOCATOR, locator_id) + device_span = self.record(DEVICE_FILE_SPAN, locator[0]) + file_record = self.record(TSFILE_RECORD, file_id) + if device_span[1] != file_id: + raise ValueError("series locator points at another TsFile") + return ( + file_record[2], + file_record[3], + locator[1], + locator[2], + locator[3], + locator[4], + device_span[2], + device_span[3], + ) + + def device_route(self, device_id: int): + """Return table id and logical-path string id for one device.""" + if self._fast_lookup is not None: + return self._fast_lookup.device_route(device_id) + device = self.record(DEVICE_RECORD, device_id) + return device[0], device[1] + + def table_name_id(self, table_id: int): + """Return the string-pool id for one table name.""" + if self._fast_lookup is not None: + return self._fast_lookup.table_name_id(table_id) + return self.record(TABLE_RECORD, table_id)[0] + + def column_name_id(self, column_id: int): + """Return the string-pool id for one column name.""" + if self._fast_lookup is not None: + return self._fast_lookup.column_name_id(column_id) + return self.record(COLUMN_SCHEMA, column_id)[1] + + def device_time_bounds(self, device_id: int): + """Return min/max timestamps for one device.""" + if self._fast_lookup is not None: + return self._fast_lookup.device_time_bounds(device_id) + device = self.record(DEVICE_RECORD, device_id) + return device[8], device[9] + def index_path_for(paths: Sequence[str]) -> str: canonical_paths = sorted(os.path.abspath(path) for path in paths) diff --git a/python/tsfile/dataset/runtime.py b/python/tsfile/dataset/runtime.py index 3ce096b1d..e06784f41 100644 --- a/python/tsfile/dataset/runtime.py +++ b/python/tsfile/dataset/runtime.py @@ -25,6 +25,7 @@ from concurrent.futures import ThreadPoolExecutor, wait from dataclasses import dataclass import os +import sys import threading from typing import Dict, Optional, Tuple @@ -57,9 +58,34 @@ from .merge import build_aligned_matrix _SERIES_DESCRIPTOR_CACHE_SIZE = 4096 +_DATACLASS_SLOTS = {"slots": True} if sys.version_info >= (3, 10) else {} -@dataclass(frozen=True) +def _descriptor_cache_size_from_env() -> int: + """Return the bounded descriptor-cache size for this runtime. + + The default remains the historical value, while allowing large random + workloads to disable or tune the cache without adding a dataset-sized + process-local map. Invalid values deliberately fall back to the default. + """ + raw_value = os.environ.get("TSFILE_DATAFRAME_DESCRIPTOR_CACHE_SIZE") + if raw_value is None: + return _SERIES_DESCRIPTOR_CACHE_SIZE + try: + return max(0, int(raw_value)) + except ValueError: + return _SERIES_DESCRIPTOR_CACHE_SIZE + + +def _descriptor_cache_stats_enabled() -> bool: + return os.environ.get("TSFILE_DATAFRAME_DESCRIPTOR_CACHE_STATS") in { + "1", + "true", + "True", + } + + +@dataclass(frozen=True, **_DATACLASS_SLOTS) class RuntimeSeriesShard: """One immutable physical fragment already expanded from the mmap route.""" @@ -72,7 +98,7 @@ class RuntimeSeriesShard: max_time: int -@dataclass(frozen=True) +@dataclass(frozen=True, **_DATACLASS_SLOTS) class RuntimeSeriesDescriptor: """Bounded process-local expansion of one logical series route.""" @@ -260,23 +286,28 @@ def __init__(self, index: MappedDatasetIndex): self._closed = False def _locator_tuple(self, file_id, locator_id): - 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: - raise ValueError("series locator points at another TsFile") + ( + file_size, + file_fingerprint_value, + locator_flags, + locator_layout, + locator_offset, + locator_length, + device_span_offset, + device_span_length, + ) = self._index.prepared_locator_metadata(file_id, locator_id) return ( id(self._index), file_id, - file_record[2], - file_record[3], + file_size, + file_fingerprint_value, locator_id, - locator[1], - locator[2], - locator[3], - locator[4], - device_span[2], - device_span[3], + locator_flags, + locator_layout, + locator_offset, + locator_length, + device_span_offset, + device_span_length, ) def get(self, file_id, locator_id, reader, time_owner=None): @@ -592,50 +623,58 @@ def __init__(self, catalog, cache_size): self._cache = OrderedDict() self._cache_lock = threading.Lock() + def _record_cache_stat(self, name): + if self._catalog._descriptor_cache_stats_enabled: + self._catalog._descriptor_cache_stats[name] += 1 + def _series_id(self, ref): device_id, field_idx = ref - device = self._catalog.index.record(DEVICE_RECORD, device_id) + table_id, _path_string_id = self._catalog.index.device_route(device_id) table_name = self._catalog.index.string( - self._catalog.index.record(TABLE_RECORD, device[0])[0] + self._catalog.index.table_name_id(table_id) ) field_name = self._catalog.table_entries[table_name].field_columns[field_idx] - column_id = self._catalog.index.find_column_id(device[0], field_name) + column_id = self._catalog.index.find_column_id(table_id, field_name) return self._catalog.index.find_series_id(device_id, column_id) def describe(self, ref, series_id=None, column_id=None): - with self._cache_lock: - result = self._cache.get(ref) - if result is not None: - self._cache.move_to_end(ref) - return result + if self._cache_size: + with self._cache_lock: + result = self._cache.get(ref) + if result is not None: + self._record_cache_stat("hits") + self._cache.move_to_end(ref) + return result + self._record_cache_stat("misses") if series_id is None: series_id = self._series_id(ref) - series = self._catalog.index.record(LOGICAL_SERIES, series_id) - if series[0] != ref[0]: + ( + series_device_id, + described_column_id, + series_min_time, + series_max_time, + count, + shard_infos, + ) = self._catalog.index.describe_series(series_id) + if series_device_id != ref[0]: raise KeyError(ref) if column_id is None: - column_id = series[1] - elif series[1] != column_id: + column_id = described_column_id + elif described_column_id != column_id: raise KeyError(ref) shards = [] - count = 0 - for span_id in range(series[2], series[2] + series[3]): - span = self._catalog.index.record(SERIES_FILE_SPAN, span_id) - locator = self._catalog.index.record(SERIES_LOCATOR, span[2]) - device_span = self._catalog.index.record(DEVICE_FILE_SPAN, locator[0]) - timeline_length = device_span[6] if device_span[4] == 1 else span[6] - count += timeline_length + for file_id, locator_id, timeline_length, min_time, max_time in shard_infos: shards.append( RuntimeSeriesShard( - self._catalog.reader_for(span[1]), - series[0], + self._catalog.reader_for(file_id), + series_device_id, column_id, - span[2], + locator_id, timeline_length, - span[4], - span[5], + min_time, + max_time, ) ) @@ -644,19 +683,21 @@ def describe(self, ref, series_id=None, column_id=None): series_id, column_id, tuple(shards), - series[4] if count else None, - series[5] if count else None, + series_min_time if count else None, + series_max_time if count else None, count, ) if self._cache_size: with self._cache_lock: existing = self._cache.get(ref) if existing is not None: + self._record_cache_stat("hits") self._cache.move_to_end(ref) return existing self._cache[ref] = result while len(self._cache) > self._cache_size: self._cache.popitem(last=False) + self._record_cache_stat("evictions") return result def __contains__(self, ref): @@ -681,7 +722,13 @@ def __init__(self, runtime: DatasetRuntime): self.runtime = runtime self.index = runtime.index self.index_identity = runtime.index.identity - self._descriptor_cache_size = _SERIES_DESCRIPTOR_CACHE_SIZE + self._descriptor_cache_size = _descriptor_cache_size_from_env() + self._descriptor_cache_stats_enabled = _descriptor_cache_stats_enabled() + self._descriptor_cache_stats = { + "hits": 0, + "misses": 0, + "evictions": 0, + } self._descriptor_cache = OrderedDict() self._descriptor_cache_lock = threading.Lock() self.table_entries = _TableMapping(self) @@ -695,11 +742,16 @@ def __init__(self, runtime: DatasetRuntime): def resolve_series_descriptor(self, table_name, tags, field_name): key = (table_name, tuple(tags), field_name) - with self._descriptor_cache_lock: - result = self._descriptor_cache.get(key) - if result is not None: - self._descriptor_cache.move_to_end(key) - return result + if self._descriptor_cache_size: + with self._descriptor_cache_lock: + result = self._descriptor_cache.get(key) + if result is not None: + if self._descriptor_cache_stats_enabled: + self._descriptor_cache_stats["hits"] += 1 + self._descriptor_cache.move_to_end(key) + return result + if self._descriptor_cache_stats_enabled: + self._descriptor_cache_stats["misses"] += 1 table_id = self.table_entries.table_id(table_name) table = self.table_entries[table_name] @@ -717,21 +769,33 @@ def resolve_series_descriptor(self, table_name, tags, field_name): with self._descriptor_cache_lock: existing = self._descriptor_cache.get(key) if existing is not None: + if self._descriptor_cache_stats_enabled: + self._descriptor_cache_stats["hits"] += 1 self._descriptor_cache.move_to_end(key) return existing self._descriptor_cache[key] = result while len(self._descriptor_cache) > self._descriptor_cache_size: self._descriptor_cache.popitem(last=False) + if self._descriptor_cache_stats_enabled: + self._descriptor_cache_stats["evictions"] += 1 return result + @property + def descriptor_cache_stats(self): + """Return a snapshot of descriptor cache activity for diagnostics.""" + stats = dict(self._descriptor_cache_stats) + stats["size"] = len(self._descriptor_cache) + stats["capacity"] = self._descriptor_cache_size + return stats + def resolve_series_descriptor_by_id(self, series_id, table_name, field_name): if series_id < 0 or series_id >= self.index.count(LOGICAL_SERIES): raise KeyError(series_id) - series = self.index.record(LOGICAL_SERIES, series_id) + device_id, column_id = self.index.series_identity(series_id) table = self.table_entries[table_name] field_idx = table.get_field_index(field_name) return self.series_shards.describe( - (series[0], field_idx), series_id=series_id, column_id=series[1] + (device_id, field_idx), series_id=series_id, column_id=column_id ) def _infer_model(self): @@ -760,8 +824,7 @@ def __len__(self): return len(self._catalog.devices) def __getitem__(self, device_id): - record = self._catalog.index.record(DEVICE_RECORD, device_id) - return record[8], record[9] + return self._catalog.index.device_time_bounds(device_id) class RuntimeSeriesReader: @@ -776,49 +839,50 @@ def _series(self, device_id, column_id): def _span(self, device_id, column_id): series_id = self._series(device_id, column_id) - series = self.runtime.index.record(LOGICAL_SERIES, series_id) - for span_id in range(series[2], series[2] + series[3]): - span = self.runtime.index.record(SERIES_FILE_SPAN, span_id) - if span[1] == self.file_id: - return span - raise KeyError((device_id, column_id, self.file_id)) + return self.runtime.index.find_series_span(series_id, self.file_id) def _identity(self, device_id, column_id): index = self.runtime.index - device = index.record(DEVICE_RECORD, device_id) - table = index.record(TABLE_RECORD, device[0]) - table_name = index.string(table[0]) - components = split_logical_series_path(index.string(device[1])) + table_id, path_string_id = index.device_route(device_id) + table_name = index.string(index.table_name_id(table_id)) + components = split_logical_series_path(index.string(path_string_id)) tags = tuple(components[1:-1]) table_entry = self.runtime.catalog.table_entries[table_name] - column_name = index.string(index.record(COLUMN_SCHEMA, column_id)[1]) + column_name = index.string(index.column_name_id(column_id)) return table_name, tags, table_entry, column_name def get_device_info(self, device_id): - record = self.runtime.index.record(DEVICE_RECORD, device_id) table_name, tags = self.runtime.catalog.devices[device_id] table = self.runtime.catalog.table_entries[table_name] + min_time, max_time = self.runtime.index.device_time_bounds(device_id) return { "table_name": table_name, "tag_columns": table.tag_columns, "tag_values": dict(zip(table.tag_columns, tags)), - "min_time": record[8], - "max_time": record[9], + "min_time": min_time, + "max_time": max_time, } def get_series_info_by_ref(self, device_id, column_id): - span = self._span(device_id, column_id) - locator = self.runtime.index.record(SERIES_LOCATOR, span[2]) - device_span = self.runtime.index.record(DEVICE_FILE_SPAN, locator[0]) + locator_id, span_min_time, span_max_time, span_length = self._span( + device_id, column_id + ) + ( + _device_span_id, + _locator_flags, + _file_id, + layout, + device_timeline_length, + ) = self.runtime.index.locator_metadata(locator_id) table_name, tags, table, column_name = self._identity(device_id, column_id) - timeline_length = device_span[6] if device_span[4] == 1 else span[6] + timeline_length = device_timeline_length if layout == 1 else span_length return { "length": timeline_length, - "min_time": span[4], - "max_time": span[5], + "min_time": span_min_time, + "max_time": span_max_time, "timeline_length": timeline_length, - "timeline_min_time": span[4], - "timeline_max_time": span[5], + "timeline_min_time": span_min_time, + "timeline_max_time": span_max_time, "table_name": table_name, "column_name": column_name, "device_id": device_id, @@ -904,7 +968,7 @@ def _query( ): span = self._span(device_id, column_id) return self._query_at_locator( - span[2], + span[0], start_time=start_time, end_time=end_time, offset=offset, @@ -950,19 +1014,17 @@ def read_device_fields_by_time_range( return np.array([], dtype=np.int64), {} spans = [self._span(device_id, column_id) for column_id in column_ids] - locators = [ - self.runtime.index.record(SERIES_LOCATOR, span[2]) for span in spans + locator_metadata = [ + self.runtime.index.locator_metadata(span[0]) for span in spans ] - device_span_ids = {locator[0] for locator in locators} + device_span_ids = {metadata[0] for metadata in locator_metadata} can_read_aligned = len(device_span_ids) == 1 if can_read_aligned: - device_span = self.runtime.index.record( - DEVICE_FILE_SPAN, next(iter(device_span_ids)) - ) + device_span = locator_metadata[0] can_read_aligned = ( - device_span[1] == self.file_id - and device_span[4] == 1 - and all(locator[1] == 1 for locator in locators) + device_span[2] == self.file_id + and device_span[3] == 1 + and all(metadata[1] == 1 for metadata in locator_metadata) ) if can_read_aligned: @@ -974,7 +1036,7 @@ def read_device_fields_by_time_range( time_owner = None for span in spans: current = self.runtime.prepared.get( - self.file_id, span[2], reader, time_owner=time_owner + self.file_id, span[0], reader, time_owner=time_owner ) prepared.append(current) if time_owner is None: From 4ac207d0b28fe14cac22a44bf2e3ae91e7f01973 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Mon, 14 Sep 2026 12:08:10 +0800 Subject: [PATCH 2/4] build(python): require Dataset Index Cython extension --- python/setup.py | 4 +- python/tests/test_dataset_index.py | 57 ++------ .../dataset/{_index_fast.pyx => _index.pyx} | 4 +- python/tsfile/dataset/index.py | 131 +++--------------- 4 files changed, 37 insertions(+), 159 deletions(-) rename python/tsfile/dataset/{_index_fast.pyx => _index.pyx} (99%) diff --git a/python/setup.py b/python/setup.py index 7f41a07f1..fe19fc91c 100644 --- a/python/setup.py +++ b/python/setup.py @@ -274,8 +274,8 @@ def finalize_options(self): exts = [ Extension("tsfile.dataset._merge", ["tsfile/dataset/_merge.pyx"], **merge_common), Extension( - "tsfile.dataset._index_fast", - ["tsfile/dataset/_index_fast.pyx"], + "tsfile.dataset._index", + ["tsfile/dataset/_index.pyx"], **merge_common, ), Extension("tsfile.tsfile_py_cpp", ["tsfile/tsfile_py_cpp.pyx"], **common), diff --git a/python/tests/test_dataset_index.py b/python/tests/test_dataset_index.py index c521958af..3cb4ac56d 100644 --- a/python/tests/test_dataset_index.py +++ b/python/tests/test_dataset_index.py @@ -124,11 +124,10 @@ def test_build_publish_map_and_lookup(tmp_path): assert index.string(file_record[0]) == str(source) -@pytest.mark.skipif( - index_module.FastIndexLookup is None, - reason="requires the built Cython Dataset Index lookup extension", -) -def test_hot_lookup_does_not_unpack_python_record_tuples(tmp_path, monkeypatch): +def test_index_lookup_is_required_and_does_not_unpack_python_record_tuples( + tmp_path, monkeypatch +): + assert index_module.IndexLookup is not None source = tmp_path / "source.tsfile" source.write_bytes(b"T" * 4096) output = tmp_path / "dataset.tsidx" @@ -136,9 +135,11 @@ def test_hot_lookup_does_not_unpack_python_record_tuples(tmp_path, monkeypatch): write_index_atomic(str(output), build_sections_from_dataframe(dataframe)) with MappedDatasetIndex(str(output), verify_sections=True) as index: + assert isinstance(index._lookup, index_module.IndexLookup) + assert not hasattr(index, "_fast_lookup") def fail_record(*_args, **_kwargs): - raise AssertionError("hot lookup should not unpack Python record tuples") + raise AssertionError("index lookup should not unpack Python record tuples") monkeypatch.setattr(index, "record", fail_record) assert index.find_device_id(0, "root.") == 0 @@ -146,13 +147,7 @@ def fail_record(*_args, **_kwargs): assert index.find_series_id(0, 0) == 0 -@pytest.mark.skipif( - index_module.FastIndexLookup is None, - reason="requires the built Cython Dataset Index lookup extension", -) -def test_fast_series_description_does_not_unpack_python_record_tuples( - tmp_path, monkeypatch -): +def test_series_description_does_not_unpack_python_record_tuples(tmp_path, monkeypatch): source = tmp_path / "source.tsfile" source.write_bytes(b"T" * 4096) output = tmp_path / "dataset.tsidx" @@ -180,13 +175,7 @@ def fail_record(*_args, **_kwargs): assert shards == [(span[1], span[2], expected_count, span[4], span[5])] -@pytest.mark.skipif( - index_module.FastIndexLookup is None, - reason="requires the built Cython Dataset Index lookup extension", -) -def test_fast_span_and_locator_metadata_do_not_unpack_python_records( - tmp_path, monkeypatch -): +def test_span_and_locator_metadata_do_not_unpack_python_records(tmp_path, monkeypatch): source = tmp_path / "source.tsfile" source.write_bytes(b"T" * 4096) output = tmp_path / "dataset.tsidx" @@ -218,11 +207,7 @@ def fail_record(*_args, **_kwargs): ) -@pytest.mark.skipif( - index_module.FastIndexLookup is None, - reason="requires the built Cython Dataset Index lookup extension", -) -def test_runtime_reader_span_lookup_uses_fast_metadata(tmp_path, monkeypatch): +def test_runtime_reader_span_lookup_uses_index_metadata(tmp_path, monkeypatch): source = tmp_path / "part.tsfile" _write_runtime_file(source, 0) @@ -242,11 +227,7 @@ def fail_record(*_args, **_kwargs): assert length >= 0 -@pytest.mark.skipif( - index_module.FastIndexLookup is None, - reason="requires the built Cython Dataset Index lookup extension", -) -def test_fast_prepared_locator_metadata_does_not_unpack_python_records( +def test_prepared_locator_metadata_does_not_unpack_python_records( tmp_path, monkeypatch ): source = tmp_path / "source.tsfile" @@ -277,13 +258,7 @@ def fail_record(*_args, **_kwargs): ) -@pytest.mark.skipif( - index_module.FastIndexLookup is None, - reason="requires the built Cython Dataset Index lookup extension", -) -def test_fast_identity_and_device_bounds_do_not_unpack_python_records( - tmp_path, monkeypatch -): +def test_identity_and_device_bounds_do_not_unpack_python_records(tmp_path, monkeypatch): source = tmp_path / "source.tsfile" source.write_bytes(b"T" * 4096) output = tmp_path / "dataset.tsidx" @@ -304,7 +279,7 @@ def fail_record(*_args, **_kwargs): raise AssertionError("identity metadata should not unpack records") def fail_string_bytes(*_args, **_kwargs): - raise AssertionError("fast string should not copy bytes") + raise AssertionError("index string lookup should not copy bytes") monkeypatch.setattr(index, "record", fail_record) monkeypatch.setattr(index, "string_bytes", fail_string_bytes) @@ -315,11 +290,7 @@ def fail_string_bytes(*_args, **_kwargs): assert index.string(expected_column_name_id) == expected_name -@pytest.mark.skipif( - index_module.FastIndexLookup is None, - reason="requires the built Cython Dataset Index lookup extension", -) -def test_fast_series_identity_does_not_unpack_python_records(tmp_path, monkeypatch): +def test_series_identity_does_not_unpack_python_records(tmp_path, monkeypatch): source = tmp_path / "source.tsfile" source.write_bytes(b"T" * 4096) output = tmp_path / "dataset.tsidx" diff --git a/python/tsfile/dataset/_index_fast.pyx b/python/tsfile/dataset/_index.pyx similarity index 99% rename from python/tsfile/dataset/_index_fast.pyx rename to python/tsfile/dataset/_index.pyx index 329d48b15..d1412f668 100644 --- a/python/tsfile/dataset/_index_fast.pyx +++ b/python/tsfile/dataset/_index.pyx @@ -17,7 +17,7 @@ # cython: boundscheck=False, wraparound=False, cdivision=True, language_level=3 -"""Cython hot lookups over the existing read-only Dataset Index mmap. +"""Cython accessors for the existing read-only Dataset Index mmap. This module deliberately owns no index-sized data structures. It receives the already-mapped byte view and section directory from ``MappedDatasetIndex`` and @@ -159,7 +159,7 @@ cdef const char* _utf8_name(object name, Py_ssize_t* length): return data -cdef class FastIndexLookup: +cdef class IndexLookup: """Typed lookup kernel over a ``MappedDatasetIndex`` memoryview.""" cdef const uint8_t[:] _view diff --git a/python/tsfile/dataset/index.py b/python/tsfile/dataset/index.py index 151c5cacb..cbcc85d52 100644 --- a/python/tsfile/dataset/index.py +++ b/python/tsfile/dataset/index.py @@ -36,11 +36,7 @@ from ..constants import ColumnCategory from .metadata import MODEL_TREE, _join_series_path - -try: - from ._index_fast import FastIndexLookup -except ImportError: # pragma: no cover - source checkouts may lack built extensions - FastIndexLookup = None +from ._index import IndexLookup MAGIC = b"TSIDX\0\0\0" VERSION_MAJOR = 1 @@ -234,7 +230,7 @@ class MappedDatasetIndex: def __init__(self, path: str, verify_sections: bool = False): self.path = path - self._fast_lookup = None + self._lookup = None self._file = open(path, "rb") try: stat = os.fstat(self._file.fileno()) @@ -247,10 +243,9 @@ def __init__(self, path: str, verify_sections: bool = False): self._mmap = mmap.mmap(self._file.fileno(), 0, access=mmap.ACCESS_READ) self._view = memoryview(self._mmap) self._entries = self._validate(verify_sections) - if FastIndexLookup is not None: - self._fast_lookup = FastIndexLookup(self._view, self._entries) + self._lookup = IndexLookup(self._view, self._entries) except Exception: - self._fast_lookup = None + self._lookup = None if getattr(self, "_view", None) is not None: self._view.release() self._view = None @@ -345,7 +340,7 @@ def _validate(self, verify_sections: bool): return entries def close(self): - self._fast_lookup = None + self._lookup = None if getattr(self, "_view", None) is not None: self._view.release() self._view = None @@ -393,9 +388,7 @@ def string_bytes(self, sid: int) -> bytes: return bytes(self._view[strings[2] + start : strings[2] + end]) def string(self, sid: int) -> str: - if self._fast_lookup is not None: - return self._fast_lookup.string(sid) - return self.string_bytes(sid).decode("utf-8") + return self._lookup.string(sid) def _equal_hash_range(self, section_type: int, hash_index: int, value_hash: int): low = 0 @@ -449,135 +442,49 @@ def _find_child( raise KeyError(name) def find_device_id(self, table_id: int, name: str) -> int: - if self._fast_lookup is not None: - return self._fast_lookup.find_device_id(table_id, name) - table = self.record(TABLE_RECORD, table_id) - return self._find_child(DEVICE_NAME_INDEX, table_id, name, table[2], table[3]) + return self._lookup.find_device_id(table_id, name) def find_column_id(self, table_id: int, name: str) -> int: - if self._fast_lookup is not None: - return self._fast_lookup.find_column_id(table_id, name) - table = self.record(TABLE_RECORD, table_id) - return self._find_child(COLUMN_NAME_INDEX, table_id, name, table[4], table[5]) + return self._lookup.find_column_id(table_id, name) def find_series_id(self, device_id: int, column_id: int) -> int: - if self._fast_lookup is not None: - return self._fast_lookup.find_series_id(device_id, column_id) - device = self.record(DEVICE_RECORD, device_id) - low, high = device[4], device[4] + device[5] - while low < high: - middle = low + (high - low) // 2 - if self.record(LOGICAL_SERIES, middle)[1] < column_id: - low = middle + 1 - else: - high = middle - if ( - low < device[4] + device[5] - and self.record(LOGICAL_SERIES, low)[1] == column_id - ): - return low - raise KeyError(column_id) + return self._lookup.find_series_id(device_id, column_id) def describe_series(self, series_id: int): - """Return scalar route metadata for one logical series. - - The optional Cython backend reads the packed mmap directly and avoids - creating intermediate ``struct.unpack`` tuples for every span. The - Python fallback preserves the same shape for source checkouts without - a compiled extension. - """ - if self._fast_lookup is not None: - return self._fast_lookup.describe_series(series_id) - - series = self.record(LOGICAL_SERIES, series_id) - shards = [] - count = 0 - for span_id in range(series[2], series[2] + series[3]): - span = self.record(SERIES_FILE_SPAN, span_id) - locator = self.record(SERIES_LOCATOR, span[2]) - device_span = self.record(DEVICE_FILE_SPAN, locator[0]) - timeline_length = device_span[6] if device_span[4] == 1 else span[6] - count += timeline_length - shards.append((span[1], span[2], timeline_length, span[4], span[5])) - return series[0], series[1], series[4], series[5], count, shards + """Return scalar route metadata for one logical series.""" + return self._lookup.describe_series(series_id) def series_identity(self, series_id: int): """Return device and column ids for one logical series.""" - if self._fast_lookup is not None: - return self._fast_lookup.series_identity(series_id) - series = self.record(LOGICAL_SERIES, series_id) - return series[0], series[1] + return self._lookup.series_identity(series_id) def find_series_span(self, series_id: int, file_id: int): """Return one series span without exposing its full record tuple.""" - if self._fast_lookup is not None: - return self._fast_lookup.find_series_span(series_id, file_id) - series = self.record(LOGICAL_SERIES, series_id) - for span_id in range(series[2], series[2] + series[3]): - span = self.record(SERIES_FILE_SPAN, span_id) - if span[1] == file_id: - return span[2], span[4], span[5], span[6] - raise KeyError((series_id, file_id)) + return self._lookup.find_series_span(series_id, file_id) def locator_metadata(self, locator_id: int): """Return locator/device-span fields needed by the runtime reader.""" - if self._fast_lookup is not None: - return self._fast_lookup.locator_metadata(locator_id) - locator = self.record(SERIES_LOCATOR, locator_id) - device_span = self.record(DEVICE_FILE_SPAN, locator[0]) - return ( - locator[0], - locator[1], - device_span[1], - device_span[4], - device_span[6], - ) + return self._lookup.locator_metadata(locator_id) def prepared_locator_metadata(self, file_id: int, locator_id: int): """Return generation and locator fields used by native prepare.""" - if self._fast_lookup is not None: - return self._fast_lookup.prepared_locator_metadata(file_id, locator_id) - locator = self.record(SERIES_LOCATOR, locator_id) - device_span = self.record(DEVICE_FILE_SPAN, locator[0]) - file_record = self.record(TSFILE_RECORD, file_id) - if device_span[1] != file_id: - raise ValueError("series locator points at another TsFile") - return ( - file_record[2], - file_record[3], - locator[1], - locator[2], - locator[3], - locator[4], - device_span[2], - device_span[3], - ) + return self._lookup.prepared_locator_metadata(file_id, locator_id) def device_route(self, device_id: int): """Return table id and logical-path string id for one device.""" - if self._fast_lookup is not None: - return self._fast_lookup.device_route(device_id) - device = self.record(DEVICE_RECORD, device_id) - return device[0], device[1] + return self._lookup.device_route(device_id) def table_name_id(self, table_id: int): """Return the string-pool id for one table name.""" - if self._fast_lookup is not None: - return self._fast_lookup.table_name_id(table_id) - return self.record(TABLE_RECORD, table_id)[0] + return self._lookup.table_name_id(table_id) def column_name_id(self, column_id: int): """Return the string-pool id for one column name.""" - if self._fast_lookup is not None: - return self._fast_lookup.column_name_id(column_id) - return self.record(COLUMN_SCHEMA, column_id)[1] + return self._lookup.column_name_id(column_id) def device_time_bounds(self, device_id: int): """Return min/max timestamps for one device.""" - if self._fast_lookup is not None: - return self._fast_lookup.device_time_bounds(device_id) - device = self.record(DEVICE_RECORD, device_id) - return device[8], device[9] + return self._lookup.device_time_bounds(device_id) def index_path_for(paths: Sequence[str]) -> str: From 01321c2fe91fba5d4d8a9e80f0e1a577a8f78489 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 17 Sep 2026 12:27:00 +0800 Subject: [PATCH 3/4] fix(python): bound Dataset Index lookups against corrupt ranges The Cython fast paths indexed the mmap with child ranges read from the index records without validating them, so a corrupt or hand-crafted .tsidx could read past the mapping and crash the interpreter (SIGSEGV/SIGBUS) instead of raising, which the previous pure-Python path did via record(). Validate the device/column/series child ranges the same way describe_series and find_series_span already do, and take ids as signed values so negative ids keep raising IndexError instead of OverflowError. Also drop the now-unused pure-Python _find_child, retarget the hash-collision test at the production Cython lookup (the old test only exercised that dead helper), and skip the dataclass-slots assertion on Python 3.9 where dataclasses do not support slots. --- python/tests/test_dataset_index.py | 138 +++++++++++++++++++++++++---- python/tsfile/dataset/_index.pyx | 60 +++++++------ python/tsfile/dataset/index.py | 25 ------ 3 files changed, 153 insertions(+), 70 deletions(-) diff --git a/python/tests/test_dataset_index.py b/python/tests/test_dataset_index.py index 3cb4ac56d..bbb53889e 100644 --- a/python/tests/test_dataset_index.py +++ b/python/tests/test_dataset_index.py @@ -17,6 +17,8 @@ from types import SimpleNamespace import os +import struct +import sys import threading import numpy as np @@ -37,6 +39,8 @@ from tsfile.dataset.index import ( COLUMN_SCHEMA, DEVICE_FILE_SPAN, + DEVICE_NAME_INDEX, + DEVICE_RECORD, DIRECTORY, HEADER, LOGICAL_SERIES, @@ -44,6 +48,7 @@ RECORDS, SERIES_FILE_SPAN, SERIES_LOCATOR, + TABLE_RECORD, TSFILE_RECORD, build_sections_from_dataframe, crc32c, @@ -90,6 +95,22 @@ def _synthetic_dataframe(source_path): ) +def _index_section_entry(blob, section_type): + return DIRECTORY.unpack_from( + blob, HEADER.size + (section_type - 1) * DIRECTORY.size + ) + + +def _patch_index_field(path, section_type, field_offset, value, fmt="= 3 + + # Same length as the row it must not match, so only the bytes can reject it. + tail = names[-1][:-1] + probe = None + for replacement in "xyz019": + value = tail[:-1] + replacement + "." + if ( + len(value) == len(names[-1]) + and value not in names + and index_module.name_hash(value.encode("utf-8")) >= rows[-1][2] + ): + probe = value + break + assert probe is not None + + # Give the last row the probe's hash: the lookup finds that row by hash and + # can only reject it by comparing the stored bytes. + with open(index_path, "r+b") as stream: + blob = bytearray(stream.read()) + entry = _index_section_entry(blob, DEVICE_NAME_INDEX) + struct.pack_into( + "= self._section_counts[4]: + if table_id < 0 or table_id >= self._section_counts[4]: raise IndexError(table_id) table_base = self._section_offsets[4] + table_id * self._section_sizes[4] first = _read_u32(view, table_base + 8) count = _read_u32(view, table_base + 12) + if first + count > self._section_counts[5]: + raise IndexError(table_id) target_hash = _fnv1a(data, length) with nogil: result = _find_child( @@ -210,7 +212,7 @@ cdef class IndexLookup: raise KeyError(name) return result - def find_column_id(self, uint32_t table_id, object name): + def find_column_id(self, Py_ssize_t table_id, object name): cdef const char* data cdef Py_ssize_t length cdef const uint8_t[:] view = self._view @@ -220,11 +222,13 @@ cdef class IndexLookup: cdef uint64_t target_hash cdef int64_t result data = _utf8_name(name, &length) - if table_id >= self._section_counts[4]: + if table_id < 0 or table_id >= self._section_counts[4]: raise IndexError(table_id) table_base = self._section_offsets[4] + table_id * self._section_sizes[4] first = _read_u32(view, table_base + 16) count = _read_u32(view, table_base + 20) + if first + count > self._section_counts[7]: + raise IndexError(table_id) target_hash = _fnv1a(data, length) with nogil: result = _find_child( @@ -245,7 +249,7 @@ cdef class IndexLookup: raise KeyError(name) return result - def find_series_id(self, uint32_t device_id, uint32_t column_id): + def find_series_id(self, Py_ssize_t device_id, Py_ssize_t column_id): cdef const uint8_t[:] view = self._view cdef uint64_t device_base cdef uint32_t first @@ -257,11 +261,13 @@ cdef class IndexLookup: cdef uint32_t row_column_id cdef uint32_t row_device_id cdef int64_t result = -1 - if device_id >= self._section_counts[6]: + if device_id < 0 or device_id >= self._section_counts[6]: raise IndexError(device_id) device_base = self._section_offsets[6] + device_id * self._section_sizes[6] first = _read_u32(view, device_base + 16) count = _read_u32(view, device_base + 20) + if first + count > self._section_counts[9]: + raise IndexError(device_id) low = first high = first + count with nogil: @@ -283,7 +289,7 @@ cdef class IndexLookup: raise KeyError(column_id) return result - def describe_series(self, uint32_t series_id): + def describe_series(self, Py_ssize_t series_id): """Expand one logical series without creating intermediate records.""" cdef const uint8_t[:] view = self._view cdef uint64_t series_base @@ -308,7 +314,7 @@ cdef class IndexLookup: cdef uint64_t count = 0 cdef list shards = [] - if series_id >= self._section_counts[9]: + if series_id < 0 or series_id >= self._section_counts[9]: raise IndexError(series_id) series_base = ( @@ -361,16 +367,16 @@ cdef class IndexLookup: return device_id, column_id, min_time, max_time, count, shards - def series_identity(self, uint32_t series_id): + def series_identity(self, Py_ssize_t series_id): """Return device and column ids for one logical series.""" cdef const uint8_t[:] view = self._view cdef uint64_t base - if series_id >= self._section_counts[9]: + if series_id < 0 or series_id >= self._section_counts[9]: raise IndexError(series_id) base = self._section_offsets[9] + series_id * self._section_sizes[9] return _read_u32(view, base), _read_u32(view, base + 4) - def find_series_span(self, uint32_t series_id, uint32_t file_id): + def find_series_span(self, Py_ssize_t series_id, Py_ssize_t file_id): """Find one series span and return locator/time/length scalars.""" cdef const uint8_t[:] view = self._view cdef uint64_t series_base @@ -380,7 +386,7 @@ cdef class IndexLookup: cdef uint32_t span_id cdef uint64_t span_base - if series_id >= self._section_counts[9]: + if series_id < 0 or series_id >= self._section_counts[9]: raise IndexError(series_id) series_base = ( self._section_offsets[9] @@ -405,14 +411,14 @@ cdef class IndexLookup: ) raise KeyError((series_id, file_id)) - def locator_metadata(self, uint32_t locator_id): + def locator_metadata(self, Py_ssize_t locator_id): """Return locator and owning device-span scalars.""" cdef const uint8_t[:] view = self._view cdef uint64_t locator_base cdef uint32_t device_span_id cdef uint64_t device_span_base - if locator_id >= self._section_counts[13]: + if locator_id < 0 or locator_id >= self._section_counts[13]: raise IndexError(locator_id) locator_base = ( self._section_offsets[13] @@ -433,7 +439,7 @@ cdef class IndexLookup: _read_u64(view, device_span_base + 24), ) - def prepared_locator_metadata(self, uint32_t file_id, uint32_t locator_id): + def prepared_locator_metadata(self, Py_ssize_t file_id, Py_ssize_t locator_id): """Return the generation and locator fields used by native prepare.""" cdef const uint8_t[:] view = self._view cdef uint64_t locator_base @@ -442,9 +448,9 @@ cdef class IndexLookup: cdef uint32_t device_span_id cdef uint32_t device_file_id - if locator_id >= self._section_counts[13]: + if locator_id < 0 or locator_id >= self._section_counts[13]: raise IndexError(locator_id) - if file_id >= self._section_counts[10]: + if file_id < 0 or file_id >= self._section_counts[10]: raise IndexError(file_id) locator_base = ( self._section_offsets[13] @@ -475,38 +481,38 @@ cdef class IndexLookup: _read_u32(view, device_span_base + 16), ) - def device_route(self, uint32_t device_id): + def device_route(self, Py_ssize_t device_id): """Return table id and logical-path string id for one device.""" cdef const uint8_t[:] view = self._view cdef uint64_t base - if device_id >= self._section_counts[6]: + if device_id < 0 or device_id >= self._section_counts[6]: raise IndexError(device_id) base = self._section_offsets[6] + device_id * self._section_sizes[6] return _read_u32(view, base), _read_u32(view, base + 4) - def table_name_id(self, uint32_t table_id): + def table_name_id(self, Py_ssize_t table_id): """Return the string-pool id for one table name.""" cdef const uint8_t[:] view = self._view cdef uint64_t base - if table_id >= self._section_counts[4]: + if table_id < 0 or table_id >= self._section_counts[4]: raise IndexError(table_id) base = self._section_offsets[4] + table_id * self._section_sizes[4] return _read_u32(view, base) - def column_name_id(self, uint32_t column_id): + def column_name_id(self, Py_ssize_t column_id): """Return the string-pool id for one column name.""" cdef const uint8_t[:] view = self._view cdef uint64_t base - if column_id >= self._section_counts[8]: + if column_id < 0 or column_id >= self._section_counts[8]: raise IndexError(column_id) base = self._section_offsets[8] + column_id * self._section_sizes[8] return _read_u32(view, base + 4) - def device_time_bounds(self, uint32_t device_id): + def device_time_bounds(self, Py_ssize_t device_id): """Return min/max timestamps for one device.""" cdef const uint8_t[:] view = self._view cdef uint64_t base - if device_id >= self._section_counts[6]: + if device_id < 0 or device_id >= self._section_counts[6]: raise IndexError(device_id) base = self._section_offsets[6] + device_id * self._section_sizes[6] return ( @@ -514,7 +520,7 @@ cdef class IndexLookup: _read_u64(view, base + 40), ) - def string(self, uint32_t string_id): + def string(self, Py_ssize_t string_id): """Decode one string pool entry without a temporary bytes slice.""" cdef const uint8_t[:] view = self._view cdef uint64_t offsets_base = self._section_offsets[1] @@ -523,7 +529,7 @@ cdef class IndexLookup: cdef uint32_t start cdef uint32_t end cdef const char* data - if string_id + 1 >= string_count: + if string_id < 0 or string_id + 1 >= string_count: raise IndexError(string_id) start = _read_u32(view, offsets_base + string_id * 4) end = _read_u32(view, offsets_base + (string_id + 1) * 4) diff --git a/python/tsfile/dataset/index.py b/python/tsfile/dataset/index.py index cbcc85d52..9b4ba539c 100644 --- a/python/tsfile/dataset/index.py +++ b/python/tsfile/dataset/index.py @@ -416,31 +416,6 @@ def find_table_ids(self, name: str) -> List[int]: if self.string_bytes(sid) == encoded ] - def _find_child( - self, section_type: int, table_id: int, name: str, first: int, count: int - ): - encoded = name.encode("utf-8") - target_hash = name_hash(encoded) - low, high = first, first + count - while low < high: - middle = low + (high - low) // 2 - row = self.record(section_type, middle) - key = (row[0], row[2], self.string_bytes(row[3])) - target = (table_id, target_hash, encoded) - if key < target: - low = middle + 1 - else: - high = middle - if low < first + count: - row = self.record(section_type, low) - if ( - row[0] == table_id - and row[2] == target_hash - and self.string_bytes(row[3]) == encoded - ): - return row[1] - raise KeyError(name) - def find_device_id(self, table_id: int, name: str) -> int: return self._lookup.find_device_id(table_id, name) From 7ac515a4cc81e443081dbdf48e1503ac6d6034fa Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 17 Sep 2026 15:33:07 +0800 Subject: [PATCH 4/4] refactor(python): drop Dataset Index descriptor cache env knobs The descriptor cache size and its hit/miss counters were exposed through TSFILE_DATAFRAME_DESCRIPTOR_CACHE_SIZE / _STATS, adding process-wide configuration surface plus a stats branch on every lookup for a bounded process-local cache that already has a sane budget. Keep the historical 4096-entry default and restore the unconditional cache reads. --- python/tests/test_dataset_index.py | 30 ------------ python/tsfile/dataset/runtime.py | 78 +++++------------------------- 2 files changed, 11 insertions(+), 97 deletions(-) diff --git a/python/tests/test_dataset_index.py b/python/tests/test_dataset_index.py index bbb53889e..4ac9f2ede 100644 --- a/python/tests/test_dataset_index.py +++ b/python/tests/test_dataset_index.py @@ -741,36 +741,6 @@ def test_runtime_descriptor_objects_use_slots(): assert not hasattr(descriptor, "__dict__") -def test_runtime_descriptor_cache_stats_report_hits_and_evictions( - tmp_path, monkeypatch -): - source = tmp_path / "devices.tsfile" - _write_runtime_devices_file(source) - monkeypatch.setattr(runtime_module, "_SERIES_DESCRIPTOR_CACHE_SIZE", 1) - monkeypatch.setenv("TSFILE_DATAFRAME_DESCRIPTOR_CACHE_STATS", "1") - - with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: - names = [str(name) for name in dataframe.list_timeseries()] - dataframe[names[0]].close() - dataframe[names[0]].close() - dataframe[names[1]].close() - - stats = dataframe._index.descriptor_cache_stats - assert stats["hits"] >= 1 - assert stats["misses"] >= 2 - assert stats["evictions"] >= 1 - - -def test_runtime_descriptor_cache_size_can_be_configured(tmp_path, monkeypatch): - source = tmp_path / "part.tsfile" - _write_runtime_file(source, 0) - monkeypatch.setenv("TSFILE_DATAFRAME_DESCRIPTOR_CACHE_SIZE", "0") - - with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: - assert dataframe._index._descriptor_cache_size == 0 - assert dataframe._index.series_shards._cache_size == 0 - - def test_reader_pool_enforces_open_file_cap(tmp_path, monkeypatch): first = tmp_path / "part1.tsfile" second = tmp_path / "part2.tsfile" diff --git a/python/tsfile/dataset/runtime.py b/python/tsfile/dataset/runtime.py index e06784f41..dca1f5993 100644 --- a/python/tsfile/dataset/runtime.py +++ b/python/tsfile/dataset/runtime.py @@ -61,30 +61,6 @@ _DATACLASS_SLOTS = {"slots": True} if sys.version_info >= (3, 10) else {} -def _descriptor_cache_size_from_env() -> int: - """Return the bounded descriptor-cache size for this runtime. - - The default remains the historical value, while allowing large random - workloads to disable or tune the cache without adding a dataset-sized - process-local map. Invalid values deliberately fall back to the default. - """ - raw_value = os.environ.get("TSFILE_DATAFRAME_DESCRIPTOR_CACHE_SIZE") - if raw_value is None: - return _SERIES_DESCRIPTOR_CACHE_SIZE - try: - return max(0, int(raw_value)) - except ValueError: - return _SERIES_DESCRIPTOR_CACHE_SIZE - - -def _descriptor_cache_stats_enabled() -> bool: - return os.environ.get("TSFILE_DATAFRAME_DESCRIPTOR_CACHE_STATS") in { - "1", - "true", - "True", - } - - @dataclass(frozen=True, **_DATACLASS_SLOTS) class RuntimeSeriesShard: """One immutable physical fragment already expanded from the mmap route.""" @@ -623,10 +599,6 @@ def __init__(self, catalog, cache_size): self._cache = OrderedDict() self._cache_lock = threading.Lock() - def _record_cache_stat(self, name): - if self._catalog._descriptor_cache_stats_enabled: - self._catalog._descriptor_cache_stats[name] += 1 - def _series_id(self, ref): device_id, field_idx = ref table_id, _path_string_id = self._catalog.index.device_route(device_id) @@ -638,14 +610,11 @@ def _series_id(self, ref): return self._catalog.index.find_series_id(device_id, column_id) def describe(self, ref, series_id=None, column_id=None): - if self._cache_size: - with self._cache_lock: - result = self._cache.get(ref) - if result is not None: - self._record_cache_stat("hits") - self._cache.move_to_end(ref) - return result - self._record_cache_stat("misses") + with self._cache_lock: + result = self._cache.get(ref) + if result is not None: + self._cache.move_to_end(ref) + return result if series_id is None: series_id = self._series_id(ref) @@ -691,13 +660,11 @@ def describe(self, ref, series_id=None, column_id=None): with self._cache_lock: existing = self._cache.get(ref) if existing is not None: - self._record_cache_stat("hits") self._cache.move_to_end(ref) return existing self._cache[ref] = result while len(self._cache) > self._cache_size: self._cache.popitem(last=False) - self._record_cache_stat("evictions") return result def __contains__(self, ref): @@ -722,13 +689,7 @@ def __init__(self, runtime: DatasetRuntime): self.runtime = runtime self.index = runtime.index self.index_identity = runtime.index.identity - self._descriptor_cache_size = _descriptor_cache_size_from_env() - self._descriptor_cache_stats_enabled = _descriptor_cache_stats_enabled() - self._descriptor_cache_stats = { - "hits": 0, - "misses": 0, - "evictions": 0, - } + self._descriptor_cache_size = _SERIES_DESCRIPTOR_CACHE_SIZE self._descriptor_cache = OrderedDict() self._descriptor_cache_lock = threading.Lock() self.table_entries = _TableMapping(self) @@ -742,16 +703,11 @@ def __init__(self, runtime: DatasetRuntime): def resolve_series_descriptor(self, table_name, tags, field_name): key = (table_name, tuple(tags), field_name) - if self._descriptor_cache_size: - with self._descriptor_cache_lock: - result = self._descriptor_cache.get(key) - if result is not None: - if self._descriptor_cache_stats_enabled: - self._descriptor_cache_stats["hits"] += 1 - self._descriptor_cache.move_to_end(key) - return result - if self._descriptor_cache_stats_enabled: - self._descriptor_cache_stats["misses"] += 1 + with self._descriptor_cache_lock: + result = self._descriptor_cache.get(key) + if result is not None: + self._descriptor_cache.move_to_end(key) + return result table_id = self.table_entries.table_id(table_name) table = self.table_entries[table_name] @@ -769,25 +725,13 @@ def resolve_series_descriptor(self, table_name, tags, field_name): with self._descriptor_cache_lock: existing = self._descriptor_cache.get(key) if existing is not None: - if self._descriptor_cache_stats_enabled: - self._descriptor_cache_stats["hits"] += 1 self._descriptor_cache.move_to_end(key) return existing self._descriptor_cache[key] = result while len(self._descriptor_cache) > self._descriptor_cache_size: self._descriptor_cache.popitem(last=False) - if self._descriptor_cache_stats_enabled: - self._descriptor_cache_stats["evictions"] += 1 return result - @property - def descriptor_cache_stats(self): - """Return a snapshot of descriptor cache activity for diagnostics.""" - stats = dict(self._descriptor_cache_stats) - stats["size"] = len(self._descriptor_cache) - stats["capacity"] = self._descriptor_cache_size - return stats - def resolve_series_descriptor_by_id(self, series_id, table_name, field_name): if series_id < 0 or series_id >= self.index.count(LOGICAL_SERIES): raise KeyError(series_id)