diff --git a/python/setup.py b/python/setup.py index c7b4cceaa..fe19fc91c 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", + ["tsfile/dataset/_index.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..4ac9f2ede 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="= 0 + assert min_time == 0 + assert max_time == 1 + assert length >= 0 + + +def test_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], + ) + + +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" + dataframe = _synthetic_dataframe(str(source)) + write_index_atomic(str(output), build_sections_from_dataframe(dataframe)) - class _Index: - @staticmethod - def record(_section_type, record_id): - return rows[record_id] + 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("index string lookup 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 + + +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" + 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) - @staticmethod - def string_bytes(sid): - return names[sid] + def fail_record(*_args, **_kwargs): + raise AssertionError("series identity should not unpack records") - assert MappedDatasetIndex._find_child(_Index(), 0, 0, "beta", 0, 3) == 11 - with pytest.raises(KeyError): - MappedDatasetIndex._find_child(_Index(), 0, 0, "missing", 0, 3) + 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(tmp_path): + """A stored hash match with different bytes must not resolve to that row.""" + source = tmp_path / "devices.tsfile" + _write_runtime_devices_file(source) + with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: + index_path = dataframe._runtime.index.path + + with MappedDatasetIndex(index_path) as index: + rows = list(index.records(DEVICE_NAME_INDEX)) + names = [index.string(row[3]) for row in rows] + assert len(rows) >= 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( + "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 IndexLookup: + """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, Py_ssize_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 < 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( + 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, Py_ssize_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 < 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( + 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, 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 + 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 < 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: + 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, 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 + 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 < 0 or 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, 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 < 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, 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 + 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 < 0 or 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, 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 < 0 or 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, 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 + 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 < 0 or locator_id >= self._section_counts[13]: + raise IndexError(locator_id) + if file_id < 0 or 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, 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 < 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, 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 < 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, 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 < 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, 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 < 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_u64(view, base + 32), + _read_u64(view, base + 40), + ) + + 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] + 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 < 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) + 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..9b4ba539c 100644 --- a/python/tsfile/dataset/index.py +++ b/python/tsfile/dataset/index.py @@ -36,6 +36,7 @@ from ..constants import ColumnCategory from .metadata import MODEL_TREE, _join_series_path +from ._index import IndexLookup MAGIC = b"TSIDX\0\0\0" VERSION_MAJOR = 1 @@ -229,6 +230,7 @@ class MappedDatasetIndex: def __init__(self, path: str, verify_sections: bool = False): self.path = path + self._lookup = None self._file = open(path, "rb") try: stat = os.fstat(self._file.fileno()) @@ -241,7 +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) + self._lookup = IndexLookup(self._view, self._entries) except Exception: + self._lookup = None if getattr(self, "_view", None) is not None: self._view.release() self._view = None @@ -336,6 +340,7 @@ def _validate(self, verify_sections: bool): return entries def close(self): + self._lookup = None if getattr(self, "_view", None) is not None: self._view.release() self._view = None @@ -383,7 +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: - 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 @@ -411,54 +416,50 @@ 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: - 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: - 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: - 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.""" + return self._lookup.describe_series(series_id) + + def series_identity(self, series_id: int): + """Return device and column ids for one logical series.""" + 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.""" + 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.""" + 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.""" + 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.""" + return self._lookup.device_route(device_id) + + def table_name_id(self, table_id: int): + """Return the string-pool id for one table name.""" + 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.""" + return self._lookup.column_name_id(column_id) + + def device_time_bounds(self, device_id: int): + """Return min/max timestamps for one device.""" + return self._lookup.device_time_bounds(device_id) def index_path_for(paths: Sequence[str]) -> str: diff --git a/python/tsfile/dataset/runtime.py b/python/tsfile/dataset/runtime.py index 3ce096b1d..dca1f5993 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,10 @@ 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) +@dataclass(frozen=True, **_DATACLASS_SLOTS) class RuntimeSeriesShard: """One immutable physical fragment already expanded from the mmap route.""" @@ -72,7 +74,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 +262,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): @@ -594,12 +601,12 @@ def __init__(self, catalog, cache_size): 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): @@ -611,31 +618,32 @@ def describe(self, ref, series_id=None, column_id=None): 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,8 +652,8 @@ 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: @@ -727,11 +735,11 @@ def resolve_series_descriptor(self, table_name, tags, field_name): 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 +768,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 +783,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 +912,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 +958,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 +980,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: