| # 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. |
| |
| from types import SimpleNamespace |
| import os |
| import struct |
| import sys |
| import threading |
| |
| import numpy as np |
| import pandas as pd |
| import pyarrow as pa |
| import pytest |
| |
| import tsfile.dataset.index as index_module |
| import tsfile.dataset.runtime as runtime_module |
| from tsfile import ( |
| ColumnCategory, |
| ColumnSchema, |
| TableSchema, |
| TsFileDataFrame, |
| TsFileTableWriter, |
| ) |
| from tsfile.constants import TSDataType |
| 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, |
| write_index_atomic, |
| ) |
| from tsfile.dataset.metadata import MetadataCatalog, SeriesStats |
| from tsfile.dataset.runtime import RuntimeSeriesReader |
| from tsfile.dataset.runtime import DatasetRuntime |
| from tsfile.tsfile_reader import TsFileReaderPy |
| |
| |
| def _synthetic_dataframe(source_path): |
| catalog = MetadataCatalog() |
| table_id = catalog.add_table("root", (), (), ("s1",)) |
| device_id = catalog.add_device(table_id, (), 1, 10) |
| catalog.series_stats_by_ref[(device_id, 0)] = SeriesStats( |
| 10, |
| 1, |
| 10, |
| 10, |
| 1, |
| 10, |
| int(TSDataType.INT64), |
| 128, |
| 16, |
| 64, |
| 16, |
| 1, |
| 1, |
| 1, |
| 1, |
| ) |
| reader = SimpleNamespace(file_path=source_path, catalog=catalog) |
| index = SimpleNamespace( |
| table_entries={"root": catalog.table_entries[0]}, |
| devices=[("root", ())], |
| series=[(0, 0)], |
| series_shards={(0, 0): [(reader, 0, 0)]}, |
| ) |
| return SimpleNamespace( |
| _index=index, |
| _readers={source_path: reader}, |
| _paths=[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 |
| assert RECORDS[COLUMN_SCHEMA].size == 32 |
| assert RECORDS[LOGICAL_SERIES].size == 32 |
| assert RECORDS[TSFILE_RECORD].size == 32 |
| assert RECORDS[DEVICE_FILE_SPAN].size == 32 |
| assert RECORDS[SERIES_FILE_SPAN].size == 40 |
| assert RECORDS[SERIES_LOCATOR].size == 24 |
| assert index_module.SECTION_COUNT == 13 |
| |
| |
| def test_crc32c_known_vector(): |
| assert crc32c(b"123456789") == 0xE3069283 |
| |
| |
| def test_build_publish_map_and_lookup(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)) |
| |
| with MappedDatasetIndex(str(output), verify_sections=True) as index: |
| table_id = index.find_table_ids("root")[0] |
| device_id = index.find_device_id(table_id, "root.") |
| column_id = index.find_column_id(table_id, "s1") |
| assert index.find_series_id(device_id, column_id) == 0 |
| assert index.count(LOGICAL_SERIES) == 1 |
| assert index.count(SERIES_LOCATOR) == 1 |
| file_record = index.record(TSFILE_RECORD, 0) |
| assert index.string(file_record[0]) == str(source) |
| |
| |
| 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)) |
| |
| 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) |
| |
| 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(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): |
| source = tmp_path / "source.tsfile" |
| source.write_bytes(b"T" * 4096) |
| output = tmp_path / "dataset.tsidx" |
| write_index_atomic( |
| str(output), build_sections_from_dataframe(_synthetic_dataframe(str(source))) |
| ) |
| with output.open("r+b") as stream: |
| stream.seek(32) |
| stream.write((1).to_bytes(8, "little")) |
| with pytest.raises(ValueError, match="header shape|checksum"): |
| MappedDatasetIndex(str(output)) |
| |
| |
| def test_reader_exposes_exact_aligned_locator_ranges(tmp_path): |
| source = tmp_path / "aligned.tsfile" |
| schema = TableSchema( |
| "weather", |
| [ |
| ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG), |
| ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD), |
| ], |
| ) |
| with TsFileTableWriter(str(source), schema) as writer: |
| writer.write_dataframe( |
| pd.DataFrame( |
| { |
| "time": [0, 1, 2], |
| "device": ["d0", "d0", "d0"], |
| "value": [1.0, 2.0, 3.0], |
| } |
| ) |
| ) |
| reader = TsFileReaderPy(str(source)) |
| try: |
| groups = reader.get_timeseries_metadata() |
| metadata = [item for group in groups.values() for item in group.timeseries] |
| assert metadata |
| assert all(item.value_metadata_length > 0 for item in metadata) |
| assert all( |
| item.value_metadata_offset + item.value_metadata_length |
| <= os.path.getsize(source) |
| for item in metadata |
| ) |
| aligned = [item for item in metadata if item.layout == 1] |
| assert aligned |
| assert all(item.time_metadata_length > 0 for item in aligned) |
| assert all( |
| item.time_chunk_meta_count == item.chunk_meta_count for item in aligned |
| ) |
| finally: |
| reader.close() |
| |
| |
| def _write_runtime_file(path, start): |
| schema = TableSchema( |
| "weather", |
| [ |
| ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG), |
| ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD), |
| ], |
| ) |
| with TsFileTableWriter(str(path), schema) as writer: |
| writer.write_dataframe( |
| pd.DataFrame( |
| { |
| "time": [start, start + 1], |
| "device": ["d0", "d0"], |
| "value": [float(start), float(start + 1)], |
| } |
| ) |
| ) |
| |
| |
| def _write_runtime_devices_file(path): |
| schema = TableSchema( |
| "weather", |
| [ |
| ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG), |
| ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD), |
| ], |
| ) |
| with TsFileTableWriter(str(path), schema) as writer: |
| writer.write_dataframe( |
| pd.DataFrame( |
| { |
| "time": [0, 1, 0, 1, 0, 1], |
| "device": ["d0", "d0", "d1", "d1", "d2", "d2"], |
| "value": [0.0, 1.0, 10.0, 11.0, 20.0, 21.0], |
| } |
| ) |
| ) |
| |
| |
| def test_hot_construction_maps_index_without_opening_readers(tmp_path, monkeypatch): |
| source = tmp_path / "part.tsfile" |
| _write_runtime_file(source, 0) |
| with TsFileDataFrame(str(source), show_progress=False, use_index=True) as first: |
| assert len(first) == 1 |
| |
| def fail_legacy_scan(*_args, **_kwargs): |
| raise AssertionError("hot construction must not build a legacy catalog") |
| |
| monkeypatch.setattr("tsfile.dataset.reader.TsFileSeriesReader", fail_legacy_scan) |
| with TsFileDataFrame(str(source), show_progress=False, use_index=True) as second: |
| assert len(second) == 1 |
| assert second._runtime.readers.open_count == 0 |
| series = second[0] |
| assert second._runtime.readers.open_count == 0 |
| assert series[0] == 0.0 |
| assert second._runtime.readers.open_count == 1 |
| assert second._runtime.prepared.size == 1 |
| assert series[1] == 1.0 |
| assert second._runtime.prepared.size == 1 |
| series.close() |
| |
| |
| def test_dataframe_does_not_use_or_create_index_by_default(tmp_path): |
| source = tmp_path / "part.tsfile" |
| _write_runtime_file(source, 0) |
| index_path = index_module.index_path_for([str(source)]) |
| |
| with TsFileDataFrame(str(source), show_progress=False) as dataframe: |
| assert dataframe._runtime is None |
| assert len(dataframe) == 1 |
| np.testing.assert_array_equal(dataframe[0][:], np.array([0.0, 1.0])) |
| aligned = dataframe.loc[0:1, [0]] |
| np.testing.assert_array_equal(aligned.timestamps, np.array([0, 1])) |
| np.testing.assert_array_equal(aligned.values, np.array([[0.0], [1.0]])) |
| |
| assert not os.path.exists(index_path) |
| |
| |
| def test_persistent_index_path_is_scoped_to_expanded_file_set(tmp_path, monkeypatch): |
| first = tmp_path / "part1.tsfile" |
| second = tmp_path / "part2.tsfile" |
| third = tmp_path / "part3.tsfile" |
| _write_runtime_file(first, 0) |
| _write_runtime_file(second, 2) |
| _write_runtime_file(third, 10) |
| |
| first_set = [str(first), str(second)] |
| second_set = [str(first), str(third)] |
| first_index = index_module.index_path_for(first_set) |
| second_index = index_module.index_path_for(second_set) |
| assert first_index != second_index |
| |
| with TsFileDataFrame(first_set, show_progress=False, use_index=True) as dataframe: |
| assert len(dataframe) == 1 |
| with TsFileDataFrame(second_set, show_progress=False, use_index=True) as dataframe: |
| assert len(dataframe) == 1 |
| assert os.path.exists(first_index) |
| assert os.path.exists(second_index) |
| |
| def fail_legacy_scan(*_args, **_kwargs): |
| raise AssertionError("matching file-set index should be reused") |
| |
| monkeypatch.setattr("tsfile.dataset.reader.TsFileSeriesReader", fail_legacy_scan) |
| with TsFileDataFrame(first_set, show_progress=False, use_index=True) as dataframe: |
| np.testing.assert_array_equal(dataframe[0][:], np.array([0.0, 1.0, 2.0, 3.0])) |
| with TsFileDataFrame(second_set, show_progress=False, use_index=True) as dataframe: |
| np.testing.assert_array_equal(dataframe[0][:], np.array([0.0, 1.0, 10.0, 11.0])) |
| |
| |
| def test_named_selection_reuses_bounded_runtime_descriptor(tmp_path, monkeypatch): |
| source = tmp_path / "part.tsfile" |
| _write_runtime_file(source, 0) |
| |
| with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: |
| name = str(dataframe.list_timeseries()[0]) |
| find_device_calls = 0 |
| original_find_device = dataframe._runtime.index.find_device_id |
| |
| def count_find_device(*args, **kwargs): |
| nonlocal find_device_calls |
| find_device_calls += 1 |
| return original_find_device(*args, **kwargs) |
| |
| def fail_series_info(*_args, **_kwargs): |
| raise AssertionError( |
| "descriptor-backed selection must not rebuild series info" |
| ) |
| |
| monkeypatch.setattr( |
| dataframe._runtime.index, "find_device_id", count_find_device |
| ) |
| monkeypatch.setattr( |
| RuntimeSeriesReader, "get_series_info_by_ref", fail_series_info |
| ) |
| |
| first = dataframe[name] |
| second = dataframe[name] |
| assert first.stats == {"start_time": 0, "end_time": 1, "count": 2} |
| np.testing.assert_array_equal(first[:], np.array([0.0, 1.0])) |
| assert second.stats == first.stats |
| assert find_device_calls == 1 |
| assert len(dataframe._index._descriptor_cache) == 1 |
| first.close() |
| second.close() |
| |
| |
| def test_listed_series_path_resolves_directly_by_snapshot_series_id( |
| 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: |
| path = dataframe.list_timeseries()[0] |
| assert isinstance(path, str) |
| assert path.series_id == 0 |
| assert path._index_identity == dataframe._runtime.index.identity |
| |
| def fail_name_lookup(*_args, **_kwargs): |
| raise AssertionError("listed SeriesPath must bypass device name lookup") |
| |
| def fail_name_rebuild(*_args, **_kwargs): |
| raise AssertionError("listed SeriesPath must not be rebuilt from the index") |
| |
| def fail_span_lookup(*_args, **_kwargs): |
| raise AssertionError("descriptor locator must bypass series span lookup") |
| |
| monkeypatch.setattr( |
| dataframe._runtime.index, "find_device_id", fail_name_lookup |
| ) |
| monkeypatch.setattr(dataframe, "_build_series_name", fail_name_rebuild) |
| monkeypatch.setattr(RuntimeSeriesReader, "_span", fail_span_lookup) |
| series = dataframe[path] |
| assert series.name is path |
| np.testing.assert_array_equal(series[:], np.array([0.0, 1.0])) |
| series.close() |
| |
| # Converting to a plain str deliberately drops the snapshot-local id. |
| assert not hasattr(str(path), "series_id") |
| |
| |
| def test_series_path_from_another_index_falls_back_to_its_name(tmp_path, monkeypatch): |
| first_dir = tmp_path / "first" |
| second_dir = tmp_path / "second" |
| first_dir.mkdir() |
| second_dir.mkdir() |
| first_source = first_dir / "part.tsfile" |
| second_source = second_dir / "part.tsfile" |
| _write_runtime_file(first_source, 0) |
| _write_runtime_file(second_source, 10) |
| |
| with TsFileDataFrame( |
| str(first_source), show_progress=False, use_index=True |
| ) as first: |
| foreign_path = first.list_timeseries()[0] |
| with TsFileDataFrame( |
| str(second_source), show_progress=False, use_index=True |
| ) as second: |
| assert foreign_path._index_identity != second._runtime.index.identity |
| find_device_calls = 0 |
| original_find_device = second._runtime.index.find_device_id |
| |
| def count_find_device(*args, **kwargs): |
| nonlocal find_device_calls |
| find_device_calls += 1 |
| return original_find_device(*args, **kwargs) |
| |
| monkeypatch.setattr( |
| second._runtime.index, "find_device_id", count_find_device |
| ) |
| series = second[foreign_path] |
| np.testing.assert_array_equal(series[:], np.array([10.0, 11.0])) |
| assert find_device_calls == 1 |
| series.close() |
| |
| |
| def test_runtime_descriptor_cache_evicts_least_recent_name(tmp_path, monkeypatch): |
| source = tmp_path / "devices.tsfile" |
| _write_runtime_devices_file(source) |
| monkeypatch.setattr(runtime_module, "_SERIES_DESCRIPTOR_CACHE_SIZE", 2) |
| |
| with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: |
| names = [str(name) for name in dataframe.list_timeseries()] |
| find_device_calls = 0 |
| original_find_device = dataframe._runtime.index.find_device_id |
| |
| def count_find_device(*args, **kwargs): |
| nonlocal find_device_calls |
| find_device_calls += 1 |
| return original_find_device(*args, **kwargs) |
| |
| monkeypatch.setattr( |
| dataframe._runtime.index, "find_device_id", count_find_device |
| ) |
| for name in names: |
| dataframe[name].close() |
| |
| assert find_device_calls == 3 |
| assert len(dataframe._index._descriptor_cache) == 2 |
| assert len(dataframe._index.series_shards._cache) == 2 |
| |
| # d0 was the least recently used name and must be resolved again. |
| dataframe[names[0]].close() |
| 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" |
| _write_runtime_file(first, 0) |
| _write_runtime_file(second, 2) |
| monkeypatch.setenv("TSFILE_DATAFRAME_MAX_OPEN_FILES", "1") |
| with TsFileDataFrame( |
| [str(first), str(second)], show_progress=False, use_index=True |
| ) as dataframe: |
| series = dataframe[0] |
| assert list(series[:]) == [0.0, 1.0, 2.0, 3.0] |
| assert dataframe._runtime.readers.open_count == 1 |
| assert dataframe._runtime.prepared.size == 2 |
| series.close() |
| |
| |
| def test_runtime_consume_concatenates_arrow_batches_without_scalar_reads(): |
| class _ArrowResult: |
| def __init__(self): |
| self._batches = iter( |
| [ |
| pa.table( |
| { |
| "time": pa.array([], type=pa.int64()), |
| "value": pa.array([], type=pa.float64()), |
| } |
| ), |
| pa.table( |
| { |
| "time": pa.array([1, 2], type=pa.int64()), |
| "value": pa.array([10.0, None], type=pa.float64()), |
| } |
| ), |
| pa.table( |
| { |
| "time": pa.array([3], type=pa.int64()), |
| "value": pa.array([30.0], type=pa.float64()), |
| } |
| ), |
| ] |
| ) |
| self.closed = False |
| |
| def __enter__(self): |
| return self |
| |
| def __exit__(self, *_): |
| self.closed = True |
| |
| def read_arrow_batch(self): |
| return next(self._batches, None) |
| |
| def next(self): |
| raise AssertionError("Runtime must not consume prepared rows one by one") |
| |
| result = _ArrowResult() |
| timestamps, values = RuntimeSeriesReader._consume(result) |
| |
| np.testing.assert_array_equal(timestamps, np.array([1, 2, 3], dtype=np.int64)) |
| np.testing.assert_allclose(values, np.array([10.0, np.nan, 30.0]), equal_nan=True) |
| assert result.closed |
| |
| |
| def test_prepared_query_reads_nullable_offset_window_in_arrow_batches(tmp_path): |
| source = tmp_path / "nullable.tsfile" |
| schema = TableSchema( |
| "weather", |
| [ |
| ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG), |
| ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD), |
| ], |
| ) |
| expected = np.arange(10, dtype=np.float64) |
| expected[2] = np.nan |
| expected[6] = np.nan |
| with TsFileTableWriter(str(source), schema) as writer: |
| writer.write_dataframe( |
| pd.DataFrame( |
| { |
| "time": np.arange(10, dtype=np.int64), |
| "device": ["d0"] * 10, |
| "value": expected, |
| } |
| ) |
| ) |
| |
| with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: |
| runtime = dataframe._runtime |
| series = runtime.index.record(LOGICAL_SERIES, 0) |
| span = runtime.index.record(SERIES_FILE_SPAN, series[2]) |
| with runtime.readers.acquire(0) as reader: |
| prepared = runtime.prepared.get(0, span[2], reader) |
| with reader.query_prepared(prepared, offset=1, limit=7) as result: |
| batches = [] |
| while True: |
| batch = result.read_arrow_batch() |
| if batch is None: |
| break |
| batches.append(batch) |
| with reader.query_prepared( |
| prepared, start_time=100, end_time=200 |
| ) as empty_result: |
| assert empty_result.read_arrow_batch() is None |
| |
| assert batches |
| table = pa.concat_tables(batches) |
| np.testing.assert_array_equal( |
| table.column("time").to_numpy(), np.arange(1, 8, dtype=np.int64) |
| ) |
| np.testing.assert_allclose( |
| table.column("value").to_numpy(zero_copy_only=False), |
| expected[1:8], |
| equal_nan=True, |
| ) |
| |
| |
| def test_prepared_locator_rejects_stale_generation_and_bad_range(tmp_path): |
| source = tmp_path / "part.tsfile" |
| _write_runtime_file(source, 0) |
| with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: |
| runtime = dataframe._runtime |
| series = runtime.index.record(LOGICAL_SERIES, 0) |
| span = runtime.index.record(SERIES_FILE_SPAN, series[2]) |
| locator = list(runtime.prepared._locator_tuple(0, span[2])) |
| with runtime.readers.acquire(0) as reader: |
| stale = list(locator) |
| stale[3] ^= 1 |
| with pytest.raises(Exception, match="prepare Dataset Index locator"): |
| reader.prepare_series(stale) |
| |
| out_of_range = list(locator) |
| out_of_range[7] = os.path.getsize(source) + 1 |
| with pytest.raises(Exception, match="prepare Dataset Index locator"): |
| reader.prepare_series(out_of_range) |
| |
| |
| def test_reader_session_revalidates_generation_when_reused(tmp_path): |
| source = tmp_path / "part.tsfile" |
| _write_runtime_file(source, 0) |
| with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: |
| pool = dataframe._runtime.readers |
| with pool.acquire(0): |
| pass |
| stat = os.stat(source) |
| os.utime(source, ns=(stat.st_atime_ns, stat.st_mtime_ns + 1_000_000)) |
| with pytest.raises(RuntimeError, match="generation changed"): |
| with pool.acquire(0): |
| pass |
| |
| |
| def test_runtime_lease_close_waits_for_query_lease(tmp_path): |
| source = tmp_path / "part.tsfile" |
| _write_runtime_file(source, 0) |
| with TsFileDataFrame(str(source), show_progress=False, use_index=True) as dataframe: |
| runtime = DatasetRuntime(str(dataframe._runtime.index.path), query_workers=1) |
| lease = runtime.lease() |
| entered = threading.Event() |
| release = threading.Event() |
| query_done = threading.Event() |
| |
| def run_query(): |
| with lease.query_lease(): |
| entered.set() |
| assert release.wait(timeout=2) |
| query_done.set() |
| |
| query_thread = threading.Thread(target=run_query) |
| query_thread.start() |
| assert entered.wait(timeout=2) |
| |
| close_done = threading.Event() |
| |
| def close_lease(): |
| lease.close() |
| close_done.set() |
| |
| close_thread = threading.Thread(target=close_lease) |
| close_thread.start() |
| assert not close_done.wait(timeout=0.05) |
| |
| release.set() |
| query_thread.join(timeout=2) |
| close_thread.join(timeout=2) |
| assert query_done.is_set() |
| assert close_done.is_set() |
