perf(python): accelerate Dataset Index metadata lookups

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",

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,13 +39,16 @@
 from tsfile.dataset.index import (
     COLUMN_SCHEMA,
     DEVICE_FILE_SPAN,
+    DEVICE_NAME_INDEX,
+    DEVICE_RECORD,
     DIRECTORY,
     HEADER,
     LOGICAL_SERIES,
     MappedDatasetIndex,
     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="<I"):
+    """Rewrite one raw index field, invalidating that section's checksum."""
+    with open(path, "r+b") as stream:
+        blob = bytearray(stream.read())
+        entry = _index_section_entry(blob, section_type)
+        struct.pack_into(fmt, blob, entry[2] + field_offset, value)
+        stream.seek(0)
+        stream.write(blob)
+
+
 def test_binary_layout_matches_cpp_v1():
     assert HEADER.size == 64
     assert DIRECTORY.size == 32
@@ -124,27 +145,288 @@ def test_build_publish_map_and_lookup(tmp_path):
         assert index.string(file_record[0]) == str(source)
 
 
-def test_child_lookup_checks_full_bytes_for_hash_collisions(monkeypatch):
-    rows = [
-        (0, 10, 42, 0, 0),
-        (0, 11, 42, 1, 0),
-        (0, 12, 42, 2, 0),
-    ]
-    names = [b"alpha", b"beta", b"gamma"]
-    monkeypatch.setattr(index_module, "name_hash", lambda _value: 42)
+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"
+    dataframe = _synthetic_dataframe(str(source))
+    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("index 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
+
+
+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"
+    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])]
+
+
+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"
+    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],
+        )
+
+
+def test_runtime_reader_span_lookup_uses_index_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
+
+
+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(
+            "<Q",
+            blob,
+            entry[2] + (entry[4] - 1) * entry[1] + 8,
+            index_module.name_hash(probe.encode("utf-8")),
+        )
+        stream.seek(0)
+        stream.write(blob)
+
+    with MappedDatasetIndex(index_path) as index:
+        untouched = list(index.records(DEVICE_NAME_INDEX))[:-1]
+        for table_id, device_id, _, sid, _ in untouched:
+            assert index.find_device_id(table_id, index.string(sid)) == device_id
+        with pytest.raises(KeyError):
+            index.find_device_id(untouched[0][0], probe)
+
+
+@pytest.mark.parametrize(
+    ("section_type", "field_offset", "lookup"),
+    [
+        (TABLE_RECORD, 12, lambda index: index.find_device_id(0, "root.")),
+        (TABLE_RECORD, 20, lambda index: index.find_column_id(0, "s1")),
+        (DEVICE_RECORD, 20, lambda index: index.find_series_id(0, 0)),
+    ],
+)
+def test_out_of_range_child_ranges_raise_instead_of_reading_out_of_bounds(
+    tmp_path, section_type, field_offset, lookup
+):
+    """Corrupt child ranges must raise, not walk past the mmap (SIGSEGV/SIGBUS)."""
+    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))
+    _patch_index_field(output, section_type, field_offset, 0xFFFFFF00)
+
+    # section checksums stay unverified here, matching DatasetRuntime's load path
+    with MappedDatasetIndex(str(output)) as index:
+        with pytest.raises(IndexError):
+            lookup(index)
+
+
+def test_lookup_rejects_negative_ids_with_index_error(tmp_path):
+    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))
+
+    calls = [
+        lambda index: index.find_device_id(-1, "root."),
+        lambda index: index.find_column_id(-1, "s1"),
+        lambda index: index.find_series_id(-1, 0),
+        lambda index: index.describe_series(-1),
+        lambda index: index.series_identity(-1),
+        lambda index: index.find_series_span(-1, 0),
+        lambda index: index.locator_metadata(-1),
+        lambda index: index.prepared_locator_metadata(-1, 0),
+        lambda index: index.device_route(-1),
+        lambda index: index.table_name_id(-1),
+        lambda index: index.column_name_id(-1),
+        lambda index: index.device_time_bounds(-1),
+        lambda index: index.string(-1),
+    ]
+    with MappedDatasetIndex(str(output)) as index:
+        for call in calls:
+            with pytest.raises(IndexError):
+                call(index)
 
 
 def test_rejects_damaged_header_checksum(tmp_path):
@@ -448,6 +730,17 @@ def count_find_device(*args, **kwargs):
         assert find_device_calls == 4
 
 
+@pytest.mark.skipif(
+    sys.version_info < (3, 10), reason="dataclasses only support slots on 3.10+"
+)
+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_reader_pool_enforces_open_file_cap(tmp_path, monkeypatch):
     first = tmp_path / "part1.tsfile"
     second = tmp_path / "part2.tsfile"