From 1ad033665590aea866801fd1782b375cc3711432 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Wed, 19 Aug 2026 19:45:09 +0800 Subject: [PATCH 1/3] feature: implement dataframe runtime index --- .gitignore | 2 + cpp/cmake/ANTLR4Dependency.cmake | 5 +- cpp/cmake/tests/ANTLR4DependencyTest.cmake | 2 + cpp/cmake/tests/TestGeneratorArguments.cmake | 4 + cpp/src/CMakeLists.txt | 2 + cpp/src/common/tsfile_common.h | 22 + cpp/src/cwrapper/tsfile_cwrapper.cc | 171 +++ cpp/src/cwrapper/tsfile_cwrapper.h | 48 + cpp/src/dataset/CMakeLists.txt | 25 + cpp/src/dataset/dataset_index.cc | 1049 +++++++++++++++++ cpp/src/dataset/dataset_index.h | 331 ++++++ cpp/src/file/read_file.cc | 49 + cpp/src/file/read_file.h | 4 + cpp/src/file/tsfile_io_reader.cc | 220 +++- cpp/src/file/tsfile_io_reader.h | 22 + .../block/prepared_series_tsblock_reader.cc | 219 ++++ .../block/prepared_series_tsblock_reader.h | 75 ++ cpp/src/reader/prepared_series.cc | 35 + cpp/src/reader/prepared_series.h | 92 ++ cpp/src/reader/qds_without_timegenerator.cc | 70 +- cpp/src/reader/qds_without_timegenerator.h | 13 +- cpp/src/reader/tsfile_executor.cc | 78 ++ cpp/src/reader/tsfile_executor.h | 16 + cpp/src/reader/tsfile_reader.cc | 33 + cpp/src/reader/tsfile_reader.h | 16 + cpp/src/reader/tsfile_series_scan_iterator.cc | 113 +- cpp/src/reader/tsfile_series_scan_iterator.h | 20 + cpp/test/CMakeLists.txt | 5 + cpp/test/dataset/dataset_index_test.cc | 298 +++++ cpp/test/reader/prepared_series_test.cc | 376 ++++++ pom.xml | 4 +- python/setup.py | 1 + python/tests/test_dataset_index.py | 582 +++++++++ python/tests/test_tsfile_dataset.py | 267 ++++- python/tsfile/dataset/_merge.pyx | 261 ++++ python/tsfile/dataset/dataframe.py | 463 ++++++-- python/tsfile/dataset/index.py | 815 +++++++++++++ python/tsfile/dataset/merge.py | 146 +-- python/tsfile/dataset/metadata.py | 55 +- python/tsfile/dataset/reader.py | 30 +- python/tsfile/dataset/runtime.py | 1013 ++++++++++++++++ python/tsfile/dataset/timeseries.py | 88 +- python/tsfile/schema.py | 7 + python/tsfile/tsfile_cpp.pxd | 41 +- python/tsfile/tsfile_py_cpp.pxd | 8 + python/tsfile/tsfile_py_cpp.pyx | 63 + python/tsfile/tsfile_reader.pyx | 111 +- 47 files changed, 7139 insertions(+), 231 deletions(-) create mode 100644 cpp/src/dataset/CMakeLists.txt create mode 100644 cpp/src/dataset/dataset_index.cc create mode 100644 cpp/src/dataset/dataset_index.h create mode 100644 cpp/src/reader/block/prepared_series_tsblock_reader.cc create mode 100644 cpp/src/reader/block/prepared_series_tsblock_reader.h create mode 100644 cpp/src/reader/prepared_series.cc create mode 100644 cpp/src/reader/prepared_series.h create mode 100644 cpp/test/dataset/dataset_index_test.cc create mode 100644 cpp/test/reader/prepared_series_test.cc create mode 100644 python/tests/test_dataset_index.py create mode 100644 python/tsfile/dataset/_merge.pyx create mode 100644 python/tsfile/dataset/index.py create mode 100644 python/tsfile/dataset/runtime.py diff --git a/.gitignore b/.gitignore index fcd5f2b6b..6b5d1d9b2 100644 --- a/.gitignore +++ b/.gitignore @@ -32,6 +32,8 @@ python/tsfile/*dll* python/tsfile/*dylib* python/tsfile/*.h python/tsfile/*.cpp +python/tsfile/**/*.cpp +python/tsfile/**/*so* python/data python/venv/* python/tests/__pycache__/* diff --git a/cpp/cmake/ANTLR4Dependency.cmake b/cpp/cmake/ANTLR4Dependency.cmake index 3d772c25c..a20a98b0b 100644 --- a/cpp/cmake/ANTLR4Dependency.cmake +++ b/cpp/cmake/ANTLR4Dependency.cmake @@ -19,7 +19,10 @@ under the License. set(TSFILE_ANTLR4_MIN_VERSION "4.9.3") set(TSFILE_ANTLR4_BUNDLED_VERSION "4.9.3") -set(TSFILE_ANTLR4_NEXT_INCOMPATIBLE_VERSION "5.0.0") +# The C++ runtime shipped by ANTLR4 4.13 and newer requires C++17. TsFile's +# public C++ baseline remains C++11, so those system runtimes must not be +# selected for the generated 4.9.3 parser. +set(TSFILE_ANTLR4_NEXT_INCOMPATIBLE_VERSION "4.13.0") set(TSFILE_ANTLR4_SYSTEM_INCLUDE_DIR "") set(_TSFILE_SYSTEM_ANTLR4_FOUND FALSE) diff --git a/cpp/cmake/tests/ANTLR4DependencyTest.cmake b/cpp/cmake/tests/ANTLR4DependencyTest.cmake index 2fc981712..5ceaecdda 100644 --- a/cpp/cmake/tests/ANTLR4DependencyTest.cmake +++ b/cpp/cmake/tests/ANTLR4DependencyTest.cmake @@ -59,6 +59,8 @@ function(_tsfile_run_antlr4_case NAME POLICY EXPECTED_SOURCE EXPECT_SUCCESS ROOT "-DCMAKE_PREFIX_PATH=${ROOT}" -DCMAKE_FIND_USE_PACKAGE_REGISTRY=FALSE -DCMAKE_FIND_USE_SYSTEM_PACKAGE_REGISTRY=FALSE + -DCMAKE_FIND_USE_CMAKE_SYSTEM_PATH=FALSE + -DCMAKE_FIND_USE_SYSTEM_ENVIRONMENT_PATH=FALSE -S "${_TSFILE_FIXTURE_SOURCE}" -B "${_TSFILE_CASE_BINARY}" RESULT_VARIABLE _TSFILE_RESULT diff --git a/cpp/cmake/tests/TestGeneratorArguments.cmake b/cpp/cmake/tests/TestGeneratorArguments.cmake index cea9db8f1..8775e6eb6 100644 --- a/cpp/cmake/tests/TestGeneratorArguments.cmake +++ b/cpp/cmake/tests/TestGeneratorArguments.cmake @@ -22,6 +22,10 @@ if (TEST_CMAKE_GENERATOR) list(APPEND _TSFILE_TEST_GENERATOR_ARGUMENTS -G "${TEST_CMAKE_GENERATOR}") endif () +if (TEST_CMAKE_MAKE_PROGRAM) + list(APPEND _TSFILE_TEST_GENERATOR_ARGUMENTS + "-DCMAKE_MAKE_PROGRAM=${TEST_CMAKE_MAKE_PROGRAM}") +endif () if (TEST_CMAKE_GENERATOR_PLATFORM) list(APPEND _TSFILE_TEST_GENERATOR_ARGUMENTS -A "${TEST_CMAKE_GENERATOR_PLATFORM}") diff --git a/cpp/src/CMakeLists.txt b/cpp/src/CMakeLists.txt index 4e34031ed..2ab2dd5d6 100644 --- a/cpp/src/CMakeLists.txt +++ b/cpp/src/CMakeLists.txt @@ -72,6 +72,7 @@ endif() add_subdirectory(common) add_subdirectory(compress) add_subdirectory(cwrapper) +add_subdirectory(dataset) add_subdirectory(encoding) add_subdirectory(file) add_subdirectory(reader) @@ -82,6 +83,7 @@ set(_TSFILE_OBJECT_TARGETS common_obj compress_obj cwrapper_obj + dataset_obj file_obj read_obj write_obj) diff --git a/cpp/src/common/tsfile_common.h b/cpp/src/common/tsfile_common.h index d763acbbd..be046a9b0 100644 --- a/cpp/src/common/tsfile_common.h +++ b/cpp/src/common/tsfile_common.h @@ -351,6 +351,8 @@ class TimeseriesIndex : public ITimeseriesIndex { TimeseriesIndex() : timeseries_meta_type_((char)255), chunk_meta_list_data_size_(0), + metadata_offset_(-1), + metadata_length_(0), measurement_name_(), data_type_(common::INVALID_DATATYPE), statistic_(nullptr), @@ -369,6 +371,8 @@ class TimeseriesIndex : public ITimeseriesIndex { { timeseries_meta_type_ = 0; chunk_meta_list_data_size_ = 0; + metadata_offset_ = -1; + metadata_length_ = 0; measurement_name_.reset(); data_type_ = common::VECTOR; chunk_meta_list_serialized_buf_.reset(); @@ -401,6 +405,16 @@ class TimeseriesIndex : public ITimeseriesIndex { FORCE_INLINE virtual common::TSDataType get_data_type() const { return data_type_; } + FORCE_INLINE void set_metadata_range(int64_t offset, uint32_t length) { + metadata_offset_ = offset; + metadata_length_ = length; + } + FORCE_INLINE int64_t get_metadata_offset() const { + return metadata_offset_; + } + FORCE_INLINE uint32_t get_metadata_length() const { + return metadata_length_; + } int init_statistic(common::TSDataType data_type) { if (statistic_ != nullptr && !statistic_from_pa_) { // clear old statistic @@ -488,6 +502,8 @@ class TimeseriesIndex : public ITimeseriesIndex { int ret = common::E_OK; timeseries_meta_type_ = that.timeseries_meta_type_; chunk_meta_list_data_size_ = that.chunk_meta_list_data_size_; + metadata_offset_ = that.metadata_offset_; + metadata_length_ = that.metadata_length_; data_type_ = that.data_type_; statistic_ = StatisticFactory::alloc_statistic_with_pa(data_type_, pa); @@ -560,6 +576,12 @@ class TimeseriesIndex : public ITimeseriesIndex { // Sum of chunk meta serialized size in List of this timeseries. uint32_t chunk_meta_list_data_size_; + // Exact byte range of this TimeseriesMetadata in the source TsFile. + // It is assigned by TsFileIOReader after deserialization and is not part + // of the on-wire TimeseriesMetadata encoding. + int64_t metadata_offset_; + uint32_t metadata_length_; + // std::string measurement_name_; common::String measurement_name_; common::TSDataType data_type_; diff --git a/cpp/src/cwrapper/tsfile_cwrapper.cc b/cpp/src/cwrapper/tsfile_cwrapper.cc index 3e27f60d0..e54afecb9 100644 --- a/cpp/src/cwrapper/tsfile_cwrapper.cc +++ b/cpp/src/cwrapper/tsfile_cwrapper.cc @@ -419,6 +419,138 @@ ERRNO tsfile_writer_write(TsFileWriter writer, Tablet tablet) { // Query +PreparedSeriesHandle tsfile_reader_prepare_series( + TsFileReader reader, const TsFilePreparedLocator* locator, + ERRNO* err_code) { + if (err_code == nullptr) { + return nullptr; + } + *err_code = common::E_INVALID_ARG; + if (reader == nullptr || locator == nullptr) { + return nullptr; + } + storage::FileGeneration generation; + generation.mapped_index_identity = locator->mapped_index_identity; + generation.file_id = locator->file_id; + generation.file_size = locator->file_size; + generation.file_fingerprint = locator->file_fingerprint; + storage::PreparedLocator native_locator; + native_locator.locator_id = locator->locator_id; + native_locator.layout = locator->layout; + native_locator.flags = locator->flags; + native_locator.value_metadata_offset = locator->value_metadata_offset; + native_locator.value_metadata_length = locator->value_metadata_length; + native_locator.time_metadata_offset = locator->time_metadata_offset; + native_locator.time_metadata_length = locator->time_metadata_length; + std::shared_ptr prepared; + *err_code = static_cast(reader)->prepare_series( + generation, native_locator, prepared); + if (*err_code != common::E_OK) { + return nullptr; + } + auto* handle = new (std::nothrow) + std::shared_ptr(std::move(prepared)); + if (handle == nullptr) { + *err_code = common::E_OOM; + } + return handle; +} + +PreparedSeriesHandle tsfile_reader_prepare_series_with_time_owner( + TsFileReader reader, const TsFilePreparedLocator* locator, + PreparedSeriesHandle aligned_time_owner, ERRNO* err_code) { + if (err_code == nullptr) { + return nullptr; + } + *err_code = common::E_INVALID_ARG; + if (reader == nullptr || locator == nullptr || + aligned_time_owner == nullptr) { + return nullptr; + } + + storage::FileGeneration generation; + generation.mapped_index_identity = locator->mapped_index_identity; + generation.file_id = locator->file_id; + generation.file_size = locator->file_size; + generation.file_fingerprint = locator->file_fingerprint; + storage::PreparedLocator native_locator; + native_locator.locator_id = locator->locator_id; + native_locator.layout = locator->layout; + native_locator.flags = locator->flags; + native_locator.value_metadata_offset = locator->value_metadata_offset; + native_locator.value_metadata_length = locator->value_metadata_length; + native_locator.time_metadata_offset = locator->time_metadata_offset; + native_locator.time_metadata_length = locator->time_metadata_length; + + auto* owner = static_cast*>( + aligned_time_owner); + std::shared_ptr prepared; + *err_code = static_cast(reader)->prepare_series( + generation, native_locator, *owner, prepared); + if (*err_code != common::E_OK) { + return nullptr; + } + auto* handle = new (std::nothrow) + std::shared_ptr(std::move(prepared)); + if (handle == nullptr) { + *err_code = common::E_OOM; + } + return handle; +} + +void tsfile_prepared_series_free(PreparedSeriesHandle prepared) { + delete static_cast*>(prepared); +} + +ResultSet tsfile_reader_query_prepared(TsFileReader reader, + PreparedSeriesHandle prepared, + Timestamp start_time, Timestamp end_time, + int offset, int limit, ERRNO* err_code) { + if (err_code == nullptr) { + return nullptr; + } + *err_code = common::E_INVALID_ARG; + if (reader == nullptr || prepared == nullptr) { + return nullptr; + } + auto* handle = + static_cast*>(prepared); + storage::ResultSet* result = nullptr; + *err_code = static_cast(reader)->query_prepared( + *handle, start_time, end_time, offset, limit, result); + return result; +} + +ResultSet tsfile_reader_query_prepared_multi( + TsFileReader reader, const PreparedSeriesHandle* prepared, + uint32_t prepared_count, Timestamp start_time, Timestamp end_time, + int offset, int limit, ERRNO* err_code) { + if (err_code == nullptr) { + return nullptr; + } + *err_code = common::E_INVALID_ARG; + if (reader == nullptr || prepared == nullptr || prepared_count == 0) { + return nullptr; + } + + std::vector> native_prepared; + native_prepared.reserve(prepared_count); + for (uint32_t i = 0; i < prepared_count; i++) { + if (prepared[i] == nullptr) { + return nullptr; + } + auto* handle = + static_cast*>(prepared[i]); + native_prepared.push_back(*handle); + } + + storage::ResultSet* result = nullptr; + *err_code = + static_cast(reader)->query_prepared_multi( + native_prepared, start_time, end_time, offset, limit, result); + return result; +} + ResultSet tsfile_query_table(TsFileReader reader, const char* table_name, char** columns, uint32_t column_num, Timestamp start_time, Timestamp end_time, @@ -1305,8 +1437,32 @@ ERRNO populate_c_metadata_map_from_cpp( aligned_idx->value_ts_idx_ != nullptr) { m.data_type = static_cast( aligned_idx->value_ts_idx_->get_data_type()); + const storage::TimeseriesIndex* value_idx = + aligned_idx->value_ts_idx_; + const storage::TimeseriesIndex* time_idx = + aligned_idx->time_ts_idx_; + if (value_idx->get_metadata_offset() >= 0) { + m.value_metadata_offset = + static_cast(value_idx->get_metadata_offset()); + m.value_metadata_length = value_idx->get_metadata_length(); + } + if (time_idx != nullptr && + time_idx->get_metadata_offset() >= 0) { + m.time_metadata_offset = + static_cast(time_idx->get_metadata_offset()); + m.time_metadata_length = time_idx->get_metadata_length(); + } + m.layout = 1; } else { m.data_type = static_cast(idx->get_data_type()); + const storage::TimeseriesIndex* value_idx = + dynamic_cast(idx.get()); + if (value_idx != nullptr && + value_idx->get_metadata_offset() >= 0) { + m.value_metadata_offset = + static_cast(value_idx->get_metadata_offset()); + m.value_metadata_length = value_idx->get_metadata_length(); + } } storage::Statistic* st = idx->get_statistic(); int32_t chunk_cnt = 0; @@ -1316,6 +1472,21 @@ ERRNO populate_c_metadata_map_from_cpp( chunk_cnt = static_cast(cl->size()); } m.chunk_meta_count = chunk_cnt; + if (chunk_cnt >= 0 && m.value_metadata_length > 0) { + m.locator_flags |= 1; + } + if (aligned_idx != nullptr) { + auto* time_chunks = idx->get_time_chunk_meta_list(); + if (time_chunks != nullptr) { + m.time_chunk_meta_count = + static_cast(time_chunks->size()); + } + if (m.time_metadata_length == 0 || + m.time_chunk_meta_count != + static_cast(chunk_cnt)) { + m.locator_flags &= ~static_cast(1); + } + } const int st_rc = fill_timeseries_statistic(st, &m.statistic); if (st_rc != common::E_OK) { for (uint32_t u = 0; u < slot; u++) { diff --git a/cpp/src/cwrapper/tsfile_cwrapper.h b/cpp/src/cwrapper/tsfile_cwrapper.h index 73cc5389a..0476f691d 100644 --- a/cpp/src/cwrapper/tsfile_cwrapper.h +++ b/cpp/src/cwrapper/tsfile_cwrapper.h @@ -199,6 +199,13 @@ typedef struct TimeseriesMetadata { int32_t chunk_meta_count; TimeseriesStatistic statistic; TimeseriesStatistic timeline_statistic; + uint64_t value_metadata_offset; + uint32_t value_metadata_length; + uint64_t time_metadata_offset; + uint32_t time_metadata_length; + uint32_t time_chunk_meta_count; + uint16_t layout; + uint16_t locator_flags; } TimeseriesMetadata; /** @@ -274,6 +281,21 @@ typedef void* TsRecord; typedef void* ResultSet; typedef void* TagFilterHandle; +typedef void* PreparedSeriesHandle; + +typedef struct TsFilePreparedLocator { + uint64_t mapped_index_identity; + uint32_t file_id; + uint64_t file_size; + uint64_t file_fingerprint; + uint32_t locator_id; + uint16_t layout; + uint16_t flags; + uint64_t value_metadata_offset; + uint32_t value_metadata_length; + uint64_t time_metadata_offset; + uint32_t time_metadata_length; +} TsFilePreparedLocator; typedef struct arrow_schema { // Array type description @@ -664,6 +686,32 @@ ERRNO tsfile_writer_write(TsFileWriter writer, Tablet tablet); /*-------------------TsFile reader query data------------------ */ +/** Deserialize one exact Dataset Index locator into a reusable series. */ +PreparedSeriesHandle tsfile_reader_prepare_series( + TsFileReader reader, const TsFilePreparedLocator* locator, ERRNO* err_code); + +/** Prepare an aligned value locator by sharing an existing parsed time index. + */ +PreparedSeriesHandle tsfile_reader_prepare_series_with_time_owner( + TsFileReader reader, const TsFilePreparedLocator* locator, + PreparedSeriesHandle aligned_time_owner, ERRNO* err_code); + +/** Release a prepared handle. Existing result sets remain independently owned. + */ +void tsfile_prepared_series_free(PreparedSeriesHandle prepared); + +/** Query a prepared series without traversing the TsFile footer index. */ +ResultSet tsfile_reader_query_prepared(TsFileReader reader, + PreparedSeriesHandle prepared, + Timestamp start_time, Timestamp end_time, + int offset, int limit, ERRNO* err_code); + +/** Query multiple aligned prepared value columns sharing one time axis. */ +ResultSet tsfile_reader_query_prepared_multi( + TsFileReader reader, const PreparedSeriesHandle* prepared, + uint32_t prepared_count, Timestamp start_time, Timestamp end_time, + int offset, int limit, ERRNO* err_code); + /** * @brief Queries time series data from a specific table within time range. * diff --git a/cpp/src/dataset/CMakeLists.txt b/cpp/src/dataset/CMakeLists.txt new file mode 100644 index 000000000..7adbdcc8f --- /dev/null +++ b/cpp/src/dataset/CMakeLists.txt @@ -0,0 +1,25 @@ +#[[ +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 + + https://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. +]] +message("Running in src/dataset directory") +set(CMAKE_POSITION_INDEPENDENT_CODE ON) +aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR} dataset_SRC_LIST) +add_library(dataset_obj OBJECT ${dataset_SRC_LIST}) + +file(GLOB HEADERS "${CMAKE_CURRENT_SOURCE_DIR}/*.h") +copy_to_dir(${HEADERS} "dataset_obj") diff --git a/cpp/src/dataset/dataset_index.cc b/cpp/src/dataset/dataset_index.cc new file mode 100644 index 000000000..f098e3aa2 --- /dev/null +++ b/cpp/src/dataset/dataset_index.cc @@ -0,0 +1,1049 @@ +/* + * 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. + */ + +#include "dataset/dataset_index.h" + +#include +#include +#include +#include +#include + +#ifdef _WIN32 +#include +#include +#include +#include +#include +#else +#include +#include +#include +#include +#include +#endif + +namespace storage { +namespace dataset { + +namespace { + +const char DATASET_INDEX_MAGIC[8] = {'T', 'S', 'I', 'D', 'X', 0, 0, 0}; + +bool is_little_endian() { + const uint16_t value = 1; + return *reinterpret_cast(&value) == 1; +} + +uint64_t align64(uint64_t value) { + if (value > std::numeric_limits::max() - 63) { + return 0; + } + return (value + 63) & ~static_cast(63); +} + +bool add_overflows(uint64_t left, uint64_t right) { + return left > std::numeric_limits::max() - right; +} + +bool multiply_overflows(uint64_t left, uint64_t right) { + return right != 0 && left > std::numeric_limits::max() / right; +} + +bool range_valid(uint32_t first, uint32_t count, uint32_t total) { + return first <= total && count <= total - first; +} + +uint32_t expected_record_size(DatasetIndexSectionType type) { + switch (type) { + case DatasetIndexSectionType::STRING_OFFSETS: + return sizeof(uint32_t); + case DatasetIndexSectionType::STRING_BYTES: + return 0; + case DatasetIndexSectionType::TABLE_NAME_INDEX: + return sizeof(TableNameIndexRecord); + case DatasetIndexSectionType::TABLE_RECORD: + return sizeof(TableRecord); + case DatasetIndexSectionType::DEVICE_NAME_INDEX: + return sizeof(DeviceNameIndexRecord); + case DatasetIndexSectionType::DEVICE_RECORD: + return sizeof(DeviceRecord); + case DatasetIndexSectionType::COLUMN_NAME_INDEX: + return sizeof(ColumnNameIndexRecord); + case DatasetIndexSectionType::COLUMN_SCHEMA: + return sizeof(ColumnSchemaRecord); + case DatasetIndexSectionType::LOGICAL_SERIES: + return sizeof(LogicalSeriesRecord); + case DatasetIndexSectionType::TSFILE_RECORD: + return sizeof(TsFileRecord); + case DatasetIndexSectionType::DEVICE_FILE_SPAN: + return sizeof(DeviceFileSpanRecord); + case DatasetIndexSectionType::SERIES_FILE_SPAN: + return sizeof(SeriesFileSpanRecord); + case DatasetIndexSectionType::SERIES_LOCATOR: + return sizeof(SeriesLocatorRecord); + } + return std::numeric_limits::max(); +} + +std::string parent_directory(const std::string& path) { + const std::string::size_type pos = path.find_last_of("/\\"); + if (pos == std::string::npos) { + return "."; + } + if (pos == 0) { + return path.substr(0, 1); + } + return path.substr(0, pos); +} + +std::string system_error_message(const char* operation) { + std::ostringstream stream; + stream << operation << " failed: " << std::strerror(errno); + return stream.str(); +} + +#ifdef _WIN32 +typedef int NativeFile; +const NativeFile INVALID_NATIVE_FILE = -1; + +NativeFile open_exclusive(const std::string& path) { + return _open(path.c_str(), _O_BINARY | _O_CREAT | _O_EXCL | _O_WRONLY, + 0666); +} + +int native_write(NativeFile file, const void* data, size_t size) { + const char* cursor = static_cast(data); + while (size > 0) { + const unsigned int chunk = static_cast( + std::min(size, std::numeric_limits::max())); + const int written = _write(file, cursor, chunk); + if (written <= 0) { + return -1; + } + cursor += written; + size -= static_cast(written); + } + return 0; +} + +int native_sync(NativeFile file) { return _commit(file); } +void native_close(NativeFile file) { _close(file); } +int process_id() { return _getpid(); } +#else +typedef int NativeFile; +const NativeFile INVALID_NATIVE_FILE = -1; + +NativeFile open_exclusive(const std::string& path) { + return ::open(path.c_str(), O_CREAT | O_EXCL | O_WRONLY, 0644); +} + +int native_write(NativeFile file, const void* data, size_t size) { + const uint8_t* cursor = static_cast(data); + while (size > 0) { + const ssize_t written = ::write(file, cursor, size); + if (written < 0) { + if (errno == EINTR) { + continue; + } + return -1; + } + if (written == 0) { + errno = EIO; + return -1; + } + cursor += written; + size -= static_cast(written); + } + return 0; +} + +int native_sync(NativeFile file) { return ::fsync(file); } +void native_close(NativeFile file) { ::close(file); } +int process_id() { return static_cast(::getpid()); } +#endif + +int write_zero_padding(NativeFile file, uint64_t length) { + static const uint8_t zeros[64] = {0}; + while (length > 0) { + const size_t chunk = + static_cast(std::min(length, sizeof(zeros))); + if (native_write(file, zeros, chunk) != 0) { + return -1; + } + length -= chunk; + } + return 0; +} + +bool string_equals(const MappedDatasetIndex& index, uint32_t sid, + const std::string& expected) { + DatasetIndexStringView value; + return index.string(sid, value) == DatasetIndexStatus::OK && + value.length == expected.size() && + (value.length == 0 || + std::memcmp(value.data, expected.data(), value.length) == 0); +} + +int compare_string_views(const DatasetIndexStringView& left, + const DatasetIndexStringView& right) { + const uint32_t common_length = std::min(left.length, right.length); + if (common_length != 0) { + const int comparison = + std::memcmp(left.data, right.data, common_length); + if (comparison != 0) { + return comparison; + } + } + return left.length < right.length ? -1 : left.length > right.length; +} + +} // namespace + +const char* dataset_index_status_name(DatasetIndexStatus status) { + switch (status) { + case DatasetIndexStatus::OK: + return "OK"; + case DatasetIndexStatus::INVALID_ARGUMENT: + return "INVALID_ARGUMENT"; + case DatasetIndexStatus::IO_ERROR: + return "IO_ERROR"; + case DatasetIndexStatus::OUT_OF_MEMORY: + return "OUT_OF_MEMORY"; + case DatasetIndexStatus::BAD_MAGIC: + return "BAD_MAGIC"; + case DatasetIndexStatus::UNSUPPORTED_VERSION: + return "UNSUPPORTED_VERSION"; + case DatasetIndexStatus::BAD_HEADER: + return "BAD_HEADER"; + case DatasetIndexStatus::BAD_DIRECTORY: + return "BAD_DIRECTORY"; + case DatasetIndexStatus::BAD_SECTION: + return "BAD_SECTION"; + case DatasetIndexStatus::BAD_CHECKSUM: + return "BAD_CHECKSUM"; + case DatasetIndexStatus::BAD_REFERENCE: + return "BAD_REFERENCE"; + case DatasetIndexStatus::NOT_FOUND: + return "NOT_FOUND"; + } + return "UNKNOWN"; +} + +uint32_t dataset_index_crc32c(const void* data, size_t length) { + const uint8_t* bytes = static_cast(data); + uint32_t crc = ~static_cast(0); + for (size_t i = 0; i < length; ++i) { + crc ^= bytes[i]; + for (int bit = 0; bit < 8; ++bit) { + const uint32_t mask = + static_cast(-(static_cast(crc & 1))); + crc = (crc >> 1) ^ (0x82F63B78U & mask); + } + } + return ~crc; +} + +uint64_t dataset_index_name_hash(const char* data, size_t length) { + uint64_t hash = 1469598103934665603ULL; + for (size_t i = 0; i < length; ++i) { + hash ^= static_cast(data[i]); + hash *= 1099511628211ULL; + } + return hash; +} + +DatasetIndexStatus DatasetIndexWriter::write_atomic( + const std::string& output_path, + const std::vector& input_sections, + std::string& error_message) { + error_message.clear(); + if (output_path.empty() || !is_little_endian()) { + error_message = "empty output path or unsupported host byte order"; + return DatasetIndexStatus::INVALID_ARGUMENT; + } + if (input_sections.size() != DATASET_INDEX_SECTION_COUNT) { + error_message = "current v1 writer requires exactly 13 sections"; + return DatasetIndexStatus::INVALID_ARGUMENT; + } + + std::vector sections(input_sections); + std::sort(sections.begin(), sections.end(), + [](const DatasetIndexSectionData& left, + const DatasetIndexSectionData& right) { + return static_cast(left.type) < + static_cast(right.type); + }); + + DatasetIndexHeader header; + std::memset(&header, 0, sizeof(header)); + std::memcpy(header.magic, DATASET_INDEX_MAGIC, sizeof(header.magic)); + header.version_major = DATASET_INDEX_VERSION_MAJOR; + header.version_minor = DATASET_INDEX_VERSION_MINOR; + header.header_size = DATASET_INDEX_HEADER_SIZE; + header.directory_offset = DATASET_INDEX_HEADER_SIZE; + header.section_count = DATASET_INDEX_SECTION_COUNT; + header.directory_entry_size = DATASET_INDEX_DIRECTORY_ENTRY_SIZE; + + std::vector directory( + DATASET_INDEX_SECTION_COUNT); + uint64_t next_offset = align64(header.directory_offset + + static_cast(directory.size()) * + sizeof(DatasetIndexDirectoryEntry)); + if (next_offset == 0) { + error_message = "directory size overflows"; + return DatasetIndexStatus::INVALID_ARGUMENT; + } + + for (uint32_t i = 0; i < DATASET_INDEX_SECTION_COUNT; ++i) { + const uint32_t expected_type = i + 1; + const uint32_t actual_type = static_cast(sections[i].type); + const uint32_t expected_size = expected_record_size(sections[i].type); + if (actual_type != expected_type || + sections[i].record_size != expected_size) { + error_message = "section type or record size does not match v1"; + return DatasetIndexStatus::INVALID_ARGUMENT; + } + if (expected_size != 0 && + (multiply_overflows(sections[i].count, expected_size) || + static_cast(sections[i].count) * expected_size != + sections[i].bytes.size())) { + error_message = "fixed-size section count does not match bytes"; + return DatasetIndexStatus::INVALID_ARGUMENT; + } + if (expected_size == 0 && + sections[i].count != sections[i].bytes.size()) { + error_message = "blob section count must equal byte length"; + return DatasetIndexStatus::INVALID_ARGUMENT; + } + DatasetIndexDirectoryEntry& entry = directory[i]; + std::memset(&entry, 0, sizeof(entry)); + entry.section_type = actual_type; + entry.record_size = expected_size; + entry.offset = next_offset; + entry.length = sections[i].bytes.size(); + entry.count = sections[i].count; + entry.crc32c = dataset_index_crc32c( + sections[i].bytes.empty() ? nullptr : sections[i].bytes.data(), + sections[i].bytes.size()); + if (add_overflows(entry.offset, entry.length)) { + error_message = "section file range overflows"; + return DatasetIndexStatus::INVALID_ARGUMENT; + } + next_offset = align64(entry.offset + entry.length); + if (next_offset == 0 && i + 1 < DATASET_INDEX_SECTION_COUNT) { + error_message = "section alignment overflows"; + return DatasetIndexStatus::INVALID_ARGUMENT; + } + } + + const DatasetIndexDirectoryEntry& last = directory.back(); + header.file_length = last.offset + last.length; + header.header_crc32c = 0; + header.header_crc32c = dataset_index_crc32c(&header, sizeof(header)); + + std::ostringstream temp_name; + temp_name << output_path << ".tmp." << process_id(); + const std::string temp_path = temp_name.str(); + NativeFile file = open_exclusive(temp_path); + if (file == INVALID_NATIVE_FILE) { + error_message = system_error_message("create temporary index"); + return DatasetIndexStatus::IO_ERROR; + } + + DatasetIndexStatus status = DatasetIndexStatus::OK; + uint64_t cursor = 0; + if (native_write(file, &header, sizeof(header)) != 0 || + native_write(file, directory.data(), + directory.size() * sizeof(directory[0])) != 0) { + error_message = system_error_message("write index header"); + status = DatasetIndexStatus::IO_ERROR; + } else { + cursor = header.directory_offset + + directory.size() * sizeof(DatasetIndexDirectoryEntry); + } + + for (uint32_t i = 0; + status == DatasetIndexStatus::OK && i < directory.size(); ++i) { + if (directory[i].offset < cursor || + write_zero_padding(file, directory[i].offset - cursor) != 0 || + native_write( + file, + sections[i].bytes.empty() ? nullptr : sections[i].bytes.data(), + sections[i].bytes.size()) != 0) { + error_message = system_error_message("write index section"); + status = DatasetIndexStatus::IO_ERROR; + break; + } + cursor = directory[i].offset + directory[i].length; + } + if (status == DatasetIndexStatus::OK && native_sync(file) != 0) { + error_message = system_error_message("fsync index"); + status = DatasetIndexStatus::IO_ERROR; + } + native_close(file); + + if (status == DatasetIndexStatus::OK) { +#ifdef _WIN32 + if (!MoveFileExA(temp_path.c_str(), output_path.c_str(), + MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH)) { + error_message = "replace index failed"; + status = DatasetIndexStatus::IO_ERROR; + } +#else + if (::rename(temp_path.c_str(), output_path.c_str()) != 0) { + error_message = system_error_message("replace index"); + status = DatasetIndexStatus::IO_ERROR; + } else { + const std::string directory_path = parent_directory(output_path); + const int directory_fd = ::open(directory_path.c_str(), O_RDONLY); + if (directory_fd >= 0) { + if (::fsync(directory_fd) != 0) { + error_message = + system_error_message("fsync index directory"); + status = DatasetIndexStatus::IO_ERROR; + } + ::close(directory_fd); + } + } +#endif + } + if (status != DatasetIndexStatus::OK) { +#ifdef _WIN32 + _unlink(temp_path.c_str()); +#else + ::unlink(temp_path.c_str()); +#endif + } + return status; +} + +MappedDatasetIndex::MappedDatasetIndex() + : mapping_(nullptr), + mapping_size_(0), + header_(nullptr), + directory_(nullptr) +#ifdef _WIN32 + , + file_handle_(INVALID_HANDLE_VALUE), + mapping_handle_(nullptr) +#else + , + fd_(-1) +#endif +{ +} + +MappedDatasetIndex::~MappedDatasetIndex() { close(); } + +DatasetIndexStatus MappedDatasetIndex::fail(DatasetIndexStatus status, + const std::string& message) { + error_message_ = message; + return status; +} + +DatasetIndexStatus MappedDatasetIndex::open(const std::string& path) { + close(); + path_ = path; + if (path.empty() || !is_little_endian()) { + return fail(DatasetIndexStatus::INVALID_ARGUMENT, + "empty path or unsupported host byte order"); + } +#ifdef _WIN32 + HANDLE file = + CreateFileA(path.c_str(), GENERIC_READ, FILE_SHARE_READ, nullptr, + OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); + if (file == INVALID_HANDLE_VALUE) { + return fail(DatasetIndexStatus::IO_ERROR, "open index failed"); + } + LARGE_INTEGER size; + if (!GetFileSizeEx(file, &size) || size.QuadPart <= 0) { + CloseHandle(file); + return fail(DatasetIndexStatus::IO_ERROR, "stat index failed"); + } + HANDLE mapping = + CreateFileMappingA(file, nullptr, PAGE_READONLY, 0, 0, nullptr); + if (mapping == nullptr) { + CloseHandle(file); + return fail(DatasetIndexStatus::IO_ERROR, "map index failed"); + } + const void* address = MapViewOfFile(mapping, FILE_MAP_READ, 0, 0, 0); + if (address == nullptr) { + CloseHandle(mapping); + CloseHandle(file); + return fail(DatasetIndexStatus::IO_ERROR, "map view failed"); + } + file_handle_ = file; + mapping_handle_ = mapping; + mapping_ = static_cast(address); + mapping_size_ = static_cast(size.QuadPart); +#else + fd_ = ::open(path.c_str(), O_RDONLY); + if (fd_ < 0) { + return fail(DatasetIndexStatus::IO_ERROR, + system_error_message("open index")); + } + struct stat stat_buffer; + if (::fstat(fd_, &stat_buffer) != 0 || stat_buffer.st_size <= 0) { + close(); + return fail(DatasetIndexStatus::IO_ERROR, + system_error_message("stat index")); + } + mapping_size_ = static_cast(stat_buffer.st_size); + void* address = + ::mmap(nullptr, mapping_size_, PROT_READ, MAP_SHARED, fd_, 0); + if (address == MAP_FAILED) { + mapping_ = nullptr; + close(); + return fail(DatasetIndexStatus::IO_ERROR, + system_error_message("mmap index")); + } + mapping_ = static_cast(address); +#endif + DatasetIndexStatus status = validate(); + if (status != DatasetIndexStatus::OK) { + const std::string message = error_message_; + close(); + error_message_ = message; + } + return status; +} + +void MappedDatasetIndex::close() { +#ifdef _WIN32 + if (mapping_ != nullptr) { + UnmapViewOfFile(mapping_); + } + if (mapping_handle_ != nullptr) { + CloseHandle(static_cast(mapping_handle_)); + } + if (file_handle_ != INVALID_HANDLE_VALUE) { + CloseHandle(static_cast(file_handle_)); + } + file_handle_ = INVALID_HANDLE_VALUE; + mapping_handle_ = nullptr; +#else + if (mapping_ != nullptr) { + ::munmap(const_cast(mapping_), mapping_size_); + } + if (fd_ >= 0) { + ::close(fd_); + } + fd_ = -1; +#endif + mapping_ = nullptr; + mapping_size_ = 0; + header_ = nullptr; + directory_ = nullptr; + path_.clear(); +} + +DatasetIndexStatus MappedDatasetIndex::section( + DatasetIndexSectionType type, DatasetIndexSectionView& result) const { + result = DatasetIndexSectionView(); + if (directory_ == nullptr) { + return DatasetIndexStatus::BAD_HEADER; + } + const uint32_t raw_type = static_cast(type); + if (raw_type == 0 || raw_type > header_->section_count) { + return DatasetIndexStatus::NOT_FOUND; + } + const DatasetIndexDirectoryEntry& entry = directory_[raw_type - 1]; + if (entry.section_type != raw_type) { + return DatasetIndexStatus::NOT_FOUND; + } + result.data = mapping_ + entry.offset; + result.length = entry.length; + result.record_size = entry.record_size; + result.count = entry.count; + return DatasetIndexStatus::OK; +} + +DatasetIndexStatus MappedDatasetIndex::string( + uint32_t sid, DatasetIndexStringView& result) const { + result = DatasetIndexStringView(); + DatasetIndexSectionView offsets; + DatasetIndexSectionView bytes; + if (section(DatasetIndexSectionType::STRING_OFFSETS, offsets) != + DatasetIndexStatus::OK || + section(DatasetIndexSectionType::STRING_BYTES, bytes) != + DatasetIndexStatus::OK || + offsets.count == 0 || sid + 1 >= offsets.count) { + return DatasetIndexStatus::NOT_FOUND; + } + const uint32_t* values = reinterpret_cast(offsets.data); + result.data = reinterpret_cast(bytes.data + values[sid]); + result.length = values[sid + 1] - values[sid]; + return DatasetIndexStatus::OK; +} + +DatasetIndexStatus MappedDatasetIndex::validate() { + if (mapping_size_ < sizeof(DatasetIndexHeader)) { + return fail(DatasetIndexStatus::BAD_HEADER, + "index is shorter than header"); + } + header_ = reinterpret_cast(mapping_); + if (std::memcmp(header_->magic, DATASET_INDEX_MAGIC, + sizeof(header_->magic)) != 0) { + return fail(DatasetIndexStatus::BAD_MAGIC, "bad dataset index magic"); + } + if (header_->version_major != DATASET_INDEX_VERSION_MAJOR || + header_->version_minor != DATASET_INDEX_VERSION_MINOR) { + return fail(DatasetIndexStatus::UNSUPPORTED_VERSION, + "unsupported dataset index version"); + } + if (header_->header_size != sizeof(DatasetIndexHeader) || + header_->directory_entry_size != sizeof(DatasetIndexDirectoryEntry) || + header_->section_count != DATASET_INDEX_SECTION_COUNT || + header_->directory_offset < sizeof(DatasetIndexHeader) || + header_->file_length != mapping_size_) { + return fail(DatasetIndexStatus::BAD_HEADER, + "header size, directory shape, or file length is invalid"); + } + for (size_t i = 0; i < sizeof(header_->reserved); ++i) { + if (header_->reserved[i] != 0) { + return fail(DatasetIndexStatus::BAD_HEADER, + "header reserved bytes are not zero"); + } + } + DatasetIndexHeader header_copy = *header_; + const uint32_t expected_header_crc = header_copy.header_crc32c; + header_copy.header_crc32c = 0; + if (dataset_index_crc32c(&header_copy, sizeof(header_copy)) != + expected_header_crc) { + return fail(DatasetIndexStatus::BAD_CHECKSUM, + "header checksum does not match"); + } + const uint64_t directory_length = + static_cast(header_->section_count) * + header_->directory_entry_size; + if (add_overflows(header_->directory_offset, directory_length) || + header_->directory_offset + directory_length > mapping_size_) { + return fail(DatasetIndexStatus::BAD_DIRECTORY, + "section directory is outside the file"); + } + directory_ = reinterpret_cast( + mapping_ + header_->directory_offset); + uint64_t minimum_section_offset = + align64(header_->directory_offset + directory_length); + uint64_t previous_end = minimum_section_offset; + for (uint32_t i = 0; i < header_->section_count; ++i) { + const DatasetIndexDirectoryEntry& entry = directory_[i]; + const DatasetIndexSectionType type = + static_cast(i + 1); + if (entry.section_type != i + 1 || + entry.record_size != expected_record_size(type)) { + return fail(DatasetIndexStatus::BAD_DIRECTORY, + "section type order or record size is invalid"); + } + if (entry.offset % DATASET_INDEX_ALIGNMENT != 0 || + entry.offset < minimum_section_offset || + entry.offset < previous_end || + add_overflows(entry.offset, entry.length) || + entry.offset + entry.length > mapping_size_) { + return fail( + DatasetIndexStatus::BAD_SECTION, + "section range is unaligned, overlapping, or out of bounds"); + } + if (entry.record_size == 0) { + if (entry.count != entry.length) { + return fail(DatasetIndexStatus::BAD_SECTION, + "blob count does not equal byte length"); + } + } else if (multiply_overflows(entry.count, entry.record_size) || + static_cast(entry.count) * entry.record_size > + entry.length) { + return fail(DatasetIndexStatus::BAD_SECTION, + "section record count exceeds section length"); + } + const uint32_t actual_crc = dataset_index_crc32c( + entry.length == 0 ? nullptr : mapping_ + entry.offset, + static_cast(entry.length)); + if (actual_crc != entry.crc32c) { + return fail(DatasetIndexStatus::BAD_CHECKSUM, + "section checksum does not match"); + } + previous_end = entry.offset + entry.length; + } + + DatasetIndexSectionView string_offsets_view; + DatasetIndexSectionView string_bytes_view; + section(DatasetIndexSectionType::STRING_OFFSETS, string_offsets_view); + section(DatasetIndexSectionType::STRING_BYTES, string_bytes_view); + if (string_offsets_view.count == 0) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "StringOffsets must contain the terminal offset"); + } + const uint32_t* string_offsets = + reinterpret_cast(string_offsets_view.data); + if (string_offsets[0] != 0) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "StringOffsets must start at zero"); + } + for (uint32_t i = 1; i < string_offsets_view.count; ++i) { + if (string_offsets[i] < string_offsets[i - 1] || + string_offsets[i] > string_bytes_view.length) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "StringOffsets is not monotonic or is out of bounds"); + } + } + if (string_offsets[string_offsets_view.count - 1] != + string_bytes_view.length) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "terminal string offset does not equal StringBytes length"); + } + const uint32_t string_count = string_offsets_view.count - 1; + + DatasetIndexSectionView table_names; + DatasetIndexSectionView tables; + DatasetIndexSectionView device_names; + DatasetIndexSectionView devices; + DatasetIndexSectionView column_names; + DatasetIndexSectionView columns; + DatasetIndexSectionView series; + DatasetIndexSectionView files; + DatasetIndexSectionView device_spans; + DatasetIndexSectionView series_spans; + DatasetIndexSectionView locators; + section(DatasetIndexSectionType::TABLE_NAME_INDEX, table_names); + section(DatasetIndexSectionType::TABLE_RECORD, tables); + section(DatasetIndexSectionType::DEVICE_NAME_INDEX, device_names); + section(DatasetIndexSectionType::DEVICE_RECORD, devices); + section(DatasetIndexSectionType::COLUMN_NAME_INDEX, column_names); + section(DatasetIndexSectionType::COLUMN_SCHEMA, columns); + section(DatasetIndexSectionType::LOGICAL_SERIES, series); + section(DatasetIndexSectionType::TSFILE_RECORD, files); + section(DatasetIndexSectionType::DEVICE_FILE_SPAN, device_spans); + section(DatasetIndexSectionType::SERIES_FILE_SPAN, series_spans); + section(DatasetIndexSectionType::SERIES_LOCATOR, locators); + +#define DATASET_RECORDS(view, type) reinterpret_cast((view).data) + const TableNameIndexRecord* table_name_records = + DATASET_RECORDS(table_names, TableNameIndexRecord); + const TableRecord* table_records = DATASET_RECORDS(tables, TableRecord); + const DeviceNameIndexRecord* device_name_records = + DATASET_RECORDS(device_names, DeviceNameIndexRecord); + const DeviceRecord* device_records = DATASET_RECORDS(devices, DeviceRecord); + const ColumnNameIndexRecord* column_name_records = + DATASET_RECORDS(column_names, ColumnNameIndexRecord); + const ColumnSchemaRecord* column_records = + DATASET_RECORDS(columns, ColumnSchemaRecord); + const LogicalSeriesRecord* series_records = + DATASET_RECORDS(series, LogicalSeriesRecord); + const TsFileRecord* file_records = DATASET_RECORDS(files, TsFileRecord); + const DeviceFileSpanRecord* device_span_records = + DATASET_RECORDS(device_spans, DeviceFileSpanRecord); + const SeriesFileSpanRecord* series_span_records = + DATASET_RECORDS(series_spans, SeriesFileSpanRecord); + const SeriesLocatorRecord* locator_records = + DATASET_RECORDS(locators, SeriesLocatorRecord); +#undef DATASET_RECORDS + + if (table_names.count != tables.count) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "TableNameIndex must contain one entry per TableRecord"); + } + std::vector seen_table_ids(tables.count, 0); + DatasetIndexStringView previous_table_name; + uint64_t previous_table_hash = 0; + bool have_previous_table_name = false; + for (uint32_t i = 0; i < table_names.count; ++i) { + const TableNameIndexRecord& value = table_name_records[i]; + if (value.name_sid >= string_count || value.table_id >= tables.count || + seen_table_ids[value.table_id] != 0 || + table_records[value.table_id].name_sid != value.name_sid) { + return fail( + DatasetIndexStatus::BAD_REFERENCE, + "TableNameIndex contains an invalid or duplicate reference"); + } + DatasetIndexStringView table_name; + if (string(value.name_sid, table_name) != DatasetIndexStatus::OK || + dataset_index_name_hash(table_name.data, table_name.length) != + value.name_hash) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "TableNameIndex name hash does not match its string"); + } + if (have_previous_table_name && + (value.name_hash < previous_table_hash || + (value.name_hash == previous_table_hash && + compare_string_views(table_name, previous_table_name) <= 0))) { + return fail( + DatasetIndexStatus::BAD_REFERENCE, + "TableNameIndex is unsorted or has duplicate table names"); + } + seen_table_ids[value.table_id] = 1; + previous_table_name = table_name; + previous_table_hash = value.name_hash; + have_previous_table_name = true; + } + for (uint32_t i = 0; i < tables.count; ++i) { + const TableRecord& value = table_records[i]; + if (value.name_sid >= string_count || value.reserved0 != 0 || + value.reserved1 != 0 || + !range_valid(value.first_device_name_index, value.device_count, + device_names.count) || + !range_valid(value.first_column_name_index, value.column_count, + column_names.count)) { + return fail( + DatasetIndexStatus::BAD_REFERENCE, + "TableRecord contains an invalid range or reserved value"); + } + } + for (uint32_t i = 0; i < device_names.count; ++i) { + const DeviceNameIndexRecord& value = device_name_records[i]; + if (value.table_id >= tables.count || + value.device_id >= devices.count || + value.name_sid >= string_count || value.reserved != 0) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "DeviceNameIndex contains an invalid reference"); + } + } + for (uint32_t i = 0; i < devices.count; ++i) { + const DeviceRecord& value = device_records[i]; + if (value.table_id >= tables.count || value.name_sid >= string_count || + value.reserved0 != 0 || value.reserved1 != 0 || + !range_valid(value.first_series_id, value.series_count, + series.count) || + !range_valid(value.first_file_span, value.file_span_count, + device_spans.count) || + (value.series_count != 0 && value.min_time > value.max_time)) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "DeviceRecord contains an invalid reference or range"); + } + } + for (uint32_t i = 0; i < column_names.count; ++i) { + const ColumnNameIndexRecord& value = column_name_records[i]; + if (value.table_id >= tables.count || + value.column_id >= columns.count || + value.name_sid >= string_count || value.reserved != 0) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "ColumnNameIndex contains an invalid reference"); + } + } + for (uint32_t i = 0; i < columns.count; ++i) { + const ColumnSchemaRecord& value = column_records[i]; + if (value.table_id >= tables.count || value.name_sid >= string_count || + value.nullable > 1 || value.reserved != 0) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "ColumnSchema contains an invalid reference or enum"); + } + } + for (uint32_t i = 0; i < series.count; ++i) { + const LogicalSeriesRecord& value = series_records[i]; + if (value.device_id >= devices.count || + value.column_id >= columns.count || + !range_valid(value.first_file_span, value.file_span_count, + series_spans.count) || + (value.file_span_count != 0 && value.min_time > value.max_time)) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "LogicalSeries contains an invalid reference or range"); + } + } + for (uint32_t i = 0; i < files.count; ++i) { + const TsFileRecord& value = file_records[i]; + if (value.path_sid >= string_count || value.reserved0 != 0 || + value.reserved1 != 0) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "TsFileRecord contains an invalid reference"); + } + } + for (uint32_t i = 0; i < device_spans.count; ++i) { + const DeviceFileSpanRecord& value = device_span_records[i]; + if (value.device_id >= devices.count || value.file_id >= files.count || + value.layout > 1 || + (value.flags & ~static_cast(1)) != 0 || + (value.layout == 0 && + (value.time_meta_offset != 0 || value.time_meta_length != 0 || + value.row_count != 0)) || + (value.layout == 1 && + (value.time_meta_length == 0 || value.row_count == 0))) { + return fail( + DatasetIndexStatus::BAD_REFERENCE, + "DeviceFileSpan contains an invalid reference or layout"); + } + if (value.layout == 1 && + (add_overflows(value.time_meta_offset, value.time_meta_length) || + value.time_meta_offset + value.time_meta_length > + file_records[value.file_id].file_size)) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "aligned time metadata range is outside its TsFile"); + } + } + for (uint32_t i = 0; i < series_spans.count; ++i) { + const SeriesFileSpanRecord& value = series_span_records[i]; + if (value.series_id >= series.count || value.file_id >= files.count || + value.locator_id >= locators.count || value.reserved != 0 || + value.min_time > value.max_time) { + return fail( + DatasetIndexStatus::BAD_REFERENCE, + "SeriesFileSpan contains an invalid reference or range"); + } + } + for (uint32_t i = 0; i < locators.count; ++i) { + const SeriesLocatorRecord& value = locator_records[i]; + if (value.device_file_span_id >= device_spans.count || + value.locator_kind > 1 || value.flags != 0 || value.padding != 0 || + value.timeseries_meta_length == 0) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "SeriesLocator contains an invalid reference or enum"); + } + const DeviceFileSpanRecord& device_span = + device_span_records[value.device_file_span_id]; + if (value.locator_kind != device_span.layout || + add_overflows(value.timeseries_meta_offset, + value.timeseries_meta_length) || + value.timeseries_meta_offset + value.timeseries_meta_length > + file_records[device_span.file_id].file_size) { + return fail(DatasetIndexStatus::BAD_REFERENCE, + "SeriesLocator metadata range is outside its TsFile"); + } + } + error_message_.clear(); + return DatasetIndexStatus::OK; +} + +DatasetIndexStatus MappedDatasetIndex::find_table_ids( + const std::string& name, std::vector& table_ids) const { + table_ids.clear(); + DatasetIndexSectionView view; + if (section(DatasetIndexSectionType::TABLE_NAME_INDEX, view) != + DatasetIndexStatus::OK) { + return DatasetIndexStatus::BAD_SECTION; + } + const TableNameIndexRecord* records = + reinterpret_cast(view.data); + const uint64_t hash = dataset_index_name_hash(name.data(), name.size()); + uint32_t low = 0; + uint32_t high = view.count; + while (low < high) { + const uint32_t middle = low + (high - low) / 2; + if (records[middle].name_hash < hash) { + low = middle + 1; + } else { + high = middle; + } + } + for (uint32_t i = low; i < view.count && records[i].name_hash == hash; + ++i) { + if (string_equals(*this, records[i].name_sid, name)) { + table_ids.push_back(records[i].table_id); + } + } + return table_ids.empty() ? DatasetIndexStatus::NOT_FOUND + : DatasetIndexStatus::OK; +} + +DatasetIndexStatus MappedDatasetIndex::find_device_id( + uint32_t table_id, const std::string& name, uint32_t& device_id) const { + const TableRecord* table = nullptr; + if (record(DatasetIndexSectionType::TABLE_RECORD, table_id, table) != + DatasetIndexStatus::OK) { + return DatasetIndexStatus::NOT_FOUND; + } + DatasetIndexSectionView view; + section(DatasetIndexSectionType::DEVICE_NAME_INDEX, view); + const DeviceNameIndexRecord* records = + reinterpret_cast(view.data); + const uint64_t hash = dataset_index_name_hash(name.data(), name.size()); + uint32_t low = table->first_device_name_index; + uint32_t high = low + table->device_count; + while (low < high) { + const uint32_t middle = low + (high - low) / 2; + if (records[middle].name_hash < hash) { + low = middle + 1; + } else { + high = middle; + } + } + const uint32_t end = table->first_device_name_index + table->device_count; + for (uint32_t i = low; i < end && records[i].name_hash == hash; ++i) { + if (records[i].table_id == table_id && + string_equals(*this, records[i].name_sid, name)) { + device_id = records[i].device_id; + return DatasetIndexStatus::OK; + } + } + return DatasetIndexStatus::NOT_FOUND; +} + +DatasetIndexStatus MappedDatasetIndex::find_column_id( + uint32_t table_id, const std::string& name, uint32_t& column_id) const { + const TableRecord* table = nullptr; + if (record(DatasetIndexSectionType::TABLE_RECORD, table_id, table) != + DatasetIndexStatus::OK) { + return DatasetIndexStatus::NOT_FOUND; + } + DatasetIndexSectionView view; + section(DatasetIndexSectionType::COLUMN_NAME_INDEX, view); + const ColumnNameIndexRecord* records = + reinterpret_cast(view.data); + const uint64_t hash = dataset_index_name_hash(name.data(), name.size()); + uint32_t low = table->first_column_name_index; + uint32_t high = low + table->column_count; + while (low < high) { + const uint32_t middle = low + (high - low) / 2; + if (records[middle].name_hash < hash) { + low = middle + 1; + } else { + high = middle; + } + } + const uint32_t end = table->first_column_name_index + table->column_count; + for (uint32_t i = low; i < end && records[i].name_hash == hash; ++i) { + if (records[i].table_id == table_id && + string_equals(*this, records[i].name_sid, name)) { + column_id = records[i].column_id; + return DatasetIndexStatus::OK; + } + } + return DatasetIndexStatus::NOT_FOUND; +} + +DatasetIndexStatus MappedDatasetIndex::find_series_id( + uint32_t device_id, uint32_t column_id, uint32_t& series_id) const { + const DeviceRecord* device = nullptr; + if (record(DatasetIndexSectionType::DEVICE_RECORD, device_id, device) != + DatasetIndexStatus::OK) { + return DatasetIndexStatus::NOT_FOUND; + } + DatasetIndexSectionView view; + section(DatasetIndexSectionType::LOGICAL_SERIES, view); + const LogicalSeriesRecord* records = + reinterpret_cast(view.data); + uint32_t low = device->first_series_id; + uint32_t high = low + device->series_count; + while (low < high) { + const uint32_t middle = low + (high - low) / 2; + if (records[middle].column_id < column_id) { + low = middle + 1; + } else { + high = middle; + } + } + if (low < device->first_series_id + device->series_count && + records[low].device_id == device_id && + records[low].column_id == column_id) { + series_id = low; + return DatasetIndexStatus::OK; + } + return DatasetIndexStatus::NOT_FOUND; +} + +} // namespace dataset +} // namespace storage diff --git a/cpp/src/dataset/dataset_index.h b/cpp/src/dataset/dataset_index.h new file mode 100644 index 000000000..faec4840e --- /dev/null +++ b/cpp/src/dataset/dataset_index.h @@ -0,0 +1,331 @@ +/* + * 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. + */ + +#ifndef DATASET_DATASET_INDEX_H +#define DATASET_DATASET_INDEX_H + +#include +#include + +#include +#include + +namespace storage { +namespace dataset { + +static const uint16_t DATASET_INDEX_VERSION_MAJOR = 1; +static const uint16_t DATASET_INDEX_VERSION_MINOR = 0; +static const uint32_t DATASET_INDEX_HEADER_SIZE = 64; +static const uint32_t DATASET_INDEX_DIRECTORY_ENTRY_SIZE = 32; +static const uint32_t DATASET_INDEX_SECTION_COUNT = 13; +static const uint64_t DATASET_INDEX_ALIGNMENT = 64; + +enum class DatasetIndexStatus { + OK = 0, + INVALID_ARGUMENT, + IO_ERROR, + OUT_OF_MEMORY, + BAD_MAGIC, + UNSUPPORTED_VERSION, + BAD_HEADER, + BAD_DIRECTORY, + BAD_SECTION, + BAD_CHECKSUM, + BAD_REFERENCE, + NOT_FOUND, +}; + +const char* dataset_index_status_name(DatasetIndexStatus status); + +enum class DatasetIndexSectionType : uint32_t { + STRING_OFFSETS = 1, + STRING_BYTES = 2, + TABLE_NAME_INDEX = 3, + TABLE_RECORD = 4, + DEVICE_NAME_INDEX = 5, + DEVICE_RECORD = 6, + COLUMN_NAME_INDEX = 7, + COLUMN_SCHEMA = 8, + LOGICAL_SERIES = 9, + TSFILE_RECORD = 10, + DEVICE_FILE_SPAN = 11, + SERIES_FILE_SPAN = 12, + SERIES_LOCATOR = 13, +}; + +#pragma pack(push, 1) +struct DatasetIndexHeader { + char magic[8]; + uint16_t version_major; + uint16_t version_minor; + uint32_t header_size; + uint64_t directory_offset; + uint32_t section_count; + uint32_t directory_entry_size; + uint64_t file_length; + uint32_t header_crc32c; + uint8_t reserved[20]; +}; + +struct DatasetIndexDirectoryEntry { + uint32_t section_type; + uint32_t record_size; + uint64_t offset; + uint64_t length; + uint32_t count; + uint32_t crc32c; +}; + +struct TableNameIndexRecord { + uint64_t name_hash; + uint32_t name_sid; + uint32_t table_id; +}; + +struct TableRecord { + uint32_t name_sid; + uint32_t reserved0; + uint32_t first_device_name_index; + uint32_t device_count; + uint32_t first_column_name_index; + uint32_t column_count; + uint64_t reserved1; +}; + +struct DeviceNameIndexRecord { + uint32_t table_id; + uint32_t device_id; + uint64_t name_hash; + uint32_t name_sid; + uint32_t reserved; +}; + +struct DeviceRecord { + uint32_t table_id; + uint32_t name_sid; + uint32_t reserved0; + uint32_t reserved1; + uint32_t first_series_id; + uint32_t series_count; + uint32_t first_file_span; + uint32_t file_span_count; + int64_t min_time; + int64_t max_time; +}; + +struct ColumnNameIndexRecord { + uint32_t table_id; + uint32_t column_id; + uint64_t name_hash; + uint32_t name_sid; + uint32_t reserved; +}; + +struct ColumnSchemaRecord { + uint32_t table_id; + uint32_t name_sid; + uint32_t column_ordinal; + uint16_t logical_type; + uint16_t physical_type; + uint16_t encoding; + uint16_t compression; + uint16_t role; + uint16_t nullable; + uint64_t reserved; +}; + +struct LogicalSeriesRecord { + uint32_t device_id; + uint32_t column_id; + uint32_t first_file_span; + uint32_t file_span_count; + int64_t min_time; + int64_t max_time; +}; + +struct TsFileRecord { + uint32_t path_sid; + uint32_t reserved0; + uint64_t file_size; + uint64_t file_fingerprint; + uint64_t reserved1; +}; + +struct DeviceFileSpanRecord { + uint32_t device_id; + uint32_t file_id; + uint64_t time_meta_offset; + uint32_t time_meta_length; + uint16_t layout; + uint16_t flags; + uint64_t row_count; +}; + +struct SeriesFileSpanRecord { + uint32_t series_id; + uint32_t file_id; + uint32_t locator_id; + uint32_t reserved; + int64_t min_time; + int64_t max_time; + uint64_t row_count; +}; + +struct SeriesLocatorRecord { + uint32_t device_file_span_id; + uint16_t locator_kind; + uint16_t flags; + uint64_t timeseries_meta_offset; + uint32_t timeseries_meta_length; + uint32_t padding; +}; +#pragma pack(pop) + +static_assert(sizeof(DatasetIndexHeader) == 64, + "DatasetIndexHeader must be 64 bytes"); +static_assert(sizeof(DatasetIndexDirectoryEntry) == 32, + "DatasetIndexDirectoryEntry must be 32 bytes"); +static_assert(sizeof(TableNameIndexRecord) == 16, + "TableNameIndexRecord must be 16 bytes"); +static_assert(sizeof(TableRecord) == 32, "TableRecord must be 32 bytes"); +static_assert(sizeof(DeviceNameIndexRecord) == 24, + "DeviceNameIndexRecord must be 24 bytes"); +static_assert(sizeof(DeviceRecord) == 48, "DeviceRecord must be 48 bytes"); +static_assert(sizeof(ColumnNameIndexRecord) == 24, + "ColumnNameIndexRecord must be 24 bytes"); +static_assert(sizeof(ColumnSchemaRecord) == 32, + "ColumnSchemaRecord must be 32 bytes"); +static_assert(sizeof(LogicalSeriesRecord) == 32, + "LogicalSeriesRecord must be 32 bytes"); +static_assert(sizeof(TsFileRecord) == 32, "TsFileRecord must be 32 bytes"); +static_assert(sizeof(DeviceFileSpanRecord) == 32, + "DeviceFileSpanRecord must be 32 bytes"); +static_assert(sizeof(SeriesFileSpanRecord) == 40, + "SeriesFileSpanRecord must be 40 bytes"); +static_assert(sizeof(SeriesLocatorRecord) == 24, + "SeriesLocatorRecord must be 24 bytes"); + +struct DatasetIndexSectionData { + DatasetIndexSectionType type; + uint32_t record_size; + uint32_t count; + std::vector bytes; +}; + +struct DatasetIndexSectionView { + const uint8_t* data; + uint64_t length; + uint32_t record_size; + uint32_t count; + + DatasetIndexSectionView() + : data(nullptr), length(0), record_size(0), count(0) {} +}; + +struct DatasetIndexStringView { + const char* data; + uint32_t length; + + DatasetIndexStringView() : data(nullptr), length(0) {} + std::string to_string() const { + return data == nullptr ? std::string() : std::string(data, length); + } +}; + +uint32_t dataset_index_crc32c(const void* data, size_t length); +uint64_t dataset_index_name_hash(const char* data, size_t length); + +class DatasetIndexWriter { + public: + static DatasetIndexStatus write_atomic( + const std::string& output_path, + const std::vector& sections, + std::string& error_message); +}; + +class MappedDatasetIndex { + public: + MappedDatasetIndex(); + ~MappedDatasetIndex(); + + DatasetIndexStatus open(const std::string& path); + void close(); + bool is_open() const { return mapping_ != nullptr; } + const std::string& error_message() const { return error_message_; } + const std::string& path() const { return path_; } + uint64_t file_length() const { return mapping_size_; } + const DatasetIndexHeader* header() const { return header_; } + + DatasetIndexStatus section(DatasetIndexSectionType type, + DatasetIndexSectionView& result) const; + DatasetIndexStatus string(uint32_t sid, + DatasetIndexStringView& result) const; + + template + DatasetIndexStatus record(DatasetIndexSectionType type, uint32_t id, + const T*& result) const { + DatasetIndexSectionView view; + DatasetIndexStatus status = section(type, view); + if (status != DatasetIndexStatus::OK) { + result = nullptr; + return status; + } + if (view.record_size != sizeof(T) || id >= view.count) { + result = nullptr; + return DatasetIndexStatus::NOT_FOUND; + } + result = reinterpret_cast( + view.data + static_cast(id) * view.record_size); + return DatasetIndexStatus::OK; + } + + DatasetIndexStatus find_table_ids(const std::string& name, + std::vector& table_ids) const; + DatasetIndexStatus find_device_id(uint32_t table_id, + const std::string& name, + uint32_t& device_id) const; + DatasetIndexStatus find_column_id(uint32_t table_id, + const std::string& name, + uint32_t& column_id) const; + DatasetIndexStatus find_series_id(uint32_t device_id, uint32_t column_id, + uint32_t& series_id) const; + + private: + DatasetIndexStatus validate(); + DatasetIndexStatus fail(DatasetIndexStatus status, + const std::string& message); + + std::string path_; + std::string error_message_; + const uint8_t* mapping_; + uint64_t mapping_size_; + const DatasetIndexHeader* header_; + const DatasetIndexDirectoryEntry* directory_; +#ifdef _WIN32 + void* file_handle_; + void* mapping_handle_; +#else + int fd_; +#endif +}; + +} // namespace dataset +} // namespace storage + +#endif // DATASET_DATASET_INDEX_H diff --git a/cpp/src/file/read_file.cc b/cpp/src/file/read_file.cc index 7d41d7095..ce1f67197 100644 --- a/cpp/src/file/read_file.cc +++ b/cpp/src/file/read_file.cc @@ -39,6 +39,55 @@ ssize_t pread(int fd, void* buf, size_t count, uint64_t offset); using namespace common; namespace storage { +namespace { +uint64_t generation_hash(uint64_t size, int64_t mtime_ns) { + uint64_t hash = 1469598103934665603ULL; + const uint64_t values[2] = {size, static_cast(mtime_ns)}; + for (size_t word = 0; word < 2; ++word) { + for (size_t byte = 0; byte < 8; ++byte) { + hash ^= static_cast(values[word] >> (byte * 8)); + hash *= 1099511628211ULL; + } + } + return hash; +} +} // namespace + +int ReadFile::generation(uint64_t& size, uint64_t& fingerprint) const { + if (fd_ < 0) { + return E_FILE_READ_ERR; + } + int64_t mtime_ns = 0; +#ifdef _WIN32 + intptr_t handle_value = _get_osfhandle(fd_); + if (handle_value == -1) { + return E_FILE_READ_ERR; + } + FILE_BASIC_INFO info; + if (!GetFileInformationByHandleEx(reinterpret_cast(handle_value), + FileBasicInfo, &info, sizeof(info))) { + return E_FILE_READ_ERR; + } + static const int64_t WINDOWS_TO_UNIX_100NS = 116444736000000000LL; + mtime_ns = (info.LastWriteTime.QuadPart - WINDOWS_TO_UNIX_100NS) * 100; +#else + struct stat info; + if (::fstat(fd_, &info) != 0) { + return E_FILE_READ_ERR; + } +#ifdef __APPLE__ + mtime_ns = static_cast(info.st_mtimespec.tv_sec) * 1000000000LL + + info.st_mtimespec.tv_nsec; +#else + mtime_ns = static_cast(info.st_mtim.tv_sec) * 1000000000LL + + info.st_mtim.tv_nsec; +#endif +#endif + size = static_cast(file_size_); + fingerprint = generation_hash(size, mtime_ns); + return E_OK; +} + void ReadFile::close() { if (fd_ >= 0) { ::close(fd_); diff --git a/cpp/src/file/read_file.h b/cpp/src/file/read_file.h index c38940098..fbfe50a8b 100644 --- a/cpp/src/file/read_file.h +++ b/cpp/src/file/read_file.h @@ -40,6 +40,10 @@ class ReadFile { FORCE_INLINE int64_t file_size() const { return file_size_; } FORCE_INLINE const std::string& file_path() const { return file_path_; } + /** Return size and the Dataset Index v1 FNV fingerprint of size+mtime_ns. + */ + int generation(uint64_t& size, uint64_t& fingerprint) const; + /* * try to reader @buf_size bytes from @offset of this file * into @buf. @read_len return the actual len reader. diff --git a/cpp/src/file/tsfile_io_reader.cc b/cpp/src/file/tsfile_io_reader.cc index e028cf902..41bcd6b93 100644 --- a/cpp/src/file/tsfile_io_reader.cc +++ b/cpp/src/file/tsfile_io_reader.cc @@ -19,7 +19,10 @@ #include "file/tsfile_io_reader.h" +#include + #include "common/allocator/alloc_base.h" +#include "reader/prepared_series.h" using namespace common; @@ -88,6 +91,196 @@ int TsFileIOReader::alloc_ssi(std::shared_ptr device_id, return ret; } +namespace { +int load_exact_timeseries_index(ReadFile* read_file, uint64_t offset, + uint32_t length, PageArena& arena, + TimeseriesIndex*& index) { + if (read_file == nullptr || length == 0 || + length > static_cast(std::numeric_limits::max()) || + offset > static_cast(std::numeric_limits::max()) || + offset > static_cast(read_file->file_size()) || + length > static_cast(read_file->file_size()) - offset) { + return E_TSFILE_CORRUPTED; + } + char* bytes = static_cast(arena.alloc(length)); + void* index_memory = arena.alloc(sizeof(TimeseriesIndex)); + if (bytes == nullptr || index_memory == nullptr) { + return E_OOM; + } + int32_t read_length = 0; + int ret = read_file->read(static_cast(offset), bytes, + static_cast(length), read_length); + if (ret != E_OK || read_length != static_cast(length)) { + return ret == E_OK ? E_TSFILE_CORRUPTED : ret; + } + ByteStream stream; + stream.wrap_from(bytes, length); + index = new (index_memory) TimeseriesIndex; + ret = index->deserialize_from(stream, &arena); + if (ret != E_OK || stream.read_pos() != length) { + index = nullptr; + return ret == E_OK ? E_TSFILE_CORRUPTED : ret; + } + index->set_metadata_range(static_cast(offset), length); + return E_OK; +} +} // namespace + +int TsFileIOReader::prepare_series(const FileGeneration& generation, + const PreparedLocator& locator, + std::shared_ptr& prepared) { + return prepare_series(generation, locator, nullptr, prepared); +} + +int TsFileIOReader::prepare_series( + const FileGeneration& generation, const PreparedLocator& locator, + const std::shared_ptr& aligned_time_owner, + std::shared_ptr& prepared) { + prepared.reset(); + uint64_t actual_size = 0; + uint64_t actual_fingerprint = 0; + if (read_file_ == nullptr || generation.file_size == 0 || + read_file_->generation(actual_size, actual_fingerprint) != E_OK || + actual_size != generation.file_size || + (generation.file_fingerprint != 0 && + actual_fingerprint != generation.file_fingerprint) || + locator.value_metadata_length == 0 || locator.layout > 1) { + return E_INVALID_ARG; + } + std::shared_ptr candidate = + std::make_shared(generation, locator); + TimeseriesIndex* value_index = nullptr; + int ret = load_exact_timeseries_index( + read_file_, locator.value_metadata_offset, + locator.value_metadata_length, candidate->arena(), value_index); + if (ret != E_OK) { + return ret; + } + if (locator.layout == 0) { + candidate->set_index(value_index); + } else { + if (locator.time_metadata_length == 0) { + return E_NOT_SUPPORT; + } + TimeseriesIndex* time_index = nullptr; + if (aligned_time_owner != nullptr) { + const FileGeneration& owner_generation = + aligned_time_owner->generation(); + const PreparedLocator& owner_locator = + aligned_time_owner->locator(); + auto* owner_index = dynamic_cast( + aligned_time_owner->index()); + if (owner_generation.mapped_index_identity != + generation.mapped_index_identity || + owner_generation.file_id != generation.file_id || + owner_generation.file_size != generation.file_size || + owner_generation.file_fingerprint != + generation.file_fingerprint || + owner_locator.layout != 1 || + owner_locator.time_metadata_offset != + locator.time_metadata_offset || + owner_locator.time_metadata_length != + locator.time_metadata_length || + owner_index == nullptr || + owner_index->time_ts_idx_ == nullptr) { + return E_INVALID_ARG; + } + time_index = owner_index->time_ts_idx_; + candidate->set_aligned_time_owner(aligned_time_owner); + } else { + ret = load_exact_timeseries_index( + read_file_, locator.time_metadata_offset, + locator.time_metadata_length, candidate->arena(), time_index); + if (ret != E_OK) { + return ret; + } + } + if (time_index->get_chunk_meta_list()->size() != + value_index->get_chunk_meta_list()->size()) { + return E_NOT_SUPPORT; + } + void* aligned_memory = + candidate->arena().alloc(sizeof(AlignedTimeseriesIndex)); + if (aligned_memory == nullptr) { + return E_OOM; + } + AlignedTimeseriesIndex* aligned = + new (aligned_memory) AlignedTimeseriesIndex; + aligned->time_ts_idx_ = time_index; + aligned->value_ts_idx_ = value_index; + candidate->set_index(aligned); + } + prepared = candidate; + return E_OK; +} + +int TsFileIOReader::alloc_prepared_ssi( + const std::shared_ptr& prepared, + TsFileSeriesScanIterator*& ssi, common::PageArena& pa, + Filter* time_filter) { + ssi = nullptr; + if (prepared == nullptr || prepared->index() == nullptr) { + return E_INVALID_ARG; + } + if (time_filter != nullptr && + !filter_stasify(prepared->index(), time_filter)) { + return E_NO_MORE_DATA; + } + void* memory = + mem_alloc(sizeof(TsFileSeriesScanIterator), MOD_TSFILE_READER); + if (memory == nullptr) { + return E_OOM; + } + ssi = new (memory) TsFileSeriesScanIterator; + int ret = ssi->init_prepared(prepared, read_file_, time_filter, pa); + if (ret == E_OK) { + ret = ssi->init_chunk_reader(); + } + if (ret != E_OK) { + ssi->destroy(); + mem_free(ssi); + ssi = nullptr; + } + return ret; +} + +int TsFileIOReader::alloc_prepared_multi_ssi( + const std::vector>& prepared, + TsFileSeriesScanIterator*& ssi, common::PageArena& pa, + Filter* time_filter) { + ssi = nullptr; + if (prepared.empty() || prepared.front() == nullptr || + prepared.front()->index() == nullptr) { + return E_INVALID_ARG; + } + auto* first_aligned = + dynamic_cast(prepared.front()->index()); + if (first_aligned == nullptr || first_aligned->time_ts_idx_ == nullptr) { + return E_NOT_SUPPORT; + } + if (time_filter != nullptr && + !filter_stasify(first_aligned->time_ts_idx_, time_filter)) { + return E_NO_MORE_DATA; + } + + void* memory = + mem_alloc(sizeof(TsFileSeriesScanIterator), MOD_TSFILE_READER); + if (memory == nullptr) { + return E_OOM; + } + ssi = new (memory) TsFileSeriesScanIterator; + int ret = ssi->init_prepared_multi(prepared, read_file_, time_filter, pa); + if (ret == E_OK) { + ret = ssi->init_chunk_reader(); + } + if (ret != E_OK) { + ssi->destroy(); + mem_free(ssi); + ssi = nullptr; + } + return ret; +} + int TsFileIOReader::alloc_multi_ssi( std::shared_ptr device_id, const std::vector& measurement_names, @@ -900,7 +1093,14 @@ int TsFileIOReader::get_time_column_metadata( return E_OOM; } ret_timeseries_index = new (buf) TimeseriesIndex; - ret_timeseries_index->deserialize_from(buffer, &pa); + if (RET_FAIL(ret_timeseries_index->deserialize_from(buffer, &pa))) { + return ret; + } + if (buffer.read_pos() > UINT32_MAX) { + return E_OVERFLOW; + } + ret_timeseries_index->set_metadata_range( + start_idx, static_cast(buffer.read_pos())); } else if (measurement_node->node_type_ == INTERNAL_MEASUREMENT) { start_idx = measurement_node->children_[0]->get_offset(); end_idx = measurement_node->children_[1]->get_offset(); @@ -947,8 +1147,18 @@ int TsFileIOReader::do_load_timeseries_index( TimeseriesIndex cur_timeseries_index; PageArena cur_timeseries_index_pa; cur_timeseries_index_pa.init(512, MOD_TSFILE_READER); // TODO 512 + const uint64_t relative_start = bs.read_pos(); if (RET_FAIL(cur_timeseries_index.deserialize_from( bs, &cur_timeseries_index_pa))) { + } else if (bs.read_pos() < relative_start || + bs.read_pos() - relative_start > UINT32_MAX) { + ret = E_OVERFLOW; + } else { + cur_timeseries_index.set_metadata_range( + start_offset + relative_start, + static_cast(bs.read_pos() - relative_start)); + } + if (RET_FAIL(ret)) { } else if (is_aligned && cur_timeseries_index.get_measurement_name().equal_to( target_measurement_name)) { @@ -1012,12 +1222,20 @@ int TsFileIOReader::do_load_all_timeseries_index( ByteStream bs; bs.wrap_from(ti_buf, read_size); while (bs.has_remaining()) { + const uint64_t relative_start = bs.read_pos(); void* buf = in_timeseries_index_pa.alloc(sizeof(TimeseriesIndex)); auto ts_idx = new (buf) TimeseriesIndex; if (RET_FAIL( ts_idx->deserialize_from(bs, &in_timeseries_index_pa))) { return ret; } + if (bs.read_pos() < relative_start || + bs.read_pos() - relative_start > UINT32_MAX) { + return E_OVERFLOW; + } + ts_idx->set_metadata_range( + start_offset + relative_start, + static_cast(bs.read_pos() - relative_start)); if (ts_idx->get_measurement_name().len_ == 0) continue; ts_indexs.push_back(ts_idx); } diff --git a/cpp/src/file/tsfile_io_reader.h b/cpp/src/file/tsfile_io_reader.h index dc9069b5c..9da3b52e9 100644 --- a/cpp/src/file/tsfile_io_reader.h +++ b/cpp/src/file/tsfile_io_reader.h @@ -20,6 +20,7 @@ #ifndef FILE_TSFILE_IO_REAER_H #define FILE_TSFILE_IO_REAER_H +#include #include #include #include @@ -33,6 +34,9 @@ #include "utils/storage_utils.h" namespace storage { class TsFileSeriesScanIterator; +class PreparedSeries; +struct FileGeneration; +struct PreparedLocator; /* * TODO: @@ -86,6 +90,24 @@ class TsFileIOReader { TsFileSeriesScanIterator*& ssi, common::PageArena& pa, Filter* time_filter = nullptr); + int prepare_series(const FileGeneration& generation, + const PreparedLocator& locator, + std::shared_ptr& prepared); + int prepare_series( + const FileGeneration& generation, const PreparedLocator& locator, + const std::shared_ptr& aligned_time_owner, + std::shared_ptr& prepared); + + int alloc_prepared_ssi(const std::shared_ptr& prepared, + TsFileSeriesScanIterator*& ssi, + common::PageArena& pa, + Filter* time_filter = nullptr); + + int alloc_prepared_multi_ssi( + const std::vector>& prepared, + TsFileSeriesScanIterator*& ssi, common::PageArena& pa, + Filter* time_filter = nullptr); + void revert_ssi(TsFileSeriesScanIterator* ssi); std::string get_file_path() const { return read_file_->file_path(); } diff --git a/cpp/src/reader/block/prepared_series_tsblock_reader.cc b/cpp/src/reader/block/prepared_series_tsblock_reader.cc new file mode 100644 index 000000000..61e5123a9 --- /dev/null +++ b/cpp/src/reader/block/prepared_series_tsblock_reader.cc @@ -0,0 +1,219 @@ +/* + * 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. + */ + +#include "reader/block/prepared_series_tsblock_reader.h" + +#include + +#include "file/tsfile_io_reader.h" +#include "reader/filter/filter.h" +#include "reader/prepared_series.h" +#include "reader/tsfile_series_scan_iterator.h" +#include "utils/errno_define.h" + +namespace storage { + +namespace { +constexpr uint32_t PREPARED_BATCH_ROWS = 65536; + +common::TSDataType prepared_value_data_type(ITimeseriesIndex* index) { + if (index == nullptr) { + return common::INVALID_DATATYPE; + } + common::TSDataType data_type = index->get_data_type(); + if (data_type == common::VECTOR) { + auto* aligned = dynamic_cast(index); + if (aligned == nullptr || aligned->value_ts_idx_ == nullptr) { + return common::INVALID_DATATYPE; + } + data_type = aligned->value_ts_idx_->get_data_type(); + } + return data_type; +} +} // namespace + +int PreparedSeriesTsBlockReader::init( + TsFileIOReader* io_reader, const std::shared_ptr& prepared, + Filter* owned_time_filter, int offset, int limit) { + owned_time_filter_ = owned_time_filter; + if (io_reader == nullptr || prepared == nullptr || + prepared->index() == nullptr || offset < 0 || closed_) { + return common::E_INVALID_ARG; + } + + io_reader_ = io_reader; + prepared_ = prepared; + remaining_limit_ = limit; + value_data_type_ = prepared_value_data_type(prepared->index()); + if (value_data_type_ == common::INVALID_DATATYPE || + value_data_type_ == common::VECTOR) { + return common::E_TYPE_NOT_SUPPORTED; + } + value_data_types_.assign(1, value_data_type_); + + pa_.init(512, common::MOD_TSFILE_READER); + if (limit == 0) { + exhausted_ = true; + return common::E_OK; + } + + int ret = io_reader_->alloc_prepared_ssi(prepared_, ssi_, pa_, + owned_time_filter_); + if (ret == common::E_NO_MORE_DATA) { + exhausted_ = true; + return common::E_OK; + } + if (ret != common::E_OK) { + return ret; + } + // The table-model multi-aligned reader applies offset through its + // chunk/page plan. Limit is enforced by sizing each native output block; + // this keeps the returned TsBlock direct and avoids RowRecord slicing. + ssi_->set_row_range(offset, -1); + return common::E_OK; +} + +int PreparedSeriesTsBlockReader::init_multi( + TsFileIOReader* io_reader, + const std::vector>& prepared, + Filter* owned_time_filter, int offset, int limit) { + owned_time_filter_ = owned_time_filter; + if (io_reader == nullptr || prepared.empty() || offset < 0 || closed_) { + return common::E_INVALID_ARG; + } + + io_reader_ = io_reader; + remaining_limit_ = limit; + value_data_types_.reserve(prepared.size()); + for (const auto& entry : prepared) { + common::TSDataType data_type = + entry == nullptr ? common::INVALID_DATATYPE + : prepared_value_data_type(entry->index()); + if (data_type == common::INVALID_DATATYPE || + data_type == common::VECTOR) { + return common::E_TYPE_NOT_SUPPORTED; + } + value_data_types_.push_back(data_type); + } + value_data_type_ = value_data_types_.front(); + + pa_.init(512, common::MOD_TSFILE_READER); + if (limit == 0) { + exhausted_ = true; + return common::E_OK; + } + + int ret = io_reader_->alloc_prepared_multi_ssi(prepared, ssi_, pa_, + owned_time_filter_); + if (ret == common::E_NO_MORE_DATA) { + exhausted_ = true; + return common::E_OK; + } + if (ret != common::E_OK) { + return ret; + } + ssi_->set_row_range(offset, -1); + return common::E_OK; +} + +int PreparedSeriesTsBlockReader::has_next(bool& has_next) { + has_next = false; + if (closed_) { + return common::E_INVALID_ARG; + } + if (block_ready_) { + has_next = true; + return common::E_OK; + } + if (exhausted_ || ssi_ == nullptr || remaining_limit_ == 0) { + exhausted_ = true; + return common::E_OK; + } + + while (true) { + const uint32_t desired_rows = + remaining_limit_ > 0 + ? std::min(PREPARED_BATCH_ROWS, + static_cast(remaining_limit_)) + : PREPARED_BATCH_ROWS; + if (block_ != nullptr) { + if (block_->get_max_row_count() != desired_rows) { + ssi_->revert_tsblock(); + block_ = nullptr; + } else { + block_->reset(); + } + } + ssi_->set_max_block_rows(desired_rows); + int ret = ssi_->get_next(block_, true); + if (ret == common::E_NO_MORE_DATA) { + exhausted_ = true; + return common::E_OK; + } + if (ret != common::E_OK) { + return ret; + } + if (block_ != nullptr && block_->get_row_count() > 0) { + if (remaining_limit_ > 0) { + remaining_limit_ -= static_cast(block_->get_row_count()); + } + block_ready_ = true; + has_next = true; + return common::E_OK; + } + } +} + +int PreparedSeriesTsBlockReader::next(common::TsBlock*& ret_block) { + ret_block = nullptr; + bool available = false; + int ret = has_next(available); + if (ret != common::E_OK) { + return ret; + } + if (!available) { + return common::E_NO_MORE_DATA; + } + ret_block = block_; + block_ready_ = false; + return common::E_OK; +} + +void PreparedSeriesTsBlockReader::close() { + if (closed_) { + return; + } + closed_ = true; + block_ready_ = false; + exhausted_ = true; + if (ssi_ != nullptr) { + ssi_->revert_tsblock(); + block_ = nullptr; + io_reader_->revert_ssi(ssi_); + ssi_ = nullptr; + } + delete owned_time_filter_; + owned_time_filter_ = nullptr; + pa_.destroy(); + prepared_.reset(); + std::vector().swap(value_data_types_); + io_reader_ = nullptr; +} + +} // namespace storage diff --git a/cpp/src/reader/block/prepared_series_tsblock_reader.h b/cpp/src/reader/block/prepared_series_tsblock_reader.h new file mode 100644 index 000000000..be416d47f --- /dev/null +++ b/cpp/src/reader/block/prepared_series_tsblock_reader.h @@ -0,0 +1,75 @@ +/* + * 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. + */ +#ifndef READER_BLOCK_PREPARED_SERIES_TSBLOCK_READER_H +#define READER_BLOCK_PREPARED_SERIES_TSBLOCK_READER_H + +#include + +#include "common/allocator/page_arena.h" +#include "reader/block/tsblock_reader.h" + +namespace storage { + +class Filter; +class PreparedSeries; +class TsFileIOReader; +class TsFileSeriesScanIterator; + +// Adapts one locator-backed PreparedSeries to the table-model batch contract. +// The underlying SSI already decodes directly into a two-column TsBlock +// (time, value), so this reader must not materialize RowRecord objects. +class PreparedSeriesTsBlockReader final : public TsBlockReader { + public: + PreparedSeriesTsBlockReader() = default; + ~PreparedSeriesTsBlockReader() override { close(); } + + int init(TsFileIOReader* io_reader, + const std::shared_ptr& prepared, + Filter* owned_time_filter, int offset, int limit); + int init_multi(TsFileIOReader* io_reader, + const std::vector>& prepared, + Filter* owned_time_filter, int offset, int limit); + + int has_next(bool& has_next) override; + int next(common::TsBlock*& ret_block) override; + void close() override; + + common::TSDataType value_data_type() const { return value_data_type_; } + const std::vector& value_data_types() const { + return value_data_types_; + } + + private: + TsFileIOReader* io_reader_ = nullptr; + std::shared_ptr prepared_; + Filter* owned_time_filter_ = nullptr; + TsFileSeriesScanIterator* ssi_ = nullptr; + common::PageArena pa_; + common::TsBlock* block_ = nullptr; + common::TSDataType value_data_type_ = common::INVALID_DATATYPE; + std::vector value_data_types_; + int remaining_limit_ = -1; + bool block_ready_ = false; + bool exhausted_ = false; + bool closed_ = false; +}; + +} // namespace storage + +#endif // READER_BLOCK_PREPARED_SERIES_TSBLOCK_READER_H diff --git a/cpp/src/reader/prepared_series.cc b/cpp/src/reader/prepared_series.cc new file mode 100644 index 000000000..473bfc014 --- /dev/null +++ b/cpp/src/reader/prepared_series.cc @@ -0,0 +1,35 @@ +/* + * 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. + */ + +#include "reader/prepared_series.h" + +namespace storage { + +PreparedSeries::PreparedSeries(const FileGeneration& generation, + const PreparedLocator& locator) + : generation_(generation), locator_(locator), arena_(), index_(nullptr) { + arena_.init(512, common::MOD_TSFILE_READER); +} + +PreparedSeries::~PreparedSeries() { + index_ = nullptr; + arena_.destroy(); +} + +} // namespace storage diff --git a/cpp/src/reader/prepared_series.h b/cpp/src/reader/prepared_series.h new file mode 100644 index 000000000..ce647966b --- /dev/null +++ b/cpp/src/reader/prepared_series.h @@ -0,0 +1,92 @@ +/* + * 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. + */ + +#ifndef READER_PREPARED_SERIES_H +#define READER_PREPARED_SERIES_H + +#include + +#include + +#include "common/allocator/page_arena.h" +#include "common/tsfile_common.h" + +namespace storage { + +struct FileGeneration { + uint64_t mapped_index_identity; + uint32_t file_id; + uint64_t file_size; + uint64_t file_fingerprint; + + FileGeneration() + : mapped_index_identity(0), + file_id(0), + file_size(0), + file_fingerprint(0) {} +}; + +struct PreparedLocator { + uint32_t locator_id; + uint16_t layout; + uint16_t flags; + uint64_t value_metadata_offset; + uint32_t value_metadata_length; + uint64_t time_metadata_offset; + uint32_t time_metadata_length; + + PreparedLocator() + : locator_id(0), + layout(0), + flags(0), + value_metadata_offset(0), + value_metadata_length(0), + time_metadata_offset(0), + time_metadata_length(0) {} +}; + +class PreparedSeries { + public: + PreparedSeries(const FileGeneration& generation, + const PreparedLocator& locator); + ~PreparedSeries(); + + const FileGeneration& generation() const { return generation_; } + const PreparedLocator& locator() const { return locator_; } + ITimeseriesIndex* index() const { return index_; } + common::PageArena& arena() { return arena_; } + void set_index(ITimeseriesIndex* index) { index_ = index; } + void set_aligned_time_owner( + const std::shared_ptr& aligned_time_owner) { + aligned_time_owner_ = aligned_time_owner; + } + + private: + FileGeneration generation_; + PreparedLocator locator_; + common::PageArena arena_; + ITimeseriesIndex* index_; + // Optional owner of the aligned time index referenced by index_. This + // lets several value PreparedSeries share one parsed time metadata arena. + std::shared_ptr aligned_time_owner_; +}; + +} // namespace storage + +#endif // READER_PREPARED_SERIES_H diff --git a/cpp/src/reader/qds_without_timegenerator.cc b/cpp/src/reader/qds_without_timegenerator.cc index 7f520f339..fcd38762d 100644 --- a/cpp/src/reader/qds_without_timegenerator.cc +++ b/cpp/src/reader/qds_without_timegenerator.cc @@ -34,6 +34,53 @@ int QDSWithoutTimeGenerator::init(TsFileIOReader* io_reader, return init_internal(io_reader, qe); } +int QDSWithoutTimeGenerator::init_prepared( + TsFileIOReader* io_reader, const std::shared_ptr& prepared, + Filter* owned_time_filter, int offset, int limit, + const std::string& column_name) { + pa_.reset(); + pa_.init(512, common::MOD_TSFILE_READER); + io_reader_ = io_reader; + qe_ = nullptr; + owned_time_filter_ = owned_time_filter; + remaining_offset_ = offset; + remaining_limit_ = limit; + is_single_path_ = true; + index_lookup_.insert({"time", 0}); + + TsFileSeriesScanIterator* ssi = nullptr; + int ret = + io_reader_->alloc_prepared_ssi(prepared, ssi, pa_, owned_time_filter_); + if (ret == E_NO_MORE_DATA) { + // Preserve the normal empty-result contract even when the global + // statistic rejects the range before an SSI is allocated. + row_record_ = new RowRecord(1); + result_set_metadata_ = std::make_shared( + std::vector(), std::vector()); + return E_OK; + } + if (ret != E_OK) { + return ret; + } + const bool table_aligned = ssi->is_multi_value(); + ssi->set_row_range(offset, table_aligned ? -1 : limit); + ssi_vec_.push_back(ssi); + tsblocks_.resize(1); + time_iters_.resize(1); + value_iters_.resize(1); + row_record_ = new RowRecord(2); + index_lookup_.insert({column_name, 1}); + load_next_tsblock(0, true); + remaining_offset_ = ssi->get_row_offset(); + if (!table_aligned) { + remaining_limit_ = ssi->get_row_limit(); + } + result_set_metadata_ = std::make_shared( + std::vector(1, column_name), + std::vector(1, ssi->get_data_type())); + return E_OK; +} + int QDSWithoutTimeGenerator::init(TsFileIOReader* io_reader, QueryExpression* qe, int offset, int limit) { remaining_offset_ = offset; @@ -56,7 +103,7 @@ int QDSWithoutTimeGenerator::init_internal(TsFileIOReader* io_reader, std::vector column_names; std::vector data_types; // Data type per valid path, captured from the timeseries index right after - // alloc_ssi — while itimeseries_index_ is still live. get_next_tsblock() + // alloc_ssi — while itimeseries_index_ is still live. load_next_tsblock() // may later destroy() the SSI (e.g. when limit==0 yields no TsBlock), which // clears the index, so this must be recorded up front. std::vector ssi_data_types; @@ -114,9 +161,9 @@ int QDSWithoutTimeGenerator::init_internal(TsFileIOReader* io_reader, value_iters_.resize(path_count); for (size_t i = 0; i < path_count; i++) { - get_next_tsblock(i, true); + load_next_tsblock(i, true); // Prefer the type carried by the value iterator, but fall back to the - // timeseries-index type captured before get_next_tsblock() when no + // timeseries-index type captured before load_next_tsblock() when no // TsBlock was produced (e.g. limit==0 skips every row, or the series is // empty). Emitting NULL_TYPE here would surface an invalid datatype // (254) to callers that map the metadata onto their own type enums. @@ -126,7 +173,7 @@ int QDSWithoutTimeGenerator::init_internal(TsFileIOReader* io_reader, data_types.push_back(col_type); } // Single-path: SSI may have consumed offset/limit by skipping chunks/pages - // during first get_next_tsblock(); sync so QDS does not double-apply. + // during first load_next_tsblock(); sync so QDS does not double-apply. if (is_single_path_) { remaining_offset_ = ssi_vec_[0]->get_row_offset(); remaining_limit_ = ssi_vec_[0]->get_row_limit(); @@ -167,6 +214,10 @@ void QDSWithoutTimeGenerator::close() { delete qe_; qe_ = nullptr; } + if (owned_time_filter_ != nullptr) { + delete owned_time_filter_; + owned_time_filter_ = nullptr; + } pa_.destroy(); } @@ -213,7 +264,7 @@ int QDSWithoutTimeGenerator::next(bool& has_next) { heap_time_.insert(std::make_pair(timev, idx)); time_iters_[idx]->next(); } else { - get_next_tsblock(idx, false); + load_next_tsblock(idx, false); } if (skip_row) { @@ -268,7 +319,7 @@ int QDSWithoutTimeGenerator::next(bool& has_next) { // Pass merge_cursor (current time) as min_time_hint // to help SSI skip chunks/pages that are entirely before // the current merge position. - get_next_tsblock_with_hint(iter->second, false, time); + load_next_tsblock_with_hint(iter->second, false, time); } std::multimap::iterator cur = iter; iter++; // cppcheck-suppress postfixOperator @@ -310,7 +361,7 @@ std::shared_ptr QDSWithoutTimeGenerator::get_metadata() { return result_set_metadata_; } -int QDSWithoutTimeGenerator::get_next_tsblock(uint32_t index, bool alloc_mem) { +int QDSWithoutTimeGenerator::load_next_tsblock(uint32_t index, bool alloc_mem) { if (tsblocks_[index] != nullptr) { delete time_iters_[index]; time_iters_[index] = nullptr; @@ -345,9 +396,8 @@ int QDSWithoutTimeGenerator::get_next_tsblock(uint32_t index, bool alloc_mem) { return ret; } -int QDSWithoutTimeGenerator::get_next_tsblock_with_hint(uint32_t index, - bool alloc_mem, - int64_t min_time_hint) { +int QDSWithoutTimeGenerator::load_next_tsblock_with_hint( + uint32_t index, bool alloc_mem, int64_t min_time_hint) { if (tsblocks_[index] != nullptr) { delete time_iters_[index]; time_iters_[index] = nullptr; diff --git a/cpp/src/reader/qds_without_timegenerator.h b/cpp/src/reader/qds_without_timegenerator.h index 1d929e575..eae2a6495 100644 --- a/cpp/src/reader/qds_without_timegenerator.h +++ b/cpp/src/reader/qds_without_timegenerator.h @@ -21,6 +21,7 @@ #define READER_QDS_WITHOUT_TIMEGENERATOR_H #include +#include #include #include "expression.h" @@ -42,11 +43,16 @@ class QDSWithoutTimeGenerator : public ResultSet { heap_time_(), remaining_offset_(0), remaining_limit_(-1), + owned_time_filter_(nullptr), is_single_path_(false) {} ~QDSWithoutTimeGenerator() { close(); } int init(TsFileIOReader* io_reader, QueryExpression* qe); int init(TsFileIOReader* io_reader, QueryExpression* qe, int offset, int limit); + int init_prepared(TsFileIOReader* io_reader, + const std::shared_ptr& prepared, + Filter* owned_time_filter, int offset, int limit, + const std::string& column_name); void close(); int next(bool& has_next); bool is_null(const std::string& column_name); @@ -56,9 +62,9 @@ class QDSWithoutTimeGenerator : public ResultSet { private: int init_internal(TsFileIOReader* io_reader, QueryExpression* qe); - int get_next_tsblock(uint32_t index, bool alloc_mem); - int get_next_tsblock_with_hint(uint32_t index, bool alloc_mem, - int64_t min_time_hint); + int load_next_tsblock(uint32_t index, bool alloc_mem); + int load_next_tsblock_with_hint(uint32_t index, bool alloc_mem, + int64_t min_time_hint); private: std::shared_ptr result_set_metadata_; @@ -72,6 +78,7 @@ class QDSWithoutTimeGenerator : public ResultSet { heap_time_; // key-->time, value-->path_index int remaining_offset_; int remaining_limit_; + Filter* owned_time_filter_; bool is_single_path_; }; diff --git a/cpp/src/reader/tsfile_executor.cc b/cpp/src/reader/tsfile_executor.cc index 15aaa161b..bdae6c863 100644 --- a/cpp/src/reader/tsfile_executor.cc +++ b/cpp/src/reader/tsfile_executor.cc @@ -22,6 +22,10 @@ #include "expression.h" #include "qds_with_timegenerator.h" #include "qds_without_timegenerator.h" +#include "reader/block/prepared_series_tsblock_reader.h" +#include "reader/filter/time_operator.h" +#include "reader/prepared_series.h" +#include "reader/table_result_set.h" using namespace common; @@ -102,6 +106,80 @@ int TsFileExecutor::execute(QueryExpression* query_expr, ResultSet*& ret_qds, return ret; } +int TsFileExecutor::prepare_series(const FileGeneration& generation, + const PreparedLocator& locator, + std::shared_ptr& prepared) { + ASSERT(is_inited_); + return io_reader_.prepare_series(generation, locator, prepared); +} + +int TsFileExecutor::prepare_series( + const FileGeneration& generation, const PreparedLocator& locator, + const std::shared_ptr& aligned_time_owner, + std::shared_ptr& prepared) { + ASSERT(is_inited_); + return io_reader_.prepare_series(generation, locator, aligned_time_owner, + prepared); +} + +int TsFileExecutor::execute_prepared( + const std::shared_ptr& prepared, int64_t start_time, + int64_t end_time, int offset, int limit, const std::string& column_name, + ResultSet*& ret_qds) { + ASSERT(is_inited_); + ret_qds = nullptr; + if (prepared == nullptr || start_time > end_time || offset < 0) { + return E_INVALID_ARG; + } + auto tsblock_reader = std::unique_ptr( + new PreparedSeriesTsBlockReader()); + int ret = tsblock_reader->init(&io_reader_, prepared, + new TimeBetween(start_time, end_time, false), + offset, limit); + if (ret != E_OK) { + return ret; + } + std::vector column_names(1, column_name); + std::vector data_types( + 1, tsblock_reader->value_data_type()); + ret_qds = new TableResultSet(std::move(tsblock_reader), column_names, + data_types, RETURN_BATCH); + return E_OK; +} + +int TsFileExecutor::execute_prepared_multi( + const std::vector>& prepared, + int64_t start_time, int64_t end_time, int offset, int limit, + ResultSet*& ret_qds) { + ASSERT(is_inited_); + ret_qds = nullptr; + if (prepared.empty() || start_time > end_time || offset < 0) { + return E_INVALID_ARG; + } + + auto tsblock_reader = std::unique_ptr( + new PreparedSeriesTsBlockReader()); + int ret = tsblock_reader->init_multi( + &io_reader_, prepared, new TimeBetween(start_time, end_time, false), + offset, limit); + if (ret != E_OK) { + return ret; + } + + std::vector column_names; + column_names.reserve(prepared.size()); + for (const auto& entry : prepared) { + column_names.push_back( + entry->index()->get_measurement_name().to_std_string()); + } + std::vector data_types = + tsblock_reader->value_data_types(); + ret_qds = + new TableResultSet(std::move(tsblock_reader), std::move(column_names), + std::move(data_types), RETURN_BATCH); + return E_OK; +} + int TsFileExecutor::execute_may_with_global_timefilter(QueryExpression* qe, ResultSet*& ret_qds) { int ret = E_OK; diff --git a/cpp/src/reader/tsfile_executor.h b/cpp/src/reader/tsfile_executor.h index 335134c89..86fe581ab 100644 --- a/cpp/src/reader/tsfile_executor.h +++ b/cpp/src/reader/tsfile_executor.h @@ -20,6 +20,7 @@ #define READER_TSFILE_EXECUTOR_H #include +#include #include "file/read_file.h" #include "query_executor.h" @@ -37,6 +38,21 @@ class TsFileExecutor // : public QueryExecutor int execute(QueryExpression* query_expr, ResultSet*& ret_qds); int execute(QueryExpression* query_expr, ResultSet*& ret_qds, int offset, int limit); + int prepare_series(const FileGeneration& generation, + const PreparedLocator& locator, + std::shared_ptr& prepared); + int prepare_series( + const FileGeneration& generation, const PreparedLocator& locator, + const std::shared_ptr& aligned_time_owner, + std::shared_ptr& prepared); + int execute_prepared(const std::shared_ptr& prepared, + int64_t start_time, int64_t end_time, int offset, + int limit, const std::string& column_name, + ResultSet*& ret_qds); + int execute_prepared_multi( + const std::vector>& prepared, + int64_t start_time, int64_t end_time, int offset, int limit, + ResultSet*& ret_qds); void destroy_query_data_set(ResultSet* qds); TsFileMeta* get_tsfile_meta() { return io_reader_.get_tsfile_meta(); } TsFileIOReader* get_tsfile_io_reader() { return &io_reader_; } diff --git a/cpp/src/reader/tsfile_reader.cc b/cpp/src/reader/tsfile_reader.cc index fb5f8fd92..6e20b2d63 100644 --- a/cpp/src/reader/tsfile_reader.cc +++ b/cpp/src/reader/tsfile_reader.cc @@ -228,6 +228,39 @@ int TsFileReader::queryByRow(std::vector& path_list, int offset, return ret; } +int TsFileReader::prepare_series(const FileGeneration& generation, + const PreparedLocator& locator, + std::shared_ptr& prepared) { + return tsfile_executor_->prepare_series(generation, locator, prepared); +} + +int TsFileReader::prepare_series( + const FileGeneration& generation, const PreparedLocator& locator, + const std::shared_ptr& aligned_time_owner, + std::shared_ptr& prepared) { + return tsfile_executor_->prepare_series(generation, locator, + aligned_time_owner, prepared); +} + +int TsFileReader::query_prepared( + const std::shared_ptr& prepared, int64_t start_time, + int64_t end_time, int offset, int limit, ResultSet*& result_set) { + if (prepared == nullptr || prepared->index() == nullptr) { + return E_INVALID_ARG; + } + return tsfile_executor_->execute_prepared( + prepared, start_time, end_time, offset, limit, + prepared->index()->get_measurement_name().to_std_string(), result_set); +} + +int TsFileReader::query_prepared_multi( + const std::vector>& prepared, + int64_t start_time, int64_t end_time, int offset, int limit, + ResultSet*& result_set) { + return tsfile_executor_->execute_prepared_multi( + prepared, start_time, end_time, offset, limit, result_set); +} + int TsFileReader::queryByRow(const std::string& table_name, const std::vector& column_names, int offset, int limit, ResultSet*& result_set, diff --git a/cpp/src/reader/tsfile_reader.h b/cpp/src/reader/tsfile_reader.h index 3ba490045..7b817a18c 100644 --- a/cpp/src/reader/tsfile_reader.h +++ b/cpp/src/reader/tsfile_reader.h @@ -24,6 +24,7 @@ #include "common/tsfile_common.h" #include "expression.h" #include "file/read_file.h" +#include "reader/prepared_series.h" #include "reader/table_query_executor.h" namespace storage { class TsFileExecutor; @@ -128,6 +129,21 @@ class TsFileReader { int queryByRow(std::vector& path_list, int offset, int limit, ResultSet*& result_set); + int prepare_series(const FileGeneration& generation, + const PreparedLocator& locator, + std::shared_ptr& prepared); + int prepare_series( + const FileGeneration& generation, const PreparedLocator& locator, + const std::shared_ptr& aligned_time_owner, + std::shared_ptr& prepared); + int query_prepared(const std::shared_ptr& prepared, + int64_t start_time, int64_t end_time, int offset, + int limit, ResultSet*& result_set); + int query_prepared_multi( + const std::vector>& prepared, + int64_t start_time, int64_t end_time, int offset, int limit, + ResultSet*& result_set); + /** * @brief Query table-model data by row with offset/limit pushdown. * diff --git a/cpp/src/reader/tsfile_series_scan_iterator.cc b/cpp/src/reader/tsfile_series_scan_iterator.cc index 79ac8d308..8c677493e 100644 --- a/cpp/src/reader/tsfile_series_scan_iterator.cc +++ b/cpp/src/reader/tsfile_series_scan_iterator.cc @@ -22,6 +22,7 @@ #include #include "common/global.h" +#include "reader/prepared_series.h" #ifdef ENABLE_THREADS #include "common/thread_pool.h" #endif @@ -30,6 +31,108 @@ using namespace common; namespace storage { +int TsFileSeriesScanIterator::init_prepared( + const std::shared_ptr& prepared, ReadFile* read_file, + Filter* time_filter, common::PageArena& data_pa) { + if (prepared == nullptr || prepared->index() == nullptr || + read_file == nullptr) { + return E_INVALID_ARG; + } + prepared_ = prepared; + itimeseries_index_ = prepared->index(); + if (auto* aligned = + dynamic_cast(itimeseries_index_)) { + // Prepared table columns must use the same multi-aligned reader as a + // normal table query. The legacy single-value aligned reader has a + // different page state machine and does not implement the table batch + // contract. A one-value MultiAlignedTimeseriesIndex is only a view; + // the PreparedSeries continues to own both exact metadata indexes. + timeseries_index_pa_.init(512, common::MOD_TSFILE_READER); + void* multi_memory = + timeseries_index_pa_.alloc(sizeof(MultiAlignedTimeseriesIndex)); + if (multi_memory == nullptr) { + return E_OOM; + } + auto* multi = new (multi_memory) MultiAlignedTimeseriesIndex; + multi->time_ts_idx_ = aligned->time_ts_idx_; + multi->value_ts_idxs_.push_back(aligned->value_ts_idx_); + itimeseries_index_ = multi; + } + measurement_name_ = + itimeseries_index_->get_measurement_name().to_std_string(); + read_file_ = read_file; + time_filter_ = time_filter; + data_pa_ = &data_pa; + return E_OK; +} + +int TsFileSeriesScanIterator::init_prepared_multi( + const std::vector>& prepared, + ReadFile* read_file, Filter* time_filter, common::PageArena& data_pa) { + if (prepared.empty() || prepared.front() == nullptr || + read_file == nullptr) { + return E_INVALID_ARG; + } + + const FileGeneration& generation = prepared.front()->generation(); + const PreparedLocator& locator = prepared.front()->locator(); + if (locator.layout != 1 || locator.time_metadata_length == 0) { + return E_NOT_SUPPORT; + } + + timeseries_index_pa_.init(512, common::MOD_TSFILE_READER); + void* multi_memory = + timeseries_index_pa_.alloc(sizeof(MultiAlignedTimeseriesIndex)); + if (multi_memory == nullptr) { + return E_OOM; + } + auto* multi = new (multi_memory) MultiAlignedTimeseriesIndex; + // Publish the placement-new object immediately so destroy() can release + // its vector if validation of a later entry fails. + itimeseries_index_ = multi; + multi->value_ts_idxs_.reserve(prepared.size()); + + for (const auto& entry : prepared) { + if (entry == nullptr || entry->index() == nullptr) { + return E_INVALID_ARG; + } + const FileGeneration& current_generation = entry->generation(); + const PreparedLocator& current_locator = entry->locator(); + if (current_generation.mapped_index_identity != + generation.mapped_index_identity || + current_generation.file_id != generation.file_id || + current_generation.file_size != generation.file_size || + current_generation.file_fingerprint != + generation.file_fingerprint || + current_locator.layout != 1 || + current_locator.time_metadata_offset != + locator.time_metadata_offset || + current_locator.time_metadata_length != + locator.time_metadata_length) { + return E_INVALID_ARG; + } + auto* aligned = dynamic_cast(entry->index()); + if (aligned == nullptr || aligned->time_ts_idx_ == nullptr || + aligned->value_ts_idx_ == nullptr) { + return E_NOT_SUPPORT; + } + if (multi->time_ts_idx_ == nullptr) { + multi->time_ts_idx_ = aligned->time_ts_idx_; + } else if (aligned->time_ts_idx_->get_chunk_meta_list()->size() != + multi->time_ts_idx_->get_chunk_meta_list()->size()) { + return E_NOT_SUPPORT; + } + multi->value_ts_idxs_.push_back(aligned->value_ts_idx_); + } + + prepared_group_ = prepared; + measurement_name_ = multi->get_measurement_name().to_std_string(); + read_file_ = read_file; + time_filter_ = time_filter; + data_pa_ = &data_pa; + return E_OK; +} + namespace { bool chunk_may_satisfy_filter(ChunkMeta* chunk_meta, Filter* filter) { return filter == nullptr || chunk_meta == nullptr || @@ -57,8 +160,6 @@ void TsFileSeriesScanIterator::destroy() { dynamic_cast(itimeseries_index_)) { std::vector().swap(multi->value_ts_idxs_); } - itimeseries_index_ = nullptr; - timeseries_index_pa_.destroy(); if (chunk_reader_ != nullptr) { // destroy() already runs manual destructors on internal members // (chunk_header_, decoders, compressor, ...), so calling @@ -69,6 +170,8 @@ void TsFileSeriesScanIterator::destroy() { common::mem_free(chunk_reader_); chunk_reader_ = nullptr; } + itimeseries_index_ = nullptr; + timeseries_index_pa_.destroy(); if (tsblock_ != nullptr) { tsblock_->~TsBlock(); tsblock_ = nullptr; @@ -82,6 +185,8 @@ void TsFileSeriesScanIterator::destroy() { std::vector::Iterator>().swap( value_chunk_meta_cursors_); device_id_.reset(); + prepared_.reset(); + std::vector>().swap(prepared_group_); std::string().swap(measurement_name_); } @@ -438,7 +543,7 @@ TsBlock* TsFileSeriesScanIterator::alloc_tsblock() { void* tsblock_buf = data_pa_->alloc(sizeof(TsBlock)); if (IS_NULL(tsblock_buf)) return nullptr; - tsblock_ = new (tsblock_buf) TsBlock(&tuple_desc_); + tsblock_ = new (tsblock_buf) TsBlock(&tuple_desc_, max_block_rows_); if (E_OK != tsblock_->init()) { tsblock_->~TsBlock(); tsblock_ = nullptr; @@ -470,7 +575,7 @@ TsBlock* TsFileSeriesScanIterator::alloc_tsblock_multi() { void* tsblock_buf = data_pa_->alloc(sizeof(TsBlock)); if (IS_NULL(tsblock_buf)) return nullptr; - tsblock_ = new (tsblock_buf) TsBlock(&tuple_desc_); + tsblock_ = new (tsblock_buf) TsBlock(&tuple_desc_, max_block_rows_); if (E_OK != tsblock_->init()) { tsblock_->~TsBlock(); tsblock_ = nullptr; diff --git a/cpp/src/reader/tsfile_series_scan_iterator.h b/cpp/src/reader/tsfile_series_scan_iterator.h index cb3832787..d6ea0e29d 100644 --- a/cpp/src/reader/tsfile_series_scan_iterator.h +++ b/cpp/src/reader/tsfile_series_scan_iterator.h @@ -21,6 +21,7 @@ #define READER_TSFILE_SERIES_SCAN_ITERATOR_H #include +#include #include #include "aligned_chunk_reader.h" @@ -34,6 +35,7 @@ namespace storage { class TsFileIOReader; +class PreparedSeries; class TsFileSeriesScanIterator { public: @@ -49,8 +51,10 @@ class TsFileSeriesScanIterator { tuple_desc_(), tsblock_(nullptr), time_filter_(nullptr), + prepared_(), is_aligned_(false), is_multi_value_(false), + max_block_rows_(0), row_offset_(0), row_limit_(-1) {} ~TsFileSeriesScanIterator() { destroy(); } @@ -65,6 +69,12 @@ class TsFileSeriesScanIterator { data_pa_ = &data_pa; return common::E_OK; } + int init_prepared(const std::shared_ptr& prepared, + ReadFile* read_file, Filter* time_filter, + common::PageArena& data_pa); + int init_prepared_multi( + const std::vector>& prepared, + ReadFile* read_file, Filter* time_filter, common::PageArena& data_pa); void destroy(); /** @@ -76,6 +86,9 @@ class TsFileSeriesScanIterator { row_offset_ = offset; row_limit_ = limit; } + void set_max_block_rows(uint32_t max_block_rows) { + max_block_rows_ = max_block_rows; + } /** Current row offset/limit after chunk/page skip; used to sync with QDS * for single-path. */ @@ -210,8 +223,15 @@ class TsFileSeriesScanIterator { common::TupleDesc tuple_desc_; common::TsBlock* tsblock_; Filter* time_filter_; + // Keeps the arena-backed index alive until the chunk reader has released + // every pointer into it. Empty for the legacy path-owned metadata arena. + std::shared_ptr prepared_; + // Multi-column prepared queries borrow one aligned value index from each + // entry. Keep every owning arena alive until all chunk readers are gone. + std::vector> prepared_group_; bool is_aligned_ = false; bool is_multi_value_ = false; + uint32_t max_block_rows_; int row_offset_; int row_limit_; }; diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt index fab471e59..f2d671998 100644 --- a/cpp/test/CMakeLists.txt +++ b/cpp/test/CMakeLists.txt @@ -152,6 +152,10 @@ enable_testing() # a platform default instead of inheriting the generator used by this build. set(_TSFILE_DEPENDENCY_TEST_CMAKE_ARGUMENTS "-DTEST_CMAKE_GENERATOR=${CMAKE_GENERATOR}") +if (CMAKE_MAKE_PROGRAM) + list(APPEND _TSFILE_DEPENDENCY_TEST_CMAKE_ARGUMENTS + "-DTEST_CMAKE_MAKE_PROGRAM=${CMAKE_MAKE_PROGRAM}") +endif () if (CMAKE_GENERATOR_PLATFORM) list(APPEND _TSFILE_DEPENDENCY_TEST_CMAKE_ARGUMENTS "-DTEST_CMAKE_GENERATOR_PLATFORM=${CMAKE_GENERATOR_PLATFORM}") @@ -224,6 +228,7 @@ file(GLOB_RECURSE TEST_SRCS "writer/*_test.cc" "cwrapper/*_test.cc" "compress/*uncompressed*_test.cc" + "dataset/*_test.cc" ) # Parser tests depend on the ANTLR4 runtime; only build them when it is enabled. diff --git a/cpp/test/dataset/dataset_index_test.cc b/cpp/test/dataset/dataset_index_test.cc new file mode 100644 index 000000000..298aa77be --- /dev/null +++ b/cpp/test/dataset/dataset_index_test.cc @@ -0,0 +1,298 @@ +/* + * 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. + */ + +#include "dataset/dataset_index.h" + +#include + +#include +#include +#include +#include +#include +#include + +#ifdef _WIN32 +#include +#else +#include +#endif + +namespace storage { +namespace dataset { +namespace { + +template +DatasetIndexSectionData fixed_section(DatasetIndexSectionType type, + const std::vector& records) { + DatasetIndexSectionData section; + section.type = type; + section.record_size = sizeof(T); + section.count = static_cast(records.size()); + section.bytes.resize(records.size() * sizeof(T)); + if (!records.empty()) { + std::memcpy(section.bytes.data(), records.data(), section.bytes.size()); + } + return section; +} + +std::vector make_minimal_sections() { + const std::string strings[] = {"table", "device", "value", "/tmp/a.tsfile"}; + std::vector offsets(1, 0); + std::vector bytes; + for (const std::string& value : strings) { + bytes.insert(bytes.end(), value.begin(), value.end()); + offsets.push_back(static_cast(bytes.size())); + } + + DatasetIndexSectionData string_bytes; + string_bytes.type = DatasetIndexSectionType::STRING_BYTES; + string_bytes.record_size = 0; + string_bytes.count = static_cast(bytes.size()); + string_bytes.bytes = bytes; + + TableNameIndexRecord table_name = { + dataset_index_name_hash(strings[0].data(), strings[0].size()), 0, 0}; + TableRecord table = {0, 0, 0, 1, 0, 1, 0}; + DeviceNameIndexRecord device_name = { + 0, 0, dataset_index_name_hash(strings[1].data(), strings[1].size()), 1, + 0}; + DeviceRecord device = {0, 1, 0, 0, 0, 1, 0, 1, 0, 9}; + ColumnNameIndexRecord column_name = { + 0, 0, dataset_index_name_hash(strings[2].data(), strings[2].size()), 2, + 0}; + ColumnSchemaRecord column = {0, 2, 0, 2, 2, 0, 0, 0, 1, 0}; + LogicalSeriesRecord series = {0, 0, 0, 1, 0, 9}; + TsFileRecord file = {3, 0, 4096, 0x5678, 0}; + DeviceFileSpanRecord device_span = {0, 0, 0, 0, 0, 0, 0}; + SeriesFileSpanRecord series_span = {0, 0, 0, 0, 0, 9, 10}; + SeriesLocatorRecord locator = {0, 0, 0, 128, 16, 0}; + + std::vector sections; + sections.push_back( + fixed_section(DatasetIndexSectionType::STRING_OFFSETS, offsets)); + sections.push_back(string_bytes); + sections.push_back( + fixed_section(DatasetIndexSectionType::TABLE_NAME_INDEX, + std::vector(1, table_name))); + sections.push_back(fixed_section(DatasetIndexSectionType::TABLE_RECORD, + std::vector(1, table))); + sections.push_back( + fixed_section(DatasetIndexSectionType::DEVICE_NAME_INDEX, + std::vector(1, device_name))); + sections.push_back(fixed_section(DatasetIndexSectionType::DEVICE_RECORD, + std::vector(1, device))); + sections.push_back( + fixed_section(DatasetIndexSectionType::COLUMN_NAME_INDEX, + std::vector(1, column_name))); + sections.push_back( + fixed_section(DatasetIndexSectionType::COLUMN_SCHEMA, + std::vector(1, column))); + sections.push_back( + fixed_section(DatasetIndexSectionType::LOGICAL_SERIES, + std::vector(1, series))); + sections.push_back(fixed_section(DatasetIndexSectionType::TSFILE_RECORD, + std::vector(1, file))); + sections.push_back( + fixed_section(DatasetIndexSectionType::DEVICE_FILE_SPAN, + std::vector(1, device_span))); + sections.push_back( + fixed_section(DatasetIndexSectionType::SERIES_FILE_SPAN, + std::vector(1, series_span))); + sections.push_back( + fixed_section(DatasetIndexSectionType::SERIES_LOCATOR, + std::vector(1, locator))); + return sections; +} + +int current_process_id() { +#ifdef _WIN32 + return _getpid(); +#else + return static_cast(getpid()); +#endif +} + +class DatasetIndexTest : public ::testing::Test { + protected: + void SetUp() override { + std::ostringstream stream; + stream << "dataset_index_test_" << current_process_id() << ".tsidx"; + path_ = stream.str(); + std::remove(path_.c_str()); + std::ostringstream temp; + temp << path_ << ".tmp." << current_process_id(); + std::remove(temp.str().c_str()); + } + + void TearDown() override { std::remove(path_.c_str()); } + + void write_valid() { + std::string error; + ASSERT_EQ(DatasetIndexStatus::OK, + DatasetIndexWriter::write_atomic( + path_, make_minimal_sections(), error)) + << error; + } + + template + void overwrite(uint64_t offset, const T& value) { + std::fstream file(path_.c_str(), + std::ios::binary | std::ios::in | std::ios::out); + ASSERT_TRUE(file.good()); + file.seekp(static_cast(offset)); + file.write(reinterpret_cast(&value), sizeof(value)); + ASSERT_TRUE(file.good()); + } + + std::string path_; +}; + +TEST_F(DatasetIndexTest, WritesMapsAndLooksUpMinimalIndex) { + write_valid(); + MappedDatasetIndex index; + ASSERT_EQ(DatasetIndexStatus::OK, index.open(path_)) + << index.error_message(); + EXPECT_EQ(DATASET_INDEX_SECTION_COUNT, index.header()->section_count); + + std::vector tables; + ASSERT_EQ(DatasetIndexStatus::OK, index.find_table_ids("table", tables)); + ASSERT_EQ(1U, tables.size()); + EXPECT_EQ(0U, tables[0]); + + uint32_t device_id = 99; + uint32_t column_id = 99; + uint32_t series_id = 99; + EXPECT_EQ(DatasetIndexStatus::OK, + index.find_device_id(0, "device", device_id)); + EXPECT_EQ(DatasetIndexStatus::OK, + index.find_column_id(0, "value", column_id)); + EXPECT_EQ(DatasetIndexStatus::OK, + index.find_series_id(device_id, column_id, series_id)); + EXPECT_EQ(0U, device_id); + EXPECT_EQ(0U, column_id); + EXPECT_EQ(0U, series_id); + + DatasetIndexStringView path; + ASSERT_EQ(DatasetIndexStatus::OK, index.string(3, path)); + EXPECT_EQ("/tmp/a.tsfile", path.to_string()); + EXPECT_EQ(DatasetIndexStatus::NOT_FOUND, + index.find_device_id(0, "missing", device_id)); +} + +TEST_F(DatasetIndexTest, RejectsDuplicateCanonicalTableNames) { + std::vector sections = make_minimal_sections(); + for (DatasetIndexSectionData& section : sections) { + if (section.type == DatasetIndexSectionType::TABLE_NAME_INDEX) { + const TableNameIndexRecord duplicate = { + dataset_index_name_hash("table", 5), 0, 1}; + const uint8_t* data = reinterpret_cast(&duplicate); + section.bytes.insert(section.bytes.end(), data, + data + sizeof(duplicate)); + ++section.count; + } else if (section.type == DatasetIndexSectionType::TABLE_RECORD) { + const TableRecord duplicate = {0, 0, 0, 0, 0, 0, 0}; + const uint8_t* data = reinterpret_cast(&duplicate); + section.bytes.insert(section.bytes.end(), data, + data + sizeof(duplicate)); + ++section.count; + } + } + std::string error; + ASSERT_EQ(DatasetIndexStatus::OK, + DatasetIndexWriter::write_atomic(path_, sections, error)) + << error; + + MappedDatasetIndex index; + EXPECT_EQ(DatasetIndexStatus::BAD_REFERENCE, index.open(path_)); + EXPECT_NE(std::string::npos, + index.error_message().find("duplicate table names")); +} + +TEST_F(DatasetIndexTest, RejectsUnsupportedVersionBeforePublishingViews) { + write_valid(); + uint16_t version = 2; + overwrite(offsetof(DatasetIndexHeader, version_major), version); + MappedDatasetIndex index; + EXPECT_EQ(DatasetIndexStatus::UNSUPPORTED_VERSION, index.open(path_)); + EXPECT_FALSE(index.is_open()); +} + +TEST_F(DatasetIndexTest, RejectsHeaderChecksumMismatch) { + write_valid(); + uint64_t wrong_length = 1; + overwrite(offsetof(DatasetIndexHeader, file_length), wrong_length); + MappedDatasetIndex index; + EXPECT_EQ(DatasetIndexStatus::BAD_HEADER, index.open(path_)); + EXPECT_FALSE(index.is_open()); +} + +TEST_F(DatasetIndexTest, RejectsSectionChecksumMismatch) { + write_valid(); + MappedDatasetIndex valid; + ASSERT_EQ(DatasetIndexStatus::OK, valid.open(path_)); + DatasetIndexSectionView strings; + ASSERT_EQ(DatasetIndexStatus::OK, + valid.section(DatasetIndexSectionType::STRING_BYTES, strings)); + const uint64_t string_offset = static_cast( + strings.data - reinterpret_cast(valid.header())); + valid.close(); + const uint8_t changed = 'X'; + overwrite(string_offset, changed); + + MappedDatasetIndex index; + EXPECT_EQ(DatasetIndexStatus::BAD_CHECKSUM, index.open(path_)); +} + +TEST_F(DatasetIndexTest, WriterRejectsMalformedSectionCount) { + std::vector sections = make_minimal_sections(); + sections.pop_back(); + std::string error; + EXPECT_EQ(DatasetIndexStatus::INVALID_ARGUMENT, + DatasetIndexWriter::write_atomic(path_, sections, error)); + EXPECT_FALSE(error.empty()); +} + +TEST(DatasetIndexChecksumTest, MatchesKnownCrc32cVector) { + const char* value = "123456789"; + EXPECT_EQ(0xE3069283U, dataset_index_crc32c(value, 9)); +} + +TEST(DatasetIndexCrossLanguageTest, OpensExternalIndexWhenConfigured) { + const char* path = std::getenv("TSFILE_DATASET_INDEX_TEST_PATH"); + if (path == nullptr || path[0] == '\0') { + GTEST_SKIP() << "TSFILE_DATASET_INDEX_TEST_PATH is not set"; + } + + MappedDatasetIndex index; + ASSERT_EQ(DatasetIndexStatus::OK, index.open(path)) + << index.error_message(); + DatasetIndexSectionView files; + DatasetIndexSectionView series; + ASSERT_EQ(DatasetIndexStatus::OK, + index.section(DatasetIndexSectionType::TSFILE_RECORD, files)); + ASSERT_EQ(DatasetIndexStatus::OK, + index.section(DatasetIndexSectionType::LOGICAL_SERIES, series)); + EXPECT_GT(files.count, 0U); + EXPECT_GT(series.count, 0U); +} + +} // namespace +} // namespace dataset +} // namespace storage diff --git a/cpp/test/reader/prepared_series_test.cc b/cpp/test/reader/prepared_series_test.cc new file mode 100644 index 000000000..0860cb88a --- /dev/null +++ b/cpp/test/reader/prepared_series_test.cc @@ -0,0 +1,376 @@ +/* + * 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. + */ + +#include "reader/prepared_series.h" + +#include +#include + +#include +#include + +#include "common/global.h" +#include "common/schema.h" +#include "common/tablet.h" +#include "file/write_file.h" +#include "reader/table_result_set.h" +#include "reader/tsfile_reader.h" +#include "writer/tsfile_table_writer.h" + +namespace storage { +namespace { + +class PagePointGuard { + public: + explicit PagePointGuard(uint32_t page_points) + : saved_(common::g_config_value_.page_writer_max_point_num_) { + common::g_config_value_.page_writer_max_point_num_ = page_points; + } + ~PagePointGuard() { + common::g_config_value_.page_writer_max_point_num_ = saved_; + } + + private: + uint32_t saved_; +}; + +class PreparedSeriesBatchTest : public ::testing::Test { + protected: + void SetUp() override { + libtsfile_init(); + const auto* test_info = + ::testing::UnitTest::GetInstance()->current_test_info(); + file_name_ = std::string("prepared_series_batch_test_") + + (test_info == nullptr ? "unknown" : test_info->name()) + + ".tsfile"; + std::remove(file_name_.c_str()); + } + + void TearDown() override { + std::remove(file_name_.c_str()); + libtsfile_destroy(); + } + + void write_nullable_table() { + PagePointGuard guard(10000); + WriteFile write_file; + int flags = O_WRONLY | O_CREAT | O_TRUNC; +#ifdef _WIN32 + flags |= O_BINARY; +#endif + ASSERT_EQ(common::E_OK, write_file.create(file_name_, flags, 0666)); + + std::vector columns = { + common::ColumnSchema("device", common::STRING, + common::ColumnCategory::TAG), + common::ColumnSchema("value", common::DOUBLE, + common::ColumnCategory::FIELD), + common::ColumnSchema("value2", common::DOUBLE, + common::ColumnCategory::FIELD), + }; + auto* schema = new TableSchema("weather", columns); + TsFileTableWriter writer(&write_file, schema); + Tablet tablet( + "weather", {"device", "value", "value2"}, + {common::STRING, common::DOUBLE, common::DOUBLE}, + {common::ColumnCategory::TAG, common::ColumnCategory::FIELD, + common::ColumnCategory::FIELD}, + 70000); + for (int row = 0; row < 70000; ++row) { + tablet.add_timestamp(row, row); + tablet.add_value(row, "device", "d0"); + if (row != 2 && row != 6) { + tablet.add_value(row, "value", static_cast(row)); + } + if (row != 4 && row != 8) { + tablet.add_value(row, "value2", static_cast(row * 10)); + } + } + ASSERT_EQ(common::E_OK, writer.write_table(tablet)); + ASSERT_EQ(common::E_OK, writer.flush()); + ASSERT_EQ(common::E_OK, writer.close()); + delete schema; + } + + std::string file_name_ = "prepared_series_batch_test.tsfile"; +}; + +TEST_F(PreparedSeriesBatchTest, + PreparedQueryReturnsDirectTableResultSetBatches) { + write_nullable_table(); + + TsFileReader reader; + ASSERT_EQ(common::E_OK, reader.open(file_name_)); + + ResultSet* fixture_result = nullptr; + ASSERT_EQ(common::E_OK, reader.query("weather", {"device", "value"}, 0, + 69999, fixture_result, 4096)); + auto* fixture_table = dynamic_cast(fixture_result); + ASSERT_NE(nullptr, fixture_table); + uint32_t fixture_row_count = 0; + common::TsBlock* fixture_block = nullptr; + while (fixture_table->get_next_tsblock(fixture_block) == common::E_OK) { + ASSERT_NE(nullptr, fixture_block); + fixture_row_count += fixture_block->get_row_count(); + } + EXPECT_EQ(70000U, fixture_row_count); + reader.destroy_query_data_set(fixture_result); + + auto metadata = reader.get_timeseries_metadata(); + AlignedTimeseriesIndex* aligned = nullptr; + for (const auto& device_entry : metadata) { + for (const auto& index : device_entry.second) { + auto* candidate = + dynamic_cast(index.get()); + if (candidate != nullptr && candidate->value_ts_idx_ != nullptr && + candidate->value_ts_idx_->get_measurement_name() + .to_std_string() == "value") { + aligned = candidate; + break; + } + } + } + ASSERT_NE(nullptr, aligned); + ASSERT_NE(nullptr, aligned->time_ts_idx_); + ASSERT_NE(nullptr, aligned->value_ts_idx_); + + FileGeneration generation; + generation.mapped_index_identity = 1; + generation.file_id = 0; + struct stat file_stat {}; + ASSERT_EQ(0, stat(file_name_.c_str(), &file_stat)); + generation.file_size = static_cast(file_stat.st_size); + generation.file_fingerprint = 0; + + PreparedLocator locator; + locator.locator_id = 0; + locator.layout = 1; + locator.flags = 1; + locator.value_metadata_offset = + aligned->value_ts_idx_->get_metadata_offset(); + locator.value_metadata_length = + aligned->value_ts_idx_->get_metadata_length(); + locator.time_metadata_offset = aligned->time_ts_idx_->get_metadata_offset(); + locator.time_metadata_length = aligned->time_ts_idx_->get_metadata_length(); + + std::shared_ptr prepared; + ASSERT_EQ(common::E_OK, + reader.prepare_series(generation, locator, prepared)); + ASSERT_NE(nullptr, prepared); + auto* prepared_aligned = + dynamic_cast(prepared->index()); + ASSERT_NE(nullptr, prepared_aligned); + ASSERT_NE(nullptr, prepared_aligned->time_ts_idx_); + ASSERT_NE(nullptr, prepared_aligned->value_ts_idx_); + EXPECT_EQ(70000, + prepared_aligned->time_ts_idx_->get_statistic()->get_count()); + EXPECT_EQ(aligned->time_ts_idx_->get_chunk_meta_list()->size(), + prepared_aligned->time_ts_idx_->get_chunk_meta_list()->size()); + + ResultSet* result = nullptr; + ASSERT_EQ(common::E_OK, + reader.query_prepared(prepared, 0, 9, 1, 7, result)); + auto* table_result = dynamic_cast(result); + ASSERT_NE(nullptr, table_result); + + std::vector timestamps; + std::vector values; + std::vector nulls; + int block_count = 0; + common::TsBlock* block = nullptr; + int ret = common::E_OK; + while ((ret = table_result->get_next_tsblock(block)) == common::E_OK) { + ASSERT_NE(nullptr, block); + ++block_count; + common::RowIterator rows(block); + while (rows.has_next()) { + uint32_t len = 0; + bool is_null = false; + const char* timestamp = rows.read(0, &len, &is_null); + ASSERT_FALSE(is_null); + timestamps.push_back(*reinterpret_cast(timestamp)); + + const char* value = rows.read(1, &len, &is_null); + nulls.push_back(is_null); + values.push_back(is_null ? 0.0 + : *reinterpret_cast(value)); + rows.next(); + } + } + EXPECT_EQ(common::E_NO_MORE_DATA, ret); + EXPECT_EQ(1, block_count); + ASSERT_EQ(7U, timestamps.size()); + for (int64_t index = 0; index < 7; ++index) { + EXPECT_EQ(index + 1, timestamps[index]); + const bool expected_null = index + 1 == 2 || index + 1 == 6; + EXPECT_EQ(expected_null, nulls[index]); + if (!expected_null) { + EXPECT_DOUBLE_EQ(static_cast(index + 1), values[index]); + } + } + reader.destroy_query_data_set(result); + + ResultSet* multi_batch = nullptr; + ASSERT_EQ(common::E_OK, + reader.query_prepared(prepared, 0, 69999, 0, 65537, multi_batch)); + auto* multi_batch_table = dynamic_cast(multi_batch); + ASSERT_NE(nullptr, multi_batch_table); + uint32_t multi_batch_rows = 0; + uint32_t multi_batch_count = 0; + block = nullptr; + while (multi_batch_table->get_next_tsblock(block) == common::E_OK) { + ASSERT_NE(nullptr, block); + ++multi_batch_count; + multi_batch_rows += block->get_row_count(); + } + EXPECT_GT(multi_batch_count, 1U); + EXPECT_EQ(65537U, multi_batch_rows); + reader.destroy_query_data_set(multi_batch); + + ResultSet* empty = nullptr; + ASSERT_EQ(common::E_OK, + reader.query_prepared(prepared, 100000, 200000, 0, -1, empty)); + auto* empty_table = dynamic_cast(empty); + ASSERT_NE(nullptr, empty_table); + block = nullptr; + EXPECT_EQ(common::E_NO_MORE_DATA, empty_table->get_next_tsblock(block)); + EXPECT_EQ(nullptr, block); + reader.destroy_query_data_set(empty); + EXPECT_EQ(common::E_OK, reader.close()); +} + +TEST_F(PreparedSeriesBatchTest, + MultiPreparedQuerySharesAlignedTimeAxisAndPreservesColumnOrder) { + write_nullable_table(); + + TsFileReader reader; + ASSERT_EQ(common::E_OK, reader.open(file_name_)); + auto metadata = reader.get_timeseries_metadata(); + AlignedTimeseriesIndex* value_index = nullptr; + AlignedTimeseriesIndex* value2_index = nullptr; + for (const auto& device_entry : metadata) { + for (const auto& index : device_entry.second) { + auto* aligned = dynamic_cast(index.get()); + if (aligned == nullptr || aligned->value_ts_idx_ == nullptr) { + continue; + } + const std::string name = + aligned->value_ts_idx_->get_measurement_name().to_std_string(); + if (name == "value") { + value_index = aligned; + } else if (name == "value2") { + value2_index = aligned; + } + } + } + ASSERT_NE(nullptr, value_index); + ASSERT_NE(nullptr, value2_index); + + FileGeneration generation; + generation.mapped_index_identity = 1; + generation.file_id = 0; + struct stat file_stat {}; + ASSERT_EQ(0, stat(file_name_.c_str(), &file_stat)); + generation.file_size = static_cast(file_stat.st_size); + + auto prepare = [&](uint32_t locator_id, AlignedTimeseriesIndex* aligned, + const std::shared_ptr& time_owner) { + PreparedLocator locator; + locator.locator_id = locator_id; + locator.layout = 1; + locator.flags = 1; + locator.value_metadata_offset = + aligned->value_ts_idx_->get_metadata_offset(); + locator.value_metadata_length = + aligned->value_ts_idx_->get_metadata_length(); + locator.time_metadata_offset = + aligned->time_ts_idx_->get_metadata_offset(); + locator.time_metadata_length = + aligned->time_ts_idx_->get_metadata_length(); + std::shared_ptr result; + EXPECT_EQ(common::E_OK, + time_owner == nullptr + ? reader.prepare_series(generation, locator, result) + : reader.prepare_series(generation, locator, time_owner, + result)); + return result; + }; + + std::shared_ptr prepared_value = + prepare(0, value_index, nullptr); + std::shared_ptr prepared_value2 = + prepare(1, value2_index, prepared_value); + ASSERT_NE(nullptr, prepared_value); + ASSERT_NE(nullptr, prepared_value2); + auto* first_aligned = + dynamic_cast(prepared_value->index()); + auto* second_aligned = + dynamic_cast(prepared_value2->index()); + ASSERT_NE(nullptr, first_aligned); + ASSERT_NE(nullptr, second_aligned); + EXPECT_EQ(first_aligned->time_ts_idx_, second_aligned->time_ts_idx_); + + ResultSet* result = nullptr; + ASSERT_EQ(common::E_OK, + reader.query_prepared_multi({prepared_value2, prepared_value}, 0, + 9, 0, -1, result)); + auto* table_result = dynamic_cast(result); + ASSERT_NE(nullptr, table_result); + auto result_metadata = table_result->get_metadata(); + ASSERT_NE(nullptr, result_metadata); + EXPECT_EQ("time", result_metadata->get_column_name(1)); + EXPECT_EQ("value2", result_metadata->get_column_name(2)); + EXPECT_EQ("value", result_metadata->get_column_name(3)); + + uint32_t row = 0; + common::TsBlock* block = nullptr; + while (table_result->get_next_tsblock(block) == common::E_OK) { + ASSERT_NE(nullptr, block); + common::RowIterator rows(block); + while (rows.has_next()) { + uint32_t len = 0; + bool is_null = false; + const char* timestamp = rows.read(0, &len, &is_null); + ASSERT_FALSE(is_null); + ASSERT_EQ(row, *reinterpret_cast(timestamp)); + + const char* value2 = rows.read(1, &len, &is_null); + EXPECT_EQ(row == 4 || row == 8, is_null); + if (!is_null) { + EXPECT_DOUBLE_EQ(static_cast(row * 10), + *reinterpret_cast(value2)); + } + + const char* value = rows.read(2, &len, &is_null); + EXPECT_EQ(row == 2 || row == 6, is_null); + if (!is_null) { + EXPECT_DOUBLE_EQ(static_cast(row), + *reinterpret_cast(value)); + } + ++row; + rows.next(); + } + } + EXPECT_EQ(10U, row); + reader.destroy_query_data_set(result); + EXPECT_EQ(common::E_OK, reader.close()); +} + +} // namespace +} // namespace storage diff --git a/pom.xml b/pom.xml index eb9380fb9..0e5988130 100644 --- a/pom.xml +++ b/pom.xml @@ -146,8 +146,8 @@ **/build/** **/.clang-format - **/tsfile/**.cpp - **/tsfile/**.h + **/tsfile/**/*.cpp + **/tsfile/**/*.h **/venv/** **/tsfile.egg-info/** diff --git a/python/setup.py b/python/setup.py index d8e9a502f..a69db1558 100644 --- a/python/setup.py +++ b/python/setup.py @@ -264,6 +264,7 @@ def finalize_options(self): ) exts = [ + Extension("tsfile.dataset._merge", ["tsfile/dataset/_merge.pyx"], **common), Extension("tsfile.tsfile_py_cpp", ["tsfile/tsfile_py_cpp.pyx"], **common), Extension("tsfile.tsfile_reader", ["tsfile/tsfile_reader.pyx"], **common), Extension("tsfile.tsfile_writer", ["tsfile/tsfile_writer.pyx"], **common), diff --git a/python/tests/test_dataset_index.py b/python/tests/test_dataset_index.py new file mode 100644 index 000000000..417b54d58 --- /dev/null +++ b/python/tests/test_dataset_index.py @@ -0,0 +1,582 @@ +# 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 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, + DIRECTORY, + HEADER, + LOGICAL_SERIES, + MappedDatasetIndex, + RECORDS, + SERIES_FILE_SPAN, + SERIES_LOCATOR, + 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 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_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) + + class _Index: + @staticmethod + def record(_section_type, record_id): + return rows[record_id] + + @staticmethod + def string_bytes(sid): + return names[sid] + + assert MappedDatasetIndex._find_child(_Index(), 0, 0, "beta", 0, 3) == 11 + with pytest.raises(KeyError): + MappedDatasetIndex._find_child(_Index(), 0, 0, "missing", 0, 3) + + +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) as first: + assert len(first) == 1 + + def fail_legacy_scan(*_args, **_kwargs): + raise AssertionError("hot construction must not build a legacy catalog") + + monkeypatch.setattr("tsfile.dataset.reader.TsFileSeriesReader", fail_legacy_scan) + with TsFileDataFrame(str(source), show_progress=False) as second: + 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_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) 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) 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) as first: + foreign_path = first.list_timeseries()[0] + with TsFileDataFrame(str(second_source), show_progress=False) 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) 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 + + +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) 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) 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) 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) 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) 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() diff --git a/python/tests/test_tsfile_dataset.py b/python/tests/test_tsfile_dataset.py index ee016d087..d24a3dafc 100644 --- a/python/tests/test_tsfile_dataset.py +++ b/python/tests/test_tsfile_dataset.py @@ -19,6 +19,7 @@ import numpy as np import pandas as pd import pytest +import threading from tsfile.dataset import dataframe as dataframe_module from tsfile import ( @@ -41,6 +42,7 @@ TsFileSeriesReader, _build_exact_tag_filter, ) +from tsfile.dataset.runtime import RuntimeSeriesReader def _write_weather_file(path, start): @@ -156,6 +158,50 @@ def _write_weather_with_extra_field_file(path, start): writer.write_dataframe(df) +def _write_weather_int_temperature_file(path, start): + schema = TableSchema( + "weather", + [ + ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG), + ColumnSchema("temperature", TSDataType.INT64, ColumnCategory.FIELD), + ColumnSchema("humidity", TSDataType.DOUBLE, ColumnCategory.FIELD), + ], + ) + with TsFileTableWriter(str(path), schema) as writer: + writer.write_dataframe( + pd.DataFrame( + { + "time": [start, start + 1], + "device": ["device_a", "device_a"], + "temperature": [20, 21], + "humidity": [50.0, 51.0], + } + ) + ) + + +def _write_weather_boolean_status_file(path, start): + schema = TableSchema( + "weather", + [ + ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG), + ColumnSchema("temperature", TSDataType.DOUBLE, ColumnCategory.FIELD), + ColumnSchema("status", TSDataType.BOOLEAN, ColumnCategory.FIELD), + ], + ) + with TsFileTableWriter(str(path), schema) as writer: + writer.write_dataframe( + pd.DataFrame( + { + "time": [start, start + 1], + "device": ["device_a", "device_a"], + "temperature": [20.0, 21.0], + "status": [True, False], + } + ) + ) + + def _write_multi_tag_file(path): schema = TableSchema( "weather", @@ -295,6 +341,129 @@ def test_dataset_loc_aligns_timestamp_union_and_preserves_requested_order(tmp_pa assert aligned.values[2, 1] == 30.0 +def test_dataset_loc_batches_aligned_fields_per_device_then_unions_devices(tmp_path): + path = tmp_path / "weather_multi_device.tsfile" + _write_weather_rows_file( + path, + { + "time": [0, 1, 1, 2], + "device": ["device_a", "device_a", "device_b", "device_b"], + "temperature": [10.0, 11.0, 20.0, 21.0], + "humidity": [100.0, 101.0, 200.0, 201.0], + }, + ) + + requested = [ + "weather.device_a.temperature", + "weather.device_a.humidity", + "weather.device_b.humidity", + ] + with TsFileDataFrame(str(path), show_progress=False) as tsdf: + aligned = tsdf.loc[0:2, requested] + + assert aligned.series_names == requested + np.testing.assert_array_equal( + aligned.timestamps, np.array([0, 1, 2], dtype=np.int64) + ) + np.testing.assert_allclose( + aligned.values, + np.array( + [ + [10.0, 100.0, np.nan], + [11.0, 101.0, 200.0], + [np.nan, np.nan, 201.0], + ] + ), + equal_nan=True, + ) + + +def test_dataset_loc_runs_independent_device_groups_concurrently(tmp_path, monkeypatch): + path = tmp_path / "weather_concurrent_devices.tsfile" + _write_weather_rows_file( + path, + { + "time": [0, 1, 0, 1], + "device": ["device_a", "device_a", "device_b", "device_b"], + "temperature": [10.0, 11.0, 20.0, 21.0], + "humidity": [100.0, 101.0, 200.0, 201.0], + }, + ) + monkeypatch.setenv("TSFILE_DATAFRAME_QUERY_WORKERS", "2") + monkeypatch.setenv("TSFILE_DATAFRAME_QUERY_PARALLEL_MIN_ROWS", "1") + + original = RuntimeSeriesReader.read_device_fields_by_time_range + rendezvous = threading.Barrier(2) + worker_ids = set() + + def observed_read(reader, device_id, column_ids, start_time, end_time): + worker_ids.add(threading.get_ident()) + rendezvous.wait(timeout=2) + return original(reader, device_id, column_ids, start_time, end_time) + + monkeypatch.setattr( + RuntimeSeriesReader, + "read_device_fields_by_time_range", + observed_read, + ) + + with TsFileDataFrame(str(path), show_progress=False) as tsdf: + aligned = tsdf.loc[ + 0:1, + [ + "weather.device_a.temperature", + "weather.device_a.humidity", + "weather.device_b.humidity", + ], + ] + + assert len(worker_ids) == 2 + np.testing.assert_allclose( + aligned.values, + np.array([[10.0, 100.0, 200.0], [11.0, 101.0, 201.0]]), + ) + + +def test_dataset_loc_keeps_small_device_groups_inline(tmp_path, monkeypatch): + path = tmp_path / "weather_inline_devices.tsfile" + _write_weather_rows_file( + path, + { + "time": [0, 1, 0, 1], + "device": ["device_a", "device_a", "device_b", "device_b"], + "temperature": [10.0, 11.0, 20.0, 21.0], + "humidity": [100.0, 101.0, 200.0, 201.0], + }, + ) + monkeypatch.setenv("TSFILE_DATAFRAME_QUERY_WORKERS", "2") + monkeypatch.setenv("TSFILE_DATAFRAME_QUERY_PARALLEL_MIN_ROWS", "8192") + + original = RuntimeSeriesReader.read_device_fields_by_time_range + caller_id = threading.get_ident() + worker_ids = set() + + def observed_read(reader, device_id, column_ids, start_time, end_time): + worker_ids.add(threading.get_ident()) + return original(reader, device_id, column_ids, start_time, end_time) + + monkeypatch.setattr( + RuntimeSeriesReader, + "read_device_fields_by_time_range", + observed_read, + ) + + with TsFileDataFrame(str(path), show_progress=False) as tsdf: + tsdf.loc[ + 0:1, + [ + "weather.device_a.temperature", + "weather.device_b.humidity", + ], + ] + + assert worker_ids == {caller_id} + + def test_dataset_reads_nullable_tag_devices_in_isolation(tmp_path): path = tmp_path / "nullable_tags.tsfile" schema = TableSchema( @@ -788,10 +957,8 @@ def test_dataset_rejects_duplicate_timestamps_across_shards(tmp_path): _write_weather_file(path1, 0) _write_weather_file(path2, 2) - with TsFileDataFrame([str(path1), str(path2)], show_progress=False) as tsdf: - series = tsdf["weather.device_a.temperature"] - with pytest.raises(ValueError, match="Duplicate timestamp"): - _ = series.timestamps + with pytest.raises(ValueError, match="Duplicate timestamp"): + TsFileDataFrame([str(path1), str(path2)], show_progress=False) def test_dataset_overlap_position_access_avoids_full_timestamp_materialization( @@ -825,6 +992,11 @@ def fail_merge(*_args, **_kwargs): monkeypatch.setattr(dataframe_module, "_merge_field_timestamps", fail_merge) + def fail_span_lookup(*_args, **_kwargs): + raise AssertionError("overlap position reads must reuse descriptor locators") + + monkeypatch.setattr(RuntimeSeriesReader, "_span", fail_span_lookup) + with TsFileDataFrame([str(path1), str(path2)], show_progress=False) as tsdf: series = tsdf["weather.device_a.temperature"] assert series[0] == 10.0 @@ -833,7 +1005,7 @@ def fail_merge(*_args, **_kwargs): np.testing.assert_array_equal(series[1:5], np.array([20.0, 30.0, 40.0, 50.0])) -def test_dataset_rejects_data_access_after_close(tmp_path): +def test_dataset_close_only_releases_current_handle(tmp_path): path = tmp_path / "weather.tsfile" _write_weather_file(path, 0) @@ -844,18 +1016,22 @@ def test_dataset_rejects_data_access_after_close(tmp_path): with pytest.raises(RuntimeError, match="TsFileDataFrame is closed"): _ = tsdf[0] - with pytest.raises(RuntimeError, match="TsFileDataFrame is closed"): + assert series[0] == 20.0 + series.close() + with pytest.raises(RuntimeError, match="Timeseries is closed"): _ = series[0] -def test_subset_close_warns_and_does_not_close_root(tmp_path): +def test_subset_close_releases_only_subset_lease(tmp_path): path = tmp_path / "weather.tsfile" _write_weather_file(path, 0) with TsFileDataFrame(str(path), show_progress=False) as tsdf: subset = tsdf[:1] - with pytest.warns(RuntimeWarning, match="no-op"): - subset.close() + subset.close() + + with pytest.raises(RuntimeError, match="TsFileDataFrame is closed"): + _ = subset[0] series = tsdf[0] assert series[0] == 20.0 @@ -871,6 +1047,79 @@ def test_dataset_rejects_incompatible_table_schemas_across_shards(tmp_path): TsFileDataFrame([str(path1), str(path2)], show_progress=False) +def test_dataset_rejects_same_field_name_with_different_type(tmp_path): + path1 = tmp_path / "double.tsfile" + path2 = tmp_path / "int64.tsfile" + _write_weather_file(path1, 0) + _write_weather_int_temperature_file(path2, 3) + + with pytest.raises(ValueError, match="Incompatible schema for table 'weather'"): + TsFileDataFrame([str(path1), str(path2)], show_progress=False) + + +def test_dataset_rejects_nonnumeric_declared_schema_difference(tmp_path): + path1 = tmp_path / "text-status.tsfile" + path2 = tmp_path / "boolean-status.tsfile" + _write_numeric_and_text_file(path1) + _write_weather_boolean_status_file(path2, 3) + + with pytest.raises(ValueError, match="Incompatible schema for table 'weather'"): + TsFileDataFrame([str(path1), str(path2)], show_progress=False) + + +def test_dataset_close_waits_for_an_active_public_query(tmp_path, monkeypatch): + path = tmp_path / "weather.tsfile" + _write_weather_file(path, 0) + entered = threading.Event() + release = threading.Event() + query_error = [] + original = RuntimeSeriesReader.read_device_fields_by_time_range + + def blocked_read(reader, device_id, column_ids, start_time, end_time): + entered.set() + assert release.wait(timeout=2) + return original(reader, device_id, column_ids, start_time, end_time) + + monkeypatch.setattr( + RuntimeSeriesReader, + "read_device_fields_by_time_range", + blocked_read, + ) + + dataframe = TsFileDataFrame(str(path), show_progress=False) + query_done = threading.Event() + + def run_query(): + try: + result = dataframe.loc[0:2, [0]] + assert result.values.shape == (3, 1) + except BaseException as exc: + query_error.append(exc) + finally: + query_done.set() + + query_thread = threading.Thread(target=run_query) + query_thread.start() + assert entered.wait(timeout=2) + + close_done = threading.Event() + + def close_dataframe(): + dataframe.close() + close_done.set() + + close_thread = threading.Thread(target=close_dataframe) + 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() + assert query_error == [] + + def test_dataset_skips_empty_tsfile_shards(tmp_path): empty_path = tmp_path / "empty.tsfile" data_path = tmp_path / "part.tsfile" diff --git a/python/tsfile/dataset/_merge.pyx b/python/tsfile/dataset/_merge.pyx new file mode 100644 index 000000000..cf26fded5 --- /dev/null +++ b/python/tsfile/dataset/_merge.pyx @@ -0,0 +1,261 @@ +# 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 + +"""Narrow typed kernels for overlapping cross-shard merges.""" + +import numpy as np +cimport numpy as cnp +from libc.stdint cimport int64_t + +cnp.import_array() + + +cdef inline bint _less( + int left, int right, int64_t[:] times, cnp.intp_t[:] starts, + cnp.intp_t[:] cursors) noexcept nogil: + cdef int64_t left_time = times[starts[left] + cursors[left]] + cdef int64_t right_time = times[starts[right] + cursors[right]] + return left_time < right_time or (left_time == right_time and left < right) + + +cdef void _push(int[:] heap, int* heap_size, int part, + int64_t[:] times, cnp.intp_t[:] starts, + cnp.intp_t[:] cursors) noexcept nogil: + cdef int pos = heap_size[0] + cdef int parent + heap_size[0] += 1 + while pos > 0: + parent = (pos - 1) // 2 + if not _less(part, heap[parent], times, starts, cursors): + break + heap[pos] = heap[parent] + pos = parent + heap[pos] = part + + +cdef int _pop(int[:] heap, int* heap_size, int64_t[:] times, + cnp.intp_t[:] starts, cnp.intp_t[:] cursors) noexcept nogil: + cdef int result = heap[0] + cdef int last + cdef int pos = 0 + cdef int child + heap_size[0] -= 1 + if heap_size[0] == 0: + return result + last = heap[heap_size[0]] + while True: + child = pos * 2 + 1 + if child >= heap_size[0]: + break + if child + 1 < heap_size[0] and _less( + heap[child + 1], heap[child], times, starts, cursors): + child += 1 + if not _less(heap[child], last, times, starts, cursors): + break + heap[pos] = heap[child] + pos = child + heap[pos] = last + return result + + +cdef tuple _flatten(list time_parts): + cdef Py_ssize_t count = len(time_parts) + cdef cnp.ndarray[cnp.intp_t, ndim=1] starts = np.empty(count, dtype=np.intp) + cdef cnp.ndarray[cnp.intp_t, ndim=1] lengths = np.empty(count, dtype=np.intp) + cdef Py_ssize_t index + cdef Py_ssize_t total = 0 + for index in range(count): + starts[index] = total + lengths[index] = len(time_parts[index]) + total += lengths[index] + return np.concatenate(time_parts).astype(np.int64, copy=False), starts, lengths, total + + +cpdef merge_timestamp_parts_overlap(list time_parts, bint deduplicate, + bint validate_unique): + cdef object flat_object + cdef cnp.ndarray[cnp.intp_t, ndim=1] starts_array + cdef cnp.ndarray[cnp.intp_t, ndim=1] lengths_array + cdef Py_ssize_t total + flat_object, starts_array, lengths_array, total = _flatten(time_parts) + cdef cnp.ndarray[cnp.int64_t, ndim=1] flat_array = flat_object + cdef int64_t[:] flat = flat_array + cdef cnp.intp_t[:] starts = starts_array + cdef cnp.intp_t[:] lengths = lengths_array + cdef cnp.ndarray[cnp.intp_t, ndim=1] cursors_array = np.zeros(len(time_parts), dtype=np.intp) + cdef cnp.intp_t[:] cursors = cursors_array + cdef cnp.ndarray[cnp.int32_t, ndim=1] heap_array = np.empty(len(time_parts), dtype=np.int32) + cdef int[:] heap = heap_array + cdef cnp.ndarray[cnp.int64_t, ndim=1] output_array = np.empty(total, dtype=np.int64) + cdef int64_t[:] output = output_array + cdef int heap_size = 0 + cdef int part + cdef Py_ssize_t out_size = 0 + cdef int64_t timestamp + cdef int64_t last_timestamp = 0 + cdef bint has_last = False + cdef bint duplicate_found = False + cdef int64_t duplicate_timestamp = 0 + cdef int part_count = len(time_parts) + with nogil: + for part in range(part_count): + if lengths[part] > 0: + _push(heap, &heap_size, part, flat, starts, cursors) + while heap_size > 0: + part = _pop(heap, &heap_size, flat, starts, cursors) + timestamp = flat[starts[part] + cursors[part]] + if has_last and timestamp == last_timestamp: + if validate_unique: + duplicate_found = True + duplicate_timestamp = timestamp + break + if not deduplicate: + output[out_size] = timestamp + out_size += 1 + else: + output[out_size] = timestamp + out_size += 1 + last_timestamp = timestamp + has_last = True + cursors[part] += 1 + if cursors[part] < lengths[part]: + _push(heap, &heap_size, part, flat, starts, cursors) + if duplicate_found: + raise ValueError( + f"Duplicate timestamp {duplicate_timestamp} found across shards." + ) + return output_array[:out_size] + + +cpdef merge_time_value_parts_overlap(list time_parts, list value_parts): + cdef object flat_object + cdef cnp.ndarray[cnp.intp_t, ndim=1] starts_array + cdef cnp.ndarray[cnp.intp_t, ndim=1] lengths_array + cdef Py_ssize_t total + flat_object, starts_array, lengths_array, total = _flatten(time_parts) + cdef cnp.ndarray[cnp.int64_t, ndim=1] flat_times_array = flat_object + cdef cnp.ndarray[cnp.float64_t, ndim=1] flat_values_array = np.concatenate(value_parts).astype(np.float64, copy=False) + cdef int64_t[:] flat_times = flat_times_array + cdef double[:] flat_values = flat_values_array + cdef cnp.intp_t[:] starts = starts_array + cdef cnp.intp_t[:] lengths = lengths_array + cdef cnp.ndarray[cnp.intp_t, ndim=1] cursors_array = np.zeros(len(time_parts), dtype=np.intp) + cdef cnp.intp_t[:] cursors = cursors_array + cdef cnp.ndarray[cnp.int32_t, ndim=1] heap_array = np.empty(len(time_parts), dtype=np.int32) + cdef int[:] heap = heap_array + cdef cnp.ndarray[cnp.int64_t, ndim=1] output_times_array = np.empty(total, dtype=np.int64) + cdef cnp.ndarray[cnp.float64_t, ndim=1] output_values_array = np.empty(total, dtype=np.float64) + cdef int64_t[:] output_times = output_times_array + cdef double[:] output_values = output_values_array + cdef int heap_size = 0 + cdef int part + cdef Py_ssize_t source + cdef Py_ssize_t out_index = 0 + cdef int part_count = len(time_parts) + cdef int64_t last_timestamp = 0 + cdef int64_t duplicate_timestamp = 0 + cdef bint has_last = False + cdef bint duplicate_found = False + with nogil: + for part in range(part_count): + if lengths[part] > 0: + _push(heap, &heap_size, part, flat_times, starts, cursors) + while heap_size > 0: + part = _pop(heap, &heap_size, flat_times, starts, cursors) + source = starts[part] + cursors[part] + if has_last and flat_times[source] == last_timestamp: + duplicate_found = True + duplicate_timestamp = flat_times[source] + break + output_times[out_index] = flat_times[source] + output_values[out_index] = flat_values[source] + last_timestamp = flat_times[source] + has_last = True + out_index += 1 + cursors[part] += 1 + if cursors[part] < lengths[part]: + _push(heap, &heap_size, part, flat_times, starts, cursors) + if duplicate_found: + raise ValueError( + f"Duplicate timestamp {duplicate_timestamp} found across shards." + ) + return output_times_array, output_values_array + + +cpdef scatter_timeline_columns(object union_timestamps, object source_timestamps, + list value_arrays, list column_indices, + object output_values): + cdef cnp.ndarray[cnp.int64_t, ndim=1] union_array = np.ascontiguousarray( + union_timestamps, dtype=np.int64 + ) + cdef cnp.ndarray[cnp.int64_t, ndim=1] source_array = np.ascontiguousarray( + source_timestamps, dtype=np.int64 + ) + cdef cnp.ndarray[cnp.float64_t, ndim=2] output_array = output_values + cdef cnp.ndarray[cnp.intp_t, ndim=1] positions_array = np.empty( + len(source_array), dtype=np.intp + ) + cdef const int64_t[:] union_view = union_array + cdef const int64_t[:] source_view = source_array + cdef cnp.intp_t[:] positions = positions_array + cdef double[:, :] output = output_array + cdef cnp.ndarray[cnp.float64_t, ndim=1] values_array + cdef const double[:] values + cdef Py_ssize_t row + cdef Py_ssize_t low + cdef Py_ssize_t high + cdef Py_ssize_t middle + cdef Py_ssize_t source_count = len(source_array) + cdef Py_ssize_t union_count = len(union_array) + cdef Py_ssize_t column_position + cdef int column_index + cdef bint missing = False + + with nogil: + for row in range(source_count): + low = 0 + high = union_count + while low < high: + middle = low + (high - low) // 2 + if union_view[middle] < source_view[row]: + low = middle + 1 + else: + high = middle + if low == union_count or union_view[low] != source_view[row]: + missing = True + break + positions[row] = low + if missing: + raise ValueError("source timestamp is absent from the aligned union") + if len(value_arrays) != len(column_indices): + raise ValueError("value arrays and column indices have different lengths") + + for column_position in range(len(value_arrays)): + values_array = np.ascontiguousarray( + value_arrays[column_position], dtype=np.float64 + ) + if len(values_array) != source_count: + raise ValueError("value array length does not match its timeline") + values = values_array + column_index = int(column_indices[column_position]) + if column_index < 0 or column_index >= output.shape[1]: + raise IndexError(column_index) + with nogil: + for row in range(source_count): + output[positions[row], column_index] = values[row] diff --git a/python/tsfile/dataset/dataframe.py b/python/tsfile/dataset/dataframe.py index a65b231bc..e3ce78372 100644 --- a/python/tsfile/dataset/dataframe.py +++ b/python/tsfile/dataset/dataframe.py @@ -19,12 +19,13 @@ """Top-level dataset accessors for TsFile shards.""" from collections import defaultdict +import contextlib from dataclasses import dataclass, field import heapq import os import sys +from types import SimpleNamespace from typing import Dict, List, Optional, Tuple, Union -import warnings import numpy as np @@ -123,20 +124,27 @@ def _series_lookup_hint(name: str) -> str: def _validate_table_schema( existing: TableEntry, incoming: TableEntry, file_path: str ) -> None: - """Reject same-name tables whose tag/field layout differs across shards.""" + """Reject same-name tables whose complete ordered schema differs.""" if ( - existing.tag_columns == incoming.tag_columns + existing.schema_columns + and incoming.schema_columns + and existing.schema_columns == incoming.schema_columns + ): + return + if ( + not existing.schema_columns + and not incoming.schema_columns + and existing.tag_columns == incoming.tag_columns and existing.tag_types == incoming.tag_types and existing.field_columns == incoming.field_columns + and existing.field_types == incoming.field_types ): return raise ValueError( f"Incompatible schema for table '{incoming.table_name}' in '{file_path}'. " - f"Expected tags={list(existing.tag_columns)}, tag_types={list(existing.tag_types)}, " - f"fields={list(existing.field_columns)} but found " - f"tags={list(incoming.tag_columns)}, tag_types={list(incoming.tag_types)}, " - f"fields={list(incoming.field_columns)}." + f"Expected columns={list(existing.schema_columns)} but found " + f"columns={list(incoming.schema_columns)}." ) @@ -169,6 +177,8 @@ def _merge_tree_table_entries(existing: TableEntry, incoming: TableEntry) -> Tab tag_columns=tag_columns, tag_types=tag_types, field_columns=tuple(field_columns), + field_types=(), + schema_columns=(), ) @@ -254,6 +264,64 @@ def _register_reader( index.series_shards[series_ref].append((reader, device_id, field_idx)) +def _validate_unique_shard_timestamps(index: _DataFrameCatalog) -> None: + """Reject overlapping shards that contain the same logical timestamp.""" + validated_timeline_pairs = set() + for series_ref in index.series: + fragments = [] + for reader, device_id, field_idx in index.series_shards[series_ref]: + stats = reader.catalog.series_stats_by_ref[(device_id, field_idx)] + timeline_identity = ( + reader.file_path, + device_id, + ( + stats.time_metadata_offset + if stats.layout + else stats.value_metadata_offset + ), + ( + stats.time_metadata_length + if stats.layout + else stats.value_metadata_length + ), + ) + fragments.append( + ( + stats.timeline_min_time, + stats.timeline_max_time, + timeline_identity, + reader, + device_id, + field_idx, + ) + ) + fragments.sort(key=lambda item: (item[0], item[1], item[2])) + for right_index, right in enumerate(fragments): + for left in fragments[:right_index]: + if left[1] < right[0]: + continue + overlap_start = max(left[0], right[0]) + overlap_end = min(left[1], right[1]) + if overlap_start > overlap_end: + continue + pair = tuple(sorted((left[2], right[2]))) + if pair in validated_timeline_pairs: + continue + left_times, _ = left[3].read_series_by_ref( + left[4], left[5], overlap_start, overlap_end + ) + right_times, _ = right[3].read_series_by_ref( + right[4], right[5], overlap_start, overlap_end + ) + duplicate = np.intersect1d(left_times, right_times, assume_unique=True) + if len(duplicate): + raise ValueError( + f"Duplicate timestamp {int(duplicate[0])} found across " + "TsFile shards while building the Dataset Index." + ) + validated_timeline_pairs.add(pair) + + def _build_runtime_series_stats(refs: List[SeriesRef]) -> dict: """Build shared-timeline series stats from native timeline metadata.""" min_time = None @@ -318,45 +386,56 @@ def _read_field_by_position( refs: List[SeriesRef], offset: int, limit: int, + cached_infos=None, + cached_locator_ids=None, ) -> Tuple[np.ndarray, np.ndarray]: """Read one logical series by global position without materializing timestamps for non-overlapping shards.""" if limit <= 0: return np.array([], dtype=np.int64), np.array([], dtype=np.float64) - infos = [] - for reader, device_id, field_idx in refs: - series_info = reader.get_series_info_by_ref(device_id, field_idx) - infos.append( - { - "length": series_info["timeline_length"], - "min_time": series_info["timeline_min_time"], - "max_time": series_info["timeline_max_time"], - "table_name": series_info["table_name"], - "column_name": series_info["column_name"], - "device_id": series_info["device_id"], - "field_idx": series_info["field_idx"], - "tag_columns": series_info["tag_columns"], - "tag_values": series_info["tag_values"], - } - ) + if cached_infos is None: + infos = [] + for reader, device_id, field_idx in refs: + series_info = reader.get_series_info_by_ref(device_id, field_idx) + infos.append( + { + "length": series_info["timeline_length"], + "min_time": series_info["timeline_min_time"], + "max_time": series_info["timeline_max_time"], + } + ) + else: + infos = cached_infos + if cached_locator_ids is None: + locator_ids = [None] * len(refs) + else: + if len(cached_locator_ids) != len(refs): + raise ValueError("cached locator count does not match series refs") + locator_ids = cached_locator_ids ordered = sorted( - zip(refs, infos), key=lambda item: (item[1]["min_time"], item[1]["max_time"]) + zip(refs, infos, locator_ids), + key=lambda item: (item[1]["min_time"], item[1]["max_time"]), ) - if _has_time_range_overlap([info for _, info in ordered]): + if _has_time_range_overlap([info for _, info, _ in ordered]): return _read_field_by_position_overlap(series_name, ordered, offset, limit) remaining_offset = offset remaining_limit = limit time_parts = [] value_parts = [] - for (reader, device_id, field_idx), info in ordered: + for (reader, device_id, field_idx), info, locator_id in ordered: shard_count = info["length"] if remaining_offset >= shard_count: remaining_offset -= shard_count continue local_limit = min(remaining_limit, shard_count - remaining_offset) - ts_arr, values = reader.read_series_by_row( - device_id, field_idx, remaining_offset, local_limit + ts_arr, values = _read_shard_by_position( + reader, + device_id, + field_idx, + locator_id, + remaining_offset, + local_limit, ) if len(ts_arr) > 0: time_parts.append(ts_arr) @@ -371,6 +450,16 @@ def _read_field_by_position( return np.concatenate(time_parts), np.concatenate(value_parts) +def _read_shard_by_position( + reader, device_id, field_idx, locator_id, offset, limit +) -> Tuple[np.ndarray, np.ndarray]: + if locator_id is not None: + read_at_locator = getattr(reader, "read_series_by_row_at_locator", None) + if read_at_locator is not None: + return read_at_locator(locator_id, offset, limit) + return reader.read_series_by_row(device_id, field_idx, offset, limit) + + def _has_time_range_overlap(infos: List[dict]) -> bool: previous_max = None for info in infos: @@ -388,12 +477,12 @@ def _has_time_range_overlap(infos: List[dict]) -> bool: def _read_field_by_position_overlap( series_name: str, - ordered: List[Tuple[SeriesRef, dict]], + ordered: List[Tuple[SeriesRef, dict, Optional[int]]], offset: int, limit: int, ) -> Tuple[np.ndarray, np.ndarray]: """Merge overlapping shard streams lazily until the requested global window is covered.""" - total_count = sum(info["length"] for _, info in ordered) + total_count = sum(info["length"] for _, info, _ in ordered) if offset >= total_count: return np.array([], dtype=np.int64), np.array([], dtype=np.float64) @@ -411,8 +500,13 @@ def fill_state(state_idx: int) -> bool: local_limit = min(chunk_size, remaining) reader, device_id, field_idx = state["ref"] - ts_arr, val_arr = reader.read_series_by_row( - device_id, field_idx, state["next_offset"], local_limit + ts_arr, val_arr = _read_shard_by_position( + reader, + device_id, + field_idx, + state["locator_id"], + state["next_offset"], + local_limit, ) state["next_offset"] += len(ts_arr) state["timestamps"] = ts_arr @@ -425,11 +519,12 @@ def fill_state(state_idx: int) -> bool: return False return True - for ref, info in ordered: + for ref, info, locator_id in ordered: state_idx = len(states) states.append( { "ref": ref, + "locator_id": locator_id, "length": info["length"], "next_offset": 0, "timestamps": np.array([], dtype=np.int64), @@ -568,9 +663,31 @@ def _query_aligned( _, table_entry, _ = self._df._get_series_components(series_ref) field_name = table_entry.field_columns[field_idx] - for reader, device_id, reader_field_idx in self._df._index.series_shards[ - series_ref - ]: + descriptor = self._df._index.series_shards.describe(series_ref) + for shard in descriptor.shards: + overlap_start = max(start_time, shard.min_time) + overlap_end = min(end_time, shard.max_time) + if shard.timeline_length <= 0 or overlap_start > overlap_end: + continue + if overlap_start <= shard.min_time and overlap_end >= shard.max_time: + estimated_rows = shard.timeline_length + elif shard.min_time == shard.max_time: + estimated_rows = 1 + else: + estimated_rows = max( + 1, + min( + shard.timeline_length, + int( + shard.timeline_length + * (overlap_end - overlap_start + 1) + / (shard.max_time - shard.min_time + 1) + ), + ), + ) + reader = shard.reader + device_id = shard.device_id + reader_field_idx = shard.column_id groups[(id(reader), device_id)].append( ( col_idx, @@ -579,22 +696,35 @@ def _query_aligned( series_names[col_idx], reader, device_id, + estimated_rows, ) ) - series_time_parts = defaultdict(list) - series_value_parts = defaultdict(list) - for entries in groups.values(): + def query_group(entries): reader = entries[0][4] device_id = entries[0][5] field_indices = list(dict.fromkeys(entry[1] for entry in entries)) ts_arr, field_vals = reader.read_device_fields_by_time_range( device_id, field_indices, start_time, end_time ) + return entries, ts_arr, field_vals + + group_entries = list(groups.values()) + group_results = self._df._runtime.map_query_groups( + query_group, + group_entries, + estimated_rows=[ + max(entry[6] for entry in entries) for entries in group_entries + ], + ) + + series_time_parts = defaultdict(list) + series_value_parts = defaultdict(list) + for entries, ts_arr, field_vals in group_results: if len(ts_arr) == 0: continue appended_series = set() - for _, _, field_name, series_name, _, _ in entries: + for _, _, field_name, series_name, _, _, _ in entries: if series_name in appended_series: continue appended_series.add(series_name) @@ -610,11 +740,12 @@ def _query_aligned( return build_aligned_matrix(series_names, series_data) def __getitem__(self, key) -> AlignedTimeseries: - start_time, end_time, series_refs, series_names = self._parse_key(key) - timestamps, values = self._query_aligned( - start_time, end_time, series_refs, series_names - ) - return AlignedTimeseries(timestamps, values, series_names) + with self._df._query_guard(): + start_time, end_time, series_refs, series_names = self._parse_key(key) + timestamps, values = self._query_aligned( + start_time, end_time, series_refs, series_names + ) + return AlignedTimeseries(timestamps, values, series_names) class TsFileDataFrame: @@ -628,6 +759,8 @@ def __init__(self, paths: Union[str, List[str]], show_progress: bool = True): self._is_view = False self._root = None self._closed = False + self._runtime = None + self._runtime_lease = None self._load_metadata() @classmethod @@ -641,40 +774,84 @@ def _from_subset( obj._paths = parent._paths obj._show_progress = parent._show_progress obj._readers = parent._readers - # Reuse the parent's full mapping but restrict the membership scope to - # the requested subset. subset_refs = list(series_refs) - parent_shards = parent._index.series_shards - subset_shards = {ref: parent_shards[ref] for ref in subset_refs} - obj._index = _DataFrameCatalog( + obj._index = SimpleNamespace( model=parent._index.model, table_entries=parent._index.table_entries, devices=parent._index.devices, device_index=parent._index.device_index, device_time_bounds=parent._index.device_time_bounds, series=subset_refs, - series_shards=subset_shards, + series_shards=parent._index.series_shards, ) + obj._runtime = parent._runtime + obj._runtime_lease = ( + parent._runtime_lease.clone() if parent._runtime_lease is not None else None + ) + obj._readers = parent._readers obj._closed = False return obj def _owner(self) -> "TsFileDataFrame": - return self._root if self._is_view else self + return self def _assert_open(self): - if self._owner()._closed: + if self._closed: raise RuntimeError("Current TsFileDataFrame is closed.") - def _load_metadata(self): - """Build the logical cross-file index and the derived per-series caches.""" - from .reader import TsFileSeriesReader - - if len(self._paths) >= 2: - self._load_metadata_parallel(TsFileSeriesReader) + @contextlib.contextmanager + def _query_guard(self): + self._assert_open() + if self._runtime_lease is None: + yield else: - self._load_metadata_serial(TsFileSeriesReader) + with self._runtime_lease.query_lease(): + yield - if not self._index.series: + def _load_metadata(self): + """Map a valid persistent index, or build it once under a file lock.""" + from .reader import TsFileSeriesReader + from .index import ( + build_index_from_dataframe, + index_matches_paths, + index_path_for, + ) + from .runtime import DatasetRuntime + + index_path = index_path_for(self._paths) + if not index_matches_paths(index_path, self._paths): + lock_path = index_path + ".lock" + os.makedirs(os.path.dirname(lock_path) or ".", exist_ok=True) + with open(lock_path, "a+b") as lock_file: + try: + import fcntl + + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) + except ImportError: + pass + if not index_matches_paths(index_path, self._paths): + if len(self._paths) >= 2: + self._load_metadata_parallel(TsFileSeriesReader) + else: + self._load_metadata_serial(TsFileSeriesReader) + + if not self._index.series: + raise ValueError( + "No valid time series found in the provided TsFile files" + ) + try: + _validate_unique_shard_timestamps(self._index) + build_index_from_dataframe(self, index_path) + finally: + for reader in self._readers.values(): + reader.close() + self._readers.clear() + + self._runtime = DatasetRuntime(index_path) + self._runtime_lease = self._runtime.lease() + self._index = self._runtime.catalog + if len(self._index.series) == 0: + self._runtime_lease.close() raise ValueError("No valid time series found in the provided TsFile files") def _show_loading_progress(self, done: int, total: int, total_series: int = None): @@ -748,14 +925,37 @@ def _get_series_components( device_key = self._index.devices[device_idx] return device_key, self._index.table_entries[device_key[0]], field_idx - def _build_series_name(self, series_ref: SeriesRefKey) -> SeriesPath: + def _build_series_name( + self, series_ref: SeriesRefKey, series_id: Optional[int] = None + ) -> SeriesPath: device_key, table_entry, field_idx = self._get_series_components(series_ref) table_name, tag_values = device_key field_name = table_entry.field_columns[field_idx] return build_logical_series_path( - table_name, tag_values, field_name, table_entry.tag_columns + table_name, + tag_values, + field_name, + table_entry.tag_columns, + index_identity=getattr(self._index, "index_identity", None), + series_id=series_id, ) + def _descriptor_from_series_path(self, series_path: SeriesPath): + identity = getattr(series_path, "_index_identity", None) + series_id = getattr(series_path, "_series_id", None) + resolver = getattr(self._index, "resolve_series_descriptor_by_id", None) + if ( + identity is None + or series_id is None + or resolver is None + or identity != getattr(self._index, "index_identity", None) + ): + return None + try: + return resolver(series_id, series_path.table, series_path.field) + except (IndexError, KeyError, ValueError): + return None + def _resolve_series_name(self, series_name) -> SeriesRefKey: """Resolve a ``SeriesPath`` or path string (``\\N`` = null tag) to a ref. @@ -763,6 +963,9 @@ def _resolve_series_name(self, series_name) -> SeriesRefKey: direct lookup -- no sparse/compressed fallback and no ambiguity. """ if isinstance(series_name, SeriesPath): + descriptor = self._descriptor_from_series_path(series_name) + if descriptor is not None: + return descriptor.ref table_name, tag_parts, field_name = ( series_name.table, list(series_name.tags), @@ -785,7 +988,15 @@ def _resolve_series_name(self, series_name) -> SeriesRefKey: except ValueError as exc: raise KeyError(_series_lookup_hint(series_name)) from exc - device_key = (table_name, _normalize_tag_values(tag_parts)) + normalized_tags = _normalize_tag_values(tag_parts) + resolver = getattr(self._index, "resolve_series_descriptor", None) + if resolver is not None: + try: + return resolver(table_name, normalized_tags, field_name).ref + except (KeyError, ValueError) as exc: + raise KeyError(_series_lookup_hint(series_name)) from exc + + device_key = (table_name, normalized_tags) device_idx = self._index.device_index.get(device_key) if device_idx is None: raise KeyError(_series_lookup_hint(series_name)) @@ -828,8 +1039,14 @@ def model(self) -> str: def list_timeseries(self, path_prefix: str = "") -> List[SeriesPath]: if not path_prefix: + series = self._index.series + series_id_at = getattr(series, "series_id", None) return [ - self._build_series_name(series_ref) for series_ref in self._index.series + self._build_series_name( + series[position], + None if series_id_at is None else series_id_at(position), + ) + for position in range(len(series)) ] try: @@ -838,7 +1055,9 @@ def list_timeseries(self, path_prefix: str = "") -> List[SeriesPath]: return [] matched = [] - for series_ref in self._index.series: + series = self._index.series + series_id_at = getattr(series, "series_id", None) + for position, series_ref in enumerate(series): device_key, table_entry, field_idx = self._get_series_components(series_ref) components = build_logical_series_components( table_entry.table_name, @@ -847,7 +1066,12 @@ def list_timeseries(self, path_prefix: str = "") -> List[SeriesPath]: table_entry.tag_columns, ) if prefix_parts == components[: len(prefix_parts)]: - matched.append(self._build_series_name(series_ref)) + matched.append( + self._build_series_name( + series_ref, + None if series_id_at is None else series_id_at(position), + ) + ) return matched def list_timeseries_metadata(self, path_prefix: str = ""): @@ -910,20 +1134,58 @@ def list_timeseries_metadata(self, path_prefix: str = ""): ordered_columns.append(extra) return df.reindex(columns=ordered_columns) - def _get_timeseries(self, series_ref: SeriesRefKey) -> Timeseries: + def _get_timeseries( + self, + series_ref: SeriesRefKey, + descriptor=None, + series_name: Optional[SeriesPath] = None, + ) -> Timeseries: self._assert_open() - series_name = self._build_series_name(series_ref) + if series_name is None: + series_name = self._build_series_name(series_ref) + if descriptor is None: + describe = getattr(self._index.series_shards, "describe", None) + if describe is not None: + descriptor = describe(series_ref) + if descriptor is None: + refs = self._index.series_shards[series_ref] + stats = _build_runtime_series_stats(refs) + cached_infos = None + cached_locator_ids = None + else: + refs = list(descriptor.refs) + stats = { + "min_time": descriptor.min_time, + "max_time": descriptor.max_time, + "count": descriptor.count, + } + cached_infos = tuple( + { + "length": shard.timeline_length, + "min_time": shard.min_time, + "max_time": shard.max_time, + } + for shard in descriptor.shards + ) + cached_locator_ids = tuple(shard.locator_id for shard in descriptor.shards) + runtime_lease = ( + self._runtime_lease.clone() if self._runtime_lease is not None else None + ) return Timeseries( series_name, - self._index.series_shards[series_ref], - _build_runtime_series_stats(self._index.series_shards[series_ref]), - self._assert_open, - lambda: _merge_field_timestamps( - series_name, self._index.series_shards[series_ref] - ), + refs, + stats, + self._assert_open if runtime_lease is None else None, + lambda: _merge_field_timestamps(series_name, refs), lambda offset, limit: _read_field_by_position( - series_name, self._index.series_shards[series_ref], offset, limit + series_name, + refs, + offset, + limit, + cached_infos, + cached_locator_ids, ), + runtime_lease=runtime_lease, ) def __getitem__(self, key): @@ -944,12 +1206,36 @@ def __getitem__(self, key): raise IndexError( f"Index {idx} out of range [0, {len(self._index.series)})" ) - return self._get_timeseries(self._index.series[idx]) + series_ref = self._index.series[idx] + descriptor = None + series_id_at = getattr(self._index.series, "series_id", None) + describe = getattr(self._index.series_shards, "describe", None) + if series_id_at is not None and describe is not None: + descriptor = describe(series_ref, series_id=series_id_at(idx)) + return self._get_timeseries(series_ref, descriptor) + + if isinstance(key, SeriesPath): + descriptor = self._descriptor_from_series_path(key) + if descriptor is not None: + return self._get_timeseries(descriptor.ref, descriptor, key) if isinstance(key, str): try: - return self._get_timeseries(self._resolve_series_name(key)) - except KeyError: + resolver = getattr(self._index, "resolve_series_descriptor", None) + if resolver is None: + return self._get_timeseries(self._resolve_series_name(key)) + + parts = split_logical_series_path(key) + if len(parts) < 2: + raise KeyError(key) + table_name, field_name = parts[0], parts[-1] + descriptor = resolver( + table_name, + _normalize_tag_values(parts[1:-1]), + field_name, + ) + return self._get_timeseries(descriptor.ref, descriptor) + except (KeyError, ValueError): pass valid_columns = {"field", "start_time", "end_time", "count"} @@ -1087,24 +1373,19 @@ def show(self, max_rows: int = 20): print(self._repr_header() + self._format_table(max_rows=max_rows)) def close(self): - if self._is_view: - warnings.warn( - "close() on a subset TsFileDataFrame is a no-op; only the root dataframe owns the readers.", - RuntimeWarning, - stacklevel=2, - ) - return if self._closed: return - for reader in self._readers.values(): - reader.close() - self._readers.clear() self._closed = True + if self._runtime_lease is not None: + self._runtime_lease.close() + else: + for reader in self._readers.values(): + reader.close() + self._readers.clear() def __del__(self): try: - if not getattr(self, "_is_view", False): - self.close() + self.close() except Exception: pass diff --git a/python/tsfile/dataset/index.py b/python/tsfile/dataset/index.py new file mode 100644 index 000000000..394860d54 --- /dev/null +++ b/python/tsfile/dataset/index.py @@ -0,0 +1,815 @@ +# 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. +# + +"""Persistent mmap-backed Dataset Index v1. + +The binary layout in this module is intentionally identical to +``cpp/src/dataset/dataset_index.h``. Python owns cold-build orchestration and +public object construction; hot-path catalog lookup reads packed records +directly from a read-only mmap without recreating the old dictionary graph. +""" + +from __future__ import annotations + +from collections import defaultdict +import contextlib +import mmap +import os +import struct +from typing import Dict, Iterable, Iterator, List, Mapping, Optional, Sequence, Tuple + +from ..constants import ColumnCategory +from .metadata import MODEL_TREE, _join_series_path + +MAGIC = b"TSIDX\0\0\0" +VERSION_MAJOR = 1 +VERSION_MINOR = 0 +HEADER_SIZE = 64 +DIRECTORY_ENTRY_SIZE = 32 +SECTION_COUNT = 13 +ALIGNMENT = 64 +INDEX_FILE_NAME = ".tsfile_dataframe_index.tsidx" + +STRING_OFFSETS = 1 +STRING_BYTES = 2 +TABLE_NAME_INDEX = 3 +TABLE_RECORD = 4 +DEVICE_NAME_INDEX = 5 +DEVICE_RECORD = 6 +COLUMN_NAME_INDEX = 7 +COLUMN_SCHEMA = 8 +LOGICAL_SERIES = 9 +TSFILE_RECORD = 10 +DEVICE_FILE_SPAN = 11 +SERIES_FILE_SPAN = 12 +SERIES_LOCATOR = 13 + +HEADER = struct.Struct("<8sHHIQIIQI20s") +DIRECTORY = struct.Struct(" int: + return (value + ALIGNMENT - 1) & ~(ALIGNMENT - 1) + + +def name_hash(value: bytes) -> int: + """Return the format-v1 FNV-1a hash (matching the C++ implementation).""" + result = 1469598103934665603 + for byte in value: + result ^= byte + result = (result * 1099511628211) & 0xFFFFFFFFFFFFFFFF + return result + + +_CRC32C_TABLE = [] +for _value in range(256): + _crc = _value + for _ in range(8): + _crc = (_crc >> 1) ^ (0x82F63B78 if _crc & 1 else 0) + _CRC32C_TABLE.append(_crc) + + +def crc32c(data) -> int: + crc = 0xFFFFFFFF + for byte in data: + crc = _CRC32C_TABLE[(crc ^ byte) & 0xFF] ^ (crc >> 8) + return crc ^ 0xFFFFFFFF + + +def file_fingerprint(path: str, stat_result=None) -> int: + """Cheap sealed-dataset generation fingerprint. + + v1 combines size and nanosecond mtime. ReaderSession rechecks the same + tuple before interpreting a locator; content replacement therefore cannot + silently reuse an active generation under the static Dataset contract. + """ + st = os.stat(path) if stat_result is None else stat_result + return name_hash(struct.pack(" int: + result = self._ids.get(value) + if result is not None: + return result + encoded = value.encode("utf-8") + result = len(self._values) + self._ids[value] = result + self._values.append(encoded) + return result + + def sections(self) -> Tuple[bytes, bytes]: + offsets = bytearray(RECORDS[STRING_OFFSETS].size * (len(self._values) + 1)) + contents = bytearray() + for index, value in enumerate(self._values): + RECORDS[STRING_OFFSETS].pack_into(offsets, index * 4, len(contents)) + contents.extend(value) + RECORDS[STRING_OFFSETS].pack_into(offsets, len(self._values) * 4, len(contents)) + if len(contents) >= 1 << 32: + raise OverflowError("Dataset Index v1 string pool exceeds 4 GiB") + return bytes(offsets), bytes(contents) + + +def _pack_records(section_type: int, rows: Iterable[tuple]) -> bytes: + record = RECORDS[section_type] + rows = list(rows) + result = bytearray(record.size * len(rows)) + for index, row in enumerate(rows): + record.pack_into(result, index * record.size, *row) + return bytes(result) + + +def write_index_atomic(path: str, section_payloads: Mapping[int, bytes]) -> None: + """Write, fsync, validate, and atomically publish one format-v1 index.""" + if set(section_payloads) != set(range(1, SECTION_COUNT + 1)): + raise ValueError("Dataset Index v1 requires exactly 13 sections") + directory_offset = HEADER_SIZE + cursor = _align64(HEADER_SIZE + SECTION_COUNT * DIRECTORY_ENTRY_SIZE) + entries = [] + for section_type in range(1, SECTION_COUNT + 1): + payload = section_payloads[section_type] + record_size = 0 if section_type == STRING_BYTES else RECORDS[section_type].size + if record_size: + if len(payload) % record_size: + raise ValueError(f"section {section_type} has a partial record") + count = len(payload) // record_size + else: + count = len(payload) + if count >= 1 << 32: + raise OverflowError(f"section {section_type} count exceeds uint32") + entries.append( + (section_type, record_size, cursor, len(payload), count, crc32c(payload)) + ) + cursor = _align64(cursor + len(payload)) + + file_length = entries[-1][2] + entries[-1][3] + header_without_crc = HEADER.pack( + MAGIC, + VERSION_MAJOR, + VERSION_MINOR, + HEADER_SIZE, + directory_offset, + SECTION_COUNT, + DIRECTORY_ENTRY_SIZE, + file_length, + 0, + b"\0" * 20, + ) + header = bytearray(header_without_crc) + struct.pack_into(" len(self._view): + raise ValueError("Dataset Index directory is out of range") + + entries = {} + previous_end = _align64(directory_end) + for index in range(section_count): + entry = DIRECTORY.unpack_from( + self._view, directory_offset + index * DIRECTORY_ENTRY_SIZE + ) + section_type, record_size, offset, length, count, checksum = entry + if section_type != STRING_BYTES and section_type not in RECORDS: + raise ValueError("invalid Dataset Index section type") + expected_size = ( + 0 if section_type == STRING_BYTES else RECORDS[section_type].size + ) + if section_type != index + 1 or record_size != expected_size: + raise ValueError("invalid Dataset Index section directory") + if ( + offset % ALIGNMENT + or offset < previous_end + or offset + length > len(self._view) + or (record_size and count * record_size > length) + or (not record_size and count != length) + ): + raise ValueError("invalid Dataset Index section range") + if ( + verify_sections + and crc32c(self._view[offset : offset + length]) != checksum + ): + raise ValueError( + f"Dataset Index section {section_type} checksum mismatch" + ) + entries[section_type] = entry + previous_end = offset + length + + offsets = entries[STRING_OFFSETS] + strings = entries[STRING_BYTES] + if offsets[4] == 0: + raise ValueError("StringOffsets lacks terminal offset") + previous = 0 + for index in range(offsets[4]): + current = RECORDS[STRING_OFFSETS].unpack_from( + self._view, offsets[2] + index * 4 + )[0] + if current < previous or current > strings[3]: + raise ValueError("invalid Dataset Index string offset") + previous = current + if previous != strings[3]: + raise ValueError("Dataset Index terminal string offset mismatch") + return entries + + def close(self): + if getattr(self, "_view", None) is not None: + self._view.release() + self._view = None + if getattr(self, "_mmap", None) is not None: + self._mmap.close() + self._mmap = None + if getattr(self, "_file", None) is not None: + self._file.close() + self._file = None + + def __enter__(self): + return self + + def __exit__(self, *_): + self.close() + + def count(self, section_type: int) -> int: + return self._entries[section_type][4] + + def record(self, section_type: int, record_id: int) -> tuple: + entry = self._entries[section_type] + if record_id < 0 or record_id >= entry[4]: + raise IndexError(record_id) + return RECORDS[section_type].unpack_from( + self._view, entry[2] + record_id * entry[1] + ) + + def records(self, section_type: int, first: int = 0, count: Optional[int] = None): + total = self.count(section_type) + end = total if count is None else first + count + if first < 0 or end < first or end > total: + raise IndexError((first, count)) + for record_id in range(first, end): + yield self.record(section_type, record_id) + + def string_bytes(self, sid: int) -> bytes: + offsets = self._entries[STRING_OFFSETS] + strings = self._entries[STRING_BYTES] + if sid < 0 or sid + 1 >= offsets[4]: + raise IndexError(sid) + start = RECORDS[STRING_OFFSETS].unpack_from(self._view, offsets[2] + sid * 4)[0] + end = RECORDS[STRING_OFFSETS].unpack_from( + self._view, offsets[2] + (sid + 1) * 4 + )[0] + return bytes(self._view[strings[2] + start : strings[2] + end]) + + def string(self, sid: int) -> str: + return self.string_bytes(sid).decode("utf-8") + + def _equal_hash_range(self, section_type: int, hash_index: int, value_hash: int): + low = 0 + high = self.count(section_type) + while low < high: + middle = low + (high - low) // 2 + if self.record(section_type, middle)[hash_index] < value_hash: + low = middle + 1 + else: + high = middle + first = low + while ( + low < self.count(section_type) + and self.record(section_type, low)[hash_index] == value_hash + ): + low += 1 + return first, low + + def find_table_ids(self, name: str) -> List[int]: + encoded = name.encode("utf-8") + first, end = self._equal_hash_range(TABLE_NAME_INDEX, 0, name_hash(encoded)) + return [ + table_id + for _, sid, table_id in self.records(TABLE_NAME_INDEX, first, end - first) + 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]) + + 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]) + + 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) + + +def index_path_for(paths: Sequence[str]) -> str: + common = os.path.commonpath([os.path.abspath(path) for path in paths]) + if not os.path.isdir(common): + common = os.path.dirname(common) + return os.path.join(common, INDEX_FILE_NAME) + + +def index_matches_paths(path: str, paths: Sequence[str]) -> bool: + """Return whether an index describes exactly the current sealed files.""" + expected = sorted(os.path.abspath(item) for item in paths) + try: + with MappedDatasetIndex(path) as index: + if index.count(TSFILE_RECORD) != len(expected): + return False + actual = [] + for file_id in range(index.count(TSFILE_RECORD)): + record = index.record(TSFILE_RECORD, file_id) + file_path = index.string(record[0]) + actual.append(file_path) + st = os.stat(file_path) + if ( + st.st_size != record[2] + or file_fingerprint(file_path, st) != record[3] + ): + return False + return actual == expected + except (OSError, ValueError, UnicodeError, IndexError): + return False + + +def build_sections_from_dataframe(dataframe) -> Mapping[int, bytes]: + """Convert one fully scanned legacy DataFrame into deterministic v1 sections. + + This is the cold-build bridge. Once published, subsequent constructions + use :class:`MappedDatasetIndex` and do not recreate this object graph. + """ + index = dataframe._index + pool = _StringPool() + file_paths = sorted(dataframe._readers) + file_ids = {path: file_id for file_id, path in enumerate(file_paths)} + + table_names = sorted(index.table_entries) + table_ids = {name: table_id for table_id, name in enumerate(table_names)} + field_types: Dict[Tuple[str, str], int] = {} + for series_ref, shards in index.series_shards.items(): + device_idx, field_idx = series_ref + table_name, _ = index.devices[device_idx] + table_entry = index.table_entries[table_name] + field_name = table_entry.field_columns[field_idx] + declared_type = ( + int(table_entry.field_types[field_idx]) + if field_idx < len(table_entry.field_types) + else None + ) + for reader, local_device, local_field in shards: + stats = reader.catalog.series_stats_by_ref[(local_device, local_field)] + if declared_type is not None and declared_type != stats.data_type: + raise ValueError( + f"physical type for {table_name}.{field_name} does not match " + "its canonical TableSchema" + ) + current = field_types.setdefault((table_name, field_name), stats.data_type) + if current != stats.data_type: + raise ValueError( + f"incompatible physical type for {table_name}.{field_name}" + ) + + columns = [] + columns_by_table = defaultdict(list) + for table_name in table_names: + table_id = table_ids[table_name] + table = index.table_entries[table_name] + if table.schema_columns: + definitions = [ + (ordinal, name, data_type, category) + for ordinal, (name, data_type, category) in enumerate( + table.schema_columns + ) + if category != int(ColumnCategory.TIME) + ] + else: + definitions = [ + (ordinal, name, int(data_type), int(ColumnCategory.TAG)) + for ordinal, (name, data_type) in enumerate( + zip(table.tag_columns, table.tag_types) + ) + ] + first_field_ordinal = len(definitions) + definitions.extend( + [ + ( + first_field_ordinal + ordinal, + name, + field_types.get((table_name, name), -1), + int(ColumnCategory.FIELD), + ) + for ordinal, name in enumerate(table.field_columns) + ] + ) + for ordinal, name, data_type, role in definitions: + column_id = len(columns) + columns.append( + ( + table_id, + pool.intern(name), + ordinal, + data_type, + data_type, + 0, + 0, + role, + 1, + 0, + ) + ) + columns_by_table[table_id].append((name, column_id)) + + device_specs = [] + old_to_new_device = {} + for old_device_id, (table_name, tags) in enumerate(index.devices): + canonical_name = _join_series_path(table_name, tags, "") + device_specs.append( + (table_ids[table_name], canonical_name, old_device_id, tags) + ) + device_specs.sort(key=lambda item: (item[0], item[1].encode("utf-8"))) + for new_device_id, (_, _, old_device_id, _) in enumerate(device_specs): + old_to_new_device[old_device_id] = new_device_id + + series_specs = [] + for old_device_id, field_idx in index.series: + table_name, _ = index.devices[old_device_id] + field_name = index.table_entries[table_name].field_columns[field_idx] + column_id = dict(columns_by_table[table_ids[table_name]])[field_name] + series_specs.append( + (old_to_new_device[old_device_id], column_id, (old_device_id, field_idx)) + ) + series_specs.sort(key=lambda item: (item[0], item[1])) + old_series_to_new = {old: new for new, (_, _, old) in enumerate(series_specs)} + + locators = [] + series_span_rows = [] + device_span_specs = {} + series_spans_by_id = defaultdict(list) + device_spans_by_device = defaultdict(list) + + for _, _, old_series_ref in series_specs: + series_id = old_series_to_new[old_series_ref] + old_device_id, _ = old_series_ref + device_id = old_to_new_device[old_device_id] + for reader, local_device, local_field in index.series_shards[old_series_ref]: + file_id = file_ids[reader.file_path] + stats = reader.catalog.series_stats_by_ref[(local_device, local_field)] + if stats.value_metadata_length <= 0: + raise ValueError( + f"missing exact TimeseriesMetadata locator in {reader.file_path}" + ) + if stats.layout and not (stats.locator_flags & 1): + raise ValueError( + f"aligned time/value chunk metadata mismatch in {reader.file_path}" + ) + span_key = (device_id, file_id) + candidate = ( + stats.time_metadata_offset if stats.layout else 0, + stats.time_metadata_length if stats.layout else 0, + stats.layout, + 1 if stats.layout and stats.locator_flags & 1 else 0, + stats.timeline_length if stats.layout else 0, + ) + previous = device_span_specs.get(span_key) + if previous is None: + device_span_specs[span_key] = candidate + else: + if previous[:4] != candidate[:4]: + raise ValueError( + "inconsistent aligned time locator for one device/file" + ) + if previous[4] != candidate[4]: + raise ValueError("inconsistent aligned timeline row count") + locator_id = len(locators) + locators.append( + ( + span_key, + stats.layout, + 0, + stats.value_metadata_offset, + stats.value_metadata_length, + 0, + ) + ) + series_spans_by_id[series_id].append( + [ + series_id, + file_id, + locator_id, + 0, + stats.timeline_min_time, + stats.timeline_max_time, + stats.length if not stats.layout else 0, + ] + ) + + device_span_rows = [] + device_span_ids = {} + for key, values in sorted(device_span_specs.items()): + device_id, file_id = key + device_span_ids[key] = len(device_span_rows) + device_span_rows.append((device_id, file_id, *values)) + device_spans_by_device[device_id].append(device_span_ids[key]) + locator_rows = [ + (device_span_ids[key], kind, flags, offset, length, padding) + for key, kind, flags, offset, length, padding in locators + ] + + logical_series_rows = [] + for series_id, (device_id, column_id, _) in enumerate(series_specs): + spans = sorted(series_spans_by_id[series_id], key=lambda row: (row[4], row[1])) + first = len(series_span_rows) + for row in spans: + series_span_rows.append(tuple(row)) + logical_series_rows.append( + ( + device_id, + column_id, + first, + len(spans), + min(row[4] for row in spans), + max(row[5] for row in spans), + ) + ) + + device_rows = [] + device_names_by_table = defaultdict(list) + series_by_device = defaultdict(list) + for series_id, row in enumerate(logical_series_rows): + series_by_device[row[0]].append(series_id) + device_span_flat = [] + # DeviceFileSpan is already globally sorted by device, so each device range + # is contiguous and its first id can be recorded directly. + for device_id, (table_id, canonical_name, _, _) in enumerate(device_specs): + name_sid = pool.intern(canonical_name) + device_names_by_table[table_id].append((canonical_name, device_id, name_sid)) + series_ids = series_by_device[device_id] + span_ids = device_spans_by_device[device_id] + first_series = series_ids[0] if series_ids else 0 + first_span = span_ids[0] if span_ids else 0 + if series_ids: + minimum = min(logical_series_rows[sid][4] for sid in series_ids) + maximum = max(logical_series_rows[sid][5] for sid in series_ids) + else: + minimum = maximum = 0 + device_rows.append( + ( + table_id, + name_sid, + 0, + 0, + first_series, + len(series_ids), + first_span, + len(span_ids), + minimum, + maximum, + ) + ) + + device_name_rows = [] + column_name_rows = [] + table_rows = [] + table_name_rows = [] + for table_name in table_names: + table_id = table_ids[table_name] + table = index.table_entries[table_name] + name_sid = pool.intern(table_name) + table_name_rows.append( + (name_hash(table_name.encode("utf-8")), name_sid, table_id) + ) + first_device = len(device_name_rows) + names = sorted( + device_names_by_table[table_id], + key=lambda item: ( + name_hash(item[0].encode("utf-8")), + item[0].encode("utf-8"), + ), + ) + for name, device_id, sid in names: + device_name_rows.append( + (table_id, device_id, name_hash(name.encode("utf-8")), sid, 0) + ) + first_column = len(column_name_rows) + column_names = sorted( + columns_by_table[table_id], + key=lambda item: ( + name_hash(item[0].encode("utf-8")), + item[0].encode("utf-8"), + ), + ) + for name, column_id in column_names: + column_name_rows.append( + ( + table_id, + column_id, + name_hash(name.encode("utf-8")), + pool.intern(name), + 0, + ) + ) + table_rows.append( + ( + name_sid, + 0, + first_device, + len(names), + first_column, + len(column_names), + 0, + ) + ) + table_name_rows.sort(key=lambda row: (row[0], pool._values[row[1]], row[2])) + + file_rows = [] + for file_id, path in enumerate(file_paths): + st = os.stat(path) + file_rows.append( + ( + pool.intern(path), + 0, + st.st_size, + file_fingerprint(path, st), + 0, + ) + ) + + string_offsets, string_bytes = pool.sections() + return { + STRING_OFFSETS: string_offsets, + STRING_BYTES: string_bytes, + TABLE_NAME_INDEX: _pack_records(TABLE_NAME_INDEX, table_name_rows), + TABLE_RECORD: _pack_records(TABLE_RECORD, table_rows), + DEVICE_NAME_INDEX: _pack_records(DEVICE_NAME_INDEX, device_name_rows), + DEVICE_RECORD: _pack_records(DEVICE_RECORD, device_rows), + COLUMN_NAME_INDEX: _pack_records(COLUMN_NAME_INDEX, column_name_rows), + COLUMN_SCHEMA: _pack_records(COLUMN_SCHEMA, columns), + LOGICAL_SERIES: _pack_records(LOGICAL_SERIES, logical_series_rows), + TSFILE_RECORD: _pack_records(TSFILE_RECORD, file_rows), + DEVICE_FILE_SPAN: _pack_records(DEVICE_FILE_SPAN, device_span_rows), + SERIES_FILE_SPAN: _pack_records(SERIES_FILE_SPAN, series_span_rows), + SERIES_LOCATOR: _pack_records(SERIES_LOCATOR, locator_rows), + } + + +def build_index_from_dataframe(dataframe, path: Optional[str] = None) -> str: + path = index_path_for(dataframe._paths) if path is None else path + write_index_atomic(path, build_sections_from_dataframe(dataframe)) + return path diff --git a/python/tsfile/dataset/merge.py b/python/tsfile/dataset/merge.py index 8d70dc552..1101a13ca 100644 --- a/python/tsfile/dataset/merge.py +++ b/python/tsfile/dataset/merge.py @@ -24,11 +24,16 @@ - duplicate timestamps for the same logical series across shards are rejected. """ -import heapq from typing import Dict, List, Tuple import numpy as np +from ._merge import ( + merge_time_value_parts_overlap, + merge_timestamp_parts_overlap, + scatter_timeline_columns, +) + def merge_timestamp_parts( time_parts: List[np.ndarray], @@ -49,39 +54,7 @@ def merge_timestamp_parts( ): return np.concatenate(parts) - total_length = sum(len(ts_part) for ts_part in parts) - merged = np.empty(total_length, dtype=np.int64) - - heap = [(int(ts_part[0]), part_idx, 0) for part_idx, ts_part in enumerate(parts)] - heapq.heapify(heap) - - out_idx = 0 - last_ts = None - while heap: - ts, part_idx, offset = heapq.heappop(heap) - - if last_ts is not None and ts == last_ts: - if validate_unique: - raise ValueError(f"Duplicate timestamp {ts} found across shards.") - if not deduplicate: - merged[out_idx] = ts - out_idx += 1 - else: - merged[out_idx] = ts - out_idx += 1 - last_ts = ts - - next_offset = offset + 1 - if next_offset < len(parts[part_idx]): - heapq.heappush( - heap, (int(parts[part_idx][next_offset]), part_idx, next_offset) - ) - - if validate_unique: - return merged[:out_idx] - if deduplicate: - return merged[:out_idx] - return merged[:out_idx] + return merge_timestamp_parts_overlap(parts, deduplicate, validate_unique) def merge_time_value_parts( @@ -90,8 +63,8 @@ def merge_time_value_parts( ) -> Tuple[np.ndarray, np.ndarray]: """Merge sorted time/value parts for one logical series. - Duplicate timestamps are validated during metadata loading, so the query - path can assume each part is already sorted and conflict-free. + Duplicate timestamps are validated during metadata loading and again by + the overlapping query merge as a defense against stale or external indexes. Fast path: if shard ranges do not overlap in time, concatenate in shard order after sorting parts by their first timestamp. @@ -117,29 +90,7 @@ def merge_time_value_parts( ): return np.concatenate(time_parts), np.concatenate(value_parts) - total_length = sum(len(ts_part) for ts_part in time_parts) - merged_ts = np.empty(total_length, dtype=np.int64) - merged_vals = np.empty(total_length, dtype=np.float64) - - heap = [ - (int(ts_part[0]), part_idx, 0) for part_idx, ts_part in enumerate(time_parts) - ] - heapq.heapify(heap) - - out_idx = 0 - while heap: - _, part_idx, offset = heapq.heappop(heap) - merged_ts[out_idx] = time_parts[part_idx][offset] - merged_vals[out_idx] = value_parts[part_idx][offset] - out_idx += 1 - - next_offset = offset + 1 - if next_offset < len(time_parts[part_idx]): - heapq.heappush( - heap, (int(time_parts[part_idx][next_offset]), part_idx, next_offset) - ) - - return merged_ts, merged_vals + return merge_time_value_parts_overlap(time_parts, value_parts) def build_aligned_matrix( @@ -150,20 +101,75 @@ def build_aligned_matrix( Each input series is assumed to already satisfy the dataset merge policy, meaning its timestamp array is unique within that logical series. """ - all_ts_arrays = [ts for ts, _ in series_data.values() if len(ts) > 0] - if not all_ts_arrays: + # Aligned measurements from one device intentionally share the same + # timestamp ndarray. Collapse those identities first, then also collapse + # exact-equal timelines from different devices. Otherwise the generic + # implementation feeds the same 2,880 timestamps into the union once per + # selected measurement and repeats np.searchsorted for every column. + timeline_groups = [] + groups_by_identity = {} + for col_idx, name in enumerate(series_names): + item = series_data.get(name) + if item is None or len(item[0]) == 0: + continue + ts_arr, val_arr = item + group = groups_by_identity.get(id(ts_arr)) + if group is None: + group = [ts_arr, []] + groups_by_identity[id(ts_arr)] = group + timeline_groups.append(group) + group[1].append((col_idx, val_arr)) + + if not timeline_groups: return np.array([], dtype=np.int64), np.empty((0, len(series_names))) - timestamps = merge_timestamp_parts(all_ts_arrays, deduplicate=True) + distinct_timelines = [] + equivalence_buckets = {} + for ts_arr, columns in timeline_groups: + last_index = len(ts_arr) - 1 + sample = tuple( + int(ts_arr[last_index * sample_index // 7]) for sample_index in range(1, 7) + ) + equivalence_key = (len(ts_arr), int(ts_arr[0]), int(ts_arr[-1]), sample) + equivalent = None + bucket = equivalence_buckets.setdefault(equivalence_key, []) + for current in bucket: + current_ts = current[0] + if np.array_equal(current_ts, ts_arr): + equivalent = current + break + if equivalent is None: + equivalent = [ts_arr, columns] + distinct_timelines.append(equivalent) + bucket.append(equivalent) + else: + equivalent[1].extend(columns) + + if len(distinct_timelines) == 1: + timestamps = distinct_timelines[0][0] + values = np.full((len(timestamps), len(series_names)), np.nan) + columns = distinct_timelines[0][1] + scatter_timeline_columns( + timestamps, + timestamps, + [val_arr for _, val_arr in columns], + [col_idx for col_idx, _ in columns], + values, + ) + return timestamps, values + + timestamps = merge_timestamp_parts( + [group[0] for group in distinct_timelines], deduplicate=True + ) values = np.full((len(timestamps), len(series_names)), np.nan) - for col_idx, name in enumerate(series_names): - if name not in series_data: - continue - ts_arr, val_arr = series_data[name] - if len(ts_arr) == 0: - continue - indices = np.searchsorted(timestamps, ts_arr) - values[indices, col_idx] = val_arr + for ts_arr, columns in distinct_timelines: + scatter_timeline_columns( + timestamps, + ts_arr, + [val_arr for _, val_arr in columns], + [col_idx for col_idx, _ in columns], + values, + ) return timestamps, values diff --git a/python/tsfile/dataset/metadata.py b/python/tsfile/dataset/metadata.py index 5eb367273..1aa665090 100644 --- a/python/tsfile/dataset/metadata.py +++ b/python/tsfile/dataset/metadata.py @@ -20,7 +20,7 @@ from dataclasses import dataclass, field import sys -from typing import Any, Dict, Iterable, Iterator, List, NamedTuple, Tuple +from typing import Any, Dict, Iterable, Iterator, List, NamedTuple, Optional, Tuple from ..constants import TSDataType @@ -41,6 +41,15 @@ class SeriesStats(NamedTuple): timeline_length: int timeline_min_time: int timeline_max_time: int + data_type: int = -1 + value_metadata_offset: int = 0 + value_metadata_length: int = 0 + time_metadata_offset: int = 0 + time_metadata_length: int = 0 + chunk_meta_count: int = 0 + time_chunk_meta_count: int = 0 + layout: int = 0 + locator_flags: int = 0 @dataclass(**_DATACLASS_SLOTS) @@ -51,6 +60,8 @@ class TableEntry: tag_columns: Tuple[str, ...] tag_types: Tuple[TSDataType, ...] field_columns: Tuple[str, ...] + field_types: Tuple[TSDataType, ...] = () + schema_columns: Tuple[Tuple[str, int, int], ...] = () _field_index_by_name: Dict[str, int] = field(init=False, repr=False) def __post_init__(self): @@ -97,6 +108,8 @@ def add_table( tag_columns: Iterable[str], tag_types: Iterable[TSDataType], field_columns: Iterable[str], + field_types: Iterable[TSDataType] = (), + schema_columns: Iterable[Tuple[str, int, int]] = (), ) -> int: table_id = len(self.table_entries) self.table_entries.append( @@ -105,6 +118,8 @@ def add_table( tag_columns=tuple(tag_columns), tag_types=tuple(tag_types), field_columns=tuple(field_columns), + field_types=tuple(field_types), + schema_columns=tuple(schema_columns), ) ) self.table_id_by_name[table_name] = table_id @@ -160,6 +175,11 @@ class SeriesPath(str): ``field`` components, where a ``None`` entry in ``tags`` means the tag is null -- unambiguously distinct from the literal string value ``"null"``. + Paths returned by ``TsFileDataFrame.list_timeseries()`` may additionally + carry a snapshot-local ``series_id``. A DataFrame backed by the same index + snapshot can use that id directly; it is only a hint and is not encoded in + the string value. + Trailing null tags are dropped (mirroring the device-id normalization), so ``tags`` keeps every interior null but not absent trailing ones. @@ -171,9 +191,20 @@ class SeriesPath(str): SeriesPath("table.\\N.sensorA.temperature") # a path string """ - __slots__ = ("_table", "_tags", "_field") + __slots__ = ( + "_table", + "_tags", + "_field", + "_index_identity", + "_series_id", + ) - def __new__(cls, *args: Any) -> "SeriesPath": + def __new__( + cls, + *args: Any, + index_identity=None, + series_id: Optional[int] = None, + ) -> "SeriesPath": if len(args) == 3: table, tags, field = args elif len(args) == 1: @@ -197,6 +228,8 @@ def __new__(cls, *args: Any) -> "SeriesPath": obj._table = table obj._tags = normalized obj._field = field + obj._index_identity = index_identity + obj._series_id = series_id return obj @property @@ -211,6 +244,11 @@ def tags(self) -> Tuple[Any, ...]: def field(self) -> str: return self._field + @property + def series_id(self) -> Optional[int]: + """Snapshot-local logical series id, when produced by a DataFrame.""" + return self._series_id + def _escape_path_component(value: Any) -> str: return ( @@ -290,8 +328,17 @@ def build_logical_series_path( tag_values: Iterable[Any], field_name: str, tag_columns: Iterable[str] = (), + *, + index_identity=None, + series_id: Optional[int] = None, ) -> SeriesPath: - return SeriesPath(table_name, tag_values, field_name) + return SeriesPath( + table_name, + tag_values, + field_name, + index_identity=index_identity, + series_id=series_id, + ) def build_logical_series_components( diff --git a/python/tsfile/dataset/reader.py b/python/tsfile/dataset/reader.py index 9b77190e1..e22467a8d 100644 --- a/python/tsfile/dataset/reader.py +++ b/python/tsfile/dataset/reader.py @@ -176,27 +176,36 @@ def _cache_metadata_table_model(self): tag_columns = [] tag_types = [] field_columns = [] + field_types = [] + schema_columns = [] for column_schema in table_schema.get_columns(): column_name = column_schema.get_column_name() column_category = column_schema.get_category() + column_type = column_schema.get_data_type() + schema_columns.append( + (column_name, int(column_type), int(column_category)) + ) if column_category == ColumnCategory.TIME: continue if column_category == ColumnCategory.TAG: tag_columns.append(column_name) - tag_types.append(column_schema.get_data_type()) + tag_types.append(column_type) # ignore fields which is not numeric, we won't use them currently. elif ( column_category == ColumnCategory.FIELD - and column_schema.get_data_type() in _NUMERIC_FIELD_TYPES + and column_type in _NUMERIC_FIELD_TYPES ): field_columns.append(column_name) - - if not field_columns: - continue + field_types.append(column_type) table_id = self._catalog.add_table( - table_name, tag_columns, tag_types, field_columns + table_name, + tag_columns, + tag_types, + field_columns, + field_types, + schema_columns, ) table_groups = [ group @@ -419,6 +428,15 @@ def _metadata_field_stats(group) -> Dict[str, SeriesStats]: timeline_length=int(timeline_statistic.row_count), timeline_min_time=int(timeline_statistic.start_time), timeline_max_time=int(timeline_statistic.end_time), + data_type=int(timeseries.data_type), + value_metadata_offset=int(timeseries.value_metadata_offset), + value_metadata_length=int(timeseries.value_metadata_length), + time_metadata_offset=int(timeseries.time_metadata_offset), + time_metadata_length=int(timeseries.time_metadata_length), + chunk_meta_count=int(timeseries.chunk_meta_count), + time_chunk_meta_count=int(timeseries.time_chunk_meta_count), + layout=int(timeseries.layout), + locator_flags=int(timeseries.locator_flags), ) return stats diff --git a/python/tsfile/dataset/runtime.py b/python/tsfile/dataset/runtime.py new file mode 100644 index 000000000..76e58e505 --- /dev/null +++ b/python/tsfile/dataset/runtime.py @@ -0,0 +1,1013 @@ +# 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. + +"""Process-local Runtime over the persistent Dataset Index.""" + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Mapping, Sequence +import contextlib +from concurrent.futures import ThreadPoolExecutor, wait +from dataclasses import dataclass +import os +import threading +from typing import Dict, Optional, Tuple + +import numpy as np + +from ..constants import TSDataType +from ..tag_filter import tag_eq, tag_is_null +from ..tsfile_reader import TsFileReaderPy +from .index import ( + COLUMN_NAME_INDEX, + COLUMN_SCHEMA, + DEVICE_FILE_SPAN, + DEVICE_RECORD, + LOGICAL_SERIES, + SERIES_FILE_SPAN, + SERIES_LOCATOR, + TABLE_RECORD, + TSFILE_RECORD, + MappedDatasetIndex, + file_fingerprint, +) +from .metadata import ( + MODEL_TABLE, + MODEL_TREE, + TableEntry, + _join_series_path, + split_logical_series_path, +) + +_SERIES_DESCRIPTOR_CACHE_SIZE = 4096 + + +@dataclass(frozen=True) +class RuntimeSeriesShard: + """One immutable physical fragment already expanded from the mmap route.""" + + reader: "RuntimeSeriesReader" + device_id: int + column_id: int + locator_id: int + timeline_length: int + min_time: int + max_time: int + + +@dataclass(frozen=True) +class RuntimeSeriesDescriptor: + """Bounded process-local expansion of one logical series route.""" + + ref: Tuple[int, int] + series_id: int + column_id: int + shards: Tuple[RuntimeSeriesShard, ...] + min_time: Optional[int] + max_time: Optional[int] + count: int + + @property + def refs(self): + return tuple( + (shard.reader, shard.device_id, shard.column_id) for shard in self.shards + ) + + +def _exact_tag_filter(tag_columns, tag_values): + result = None + for index, name in enumerate(tag_columns): + value = tag_values[index] if index < len(tag_values) else None + current = tag_is_null(name) if value is None else tag_eq(name, str(value)) + result = current if result is None else result & current + return result + + +class RuntimeLease: + def __init__(self, runtime: "DatasetRuntime"): + self._runtime = runtime + self._lock = threading.Lock() + self._closed = False + runtime._acquire_object() + + def clone(self): + with self._lock: + if self._closed: + raise RuntimeError("Runtime lease is closed") + return RuntimeLease(self._runtime) + + def query_lease(self): + with self._lock: + if self._closed: + raise RuntimeError("Runtime lease is closed") + self._runtime._acquire_query() + return _QueryLease(self._runtime, acquired=True) + + def close(self): + with self._lock: + if self._closed: + return + self._closed = True + self._runtime._release_object() + + +class _QueryLease: + def __init__(self, runtime: "DatasetRuntime", acquired: bool = False): + self._runtime = runtime + self._acquired = acquired + + def __enter__(self): + if not self._acquired: + self._runtime._acquire_query() + self._acquired = True + return self + + def __exit__(self, *_): + self.close() + + def close(self): + if self._acquired: + self._acquired = False + self._runtime._release_query() + + +class _ReaderSession: + def __init__(self, file_id: int, path: str, expected_size: int, fingerprint: int): + self.file_id = file_id + self.path = path + self.expected_size = expected_size + self.fingerprint = fingerprint + self._validate_generation() + self.reader = TsFileReaderPy(path) + self.active_uses = 0 + + def _validate_generation(self): + st = os.stat(self.path) + if ( + st.st_size != self.expected_size + or file_fingerprint(self.path, st) != self.fingerprint + ): + raise RuntimeError( + "TsFile generation changed after Dataset Index publication: " + f"{self.path}" + ) + + def close(self): + self.reader.close() + + +class ReaderSessionPool: + """Per-Runtime LRU pool with a hard cap on simultaneously open Readers.""" + + def __init__(self, index: MappedDatasetIndex, max_open_files: int): + self._index = index + self.max_open_files = max(1, int(max_open_files)) + self._sessions: "OrderedDict[int, _ReaderSession]" = OrderedDict() + self._condition = threading.Condition() + self._closed = False + + def _new_session(self, file_id: int): + record = self._index.record(TSFILE_RECORD, file_id) + return _ReaderSession( + file_id, + self._index.string(record[0]), + record[2], + record[3], + ) + + @contextlib.contextmanager + def acquire(self, file_id: int): + with self._condition: + while True: + if self._closed: + raise RuntimeError("ReaderSessionPool is closed") + session = self._sessions.get(file_id) + if session is not None: + session._validate_generation() + self._sessions.move_to_end(file_id) + session.active_uses += 1 + break + if len(self._sessions) < self.max_open_files: + session = self._new_session(file_id) + session.active_uses = 1 + self._sessions[file_id] = session + break + idle_id = next( + ( + key + for key, value in self._sessions.items() + if value.active_uses == 0 + ), + None, + ) + if idle_id is None: + self._condition.wait() + continue + idle = self._sessions.pop(idle_id) + idle.close() + session = self._new_session(file_id) + session.active_uses = 1 + self._sessions[file_id] = session + break + try: + yield session.reader + finally: + with self._condition: + session.active_uses -= 1 + self._condition.notify_all() + + def close(self): + with self._condition: + self._closed = True + while any(session.active_uses for session in self._sessions.values()): + self._condition.wait() + sessions = list(self._sessions.values()) + self._sessions.clear() + for session in sessions: + session.close() + + @property + def open_count(self): + with self._condition: + return len(self._sessions) + + +class PreparedSeriesCache: + """Runtime-wide single-flight cache of native exact-locator metadata.""" + + def __init__(self, index: MappedDatasetIndex): + self._index = index + self._condition = threading.Condition() + self._entries = {} + self._loading = set() + 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") + return ( + id(self._index), + file_id, + file_record[2], + file_record[3], + locator_id, + locator[1], + locator[2], + locator[3], + locator[4], + device_span[2], + device_span[3], + ) + + def get(self, file_id, locator_id, reader, time_owner=None): + key = (id(self._index), file_id, locator_id) + with self._condition: + while True: + if self._closed: + raise RuntimeError("PreparedSeriesCache is closed") + result = self._entries.get(key) + if result is not None: + return result + if key not in self._loading: + self._loading.add(key) + break + self._condition.wait() + try: + result = reader.prepare_series( + self._locator_tuple(file_id, locator_id), time_owner=time_owner + ) + except Exception: + with self._condition: + self._loading.remove(key) + self._condition.notify_all() + raise + with self._condition: + if self._closed: + result.close() + self._loading.remove(key) + self._condition.notify_all() + raise RuntimeError("PreparedSeriesCache is closed") + self._entries[key] = result + self._loading.remove(key) + self._condition.notify_all() + return result + + def close(self): + with self._condition: + self._closed = True + while self._loading: + self._condition.wait() + entries = list(self._entries.values()) + self._entries.clear() + for prepared in entries: + prepared.close() + + @property + def size(self): + with self._condition: + return len(self._entries) + + +class DatasetRuntime: + def __init__( + self, + path: str, + max_open_files: Optional[int] = None, + query_workers: Optional[int] = None, + query_parallel_min_rows: Optional[int] = None, + ): + self.index = MappedDatasetIndex(path) + maximum = ( + int(os.environ.get("TSFILE_DATAFRAME_MAX_OPEN_FILES", "16")) + if max_open_files is None + else max_open_files + ) + workers = ( + int( + os.environ.get( + "TSFILE_DATAFRAME_QUERY_WORKERS", + str(min(4, os.cpu_count() or 1)), + ) + ) + if query_workers is None + else query_workers + ) + self.query_workers = max(1, int(workers)) + minimum_rows = ( + int(os.environ.get("TSFILE_DATAFRAME_QUERY_PARALLEL_MIN_ROWS", "8192")) + if query_parallel_min_rows is None + else query_parallel_min_rows + ) + self.query_parallel_min_rows = max(1, int(minimum_rows)) + self._query_executor = ( + ThreadPoolExecutor( + max_workers=self.query_workers, + thread_name_prefix="tsfile-dataframe-query", + ) + if self.query_workers > 1 + else None + ) + self.readers = ReaderSessionPool(self.index, maximum) + self.prepared = PreparedSeriesCache(self.index) + self._condition = threading.Condition() + self._object_leases = 0 + self._query_leases = 0 + self._accepting = True + self._torn_down = False + self.catalog = MappedDataFrameCatalog(self) + + def lease(self): + return RuntimeLease(self) + + def query_lease(self): + return _QueryLease(self) + + def map_query_groups(self, function, groups, estimated_rows=None): + """Run independent query groups under the caller's query lease.""" + groups = list(groups) + if not groups: + return [] + estimates = list(estimated_rows) if estimated_rows is not None else [] + large_enough = not estimates or max(estimates, default=0) >= ( + self.query_parallel_min_rows + ) + if self._query_executor is None or len(groups) == 1 or not large_enough: + return [function(group) for group in groups] + + futures = [self._query_executor.submit(function, group) for group in groups] + try: + # Preserve group order so merge behavior stays deterministic. + return [future.result() for future in futures] + except BaseException: + for future in futures: + future.cancel() + wait(futures) + raise + + def _acquire_object(self): + with self._condition: + if not self._accepting: + raise RuntimeError("Dataset Runtime is closing") + self._object_leases += 1 + + def _release_object(self): + teardown = False + with self._condition: + self._object_leases -= 1 + if self._object_leases == 0: + self._accepting = False + while self._query_leases: + self._condition.wait() + teardown = not self._torn_down + self._torn_down = True + if teardown: + if self._query_executor is not None: + self._query_executor.shutdown(wait=True, cancel_futures=True) + self.prepared.close() + self.readers.close() + self.index.close() + + def _acquire_query(self): + with self._condition: + if not self._accepting: + raise RuntimeError("Dataset Runtime is closing") + self._query_leases += 1 + + def _release_query(self): + with self._condition: + self._query_leases -= 1 + self._condition.notify_all() + + +class _TableMapping(Mapping): + def __init__(self, catalog: "MappedDataFrameCatalog"): + self._catalog = catalog + self._names = {} + self._cache = {} + for table_id in range(catalog.index.count(TABLE_RECORD)): + record = catalog.index.record(TABLE_RECORD, table_id) + name = catalog.index.string(record[0]) + if name in self._names: + raise ValueError( + f"Dataset Index contains duplicate canonical table '{name}'" + ) + self._names[name] = table_id + + def __getitem__(self, name): + result = self._cache.get(name) + if result is not None: + return result + table_id = self._names[name] + table = self._catalog.index.record(TABLE_RECORD, table_id) + tags = [] + tag_types = [] + fields = [] + field_types = [] + schema_columns = [] + for name_index in self._catalog.index.records( + COLUMN_NAME_INDEX, table[4], table[5] + ): + column = self._catalog.index.record(COLUMN_SCHEMA, name_index[1]) + column_name = self._catalog.index.string(column[1]) + schema_columns.append( + (column[2], column_name, int(column[3]), int(column[7])) + ) + if column[7] == 0: + tags.append((column[2], column_name, TSDataType(column[3]))) + elif column[7] == 1 and column[3] in { + int(TSDataType.INT32), + int(TSDataType.INT64), + int(TSDataType.FLOAT), + int(TSDataType.DOUBLE), + }: + fields.append((column[2], column_name, TSDataType(column[3]))) + tags.sort() + fields.sort() + schema_columns.sort() + result = TableEntry( + name, + tuple(item[1] for item in tags), + tuple(item[2] for item in tags), + tuple(item[1] for item in fields), + tuple(item[2] for item in fields), + tuple((item[1], item[2], item[3]) for item in schema_columns), + ) + self._cache[name] = result + return result + + def __iter__(self): + return iter(self._names) + + def __len__(self): + return len(self._names) + + def table_id(self, name): + return self._names[name] + + +class _DeviceSequence(Sequence): + def __init__(self, catalog): + self._catalog = catalog + + def __len__(self): + return self._catalog.index.count(DEVICE_RECORD) + + def __getitem__(self, device_id): + if isinstance(device_id, slice): + return [self[index] for index in range(*device_id.indices(len(self)))] + record = self._catalog.index.record(DEVICE_RECORD, device_id) + table = self._catalog.index.record(TABLE_RECORD, record[0]) + table_name = self._catalog.index.string(table[0]) + components = split_logical_series_path(self._catalog.index.string(record[1])) + return table_name, tuple(components[1:-1]) + + +class _DeviceIndexMapping(Mapping): + def __init__(self, catalog): + self._catalog = catalog + + def get(self, key, default=None): + table_name, tags = key + try: + table_id = self._catalog.table_entries.table_id(table_name) + except KeyError: + return default + + try: + return self._catalog.index.find_device_id( + table_id, _join_series_path(table_name, tags, "") + ) + except KeyError: + return default + + def __getitem__(self, key): + result = self.get(key) + if result is None: + raise KeyError(key) + return result + + def __iter__(self): + return iter(self._catalog.devices) + + def __len__(self): + return len(self._catalog.devices) + + +class _SeriesSequence(Sequence): + def __init__(self, catalog, ids=None): + self._catalog = catalog + self._ids = ids + + def __len__(self): + return ( + self._catalog.index.count(LOGICAL_SERIES) + if self._ids is None + else len(self._ids) + ) + + def series_id(self, position): + return position if self._ids is None else self._ids[position] + + def __getitem__(self, position): + if isinstance(position, slice): + return [self[index] for index in range(*position.indices(len(self)))] + series_id = self.series_id(position) + series = self._catalog.index.record(LOGICAL_SERIES, series_id) + device = self._catalog.index.record(DEVICE_RECORD, series[0]) + table_name = self._catalog.index.string( + self._catalog.index.record(TABLE_RECORD, device[0])[0] + ) + table = self._catalog.table_entries[table_name] + column_name = self._catalog.index.string( + self._catalog.index.record(COLUMN_SCHEMA, series[1])[1] + ) + return series[0], table.get_field_index(column_name) + + def __iter__(self): + for index in range(len(self)): + yield self[index] + + +class _RouteMapping(Mapping): + def __init__(self, catalog, cache_size): + self._catalog = catalog + self._cache_size = cache_size + self._cache = OrderedDict() + self._cache_lock = threading.Lock() + + def _series_id(self, ref): + device_id, field_idx = ref + device = self._catalog.index.record(DEVICE_RECORD, device_id) + table_name = self._catalog.index.string( + self._catalog.index.record(TABLE_RECORD, device[0])[0] + ) + field_name = self._catalog.table_entries[table_name].field_columns[field_idx] + column_id = self._catalog.index.find_column_id(device[0], 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 series_id is None: + series_id = self._series_id(ref) + series = self._catalog.index.record(LOGICAL_SERIES, series_id) + if series[0] != ref[0]: + raise KeyError(ref) + if column_id is None: + column_id = series[1] + elif series[1] != 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 + shards.append( + RuntimeSeriesShard( + self._catalog.reader_for(span[1]), + series[0], + column_id, + span[2], + timeline_length, + span[4], + span[5], + ) + ) + + result = RuntimeSeriesDescriptor( + ref, + series_id, + column_id, + tuple(shards), + series[4] if count else None, + series[5] if count else None, + count, + ) + if self._cache_size: + with self._cache_lock: + existing = self._cache.get(ref) + if existing is not None: + self._cache.move_to_end(ref) + return existing + self._cache[ref] = result + while len(self._cache) > self._cache_size: + self._cache.popitem(last=False) + return result + + def __contains__(self, ref): + try: + self.describe(ref) + return True + except KeyError: + return False + + def __getitem__(self, ref): + return list(self.describe(ref).refs) + + def __iter__(self): + return iter(self._catalog.series) + + def __len__(self): + return len(self._catalog.series) + + +class MappedDataFrameCatalog: + 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 = OrderedDict() + self._descriptor_cache_lock = threading.Lock() + self.table_entries = _TableMapping(self) + self.devices = _DeviceSequence(self) + self.device_index = _DeviceIndexMapping(self) + self.device_time_bounds = _DeviceTimeBounds(self) + self.series = _SeriesSequence(self) + self.series_shards = _RouteMapping(self, self._descriptor_cache_size) + self._readers = {} + self.model = self._infer_model() + + 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 + + table_id = self.table_entries.table_id(table_name) + table = self.table_entries[table_name] + field_idx = table.get_field_index(field_name) + device_id = self.index.find_device_id( + table_id, _join_series_path(table_name, tags, "") + ) + column_id = self.index.find_column_id(table_id, field_name) + series_id = self.index.find_series_id(device_id, column_id) + result = self.series_shards.describe( + (device_id, field_idx), series_id=series_id, column_id=column_id + ) + + if self._descriptor_cache_size: + with self._descriptor_cache_lock: + existing = self._descriptor_cache.get(key) + if existing is not None: + 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) + return result + + 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) + 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] + ) + + def _infer_model(self): + if len(self.table_entries) != 1: + return MODEL_TABLE + table = next(iter(self.table_entries.values())) + if table.tag_columns == tuple( + f"_col_{index + 1}" for index in range(len(table.tag_columns)) + ): + return MODEL_TREE + return MODEL_TABLE + + def reader_for(self, file_id): + reader = self._readers.get(file_id) + if reader is None: + reader = RuntimeSeriesReader(self.runtime, file_id) + self._readers[file_id] = reader + return reader + + +class _DeviceTimeBounds(Sequence): + def __init__(self, catalog): + self._catalog = catalog + + 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] + + +class RuntimeSeriesReader: + """File-specific facade whose metadata comes from mmap, not a Python catalog.""" + + def __init__(self, runtime: DatasetRuntime, file_id: int): + self.runtime = runtime + self.file_id = file_id + + def _series(self, device_id, column_id): + return self.runtime.index.find_series_id(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)) + + 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])) + 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]) + 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] + 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], + } + + 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]) + 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] + return { + "length": timeline_length, + "min_time": span[4], + "max_time": span[5], + "timeline_length": timeline_length, + "timeline_min_time": span[4], + "timeline_max_time": span[5], + "table_name": table_name, + "column_name": column_name, + "device_id": device_id, + "field_idx": column_id, + "tag_columns": table.tag_columns, + "tag_values": dict(zip(table.tag_columns, tags)), + } + + @staticmethod + def _consume(result): + timestamp_parts = [] + value_parts = [] + with result: + read_arrow = getattr(result, "read_arrow_record_batch", None) + if read_arrow is None: + read_arrow = result.read_arrow_batch + while True: + arrow_batch = read_arrow() + if arrow_batch is None: + break + if arrow_batch.num_rows == 0: + continue + timestamp_parts.append( + np.asarray( + arrow_batch.column(0).to_numpy(zero_copy_only=False), + dtype=np.int64, + ) + ) + value_parts.append( + np.asarray( + arrow_batch.column(1).to_numpy(zero_copy_only=False), + dtype=np.float64, + ) + ) + if not timestamp_parts: + return np.array([], dtype=np.int64), np.array([], dtype=np.float64) + if len(timestamp_parts) == 1: + return timestamp_parts[0], value_parts[0] + return np.concatenate(timestamp_parts), np.concatenate(value_parts) + + @staticmethod + def _consume_multi(result, column_names): + timestamp_parts = [] + value_parts = {name: [] for name in column_names} + with result: + read_arrow = getattr(result, "read_arrow_record_batch", None) + if read_arrow is None: + read_arrow = result.read_arrow_batch + while True: + arrow_batch = read_arrow() + if arrow_batch is None: + break + if arrow_batch.num_rows == 0: + continue + timestamp_parts.append( + np.asarray( + arrow_batch.column(0).to_numpy(zero_copy_only=False), + dtype=np.int64, + ) + ) + for column_index, name in enumerate(column_names, start=1): + value_parts[name].append( + np.asarray( + arrow_batch.column(column_index).to_numpy( + zero_copy_only=False + ), + dtype=np.float64, + ) + ) + if not timestamp_parts: + return np.array([], dtype=np.int64), { + name: np.array([], dtype=np.float64) for name in column_names + } + timestamps = ( + timestamp_parts[0] + if len(timestamp_parts) == 1 + else np.concatenate(timestamp_parts) + ) + values = { + name: parts[0] if len(parts) == 1 else np.concatenate(parts) + for name, parts in value_parts.items() + } + return timestamps, values + + def _query( + self, + device_id, + column_id, + start_time=None, + end_time=None, + offset=None, + limit=None, + ): + span = self._span(device_id, column_id) + return self._query_at_locator( + span[2], + start_time=start_time, + end_time=end_time, + offset=offset, + limit=limit, + ) + + def _query_at_locator( + self, + locator_id, + start_time=None, + end_time=None, + offset=None, + limit=None, + ): + with self.runtime.readers.acquire(self.file_id) as reader: + prepared = self.runtime.prepared.get(self.file_id, locator_id, reader) + if offset is None: + result = reader.query_prepared( + prepared, start_time=start_time, end_time=end_time + ) + else: + result = reader.query_prepared(prepared, offset=offset, limit=limit) + return self._consume(result) + + def read_series_by_ref(self, device_id, column_id, start_time, end_time): + return self._query(device_id, column_id, start_time, end_time) + + def read_series_by_row(self, device_id, column_id, offset, limit): + if limit <= 0: + return np.array([], dtype=np.int64), np.array([], dtype=np.float64) + return self._query(device_id, column_id, offset=offset, limit=limit) + + def read_series_by_row_at_locator(self, locator_id, offset, limit): + """Read by a locator already validated against this Runtime snapshot.""" + if limit <= 0: + return np.array([], dtype=np.int64), np.array([], dtype=np.float64) + return self._query_at_locator(locator_id, offset=offset, limit=limit) + + def read_device_fields_by_time_range( + self, device_id, column_ids, start_time, end_time + ): + if not column_ids: + 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 + ] + device_span_ids = {locator[0] for locator in locators} + 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)) + ) + can_read_aligned = ( + device_span[1] == self.file_id + and device_span[4] == 1 + and all(locator[1] == 1 for locator in locators) + ) + + if can_read_aligned: + column_names = [ + self._identity(device_id, column_id)[3] for column_id in column_ids + ] + with self.runtime.readers.acquire(self.file_id) as reader: + prepared = [] + time_owner = None + for span in spans: + current = self.runtime.prepared.get( + self.file_id, span[2], reader, time_owner=time_owner + ) + prepared.append(current) + if time_owner is None: + time_owner = current + result = reader.query_prepared_multi( + prepared, start_time=start_time, end_time=end_time + ) + return self._consume_multi(result, column_names) + + parts = [ + self.read_series_by_ref(device_id, column_id, start_time, end_time) + for column_id in column_ids + ] + if not parts: + return np.array([], dtype=np.int64), {} + # Existing dataframe merge aligns separate field parts across files; + # this method is only a compatibility surface for a single file. + timestamps = parts[0][0] + values = {} + for column_id, (current_timestamps, current_values) in zip(column_ids, parts): + if not np.array_equal(current_timestamps, timestamps): + raise ValueError("single-file fields do not share one timeline") + _, _, _, name = self._identity(device_id, column_id) + values[name] = current_values + return timestamps, values diff --git a/python/tsfile/dataset/timeseries.py b/python/tsfile/dataset/timeseries.py index 39d43b442..f85bde139 100644 --- a/python/tsfile/dataset/timeseries.py +++ b/python/tsfile/dataset/timeseries.py @@ -18,6 +18,7 @@ """Timeseries handles returned by the dataset package.""" +import contextlib from typing import Callable, List, Optional, Tuple import numpy as np @@ -76,9 +77,10 @@ def __init__( name: str, series_refs: list, stats: dict, - ensure_open: Callable[[], None], + ensure_open: Optional[Callable[[], None]], load_timestamps: Callable[[], np.ndarray], read_by_position: Callable[[int, int], Tuple[np.ndarray, np.ndarray]], + runtime_lease=None, ): self._name = name self._series_refs = series_refs @@ -87,6 +89,23 @@ def __init__( self._load_timestamps = load_timestamps self._read_by_position = read_by_position self._timestamps = None + self._runtime_lease = runtime_lease + self._closed = False + + def _assert_open(self): + if self._closed: + raise RuntimeError("Current Timeseries is closed.") + if self._ensure_open is not None: + self._ensure_open() + + @contextlib.contextmanager + def _query_guard(self): + self._assert_open() + if self._runtime_lease is None: + yield + else: + with self._runtime_lease.query_lease(): + yield @property def name(self) -> str: @@ -94,10 +113,10 @@ def name(self) -> str: @property def timestamps(self) -> np.ndarray: - self._ensure_open() - if self._timestamps is None: - self._timestamps = self._load_timestamps() - return self._timestamps + with self._query_guard(): + if self._timestamps is None: + self._timestamps = self._load_timestamps() + return self._timestamps @property def stats(self) -> dict: @@ -111,7 +130,10 @@ def __len__(self) -> int: return self._stats["count"] def __getitem__(self, key): - self._ensure_open() + with self._query_guard(): + return self._getitem_open(key) + + def _getitem_open(self, key): length = len(self) if isinstance(key, int): @@ -148,23 +170,43 @@ def __getitem__(self, key): def _query_time_range( self, start_time: int, end_time: int ) -> Tuple[np.ndarray, np.ndarray]: - self._ensure_open() - time_parts = [] - value_parts = [] - for reader, device_id, field_idx in self._series_refs: - device_info = reader.get_device_info(device_id) - if ( - device_info["max_time"] < start_time - or device_info["min_time"] > end_time - ): - continue - ts_arr, val_arr = reader.read_series_by_ref( - device_id, field_idx, start_time, end_time - ) - if len(ts_arr) > 0: - time_parts.append(ts_arr) - value_parts.append(val_arr) - return merge_time_value_parts(time_parts, value_parts) + with self._query_guard(): + time_parts = [] + value_parts = [] + for reader, device_id, field_idx in self._series_refs: + device_info = reader.get_device_info(device_id) + if ( + device_info["max_time"] < start_time + or device_info["min_time"] > end_time + ): + continue + ts_arr, val_arr = reader.read_series_by_ref( + device_id, field_idx, start_time, end_time + ) + if len(ts_arr) > 0: + time_parts.append(ts_arr) + value_parts.append(val_arr) + return merge_time_value_parts(time_parts, value_parts) + + def close(self): + if self._closed: + return + self._closed = True + if self._runtime_lease is not None: + self._runtime_lease.close() + + def __enter__(self): + self._assert_open() + return self + + def __exit__(self, *_): + self.close() + + def __del__(self): + try: + self.close() + except Exception: + pass def __repr__(self): stats = self.stats diff --git a/python/tsfile/schema.py b/python/tsfile/schema.py index fcfee5e99..6a8bdaa9b 100644 --- a/python/tsfile/schema.py +++ b/python/tsfile/schema.py @@ -100,6 +100,13 @@ class TimeseriesMetadata: chunk_meta_count: int statistic: TimeseriesStatisticType timeline_statistic: TimeseriesStatisticType + value_metadata_offset: int = 0 + value_metadata_length: int = 0 + time_metadata_offset: int = 0 + time_metadata_length: int = 0 + time_chunk_meta_count: int = 0 + layout: int = 0 + locator_flags: int = 0 @dataclass(frozen=True) diff --git a/python/tsfile/tsfile_cpp.pxd b/python/tsfile/tsfile_cpp.pxd index 5979e293e..ed750c79b 100644 --- a/python/tsfile/tsfile_cpp.pxd +++ b/python/tsfile/tsfile_cpp.pxd @@ -17,7 +17,7 @@ # #cython: language_level=3 -from libc.stdint cimport uint32_t, int32_t, int64_t, uint64_t, uint8_t +from libc.stdint cimport uint16_t, uint32_t, int32_t, int64_t, uint64_t, uint8_t ctypedef int32_t ErrorCode @@ -38,6 +38,7 @@ cdef extern from "cwrapper/tsfile_cwrapper.h": ctypedef void * Tablet ctypedef void * TsRecord ctypedef void * ResultSet + ctypedef void * PreparedSeriesHandle # enum types ctypedef enum TSDataType: @@ -172,6 +173,13 @@ cdef extern from "cwrapper/tsfile_cwrapper.h": int32_t chunk_meta_count TimeseriesStatistic statistic TimeseriesStatistic timeline_statistic + uint64_t value_metadata_offset + uint32_t value_metadata_length + uint64_t time_metadata_offset + uint32_t time_metadata_length + uint32_t time_chunk_meta_count + uint16_t layout + uint16_t locator_flags ctypedef struct DeviceID: char * path @@ -200,6 +208,19 @@ cdef extern from "cwrapper/tsfile_cwrapper.h": TSDataType * data_types int column_num + ctypedef struct TsFilePreparedLocator: + uint64_t mapped_index_identity + uint32_t file_id + uint64_t file_size + uint64_t file_fingerprint + uint32_t locator_id + uint16_t layout + uint16_t flags + uint64_t value_metadata_offset + uint32_t value_metadata_length + uint64_t time_metadata_offset + uint32_t time_metadata_length + # Function Declarations ctypedef void * TagFilterHandle @@ -276,6 +297,22 @@ cdef extern from "cwrapper/tsfile_cwrapper.h": const char** columns, uint32_t column_num, int64_t start_time, int64_t end_time, ErrorCode *err_code) + PreparedSeriesHandle tsfile_reader_prepare_series( + TsFileReader reader, const TsFilePreparedLocator * locator, + ErrorCode * err_code) nogil + PreparedSeriesHandle tsfile_reader_prepare_series_with_time_owner( + TsFileReader reader, const TsFilePreparedLocator * locator, + PreparedSeriesHandle aligned_time_owner, ErrorCode * err_code) nogil + void tsfile_prepared_series_free(PreparedSeriesHandle prepared) + ResultSet tsfile_reader_query_prepared( + TsFileReader reader, PreparedSeriesHandle prepared, + int64_t start_time, int64_t end_time, int offset, int limit, + ErrorCode * err_code) nogil + ResultSet tsfile_reader_query_prepared_multi( + TsFileReader reader, const PreparedSeriesHandle * prepared, + uint32_t prepared_count, int64_t start_time, int64_t end_time, + int offset, int limit, ErrorCode * err_code) nogil + ResultSet tsfile_query_table_on_tree(TsFileReader reader, char** columns, uint32_t column_num, int64_t start_time, int64_t end_time, @@ -432,7 +469,7 @@ cdef extern from "cwrapper/tsfile_cwrapper.h": # Arrow batch reading function ErrorCode tsfile_result_set_get_next_tsblock_as_arrow(ResultSet result_set, ArrowArray* out_array, - ArrowSchema* out_schema); + ArrowSchema* out_schema) nogil # Arrow batch writing function ErrorCode _tsfile_writer_write_arrow_table(TsFileWriter writer, diff --git a/python/tsfile/tsfile_py_cpp.pxd b/python/tsfile/tsfile_py_cpp.pxd index 0e2f91cbc..adfe30939 100644 --- a/python/tsfile/tsfile_py_cpp.pxd +++ b/python/tsfile/tsfile_py_cpp.pxd @@ -43,6 +43,14 @@ cdef public api void free_c_tablet(Tablet tablet) cdef public api void free_c_row_record(TsRecord record) cdef public api TsFileWriter tsfile_writer_new_c(object pathname, uint64_t memory_threshold) except NULL cdef public api TsFileReader tsfile_reader_new_c(object pathname) except NULL +cdef public api PreparedSeriesHandle tsfile_reader_prepare_series_c( + TsFileReader reader, object locator) except NULL +cdef public api PreparedSeriesHandle tsfile_reader_prepare_series_with_time_owner_c( + TsFileReader reader, object locator, + PreparedSeriesHandle aligned_time_owner) except NULL +cdef public api ResultSet tsfile_reader_query_prepared_c( + TsFileReader reader, PreparedSeriesHandle prepared, int64_t start_time, + int64_t end_time, int offset, int limit) cdef public api ErrorCode tsfile_writer_register_device_py_cpp(TsFileWriter writer, DeviceSchema *schema) cdef public api ErrorCode tsfile_writer_register_timeseries_py_cpp(TsFileWriter writer, object device_name, TimeseriesSchema *schema) diff --git a/python/tsfile/tsfile_py_cpp.pyx b/python/tsfile/tsfile_py_cpp.pyx index 8406c8519..54a48a8c0 100644 --- a/python/tsfile/tsfile_py_cpp.pyx +++ b/python/tsfile/tsfile_py_cpp.pyx @@ -789,6 +789,62 @@ cdef TsFileReader tsfile_reader_new_c(object pathname) except NULL: check_error(errno) return reader +cdef PreparedSeriesHandle tsfile_reader_prepare_series_c( + TsFileReader reader, object locator) except NULL: + cdef TsFilePreparedLocator native + cdef ErrorCode code = 0 + native.mapped_index_identity = locator[0] + native.file_id = locator[1] + native.file_size = locator[2] + native.file_fingerprint = locator[3] + native.locator_id = locator[4] + native.layout = locator[5] + native.flags = locator[6] + native.value_metadata_offset = locator[7] + native.value_metadata_length = locator[8] + native.time_metadata_offset = locator[9] + native.time_metadata_length = locator[10] + cdef PreparedSeriesHandle prepared + with nogil: + prepared = tsfile_reader_prepare_series(reader, &native, &code) + check_error(code, b"Failed to prepare Dataset Index locator") + return prepared + +cdef PreparedSeriesHandle tsfile_reader_prepare_series_with_time_owner_c( + TsFileReader reader, object locator, + PreparedSeriesHandle aligned_time_owner) except NULL: + cdef TsFilePreparedLocator native + cdef ErrorCode code = 0 + native.mapped_index_identity = locator[0] + native.file_id = locator[1] + native.file_size = locator[2] + native.file_fingerprint = locator[3] + native.locator_id = locator[4] + native.layout = locator[5] + native.flags = locator[6] + native.value_metadata_offset = locator[7] + native.value_metadata_length = locator[8] + native.time_metadata_offset = locator[9] + native.time_metadata_length = locator[10] + cdef PreparedSeriesHandle prepared + with nogil: + prepared = tsfile_reader_prepare_series_with_time_owner( + reader, &native, aligned_time_owner, &code + ) + check_error(code, b"Failed to prepare aligned Dataset Index locator") + return prepared + +cdef ResultSet tsfile_reader_query_prepared_c( + TsFileReader reader, PreparedSeriesHandle prepared, + int64_t start_time, int64_t end_time, int offset, int limit): + cdef ErrorCode code = 0 + cdef ResultSet result + with nogil: + result = tsfile_reader_query_prepared( + reader, prepared, start_time, end_time, offset, limit, &code) + check_error(code, b"Failed to query prepared series") + return result + cpdef object get_tsfile_config(): return { "tsblock_mem_inc_step_size_": g_config_value_.tsblock_mem_inc_step_size_, @@ -1263,6 +1319,13 @@ cdef object timeseries_metadata_c_to_py(TimeseriesMetadata* m): int(m.chunk_meta_count), stat, timeline_stat, + int(m.value_metadata_offset), + int(m.value_metadata_length), + int(m.time_metadata_offset), + int(m.time_metadata_length), + int(m.time_chunk_meta_count), + int(m.layout), + int(m.locator_flags), ) cdef tuple c_device_segments_to_tuple(char** segs, uint32_t n): diff --git a/python/tsfile/tsfile_reader.pyx b/python/tsfile/tsfile_reader.pyx index a2e8fe263..c4354799c 100644 --- a/python/tsfile/tsfile_reader.pyx +++ b/python/tsfile/tsfile_reader.pyx @@ -23,6 +23,7 @@ from typing import List, Optional, Dict import pandas as pd from libc.string cimport strlen +from libc.stdlib cimport free, malloc from cpython.bytes cimport PyBytes_FromStringAndSize from libc.string cimport memset import pyarrow as pa @@ -153,7 +154,8 @@ cdef class ResultSetPy: df = df.astype(data_type_dict) return df - def read_arrow_batch(self): + def read_arrow_record_batch(self): + """Read one native TsBlock as an Arrow RecordBatch.""" self.check_result_set_invalid() cdef ArrowArray arrow_array @@ -164,7 +166,9 @@ cdef class ResultSetPy: memset(&arrow_array, 0, sizeof(ArrowArray)) memset(&arrow_schema, 0, sizeof(ArrowSchema)) - code = tsfile_result_set_get_next_tsblock_as_arrow(self.result, &arrow_array, &arrow_schema) + with nogil: + code = tsfile_result_set_get_next_tsblock_as_arrow( + self.result, &arrow_array, &arrow_schema) if code == RET_NO_MORE_DATA: return None @@ -177,9 +181,7 @@ cdef class ResultSetPy: try: schema_ptr = &arrow_schema array_ptr = &arrow_array - batch = pa.RecordBatch._import_from_c(array_ptr, schema_ptr) - table = pa.Table.from_batches([batch]) - return table + return pa.RecordBatch._import_from_c(array_ptr, schema_ptr) except Exception as e: if arrow_array.release != NULL: arrow_array.release(&arrow_array) @@ -187,6 +189,13 @@ cdef class ResultSetPy: arrow_schema.release(&arrow_schema) raise e + def read_arrow_batch(self): + """Read one native TsBlock as an Arrow Table.""" + batch = self.read_arrow_record_batch() + if batch is None: + return None + return pa.Table.from_batches([batch]) + def get_value_by_index(self, index : int): """ Get value by index from query result set. @@ -309,6 +318,30 @@ cdef class ResultSetPy: def __exit__(self, exc_type, exc_val, exc_tb): self.close() +cdef class PreparedSeriesPy: + """Reusable native metadata parsed from one exact Dataset Index locator.""" + cdef PreparedSeriesHandle prepared + + def __cinit__(self): + self.prepared = NULL + + cdef init_c(self, PreparedSeriesHandle prepared): + self.prepared = prepared + + def close(self): + if self.prepared != NULL: + tsfile_prepared_series_free(self.prepared) + self.prepared = NULL + + def __dealloc__(self): + self.close() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + cdef class TsFileReaderPy: """ Cython wrapper class for interacting with TsFileReader C implementation. @@ -470,6 +503,74 @@ cdef class TsFileReaderPy: self.activate_result_set_list.add(pyresult) return pyresult + def prepare_series(self, locator, PreparedSeriesPy time_owner=None) -> PreparedSeriesPy: + """Prepare an 11-field native locator tuple for repeated queries.""" + if len(locator) != 11: + raise ValueError("prepared locator must contain exactly 11 fields") + cdef PreparedSeriesHandle prepared = NULL + if time_owner is None: + prepared = tsfile_reader_prepare_series_c(self.reader, locator) + else: + if time_owner.prepared == NULL: + raise RuntimeError("PreparedSeries time owner is closed") + prepared = tsfile_reader_prepare_series_with_time_owner_c( + self.reader, locator, time_owner.prepared) + py_prepared = PreparedSeriesPy() + py_prepared.init_c(prepared) + return py_prepared + + def query_prepared(self, PreparedSeriesPy prepared, + start_time : int = INT64_MIN, + end_time : int = INT64_MAX, + offset : int = 0, limit : int = -1) -> ResultSetPy: + if prepared.prepared == NULL: + raise RuntimeError("PreparedSeries is closed") + cdef ResultSet result = tsfile_reader_query_prepared_c( + self.reader, prepared.prepared, start_time, end_time, offset, limit) + pyresult = ResultSetPy(self, True) + pyresult.init_c(result, "prepared") + self.activate_result_set_list.add(pyresult) + return pyresult + + def query_prepared_multi(self, prepared_list, + start_time : int = INT64_MIN, + end_time : int = INT64_MAX, + offset : int = 0, limit : int = -1) -> ResultSetPy: + """Query aligned prepared value columns through one shared time reader.""" + cdef Py_ssize_t count = len(prepared_list) + cdef Py_ssize_t i + cdef PreparedSeriesPy prepared + cdef PreparedSeriesHandle* handles = NULL + cdef ResultSet result = NULL + cdef ErrorCode code = 0 + cdef int64_t c_start_time = start_time + cdef int64_t c_end_time = end_time + cdef int c_offset = offset + cdef int c_limit = limit + if count <= 0: + raise ValueError("prepared_list must not be empty") + handles = malloc( + count * sizeof(PreparedSeriesHandle)) + if handles == NULL: + raise MemoryError() + try: + for i in range(count): + prepared = prepared_list[i] + if prepared.prepared == NULL: + raise RuntimeError("PreparedSeries is closed") + handles[i] = prepared.prepared + with nogil: + result = tsfile_reader_query_prepared_multi( + self.reader, handles, count, + c_start_time, c_end_time, c_offset, c_limit, &code) + check_error(code, b"Failed to query aligned prepared series") + finally: + free(handles) + pyresult = ResultSetPy(self, True) + pyresult.init_c(result, "prepared") + self.activate_result_set_list.add(pyresult) + return pyresult + def notify_result_set_discard(self, result_set: ResultSetPy): """ Remove activate result set from activate_result_set_list, called when a result set close. From 3cb8342d8bc20c3aa50850270b71a358053d3744 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 20 Aug 2026 11:36:47 +0800 Subject: [PATCH 2/3] fix: resolve C++ CI failures --- .github/workflows/wheels.yml | 3 +++ cpp/src/reader/aligned_chunk_reader.cc | 5 ++++- cpp/test/dataset/dataset_index_test.cc | 15 +++++++++------ cpp/test/writer/tsfile_writer_test.cc | 21 ++++++++++++++++++++- python/setup.py | 10 +++++++++- 5 files changed, 45 insertions(+), 9 deletions(-) diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml index d95e58edf..64b5b6fe5 100644 --- a/.github/workflows/wheels.yml +++ b/.github/workflows/wheels.yml @@ -96,6 +96,7 @@ jobs: ./mvnw -Pwith-cpp clean package \ -DskipTests -Dspotless.check.skip=true -Dspotless.apply.skip=true \ -Dbuild.test=OFF \ + -Dtsfile.dependency.source=BUNDLED \ -Denable.lzma2=ON \ -Dcmake.args="-DCMAKE_OSX_DEPLOYMENT_TARGET=12.0" otool -l cpp/target/build/lib/libtsfile*.dylib | grep -A2 LC_VERSION_MIN_MACOSX || true @@ -138,6 +139,7 @@ jobs: chmod +x mvnw || true ./mvnw -Pwith-cpp clean package \ -DskipTests -Dbuild.test=OFF \ + -Dtsfile.dependency.source=BUNDLED \ -Denable.lzma2=ON \ -Dspotless.check.skip=true -Dspotless.apply.skip=true test -d cpp/target/build/lib && test -d cpp/target/build/include @@ -205,6 +207,7 @@ jobs: chmod +x mvnw || true ./mvnw -Pwith-cpp clean package \ -DskipTests -Dbuild.test=OFF \ + -Dtsfile.dependency.source=BUNDLED \ -Denable.lzma2=ON \ -Dspotless.check.skip=true -Dspotless.apply.skip=true test -d cpp/target/build/lib diff --git a/cpp/src/reader/aligned_chunk_reader.cc b/cpp/src/reader/aligned_chunk_reader.cc index a140c1e98..1eb9f4833 100644 --- a/cpp/src/reader/aligned_chunk_reader.cc +++ b/cpp/src/reader/aligned_chunk_reader.cc @@ -144,6 +144,10 @@ void AlignedChunkReader::destroy() { value_compressor_->after_uncompress(value_uncompressed_buf_); value_uncompressed_buf_ = nullptr; } + // Multi-value readers keep the current page's decompressed buffers in + // ValueColumnState. Release them while the column compressors are still + // alive; the columns are deleted below and cannot release them afterwards. + release_current_page_state(); value_page_col_notnull_bitmap_.clear(); value_page_col_notnull_bitmap_.shrink_to_fit(); if (time_decoder_ != nullptr) { @@ -213,7 +217,6 @@ void AlignedChunkReader::destroy() { // vector to actually release the storage, matching the chunk_pages_ / // page_all_times_ handling above. std::vector().swap(value_columns_); - release_current_page_state(); std::vector>().swap(per_page_times_); #ifdef ENABLE_THREADS decode_pool_ = nullptr; // borrowed, not owned diff --git a/cpp/test/dataset/dataset_index_test.cc b/cpp/test/dataset/dataset_index_test.cc index 298aa77be..697e86049 100644 --- a/cpp/test/dataset/dataset_index_test.cc +++ b/cpp/test/dataset/dataset_index_test.cc @@ -52,6 +52,13 @@ DatasetIndexSectionData fixed_section(DatasetIndexSectionType type, return section; } +template +void append_fixed_record(DatasetIndexSectionData& section, const T& record) { + const size_t offset = section.bytes.size(); + section.bytes.resize(offset + sizeof(T)); + std::memcpy(section.bytes.data() + offset, &record, sizeof(T)); +} + std::vector make_minimal_sections() { const std::string strings[] = {"table", "device", "value", "/tmp/a.tsfile"}; std::vector offsets(1, 0); @@ -202,15 +209,11 @@ TEST_F(DatasetIndexTest, RejectsDuplicateCanonicalTableNames) { if (section.type == DatasetIndexSectionType::TABLE_NAME_INDEX) { const TableNameIndexRecord duplicate = { dataset_index_name_hash("table", 5), 0, 1}; - const uint8_t* data = reinterpret_cast(&duplicate); - section.bytes.insert(section.bytes.end(), data, - data + sizeof(duplicate)); + append_fixed_record(section, duplicate); ++section.count; } else if (section.type == DatasetIndexSectionType::TABLE_RECORD) { const TableRecord duplicate = {0, 0, 0, 0, 0, 0, 0}; - const uint8_t* data = reinterpret_cast(&duplicate); - section.bytes.insert(section.bytes.end(), data, - data + sizeof(duplicate)); + append_fixed_record(section, duplicate); ++section.count; } } diff --git a/cpp/test/writer/tsfile_writer_test.cc b/cpp/test/writer/tsfile_writer_test.cc index 62d5167f3..401741c98 100644 --- a/cpp/test/writer/tsfile_writer_test.cc +++ b/cpp/test/writer/tsfile_writer_test.cc @@ -20,10 +20,17 @@ #include +#include #include #include #include +#ifdef _WIN32 +#include +#else +#include +#endif + #include "common/path.h" #include "common/record.h" #include "common/schema.h" @@ -78,6 +85,18 @@ class TsFileWriterTest : public ::testing::Test { for (int i = 0; i < length; ++i) { random_string += chars[dis(gen)]; } + + // CTest runs writer tests in separate processes. A clock-seeded + // generator can produce the same name when two processes start in + // the same clock tick, allowing one test to remove the other's file. +#ifdef _WIN32 + const auto process_id = static_cast(_getpid()); +#else + const auto process_id = static_cast(getpid()); +#endif + static std::atomic counter{0}; + random_string += "_" + std::to_string(process_id) + "_" + + std::to_string(counter.fetch_add(1)); return random_string; } @@ -1692,4 +1711,4 @@ TEST_F(TsFileWriterTest, WriterReuseAfterDestroyProducesValidSecondFile) { // wf was passed to init() but init() did not take ownership. delete wf; remove(second_path.c_str()); -} \ No newline at end of file +} diff --git a/python/setup.py b/python/setup.py index a69db1558..e74c3296f 100644 --- a/python/setup.py +++ b/python/setup.py @@ -263,8 +263,16 @@ def finalize_options(self): runtime_library_dirs=runtime_library_dirs, ) +merge_common = common.copy() +merge_common.update( + library_dirs=[], + libraries=[], + extra_link_args=[], + runtime_library_dirs=[], +) + exts = [ - Extension("tsfile.dataset._merge", ["tsfile/dataset/_merge.pyx"], **common), + Extension("tsfile.dataset._merge", ["tsfile/dataset/_merge.pyx"], **merge_common), Extension("tsfile.tsfile_py_cpp", ["tsfile/tsfile_py_cpp.pyx"], **common), Extension("tsfile.tsfile_reader", ["tsfile/tsfile_reader.pyx"], **common), Extension("tsfile.tsfile_writer", ["tsfile/tsfile_writer.pyx"], **common), From 33a15fbd70ba2caf68585d868bdd2904b1f75126 Mon Sep 17 00:00:00 2001 From: ColinLee Date: Thu, 20 Aug 2026 18:57:15 +0800 Subject: [PATCH 3/3] fix(dataset): unify numeric field whitelist and union-align non-aligned loc reads --- python/tests/test_tsfile_dataset.py | 68 +++++++++++++++++++++++++++++ python/tsfile/constants.py | 14 ++++++ python/tsfile/dataset/reader.py | 11 +---- python/tsfile/dataset/runtime.py | 36 +++++++-------- 4 files changed, 100 insertions(+), 29 deletions(-) diff --git a/python/tests/test_tsfile_dataset.py b/python/tests/test_tsfile_dataset.py index d24a3dafc..55115987e 100644 --- a/python/tests/test_tsfile_dataset.py +++ b/python/tests/test_tsfile_dataset.py @@ -1067,6 +1067,23 @@ def test_dataset_rejects_nonnumeric_declared_schema_difference(tmp_path): TsFileDataFrame([str(path1), str(path2)], show_progress=False) +def test_dataset_table_model_omits_boolean_fields(tmp_path): + # The numeric dataset surface must drop BOOLEAN fields (they are not + # numeric) rather than surfacing them as a series that crashes the runtime + # read path. Regression for the reader/runtime field-type whitelist drift. + path = tmp_path / "weather_boolean.tsfile" + _write_weather_boolean_status_file(path, 0) + + with TsFileDataFrame(str(path), show_progress=False) as tsdf: + assert tsdf.model == "table" + assert tsdf.list_timeseries() == ["weather.device_a.temperature"] + with pytest.raises(KeyError): + tsdf["weather.device_a.status"] + np.testing.assert_array_equal( + tsdf["weather.device_a.temperature"][:], np.array([20.0, 21.0]) + ) + + def test_dataset_close_waits_for_an_active_public_query(tmp_path, monkeypatch): path = tmp_path / "weather.tsfile" _write_weather_file(path, 0) @@ -2075,3 +2092,54 @@ def test_dataset_tree_model_omits_non_numeric_measurements(tmp_path): np.testing.assert_array_equal( tsdf["root.a.b.temp"][:], np.array([0.5, 1.5, 2.5]) ) + + +def test_dataset_tree_model_loc_aligns_sparse_non_aligned_fields(tmp_path): + # Non-aligned (tree) device whose two measurements are sampled at + # different timestamps must produce a timestamp union with NaN fill, + # instead of raising when the per-field timelines differ. Regression for + # read_device_fields_by_time_range on non-aligned devices. + from tsfile import Field, RowRecord, TimeseriesSchema, TsFileWriter + + path = tmp_path / "sparse_tree.tsfile" + writer = TsFileWriter(str(path)) + writer.register_timeseries("root.a.b", TimeseriesSchema("m1", TSDataType.DOUBLE)) + writer.register_timeseries("root.a.b", TimeseriesSchema("m2", TSDataType.DOUBLE)) + writer.write_row_record( + RowRecord("root.a.b", 0, [Field("m1", 0.5, TSDataType.DOUBLE)]) + ) + writer.write_row_record( + RowRecord("root.a.b", 1, [Field("m2", 10.5, TSDataType.DOUBLE)]) + ) + writer.write_row_record( + RowRecord("root.a.b", 2, [Field("m1", 2.5, TSDataType.DOUBLE)]) + ) + writer.write_row_record( + RowRecord("root.a.b", 3, [Field("m2", 30.5, TSDataType.DOUBLE)]) + ) + writer.write_row_record( + RowRecord("root.a.b", 4, [Field("m1", 4.5, TSDataType.DOUBLE)]) + ) + writer.close() + + with TsFileDataFrame(str(path), show_progress=False) as tsdf: + assert tsdf.model == "tree" + aligned = tsdf.loc[0:5, ["root.a.b.m1", "root.a.b.m2"]] + assert isinstance(aligned, AlignedTimeseries) + assert aligned.series_names == ["root.a.b.m1", "root.a.b.m2"] + np.testing.assert_array_equal( + aligned.timestamps, np.array([0, 1, 2, 3, 4], dtype=np.int64) + ) + np.testing.assert_allclose( + aligned.values, + np.array( + [ + [0.5, np.nan], + [np.nan, 10.5], + [2.5, np.nan], + [np.nan, 30.5], + [4.5, np.nan], + ] + ), + equal_nan=True, + ) diff --git a/python/tsfile/constants.py b/python/tsfile/constants.py index 17b3ba34d..659988c5b 100644 --- a/python/tsfile/constants.py +++ b/python/tsfile/constants.py @@ -158,6 +158,20 @@ def from_pandas_datatype(cls, dtype): return cls.STRING +# Field data types exposed by the numeric dataset surface. The dataset reads +# value columns as float64 (missing values become NaN), so only numeric field +# types are exposed. This must stay in sync with the runtime read path in +# python/tsfile/dataset/runtime.py, which decodes the same set from the index. +NUMERIC_DATASET_FIELD_TYPES = frozenset( + { + TSDataType.INT32, + TSDataType.INT64, + TSDataType.FLOAT, + TSDataType.DOUBLE, + } +) + + _TSDATATYPE_COMPATIBLE_SOURCES = { TSDataType.INT64: (TSDataType.INT32, TSDataType.TIMESTAMP), TSDataType.STRING: (TSDataType.TEXT,), diff --git a/python/tsfile/dataset/reader.py b/python/tsfile/dataset/reader.py index e22467a8d..f88049cce 100644 --- a/python/tsfile/dataset/reader.py +++ b/python/tsfile/dataset/reader.py @@ -24,7 +24,7 @@ import numpy as np -from ..constants import ColumnCategory, TSDataType +from ..constants import ColumnCategory, NUMERIC_DATASET_FIELD_TYPES, TSDataType from ..tag_filter import tag_eq, tag_is_null from ..tsfile_reader import TsFileReaderPy from .metadata import ( @@ -36,14 +36,7 @@ resolve_series_path, ) -_NUMERIC_FIELD_TYPES = { - TSDataType.BOOLEAN, - TSDataType.INT32, - TSDataType.INT64, - TSDataType.FLOAT, - TSDataType.DOUBLE, - TSDataType.TIMESTAMP, -} +_NUMERIC_FIELD_TYPES = NUMERIC_DATASET_FIELD_TYPES def _to_python_scalar(value): diff --git a/python/tsfile/dataset/runtime.py b/python/tsfile/dataset/runtime.py index 76e58e505..a44bbf49f 100644 --- a/python/tsfile/dataset/runtime.py +++ b/python/tsfile/dataset/runtime.py @@ -30,7 +30,7 @@ import numpy as np -from ..constants import TSDataType +from ..constants import NUMERIC_DATASET_FIELD_TYPES, TSDataType from ..tag_filter import tag_eq, tag_is_null from ..tsfile_reader import TsFileReaderPy from .index import ( @@ -53,6 +53,7 @@ _join_series_path, split_logical_series_path, ) +from .merge import build_aligned_matrix _SERIES_DESCRIPTOR_CACHE_SIZE = 4096 @@ -473,10 +474,7 @@ def __getitem__(self, name): if column[7] == 0: tags.append((column[2], column_name, TSDataType(column[3]))) elif column[7] == 1 and column[3] in { - int(TSDataType.INT32), - int(TSDataType.INT64), - int(TSDataType.FLOAT), - int(TSDataType.DOUBLE), + int(data_type) for data_type in NUMERIC_DATASET_FIELD_TYPES }: fields.append((column[2], column_name, TSDataType(column[3]))) tags.sort() @@ -995,19 +993,17 @@ def read_device_fields_by_time_range( ) return self._consume_multi(result, column_names) - parts = [ - self.read_series_by_ref(device_id, column_id, start_time, end_time) - for column_id in column_ids - ] - if not parts: - return np.array([], dtype=np.int64), {} - # Existing dataframe merge aligns separate field parts across files; - # this method is only a compatibility surface for a single file. - timestamps = parts[0][0] - values = {} - for column_id, (current_timestamps, current_values) in zip(column_ids, parts): - if not np.array_equal(current_timestamps, timestamps): - raise ValueError("single-file fields do not share one timeline") + # Non-aligned device (or fields spanning different device spans): each + # field carries its own timeline, so align them onto a single timestamp + # union with NaN for missing values. This matches the cross-file merge + # semantics of build_aligned_matrix rather than requiring identical + # per-field timelines. + parts = {} + for column_id in column_ids: _, _, _, name = self._identity(device_id, column_id) - values[name] = current_values - return timestamps, values + parts[name] = self.read_series_by_ref( + device_id, column_id, start_time, end_time + ) + names = list(parts) + timestamps, matrix = build_aligned_matrix(names, parts) + return timestamps, {name: matrix[:, index] for index, name in enumerate(names)}