From 5a46c0e579e4b78a4e2aef4e0487c652cb694835 Mon Sep 17 00:00:00 2001 From: lgqss Date: Tue, 10 Feb 2026 16:00:59 +0800 Subject: [PATCH 1/5] merge_2025_12_05 --- include/column/column_record.h | 8 +- include/column/file_manager.h | 28 +- include/column/parquet_cache.h | 3 +- include/column/row2column.h | 118 ++- include/common/baikal_heartbeat.h | 1 + include/common/common.h | 72 +- include/common/data_stream_manager.h | 1 + include/common/datetime.h | 2 + include/common/expr_value.h | 2 +- include/common/file_system.h | 31 +- include/common/lru_cache.h | 1 + include/common/lru_cache.hpp | 17 + include/common/mut_table_key.h | 17 + include/common/schema_factory.h | 47 +- include/common/statistics.h | 47 +- include/common/type_utils.h | 1 + include/engine/external_filesystem.h | 6 +- include/engine/qos.h | 2 - include/engine/rocks_wrapper.h | 8 +- include/engine/rocksdb_compaction_service.h | 196 +++- include/engine/rocksdb_merge_operator.h | 5 +- include/engine/split_compaction_filter.h | 14 +- include/engine/sst_file_writer.h | 6 +- include/engine/table_iterator.h | 2 +- include/engine/transaction.h | 11 +- include/engine/transaction_pool.h | 12 +- include/exec/access_path.h | 1 + include/exec/agg_node.h | 6 + include/exec/exchange_receiver_node.h | 16 +- include/exec/exchange_sender_node.h | 28 +- include/exec/exec_node.h | 3 +- include/exec/fetcher_store.h | 10 + include/exec/filter_node.h | 4 + include/exec/index_ddl_manager_node.h | 29 + include/exec/join_node.h | 56 +- include/exec/joiner.h | 4 +- include/exec/mpp_property.h | 3 +- include/exec/packet_node.h | 3 +- include/exec/parquet_scan_node.h | 39 +- include/exec/property.h | 1 + include/exec/rocksdb_scan_node.h | 10 + include/exec/scan_node.h | 9 +- include/exec/select_manager_node.h | 40 +- include/exec/sort_node.h | 6 +- include/exec/vectorize_helpper.h | 444 ++++++++- include/expr/agg_fn_call.h | 5 + include/expr/arrow_function.h | 29 +- include/expr/predicate.h | 11 +- include/expr/row_expr.h | 2 +- include/expr/scalar_fn_call.h | 14 +- include/expr/window_fn_call.h | 13 +- include/logical_plan/ddl_planner.h | 1 + include/logical_plan/query_context.h | 9 + include/meta_server/region_manager.h | 14 +- include/meta_server/table_manager.h | 46 +- include/physical_plan/fragment.h | 2 + include/physical_plan/index_selector.h | 25 +- include/physical_plan/mpp_analyzer.h | 3 + include/physical_plan/physical_planner.h | 1 + include/physical_plan/plan_router.h | 4 +- include/protocol/db_service.h | 42 + include/protocol/show_helper.h | 1 + .../reverse/boolean_engine/boolean_executor.h | 17 +- .../reverse/boolean_engine/logical_query.h | 21 +- .../reverse/boolean_engine/logical_query.hpp | 114 ++- include/reverse/reverse_arrow.h | 16 +- include/reverse/reverse_common.h | 10 +- include/reverse/reverse_common.hpp | 8 +- include/reverse/reverse_index.h | 10 +- include/reverse/reverse_index.hpp | 4 +- include/reverse/reverse_interface.h | 25 +- include/reverse/reverse_interface.hpp | 56 +- include/runtime/arrow_io_excutor.h | 4 + include/runtime/row_batch.h | 7 +- include/runtime/runtime_state.h | 4 +- include/session/user_info.h | 2 +- include/sqlparser/gen_source.sh | 3 +- include/sqlparser/sql_parse.y | 7 + include/store/meta_writer.h | 7 +- include/store/region.h | 121 ++- include/store/store.h | 6 +- include/vector_index/vector_index.h | 56 +- proto/common.proto | 1 + proto/compaction.interface.proto | 2 +- proto/dms.proto | 22 + proto/expr.proto | 1 + proto/meta.interface.proto | 29 + proto/optype.proto | 2 +- proto/plan.proto | 11 +- proto/store.interface.proto | 3 +- src/column/column_record.cpp | 307 ++----- src/column/file_manager.cpp | 149 ++- src/column/parquet_cache.cpp | 4 + src/column/region_column.cpp | 141 +-- src/column/row2column.cpp | 248 ++++- src/common/baikal_heartbeat.cpp | 21 +- src/common/common.cpp | 8 +- src/common/datetime.cpp | 315 ++++--- src/common/default_room_define.cpp | 4 +- src/common/expr_value.cpp | 2 +- src/common/file_system.cpp | 102 ++- src/common/information_schema.cpp | 111 ++- src/common/memory_profile.cpp | 5 + src/common/meta_server_interact.cpp | 2 +- src/common/schema_factory.cpp | 121 ++- src/engine/external_filesystem.cpp | 14 +- src/engine/olap_pre_split.cpp | 2 +- src/engine/qos.cpp | 2 - src/engine/rocks_wrapper.cpp | 79 +- src/engine/rocksdb_compaction_service.cpp | 6 +- src/engine/rocksdb_merge_operator.cpp | 122 ++- src/engine/table_iterator.cpp | 14 +- src/engine/transaction.cpp | 29 +- src/engine/transaction_pool.cpp | 4 +- src/exec/access_path.cpp | 58 +- src/exec/apply_node.cpp | 2 +- src/exec/dml_node.cpp | 7 +- src/exec/exchange_receiver_node.cpp | 104 ++- src/exec/exchange_sender_node.cpp | 65 +- src/exec/exec_node.cpp | 29 +- src/exec/fetcher_store.cpp | 129 ++- src/exec/filter_node.cpp | 31 +- src/exec/full_export_node.cpp | 1 - src/exec/index_ddl_manager_node.cpp | 131 ++- src/exec/join_node.cpp | 484 ++++++++-- src/exec/joiner.cpp | 59 +- src/exec/mysql_scan_node.cpp | 30 +- src/exec/packet_node.cpp | 83 +- src/exec/parquet_scan_node.cpp | 233 +++-- src/exec/rocksdb_scan_node.cpp | 101 +- src/exec/scan_node.cpp | 12 +- src/exec/select_manager_node.cpp | 168 +++- src/exec/sort_node.cpp | 9 + src/exec/update_manager_node.cpp | 5 +- src/exec/window_node.cpp | 6 +- src/expr/agg_fn_call.cpp | 91 +- src/expr/arrow_agg_function.cpp | 138 +++ src/expr/arrow_function.cpp | 286 +++++- src/expr/arrow_string_function.cpp | 100 ++ src/expr/arrow_time_function.cpp | 115 +-- src/expr/expr_node.cpp | 4 +- src/expr/fn_manager.cpp | 2 +- src/expr/internal_functions.cpp | 4 +- src/expr/predicate.cpp | 10 +- src/expr/scalar_fn_call.cpp | 10 +- src/expr/window_fn_call.cpp | 150 ++- src/logical_plan/ddl_planner.cpp | 199 +++- src/logical_plan/ddl_work_planner.cpp | 1 - src/logical_plan/kill_planner.cpp | 8 +- src/logical_plan/logical_planner.cpp | 25 +- src/logical_plan/select_planner.cpp | 2 +- src/logical_plan/update_planner.cpp | 2 +- src/meta_server/common_state_machine.cpp | 10 +- src/meta_server/ddl_manager.cpp | 24 +- src/meta_server/meta_server.cpp | 4 + src/meta_server/meta_state_machine.cpp | 4 +- src/meta_server/region_manager.cpp | 9 +- src/meta_server/schema_manager.cpp | 99 +- src/meta_server/table_manager.cpp | 850 ++++++++++++++--- src/meta_server/tso_state_machine.cpp | 2 +- src/physical_plan/exec_type_analyzer.cpp | 40 +- src/physical_plan/index_selector.cpp | 289 +++++- src/physical_plan/limit_calc.cpp | 4 + src/physical_plan/mpp_analyzer.cpp | 137 ++- src/physical_plan/physical_planner.cpp | 70 +- src/physical_plan/plan_router.cpp | 18 +- src/physical_plan/separate.cpp | 3 +- src/protocol/db_service.cpp | 52 +- src/protocol/handle_helper.cpp | 23 +- src/protocol/main.cpp | 3 +- src/protocol/network_server.cpp | 64 +- src/protocol/show_helper.cpp | 102 ++- src/protocol/state_machine.cpp | 13 +- src/raft/log_entry_reader.cpp | 3 + src/raft/my_raft_log_storage.cpp | 1 + src/reverse/reverse_common.cpp | 12 +- src/reverse/reverse_interface.cpp | 2 +- src/runtime/arrow_io_excutor.cpp | 25 +- src/runtime/chunk.cpp | 18 +- src/session/binlog_context.cpp | 2 +- src/store/backup.cpp | 2 +- src/store/main.cpp | 7 +- src/store/meta_writer.cpp | 49 +- src/store/region.cpp | 520 +++++++++-- src/store/region_binlog.cpp | 144 ++- src/store/region_control.cpp | 5 +- src/store/region_olap.cpp | 2 +- src/store/store.cpp | 65 +- src/vector_index/vector_index.cpp | 477 +++++++--- test/test_arrow_compute.cpp | 130 ++- test/test_common.cpp | 25 +- test/test_date_time.cpp | 34 + test/test_dms.cpp | 1 - test/test_partition_meta.cpp | 2 + test/test_privilege_manager.cpp | 8 + test/test_region_manager.cpp | 9 +- test/test_rocksdb.cpp | 867 +++++++++++++++++- test/test_schema_manager.cpp | 4 +- test/test_table_manager.cpp | 11 +- 199 files changed, 8891 insertions(+), 2221 deletions(-) diff --git a/include/column/column_record.h b/include/column/column_record.h index 2b250e0a9..2d475b310 100644 --- a/include/column/column_record.h +++ b/include/column/column_record.h @@ -77,11 +77,11 @@ class ColumnRecord { const std::vector& fields, int64_t& userid) { return encode_row_key(record_batch, record_batch->num_rows() - 1, fields, userid); } - - static std::shared_ptr make_schema(const std::string& name, arrow::Type::type type); static ExprValue get_vectorized_value(const std::shared_ptr& array, int row_idx); - static std::shared_ptr make_array_from_exprvalue( - const pb::PrimitiveType type, const ExprValue& expr_value, const int length); + + static std::shared_ptr make_column_schema(int64_t tableid, + SmartTable table_info, SmartIndex pri_info, + const std::unordered_map& field_id2info_map); int init(); diff --git a/include/column/file_manager.h b/include/column/file_manager.h index cd60a0e91..a52a40e8e 100644 --- a/include/column/file_manager.h +++ b/include/column/file_manager.h @@ -187,7 +187,7 @@ class ParquetFile { // 获取每个parquet文件中符合条件的rowgroup和rowranges int get_qualified_rowgroup_and_rowranges( - const std::vector& key_ranges, + const pb::PossibleIndex& possible_index, std::vector& rowgroup_indices, std::vector>>& rowranges); @@ -229,7 +229,8 @@ class ParquetFile { ::arrow::Status GetRecordBatchReader(std::unique_ptr<::arrow::RecordBatchReader>* out); static bool check_interval_overlapped( - const pb::PossibleIndex::Range& index_range, const std::string& file_start_key, const std::string& file_end_key); + const pb::PossibleIndex::Range& index_range, bool is_eq, bool is_left_open, bool is_right_open, + const std::string& file_start_key, const std::string& file_end_key); private: std::shared_ptr _file_info; @@ -243,22 +244,19 @@ class ParquetFile { }; struct ParquetFileReaderOptions { - bool need_order_info = true; - int64_t raftindex = 0; - std::shared_ptr schema_info = nullptr; - std::shared_ptr file_info = nullptr; + pb::PossibleIndex* pos_index = nullptr; + std::map lower_short_name_fields; + std::shared_ptr schema = nullptr; }; class ParquetFileReader : public ::arrow::RecordBatchReader { public: - ParquetFileReader(ParquetFileReaderOptions& options) : _options(options) { - _parquet_file = std::make_shared(_options.file_info); - } + ParquetFileReader(ParquetFileReaderOptions& options, std::shared_ptr file) : _options(options), _parquet_file(file) { } virtual ~ParquetFileReader() { } int init(); - std::shared_ptr schema() const override { return nullptr; } + std::shared_ptr schema() const override { return _options.schema; } virtual ::arrow::Status ReadNext(std::shared_ptr<::arrow::RecordBatch>* batch) override; private: @@ -266,6 +264,9 @@ class ParquetFileReader : public ::arrow::RecordBatchReader { std::shared_ptr _parquet_file; std::unique_ptr<::arrow::RecordBatchReader> _reader; bool _init = false; + int64_t _raftindex = 0; + std::shared_ptr _read_contents = nullptr; + }; class ParquetFileManager { @@ -275,6 +276,11 @@ class ParquetFileManager { return &instance; } + void close() { + std::unique_lock l(_mutex); + _lru_cache.clear(); + } + bool link_file(const std::string& old_path, const std::string& new_path) { return ::link(old_path.c_str(), new_path.c_str()) == 0; } @@ -335,7 +341,7 @@ class ColumnFileManager { int load_snapshot(bool restart); int pick_minor_compact_file(int64_t applied_index, int64_t& start_version); int pick_major_compact_file(std::vector>& file_infos); - int pick_base_compact_file(std::vector>& file_infos); + int pick_base_compact_file(std::vector>& file_infos, bool only_read_base); int finish_minor_compact(const std::shared_ptr& new_file, int64_t last_max_version); int finish_major_compact(const std::vector>& old_files, const std::vector>& new_files, bool is_base); diff --git a/include/column/parquet_cache.h b/include/column/parquet_cache.h index 1465708a3..b1156dc2d 100644 --- a/include/column/parquet_cache.h +++ b/include/column/parquet_cache.h @@ -16,6 +16,7 @@ #include "file_system.h" #include "column_record.h" #include "rocksdb_filesystem.h" +#include "arrow_io_excutor.h" namespace baikaldb { DECLARE_int64(parquet_cache_size_mb); @@ -238,7 +239,7 @@ class ParquetArrowReader { class ParquetArrowReadableFile : public ::arrow::io::RandomAccessFile { public: - ParquetArrowReadableFile(const std::shared_ptr& reader, int64_t size, ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()) : + ParquetArrowReadableFile(const std::shared_ptr& reader, int64_t size, ::arrow::MemoryPool* pool = GetMemoryPoolForRead()) : _file_reader(reader), _file_size(size), _pool(pool) { } ~ParquetArrowReadableFile() override { diff --git a/include/column/row2column.h b/include/column/row2column.h index 1a156acef..95189aa57 100644 --- a/include/column/row2column.h +++ b/include/column/row2column.h @@ -77,39 +77,143 @@ class RocksdbBaseReader : public Row2ColumnReader { bool _init = false; }; +struct RaftLogCacheIter { + TimeCost begin_time; + int64_t commit_index = -1; + std::map log_index_req_map; +}; + +struct RaftLogCache { + bool empty() { + return txn_id_raft_log_map.empty(); + } + + int64_t time() { + int64_t t = 0; + for (auto iter : txn_id_raft_log_map) { + if (t < iter.second->begin_time.get_time()) { + t = iter.second->begin_time.get_time(); + } + } + + return t; + } + + std::map> txn_id_raft_log_map; +}; + +class RaftLogMgr { +public: + ~RaftLogMgr() {} + + static RaftLogMgr* get_instance() { + static RaftLogMgr _instance; + return &_instance; + } + + std::shared_ptr get_raft_log_cache(int64_t region_id) { + std::unique_lock l(_lock); + auto iter = _region_id_raft_log.find(region_id); + if (iter == _region_id_raft_log.end()) { + auto cache = std::make_shared(); + return cache; + } else { + auto cache = iter->second; + _region_id_raft_log.erase(iter); + return cache; + } + } + + void release_raft_log_cache(int64_t region_id, std::shared_ptr cache) { + if (!cache->empty()) { + std::unique_lock l(_lock); + if (_region_id_raft_log.count(region_id) > 0) { + DB_COLUMN_FATAL("region_id: %ld, cache not empty", region_id); + } + _region_id_raft_log[region_id] = cache; + } + } + + +private: + bthread::Mutex _lock; + std::map> _region_id_raft_log; + +private: + RaftLogMgr() {} + DISALLOW_COPY_AND_ASSIGN(RaftLogMgr); +}; + class RaftLogReader : public Row2ColumnReader { public: - RaftLogReader(const Row2ColOptions& options) : Row2ColumnReader(options) { + RaftLogReader(const Row2ColOptions& options) : Row2ColumnReader(options), _region_id(options.region_id) { _column_record = std::make_shared(_schema_info->schema_with_order_info, _options.read_batch_size); + _raft_log_cache = RaftLogMgr::get_instance()->get_raft_log_cache(_region_id); } - virtual ~RaftLogReader() {} + + virtual ~RaftLogReader() { + int64_t time_cost = 0; + if (!_raft_log_cache->txn_id_raft_log_map.empty()) { + time_cost = _raft_log_cache->time(); + if (time_cost > 15 * 60 * 1000 * 1000ULL) { + DB_COLUMN_FATAL("region_id: %ld, time_cost: %ld", _region_id, time_cost); + } + } + + RaftLogMgr::get_instance()->release_raft_log_cache(_region_id, _raft_log_cache); + } + int init(); + + int get_raft_log(int64_t start_index, int64_t end_index, uint64_t txn_id, std::map& pre_reqs); + arrow::Status ReadNext(std::shared_ptr* out) override; + std::shared_ptr<::arrow::Schema> schema() const override { return _schema_info->schema_with_order_info; } + + int delete_column_txn_log_index() { + return MetaWriter::get_instance()->delete_column_txn_log_index(_region_id, _txn_ids); + } + int64_t get_last_raft_index() const { return _last_index; } + int64_t row_count() const { return _total_row_nums; } + int64_t put_count() const { return _put_count; } + int64_t delete_count() const { return _delete_count; } + void commit(int64_t txn_id, int64_t raft_index); + + void rollback(int64_t txn_id, int64_t raft_index); + + void insert(int64_t txn_id, int64_t raft_index, pb::StoreReq& request); + + int get(int64_t txn_id, std::map& log_index_req_map); + private: - int64_t _first_index = -1; - int64_t _last_index = -1; - int64_t _skip_count = 0; - int64_t _put_count = 0; + int64_t _first_index = -1; + int64_t _last_index = -1; + int64_t _skip_count = 0; + int64_t _put_count = 0; int64_t _delete_count = 0; - int64_t _merge_count = 0; + int64_t _merge_count = 0; bool _init = false; std::vector> _batchs; + int64_t _region_id = 0; + int64_t _commited_txn_id = -1; + std::shared_ptr _raft_log_cache; + std::vector _txn_ids; }; } // baikaldb \ No newline at end of file diff --git a/include/common/baikal_heartbeat.h b/include/common/baikal_heartbeat.h index 464edc2a0..34bd4e784 100755 --- a/include/common/baikal_heartbeat.h +++ b/include/common/baikal_heartbeat.h @@ -24,6 +24,7 @@ struct HeartBeatTableName { std::string namespace_name; std::string database; std::string table_name; + std::set partition_ids; }; struct SubTableNames { diff --git a/include/common/common.h b/include/common/common.h index 45e857b23..ea82dace4 100644 --- a/include/common/common.h +++ b/include/common/common.h @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -54,6 +55,7 @@ #include "brpc/reloadable_flags.h" #endif #include +#include #include #include "log.h" #include "proto/common.pb.h" @@ -1150,11 +1152,12 @@ struct BvarMap { int64_t affected_rows, int64_t scan_rows, int64_t read_disk_size, int64_t filter_rows, int64_t region_count, const std::map& field_range_type_, - const uint64_t sign_, const std::set& subquery_signs_) + const uint64_t sign_, const std::set& subquery_signs_, + const std::string& tag) : table_id(table_id), sum(sum), err_sum(err_sum), count(count), err_count(err_count), affected_rows(affected_rows), scan_rows(scan_rows), read_disk_size(read_disk_size), filter_rows(filter_rows), - region_count(region_count) { + region_count(region_count), resource_tag(tag) { field_range_type = field_range_type_; parent_sign = sign_; if (subquery_signs_.size() > 0) { @@ -1177,6 +1180,9 @@ struct BvarMap { subquery_signs = other.subquery_signs; } } + if (resource_tag.empty()) { + resource_tag = other.resource_tag; + } sum += other.sum; err_sum += other.err_sum; @@ -1215,6 +1221,7 @@ struct BvarMap { std::map field_range_type; uint64_t parent_sign; std::set subquery_signs; + std::string resource_tag; }; public: BvarMap() {} @@ -1222,9 +1229,10 @@ struct BvarMap { int64_t affected_rows, int64_t scan_rows, int64_t read_disk_size, int64_t filter_rows, int64_t region_count, const std::map& field_range_type_, int64_t err_count, - uint64_t parent_sign, std::set& subquery_signs) { + uint64_t parent_sign, std::set& subquery_signs, + const std::string& tag) { internal_map[key][index_id] = SumCount(table_id, cost, err_cost, 1, err_count, affected_rows, - scan_rows, read_disk_size, filter_rows, region_count, field_range_type_, parent_sign, subquery_signs); + scan_rows, read_disk_size, filter_rows, region_count, field_range_type_, parent_sign, subquery_signs, tag); } BvarMap& operator+=(const BvarMap& other) { @@ -1376,6 +1384,9 @@ struct RocksdbVars { // 统计未提交的binlog最大时间 bvar::Maxer binlog_not_commit_max_cost; bvar::Window> binlog_not_commit_max_cost_minute; + // 对于vector memtable, 只要读取就需要排序，因此只记录读取次数而非扫描量 + bvar::Adder raft_log_scan_times_count; + bvar::PerSecond> raft_log_scan_count_qps; private: RocksdbVars(): rocksdb_put_time_cost_latency("rocksdb_put_time_cost_latency", &rocksdb_put_time, -1), @@ -1392,7 +1403,9 @@ struct RocksdbVars { qos_fetch_tokens_wait_count("qos_fetch_tokens_wait_count"), qos_fetch_tokens_qps("qos_fetch_tokens_qps", &qos_fetch_tokens_count), qos_token_waste_qps("qos_token_waste_qps", &qos_token_waste_count), - binlog_not_commit_max_cost_minute("binlog_not_commit_max_cost_minute", &binlog_not_commit_max_cost, 60) { + binlog_not_commit_max_cost_minute("binlog_not_commit_max_cost_minute", &binlog_not_commit_max_cost, 60), + raft_log_scan_times_count("raft_log_scan_times_count"), + raft_log_scan_count_qps("raft_log_scan_count_qps", &raft_log_scan_times_count, 60) { } }; @@ -1643,6 +1656,55 @@ class Singleton { } }; +template +class BlockingQueue { +public: + explicit BlockingQueue(size_t capacity = 0) : _capacity(capacity) {} + ~BlockingQueue() {} + + bool blocking_get(T* out) { + std::unique_lock l(_mutex); + while (!_shutdown && _queue.empty()) { + _not_empty_cv.wait(l); + } + if (!_queue.empty()) { + *out = _queue.front(); + _queue.pop(); + _not_full_cv.notify_one(); + return true; + } + return false; + } + + bool blocking_put(const T& val) { + std::unique_lock l(_mutex); + while (!_shutdown && _capacity != 0 && _queue.size() >= _capacity) { + _not_full_cv.wait(l); + } + if (!_shutdown) { + _queue.push(val); + _not_empty_cv.notify_one(); + return true; + } + return false; + } + + void shutdown() { + std::unique_lock l(_mutex); + _shutdown = true; + _not_empty_cv.notify_all(); + _not_full_cv.notify_all(); + } + +private: + const size_t _capacity = 0; // 0表示无界队列 + bool _shutdown = false; + bthread::Mutex _mutex; + bthread::ConditionVariable _not_empty_cv; + bthread::ConditionVariable _not_full_cv; + std::queue _queue; +}; + enum class IconvOnError { ABORT = 0, // 失败放弃 IGNORE = 1, // 忽略跳过 diff --git a/include/common/data_stream_manager.h b/include/common/data_stream_manager.h index 4d6352331..87ff80a0d 100644 --- a/include/common/data_stream_manager.h +++ b/include/common/data_stream_manager.h @@ -65,6 +65,7 @@ class DataStreamManager { std::hash{}(key.node_id); } }; + arrow::compute::Expression* _condition = nullptr; bthread::Mutex _mtx; std::unordered_map, ReceiverHasher> _receiver_map; }; diff --git a/include/common/datetime.h b/include/common/datetime.h index 6db1bcade..0b40814a5 100755 --- a/include/common/datetime.h +++ b/include/common/datetime.h @@ -17,6 +17,8 @@ #include "common.h" namespace baikaldb { +time_t mktime_fixed_r(const struct tm* tm, int tz_offset_hours = 8); +struct tm *localtime_fixed_r(const time_t *timep, struct tm *result, int tz_offset_hours = 8); extern std::string timestamp_to_str(time_t timestamp, bool is_utc = false); extern time_t str_to_timestamp(const char* str_time); diff --git a/include/common/expr_value.h b/include/common/expr_value.h index 5f32a058e..a94d31e0b 100644 --- a/include/common/expr_value.h +++ b/include/common/expr_value.h @@ -1013,7 +1013,7 @@ struct ExprValue { } bool is_string() const { - return type == pb::STRING || type == pb::HEX || type == pb::BITMAP || type == pb::HLL || type == pb::TDIGEST || type == pb::JSON; + return ::baikaldb::is_string(type); } bool is_double() const { diff --git a/include/common/file_system.h b/include/common/file_system.h index 1d729a042..17512f7d0 100644 --- a/include/common/file_system.h +++ b/include/common/file_system.h @@ -33,9 +33,11 @@ #include #include #include + #include "baidu/inf/afs-api/client/afs_filesystem.h" #include "baidu/inf/afs-api/common/afs_common.h" #include "baidu/inf/afs-api/client/afs_impl.h" + #else #include #include @@ -53,6 +55,14 @@ DECLARE_int32(file_buffer_size); typedef boost::filesystem::directory_iterator dir_iter; +extern const char* AFS_CLIENT_CONF_PATH; + +struct MemBuf : std::streambuf { + MemBuf(char* begin, char* end) { + this->setg(begin, begin, end); + } +}; + enum class FileMode { I_FILE, I_DIR, @@ -74,7 +84,8 @@ class FileWriter { virtual bool sync() = 0; }; -#ifdef BAIDU_INTERNAL + +#if defined(BAIDU_INTERNAL) class AfsFileWriter : public FileWriter { public: AfsFileWriter(afs::Writer* writer) : _writer(writer) {} @@ -132,7 +143,8 @@ class PosixFileReader : public FileReader { bool _error = false; }; -#ifdef BAIDU_INTERNAL + +#if defined(BAIDU_INTERNAL) class AfsFileReader : public FileReader { public: AfsFileReader(afs::Reader* reader) : _reader(reader) {} @@ -151,6 +163,11 @@ class AfsFileReader : public FileReader { }; #endif +struct DirEntry { + FileMode mode; + std::string path; +}; + class FileSystem { public: FileSystem(bool is_posix) : _is_posix(is_posix) {} @@ -162,6 +179,7 @@ class FileSystem { virtual std::shared_ptr open_writer(const std::string& path, bool is_create) = 0; virtual int close_writer(std::shared_ptr file_writer) = 0; virtual int read_dir(const std::string& path, std::vector& direntrys) = 0; + virtual int read_dir(const std::string& path, std::vector& direntrys) = 0; virtual int get_file_info(const std::string& path, FileInfo& file_info, std::string* err_msg) = 0; virtual int destroy() = 0; virtual int delete_path(const std::string& path, bool recursive = false) = 0; @@ -206,6 +224,7 @@ class PosixFileSystem : public FileSystem { } virtual int read_dir(const std::string& path, std::vector& direntrys) override; + virtual int read_dir(const std::string& path, std::vector& direntrys) override; virtual int get_file_info(const std::string& path, FileInfo& file_info, std::string* err_msg) override; // Parquet读取使用 @@ -213,7 +232,8 @@ class PosixFileSystem : public FileSystem { virtual ::arrow::Status close_arrow_reader(std::shared_ptr<::arrow::io::RandomAccessFile> arrow_reader) override; }; -#ifdef BAIDU_INTERNAL + +#if defined(BAIDU_INTERNAL) class AfsFileSystem : public FileSystem { public: AfsFileSystem(const std::string& afs_uri, @@ -238,6 +258,7 @@ class AfsFileSystem : public FileSystem { virtual int delete_path(const std::string& path, bool recursive = false) override; virtual int read_dir(const std::string& path, std::vector& direntrys) override; + virtual int read_dir(const std::string& path, std::vector& direntrys) override; virtual int get_file_info(const std::string& path, FileInfo& file_info, std::string* err_msg) override; // Parquet使用 @@ -285,6 +306,10 @@ class ReadDirImpl { // return 0 : success; entry is valid // return 1 : finish; entry is not valid int next_entry(std::string& entry); + // 获取该目录下所有的文件 + static int get_all_files(FileSystem* fs, const std::string& path, std::vector& files); + // 获取该目录下所有的子目录 + static int get_all_dirs(FileSystem* fs, const std::string& path, std::vector& dirs); private: size_t _idx = 0; diff --git a/include/common/lru_cache.h b/include/common/lru_cache.h index 1367e2a89..b08d3383f 100644 --- a/include/common/lru_cache.h +++ b/include/common/lru_cache.h @@ -49,6 +49,7 @@ class Cache { int find(const ItemKey& key, ItemType* value); int add(const ItemKey& key, const ItemType& value); int del(const ItemKey& key); + int clear(); private: //双链表，从尾部插入数据，超过阈值数据从头部删除 butil::LinkedList> _lru_list; diff --git a/include/common/lru_cache.hpp b/include/common/lru_cache.hpp index ef02807d7..e804d3c31 100644 --- a/include/common/lru_cache.hpp +++ b/include/common/lru_cache.hpp @@ -97,6 +97,23 @@ int Cache::del(const ItemKey& key) { return 0; } +template +int Cache::clear() { + std::lock_guard lock(_mutex); + std::vector*> nodes; + nodes.reserve(_lru_map.size()); + for (auto& pair : _lru_map) { + LruNode* node = pair.second; + node->RemoveFromList(); + nodes.emplace_back(node); + } + _lru_map.clear(); + for (auto& node : nodes) { + delete node; + } + return 0; +} + } /* vim: set ts=4 sw=4 sts=4 tw=100 */ diff --git a/include/common/mut_table_key.h b/include/common/mut_table_key.h index 771df6add..438985430 100644 --- a/include/common/mut_table_key.h +++ b/include/common/mut_table_key.h @@ -37,6 +37,23 @@ class MutTableKey { _full(full), _data(key.data(), key.size()) { } + MutTableKey(MutTableKey&& key) noexcept : + _full(key.get_full()), + _data(std::move(key._data)) {} + + MutTableKey(const MutTableKey& key) : + _full(key.get_full()), + _data(key._data) {} + + MutTableKey& operator=(const MutTableKey& key) { + if (this == &key) { + return *this; + } + _full = key.get_full(); + _data = key._data; + return *this; + } + MutTableKey(const TableKey& key); MutTableKey& append_i8(int8_t val) { diff --git a/include/common/schema_factory.h b/include/common/schema_factory.h index 2f55e54a0..46a575bc4 100644 --- a/include/common/schema_factory.h +++ b/include/common/schema_factory.h @@ -165,6 +165,7 @@ struct DistInfo { struct TTLInfo { TTLInfo() { } int64_t ttl_duration_s = 0; // >0表示配置有ttl，单位s + std::shared_ptr ttl_field = nullptr; // ttl指定的列，为nullptr表示按照修改时间ttl int64_t online_ttl_expire_time_us = 0; // online ttl 过期时间 }; @@ -214,6 +215,7 @@ struct TableInfo { bool has_fulltext = false; bool has_rollup_index = false; bool has_vector_index = false; + std::set rollup_indexs; // 该表是否已和 binlog 表关联 bool is_linked = false; bool is_binlog = false; @@ -261,6 +263,19 @@ struct TableInfo { return nullptr; } + std::shared_ptr get_ttl_field() { + return ttl_info.ttl_field; + } + + FieldInfo* get_field_ptr(const std::string& short_name) { + for (auto& info : fields) { + if (info.short_name == short_name) { + return &info; + } + } + return nullptr; + } + int32_t get_field_id_by_short_name(const std::string& short_name) { for (auto& info : fields) { if (info.short_name == short_name) { @@ -269,7 +284,6 @@ struct TableInfo { } return -1; } - }; struct IndexInfo { @@ -290,6 +304,7 @@ struct IndexInfo { // all pk fields not in index fields // empty if this is a pk index + // 索引或主键包含变长字段时为完整pk std::vector pk_fields; // 索引定长(length>0)且overlap时, 主键字段在索引中的位置 @@ -309,15 +324,20 @@ struct IndexInfo { int64_t restore_time = -1; int64_t disable_time = -1; int32_t max_field_id = 0; + + // rollup + pb::RollupType rollup_type = pb::SUM; + int32_t publish_timestamp; + //vector index int32_t dimension = 0; - int32_t nprobe = 5; - pb::RollupType rollup_type = pb::SUM; - int32_t efsearch = 16; - int32_t efconstruction = 40; std::string vector_description; pb::MetricType metric_type = pb::METRIC_L2; - int32_t publish_timestamp; + // -- HNSW + int32_t efsearch = 16; + int32_t efconstruction = 40; + // -- IVF + int32_t nprobe = 5; }; struct DatabaseInfo { @@ -1254,6 +1274,8 @@ typedef ::google::protobuf::RepeatedPtrField StatisticsVec; std::map* region_infos); int get_partition_regions(int64_t table_id, const std::vector& partition_ids, std::map& region_infos); + int get_all_partition_ids( + int64_t table_id, std::vector& partition_ids); int check_region_ranges_consecutive(int64_t table_id); int get_region_by_key(int64_t main_table_id, IndexInfo& index, @@ -1757,10 +1779,18 @@ typedef ::google::protobuf::RepeatedPtrField StatisticsVec; void update_valiable_addresses(const std::vector& baikal_status_vec); void set_db_unavailable(const std::string& db_address); void update_meta_map(const std::string& meta_name); - + bool table_suitable_for_broadcast_join(const int64_t table_id); std::string get_address() const { return _my_address; } + void get_special_signs(std::unordered_map& special_signs) { + BAIDU_SCOPED_LOCK(_special_signs_mutex); + special_signs = _special_signs; + } + void set_special_signs(std::unordered_map& special_signs) { + BAIDU_SCOPED_LOCK(_special_signs_mutex); + _special_signs = special_signs; + } private: SchemaFactory() { _is_inited = false; @@ -1827,6 +1857,9 @@ typedef ::google::protobuf::RepeatedPtrField StatisticsVec; DoublBufferedDBAddresses _doubly_valiable_addresses; // ip:port std::string _my_address; + + bthread::Mutex _special_signs_mutex; + std::unordered_map _special_signs; }; } diff --git a/include/common/statistics.h b/include/common/statistics.h index 063bb91a6..10eeb6527 100644 --- a/include/common/statistics.h +++ b/include/common/statistics.h @@ -29,12 +29,15 @@ class Statistics { if (statistics.has_histogram()) { init_histogram(statistics.histogram()); _total_rows = statistics.histogram().total_rows(); + _is_valid = true; } if (statistics.has_cmsketch()) { init_cmsketch(statistics.cmsketch()); + _is_valid = true; } if (statistics.has_hll()) { - init_hll(statistics.hll()); + init_hll(statistics.hll()); + _is_valid = true; } } @@ -50,6 +53,10 @@ class Statistics { return _total_rows; } + bool is_valid() const { + return _is_valid; + } + std::shared_ptr get_cmsketchcolumn_ptr(int field_id) { if (field_id <= 0) { return nullptr; @@ -84,19 +91,21 @@ class Statistics { } void histogram_to_string(std::vector>& rows, std::vector& fields) { - if (_field_histogram.size() != fields.size()) { - DB_FATAL("use select * from table_name"); - return; - } - int i = 0; - for (auto iter = _field_histogram.begin(); iter != _field_histogram.end(); iter++) { + auto statistics_fields = get_field_names(); + for (const auto& field : fields) { + std::string field_name = field.name; + auto field_name_id = statistics_fields.find(field_name); + if (field_name_id == statistics_fields.end()) { + continue; + } + auto iter = _field_histogram.find(field_name_id->second); std::vector row; - row.push_back(std::to_string(iter->first)); - row.push_back(fields[i++].name); - row.push_back(std::to_string(iter->second->get_distinct_cnt())); - row.push_back(std::to_string(iter->second->get_null_value_cnt())); - row.push_back(std::to_string(iter->second->get_bucket_count())); - rows.push_back(row); + row.emplace_back(std::to_string(field_name_id->second)); + row.emplace_back(field_name); + row.emplace_back(std::to_string(iter->second->get_distinct_cnt())); + row.emplace_back(std::to_string(iter->second->get_null_value_cnt())); + row.emplace_back(std::to_string(iter->second->get_bucket_count())); + rows.emplace_back(row); } } @@ -132,7 +141,7 @@ class Statistics { //get_histogram_count返-2时说明超出取值范围时，根据need_mapping标记判断是否映射到已存在的范围，默认进行映射 double get_histogram_ratio(const int field_id, const ExprValue& lower, const ExprValue& upper, bool need_mapping = true) { - if (_sample_rows == 0) { + if (_sample_rows == 0 || _field_histogram.find(field_id) == _field_histogram.end()) { return 1.0; } @@ -168,7 +177,7 @@ class Statistics { } double get_cmsketch_ratio(const int field_id, const ExprValue& value) { - if (_total_rows == 0) { + if (_total_rows == 0 || _field_histogram.find(field_id) == _field_histogram.end()) { return 1.0; } @@ -194,8 +203,9 @@ class Statistics { return iter->second->get_distinct_cnt(); } - bool is_cms_exist() { - return !_field_cmsketch.empty(); + bool is_cms_exist(int field_id) { + // return !_field_cmsketch.empty(); + return _field_histogram.find(field_id) != _field_histogram.end(); } private: @@ -228,11 +238,14 @@ class Statistics { } } + std::map get_field_names() const; + private: int64_t _table_id = 0; int64_t _version = 0; int64_t _sample_rows = 0; int64_t _total_rows = 0; + bool _is_valid = false; // 至少有一种统计数据时为真 std::map> _field_histogram; std::map> _field_cmsketch; std::map> _field_hll; diff --git a/include/common/type_utils.h b/include/common/type_utils.h index e24fad62f..9ce6f1cee 100644 --- a/include/common/type_utils.h +++ b/include/common/type_utils.h @@ -263,6 +263,7 @@ inline bool is_string(pb::PrimitiveType type) { case pb::BITMAP: case pb::HLL: case pb::TDIGEST: + case pb::JSON: return true; default: return false; diff --git a/include/engine/external_filesystem.h b/include/engine/external_filesystem.h index 2406018cd..6d27f21b4 100644 --- a/include/engine/external_filesystem.h +++ b/include/engine/external_filesystem.h @@ -146,7 +146,7 @@ class CompactionSstCache { Cache _cache; }; -#ifdef BAIDU_INTERNAL +#if defined(BAIDU_INTERNAL) struct AfsRWInfo { std::string uri; std::string absolute_path; @@ -268,7 +268,6 @@ class AfsExtFileWriter : public ExtFileWriter { private: std::vector _afs_rw_infos; }; - #endif class CompactionExtFileReader : public ExtFileReader { @@ -389,7 +388,8 @@ class ExtFileSystem { DISALLOW_COPY_AND_ASSIGN(ExtFileSystem); }; -#ifdef BAIDU_INTERNAL +#if defined(BAIDU_INTERNAL) + class AfsExtFileSystem : public ExtFileSystem { public: struct AfsUgi { diff --git a/include/engine/qos.h b/include/engine/qos.h index 907681cfd..47df180ff 100644 --- a/include/engine/qos.h +++ b/include/engine/qos.h @@ -18,7 +18,6 @@ #include "concurrency.h" namespace baikaldb { -DECLARE_int64(sql_token_bucket_timeout_min); DECLARE_int64(min_global_extended_percent); DECLARE_int64(qps_statistics_minutes_ago); DECLARE_int64(max_tokens_per_second); @@ -29,7 +28,6 @@ DECLARE_int64(get_token_weight); DECLARE_int64(qos_reject_interval_s); DECLARE_int64(qos_reject_ratio); DECLARE_int64(qos_reject_timeout_s); -DECLARE_int64(qos_reject_max_scan_ratio); DECLARE_int64(qos_reject_growth_multiple); DECLARE_int64(qos_need_reject); DECLARE_int64(sign_concurrency); diff --git a/include/engine/rocks_wrapper.h b/include/engine/rocks_wrapper.h index f95d08b86..a084d786a 100644 --- a/include/engine/rocks_wrapper.h +++ b/include/engine/rocks_wrapper.h @@ -39,7 +39,6 @@ DECLARE_int32(key_point_collector_interval); DECLARE_bool(rocks_use_partitioned_index_filters); DECLARE_bool(rocks_use_ribbon_filter); DECLARE_bool(olap_table_only); -DECLARE_bool(olap_import_mode); DECLARE_bool(rocks_use_sst_partitioner_fixed_prefix); DECLARE_int64(rocks_block_cache_size_mb); DECLARE_double(rocks_high_pri_pool_ratio); @@ -414,9 +413,16 @@ class RocksWrapper { } void close() { delete _txn_db; + _txn_db = nullptr; if (_cold_txn_db != nullptr) { delete _cold_txn_db; + _cold_txn_db = nullptr; } + _is_init = false; + _column_families.clear(); + _cold_column_family = nullptr; + _cold_binlog_cf = nullptr; + _old_binlog_cf = nullptr; } bool is_any_stall() { uint64_t value = 0; diff --git a/include/engine/rocksdb_compaction_service.h b/include/engine/rocksdb_compaction_service.h index e3f424fdf..5499c0862 100644 --- a/include/engine/rocksdb_compaction_service.h +++ b/include/engine/rocksdb_compaction_service.h @@ -14,7 +14,7 @@ #endif namespace baikaldb { - +DECLARE_bool(enable_remote_compaction); extern void set_rocksdb_flags(pb::RocksdbGFLAGS* rocksdb_gflags); class MyCompactionService : public rocksdb::CompactionService { @@ -78,4 +78,198 @@ class MyCompactionService : public rocksdb::CompactionService { bthread::Mutex _mutex; }; +// struct CompactionServiceJobInfo { +// std::string db_name; +// std::string db_id; +// std::string db_session_id; +// uint64_t job_id; // job_id is only unique within the current DB and session, +// // restart DB will reset the job_id. `db_id` and +// // `db_session_id` could help you build unique id across +// // different DBs and sessions. + +// Env::Priority priority; + +// // Additional Compaction Details that can be useful in the CompactionService +// CompactionReason compaction_reason; +// bool is_full_compaction; +// bool is_manual_compaction; +// bool bottommost_level; +// bool is_l0_compaction; +// CompactionServiceJobInfo(std::string db_name_, std::string db_id_, +// std::string db_session_id_, uint64_t job_id_, +// Env::Priority priority_, +// CompactionReason compaction_reason_, +// bool is_full_compaction_, bool is_manual_compaction_, +// bool bottommost_level_, bool is_l0_compaction_) +// : db_name(std::move(db_name_)), +// db_id(std::move(db_id_)), +// db_session_id(std::move(db_session_id_)), +// job_id(job_id_), +// priority(priority_), +// compaction_reason(compaction_reason_), +// is_full_compaction(is_full_compaction_), +// is_manual_compaction(is_manual_compaction_), +// bottommost_level(bottommost_level_), +// is_l0_compaction(is_l0_compaction_) {} +// }; + +// struct CompactionServiceInput { +// std::string cf_name; + +// std::vector snapshots; + +// // SST files for compaction, it should already be expended to include all the +// // files needed for this compaction, for both input level files and output +// // level files. +// std::vector input_files; +// int output_level; + +// // db_id is used to generate unique id of sst on the remote compactor +// std::string db_id; + +// // information for subcompaction +// bool has_begin = false; +// std::string begin; +// bool has_end = false; +// std::string end; + +// uint64_t options_file_number; + +// // serialization interface to read and write the object +// static Status Read(const std::string& data_str, CompactionServiceInput* obj); +// Status Write(std::string* output); + +// bool is_l0_compaction = false; +// }; +#ifdef REMOTE_COMPACTION + +class TESTCompactionService : public rocksdb::CompactionService { +public: + TESTCompactionService(const std::string& db_path, const std::string& address) { + _db_path = db_path; + _address = address; + } + + static const char* kClassName() { return "TESTCompactionService"; } + + const char* Name() const override { return kClassName(); } + + rocksdb::CompactionServiceScheduleResponse Schedule(const rocksdb::CompactionServiceJobInfo& info, + const std::string& compaction_service_input) override { + rocksdb::CompactionServiceInput compaction_input; + rocksdb::Status s = rocksdb::CompactionServiceInput::Read(compaction_service_input, &compaction_input); + if (!s.ok()) { + DB_WARNING("CompactionServiceInput Read failed, error: %s", s.ToString().c_str()); + return rocksdb::CompactionServiceScheduleResponse("", rocksdb::CompactionServiceJobStatus::kUseLocal); + } + + std::ostringstream oss; + oss << "CompactionServiceJobInfo: db_name: " << info.db_name; + oss << ", db_id: " << info.db_id; + oss << ", db_session_id: " << info.db_session_id; + oss << ", job_id: " << info.job_id; + oss << ", priority: " << info.priority; + // oss << ", compaction_reason: " << info.compaction_reason; + oss << ", is_full_compaction: " << info.is_full_compaction; + oss << ", is_manual_compaction: " << info.is_manual_compaction; + oss << ", bottommost_level: " << info.bottommost_level; + oss << ", is_l0_compaction: " << info.is_l0_compaction; + oss << ", cf_name: " << compaction_input.cf_name; + oss << ", output_level: " << compaction_input.output_level; + oss << ", db_id: " << compaction_input.db_id; + oss << ", has_begin: " << compaction_input.has_begin; + oss << ", begin: " << compaction_input.begin; + oss << ", has_end: " << compaction_input.has_end; + oss << ", end: " << compaction_input.end; + oss << ", options_file_number: " << compaction_input.options_file_number; + oss << ", is_l0_compaction: " << compaction_input.is_l0_compaction; + oss << " snapshots: "; + for (const auto s : compaction_input.snapshots) { + oss << s << " "; + } + oss << ", input_files: "; + for (const auto f : compaction_input.input_files) { + oss << f << " "; + } + + DB_WARNING("%s", oss.str().c_str()); + if (!FLAGS_enable_remote_compaction) { + DB_WARNING("remote compaction is not enabled"); + return rocksdb::CompactionServiceScheduleResponse("", rocksdb::CompactionServiceJobStatus::kUseLocal); + } + return rocksdb::CompactionServiceScheduleResponse("", rocksdb::CompactionServiceJobStatus::kSuccess); + } + + +// struct CompactionServiceOutputFile { +// std::string file_name; +// SequenceNumber smallest_seqno; +// SequenceNumber largest_seqno; +// std::string smallest_internal_key; +// std::string largest_internal_key; +// uint64_t oldest_ancester_time = kUnknownOldestAncesterTime; +// uint64_t file_creation_time = kUnknownFileCreationTime; +// uint64_t epoch_number = kUnknownEpochNumber; +// std::string file_checksum = kUnknownFileChecksum; +// std::string file_checksum_func_name = kUnknownFileChecksumFuncName; +// uint64_t paranoid_hash; +// bool marked_for_compaction; +// UniqueId64x2 unique_id{}; +// }; + +// struct CompactionServiceResult { +// Status status; +// std::vector output_files; +// int output_level; + +// // location of the output files +// std::string output_path; + +// uint64_t bytes_read = 0; +// uint64_t bytes_written = 0; +// CompactionJobStats stats; + +// // serialization interface to read and write the object +// static Status Read(const std::string& data_str, CompactionServiceResult* obj); +// Status Write(std::string* output); +// }; + rocksdb::CompactionServiceJobStatus Wait(const std::string& scheduled_job_id, std::string* result) override { + rocksdb::CompactionServiceResult compaction_result; + compaction_result.status = rocksdb::Status::OK(); + compaction_result.output_level = 6; + compaction_result.output_path = "./"; + rocksdb::CompactionServiceOutputFile output_file; + output_file.file_name = "abc.sst"; + output_file.smallest_seqno = 0; + output_file.largest_seqno = 0; + output_file.epoch_number = 1; + output_file.paranoid_hash = 0; + output_file.marked_for_compaction = false; + rocksdb::InternalKey smallest_ikey("0", 0, rocksdb::ValueType::kTypeValue); + rocksdb::InternalKey largest_ikey("9", 0, rocksdb::ValueType::kTypeValue); + output_file.smallest_internal_key = smallest_ikey.Encode().ToString(); + output_file.largest_internal_key = largest_ikey.Encode().ToString(); + compaction_result.output_files.emplace_back(output_file); + + DB_WARNING("scheduled_job_id: %s, compaction_result serialize", scheduled_job_id.c_str()); + auto s = compaction_result.Write(result); + if (!s.ok()) { + DB_FATAL("scheduled_job_id: %s, compaction_result serialize fail reason: %s", + scheduled_job_id.c_str(), s.ToString().c_str()); + return rocksdb::CompactionServiceJobStatus::kFailure; + } + return rocksdb::CompactionServiceJobStatus::kSuccess; + } + + void OnInstallation(const std::string& scheduled_job_id, + rocksdb::CompactionServiceJobStatus status) override { + DB_WARNING("scheduled_job_id: %s, compaction_result install", + scheduled_job_id.c_str()); + } +private: + std::string _db_path; + std::string _address; +}; +#endif + } // namespace baikaldb diff --git a/include/engine/rocksdb_merge_operator.h b/include/engine/rocksdb_merge_operator.h index 30d499e64..50029b2a6 100644 --- a/include/engine/rocksdb_merge_operator.h +++ b/include/engine/rocksdb_merge_operator.h @@ -26,9 +26,8 @@ class OLAPMergeOperator : public rocksdb::MergeOperator { bool PartialMerge(const rocksdb::Slice& key, const rocksdb::Slice& left_operand, const rocksdb::Slice& right_operand, std::string* new_value, - rocksdb::Logger* /*logger*/) const override { - return false; - } + rocksdb::Logger* /*logger*/) const override; + const char* Name() const override { return "OLAPMergeOperator"; } }; } // namespace baikaldb \ No newline at end of file diff --git a/include/engine/split_compaction_filter.h b/include/engine/split_compaction_filter.h index f57e0c196..5793adba7 100644 --- a/include/engine/split_compaction_filter.h +++ b/include/engine/split_compaction_filter.h @@ -25,9 +25,9 @@ namespace baikaldb { DECLARE_int32(rocks_binlog_ttl_days); class SplitCompactionFilter : public rocksdb::CompactionFilter { struct FilterRegionInfo { - FilterRegionInfo(bool use_ttl, const std::string& end_key, int64_t online_ttl_base_expire_time_us) : - use_ttl(use_ttl), end_key(end_key), online_ttl_base_expire_time_us(online_ttl_base_expire_time_us) {} - bool use_ttl = false; + FilterRegionInfo(bool use_normal_ttl, const std::string& end_key, int64_t online_ttl_base_expire_time_us) : + use_normal_ttl(use_normal_ttl), end_key(end_key), online_ttl_base_expire_time_us(online_ttl_base_expire_time_us) {} + bool use_normal_ttl = false; std::string end_key; int64_t online_ttl_base_expire_time_us = 0; }; @@ -96,7 +96,7 @@ typedef DoubleBuffer DoubleBufBinlog; rocksdb::Slice key_slice(key); key_slice.remove_prefix(sizeof(int64_t) * 2); rocksdb::Slice value_slice(value); - if (filter_info->use_ttl) { + if (filter_info->use_normal_ttl) { ttl_decode(value_slice, index_info, filter_info->online_ttl_base_expire_time_us); } return !Transaction::fits_region_range(key_slice, value_slice, @@ -106,14 +106,14 @@ typedef DoubleBuffer DoubleBufBinlog; } void set_filter_region_info(int64_t region_id, const std::string& end_key, - bool use_ttl, int64_t online_ttl_base_expire_time_us) { + bool use_normal_ttl, int64_t online_ttl_base_expire_time_us) { FilterRegionInfo* old = get_filter_region_info(region_id); // 已存在不更新 if (old != nullptr && old->end_key == end_key) { return; } - auto call = [region_id, end_key, use_ttl, online_ttl_base_expire_time_us](KeyMap& key_map) { - FilterRegionInfo* new_info = new FilterRegionInfo(use_ttl, end_key, online_ttl_base_expire_time_us); + auto call = [region_id, end_key, use_normal_ttl, online_ttl_base_expire_time_us](KeyMap& key_map) { + FilterRegionInfo* new_info = new FilterRegionInfo(use_normal_ttl, end_key, online_ttl_base_expire_time_us); key_map[region_id] = new_info; }; _range_key_map.modify(call); diff --git a/include/engine/sst_file_writer.h b/include/engine/sst_file_writer.h index 7237c92c2..7d3788c0b 100644 --- a/include/engine/sst_file_writer.h +++ b/include/engine/sst_file_writer.h @@ -21,11 +21,7 @@ namespace baikaldb { class SstFileWriter { public: - SstFileWriter(const rocksdb::Options& options, bool force_lz4 = true) : _options(options) { - if (force_lz4) { - _options.bottommost_compression = rocksdb::kLZ4Compression; - _options.bottommost_compression_opts = rocksdb::CompressionOptions(); - } + SstFileWriter(const rocksdb::Options& options) : _options(options) { //ingest到L6后会保留filter，因此需要去除 bool skip_filters = _options.optimize_filters_for_hits; _sst_writer.reset(new rocksdb::SstFileWriter(rocksdb::EnvOptions(), _options, nullptr, diff --git a/include/engine/table_iterator.h b/include/engine/table_iterator.h index 523ef9723..cde22895c 100644 --- a/include/engine/table_iterator.h +++ b/include/engine/table_iterator.h @@ -158,7 +158,7 @@ class Iterator { int _upper_sufix = 0; bool _valid = true; - bool _use_ttl = false; + bool _use_normal_ttl = false; bool _is_cstore = false; bool _is_offline_binlog = false; int64_t _read_ttl_timestamp_us = 0; diff --git a/include/engine/transaction.h b/include/engine/transaction.h index a5dcc9eae..d5af0b4cc 100644 --- a/include/engine/transaction.h +++ b/include/engine/transaction.h @@ -31,6 +31,7 @@ namespace baikaldb { DECLARE_bool(disable_wal); DECLARE_bool(leader_merge_in_raft); +DECLARE_bool(enable_column_engine); class ReverseIndexBase; typedef std::map CachePlanMap; @@ -459,8 +460,8 @@ class Transaction { _watt_stats_version = version; } - void set_use_ttl(bool use_ttl) { _use_ttl = use_ttl; } - bool use_ttl() const { return _use_ttl; } + void set_use_normal_ttl(bool use_normal_ttl) { _use_normal_ttl = use_normal_ttl; } + bool use_normal_ttl() const { return _use_normal_ttl; } bool use_cold_db() const { return _use_cold_db; } static int get_full_primary_key( rocksdb::Slice index_bytes, @@ -568,6 +569,10 @@ class Transaction { _reverse_set.insert(base); } + bool has_column_engine() { + return FLAGS_enable_column_engine && _table_info->schema_conf.enable_column_engine(); + } + public: int64_t num_increase_rows = 0; int64_t last_active_time = 0; @@ -672,7 +677,7 @@ class Transaction { std::set _pri_field_ids; // for cstore bthread_mutex_t _txn_mutex; - bool _use_ttl = false; + bool _use_normal_ttl = false; bool _is_separate = false; int64_t _read_ttl_timestamp_us = 0; //ttl读取时间 int64_t _write_ttl_timestamp_us = 0; //ttl写入时间 diff --git a/include/engine/transaction_pool.h b/include/engine/transaction_pool.h index 4b0e5e375..902cfa721 100644 --- a/include/engine/transaction_pool.h +++ b/include/engine/transaction_pool.h @@ -36,7 +36,7 @@ class TransactionPool { TransactionPool() : _num_prepared_txn(0), _txn_count(0) {} - int init(int64_t region_id, bool use_ttl, int64_t online_ttl_base_expire_time_us); + int init(int64_t region_id, bool use_normal_ttl, int64_t online_ttl_base_expire_time_us); // -1 means insert error (already exists) int begin_txn(uint64_t txn_id, SmartTransaction& txn, int64_t primary_region_id, @@ -96,16 +96,16 @@ class TransactionPool { --_has_write_txn_count; } - bool use_ttl() const { - return _use_ttl; + bool use_normal_ttl() const { + return _use_normal_ttl; } int64_t online_ttl_base_expire_time_us() const { return _online_ttl_base_expire_time_us; } - void update_ttl_info(bool use_ttl, int64_t online_ttl_base_expire_time_us) { - _use_ttl = use_ttl; + void update_ttl_info(bool use_normal_ttl, int64_t online_ttl_base_expire_time_us) { + _use_normal_ttl = use_normal_ttl; _online_ttl_base_expire_time_us = online_ttl_base_expire_time_us; } @@ -150,7 +150,7 @@ class TransactionPool { private: int64_t _region_id = 0; int64_t _latest_has_write_txn_ts = 0; - bool _use_ttl = false; + bool _use_normal_ttl = false; // 是否是普通ttl int64_t _online_ttl_base_expire_time_us = 0; // txn_id => txn handler mapping diff --git a/include/exec/access_path.h b/include/exec/access_path.h index 890936754..a5ed738b2 100755 --- a/include/exec/access_path.h +++ b/include/exec/access_path.h @@ -74,6 +74,7 @@ enum IndexHint { void fetch_field_ids() { if (index_type == pb::I_KEY || index_type == pb::I_UNIQ || index_type == pb::I_PRIMARY) { + // TODO: 普通索引支持使用主键字段这边暂时不修改，后续可能需要修改I_KEY相关逻辑 for (auto& field : index_info_ptr->fields) { index_field_ids.insert(field.id); } diff --git a/include/exec/agg_node.h b/include/exec/agg_node.h index 428c274b7..feae0d902 100644 --- a/include/exec/agg_node.h +++ b/include/exec/agg_node.h @@ -96,6 +96,12 @@ class AggNode : public ExecNode { _node_type = pb::MERGE_AGG_NODE; _is_merger = true; } + void transfer_to_agg() { + _pb_node.set_node_type(pb::AGG_NODE); + _node_type = pb::AGG_NODE; + _is_merger = false; + _has_merger = false; + } int delete_self() { if (_parent == nullptr) { DB_WARNING("_parent is nullptr"); diff --git a/include/exec/exchange_receiver_node.h b/include/exec/exchange_receiver_node.h index ac52f75aa..4cce8ca2e 100644 --- a/include/exec/exchange_receiver_node.h +++ b/include/exec/exchange_receiver_node.h @@ -67,6 +67,9 @@ class DataStreamReceiver { RuntimeState* runtime_state() { return _state; } + void set_conditions(arrow::compute::Expression* condition) { + _condition = condition; + } public: class SenderQueue; @@ -115,6 +118,8 @@ class DataStreamReceiver { RuntimeState* _state = nullptr; ExchangeReceiverNode* _exchange_receiver_node = nullptr; + arrow::compute::Expression* _condition = nullptr; + bthread::Mutex _mtx; // ，需要接收region的信息 std::unordered_map _region_info_map; @@ -167,12 +172,7 @@ class ExchangeReceiverNode : public ExecNode { return true; } - int init_sort_info(SortNode* sort_node) { - _slot_order_exprs = sort_node->slot_order_exprs(); - _is_asc = sort_node->is_asc(); - _is_null_first = sort_node->is_null_first(); - return 0; - } + int init_condition_and_sort_info(ExecNode* select_manager); std::shared_ptr get_data_stream_receiver() const { return _data_stream_receiver; @@ -251,6 +251,10 @@ class ExchangeReceiverNode : public ExecNode { std::vector _is_asc; std::vector _is_null_first; + // filter + std::vector _conditions; + arrow::compute::Expression _vectorize_conditions; + // 已下只有主db赋值 ExecNode* _exchange_sender_node = nullptr; ExecNode* _relate_select_manager_node = nullptr; diff --git a/include/exec/exchange_sender_node.h b/include/exec/exchange_sender_node.h index a0cff656f..c4bd3bb22 100644 --- a/include/exec/exchange_sender_node.h +++ b/include/exec/exchange_sender_node.h @@ -13,7 +13,7 @@ // limitations under the License. #pragma once -#include "exec_node.h" +#include "scan_node.h" #include "proto/db.interface.pb.h" #include @@ -235,6 +235,25 @@ class ExchangeSenderNode : public ExecNode { arrow::Status repartition(std::shared_ptr batch, std::unordered_map>& hash_batch_map); + bool is_db_fragment() { + bool has_dblink_scan_node = false; + std::vector scan_nodes; + get_node_pass_subquery(pb::SCAN_NODE, scan_nodes); + for (auto* scan_node : scan_nodes) { + if (scan_node != nullptr) { + if (static_cast(scan_node)->is_file_scan_node()) { + has_dblink_scan_node = true; + break; + } + if (static_cast(scan_node)->is_mysql_scan_node()) { + has_dblink_scan_node = true; + break; + } + } + } + return (get_node_pass_subquery(pb::EXCHANGE_RECEIVER_NODE) != nullptr) || has_dblink_scan_node; + } + uint64_t _log_id = 0; int32_t _fragment_id = 0; int32_t _receiver_fragment_id = 0; @@ -261,7 +280,12 @@ class ExchangeSenderNode : public ExecNode { std::atomic _repartition_cost_us {0}; std::atomic _repartition_rows {0}; - StockRecordBatchKeeper _broadcast_record_batch_keepper; // for broadcast type + // for broadcast type + StockRecordBatchKeeper _broadcast_record_batch_keepper; + + // for random type + int _target_channel_idx_for_random_type = -1; + std::shared_ptr _empty_recordbatch_buffer; }; } // naemspace baikaldb diff --git a/include/exec/exec_node.h b/include/exec/exec_node.h index 35879acf2..35c7116d2 100644 --- a/include/exec/exec_node.h +++ b/include/exec/exec_node.h @@ -386,7 +386,7 @@ class ExecNode { bool is_get_keypoint() { return _is_get_keypoint; } - bool set_has_optimized(bool has_optimized) { + void set_has_optimized(bool has_optimized) { _has_optimized = has_optimized; } bool has_optimized() { @@ -466,6 +466,7 @@ class ExecNode { _partition_property.hash_partition_propertys = partition_property->hash_partition_propertys; _partition_property.add_need_cast_string_columns(partition_property->need_cast_string_columns); } + void set_child_node_property_any_type(); NodePartitionProperty* partition_property() { return &_partition_property; diff --git a/include/exec/fetcher_store.h b/include/exec/fetcher_store.h index e77e28dd4..1f1fd85ba 100755 --- a/include/exec/fetcher_store.h +++ b/include/exec/fetcher_store.h @@ -40,6 +40,7 @@ enum ErrorType { struct RegionReturnData { std::shared_ptr row_data = nullptr; std::shared_ptr arrow_data = nullptr; + int64_t partition_id = 0; // 只有分区表全局索引使用 void set_row_data(std::shared_ptr& batch) { row_data = batch; arrow_data = nullptr; @@ -48,6 +49,9 @@ struct RegionReturnData { arrow_data = batch; row_data = nullptr; } + void set_partition_id(int64_t p_id) { + partition_id = p_id; + } }; struct RegionInfoData { @@ -121,6 +125,8 @@ class OnRPCDone : public google::protobuf::Closure { std::shared_ptr single_response); ErrorType check_status(); + bool need_copy(MemRow* row); + void clear_request() { if (_need_check_memory && _has_multi_plan && _op_type == pb::OP_SELECT) { // 每个region一个plan, 没有复用 @@ -771,6 +777,10 @@ class FetcherStore { std::shared_ptr shared_plan; bool is_full_export = false; bool need_check_memory = false; + std::vector* conditions = nullptr; + arrow::compute::Expression* vectorize_conditions = nullptr; + + bool broadcast_all_peer_without_raft = false; }; template diff --git a/include/exec/filter_node.h b/include/exec/filter_node.h index d6c453d33..5c5d456c6 100644 --- a/include/exec/filter_node.h +++ b/include/exec/filter_node.h @@ -112,6 +112,10 @@ class FilterNode : public ExecNode { } virtual int show_explain(QueryContext* ctx, std::vector>& output, int& next_id, int display_id); + bool has_huge_in_condition(); + + void cut_huge_in_condition(std::vector& huge_in_conditions); + void reset(RuntimeState* state) override { _child_eos = false; _child_row_idx = 0; diff --git a/include/exec/index_ddl_manager_node.h b/include/exec/index_ddl_manager_node.h index c69342bef..8c42c9e51 100755 --- a/include/exec/index_ddl_manager_node.h +++ b/include/exec/index_ddl_manager_node.h @@ -15,6 +15,8 @@ #pragma once #include "dml_manager_node.h" #include "fetcher_store.h" +#include "lock_primary_node.h" +#include "lock_secondary_node.h" namespace baikaldb { @@ -39,11 +41,38 @@ class IndexDDLManagerNode : public DmlManagerNode { _is_global_index = flags; } + int create_lock_primary_node(int64_t table_id, + std::unique_ptr& lock_primary_node) const; + + int create_lock_secondary_node(int64_t table_id, + std::vector>& lock_secondary_nodes) const; + + int init_lock_nodes_if_not_exist() { + if (_lock_primary != nullptr) { + // 已经初始化完成 + return 1; + } + int ret = create_lock_primary_node(_table_id, _lock_primary); + if (ret < 0) { + return ret; + } + ret = create_lock_secondary_node(_table_id, _lock_secondaries); + if (ret < 0) { + return ret; + } + _op_type = pb::OP_DELETE; + return 0; + } + + static bool is_ddl_delete(const pb::ScanNode& scan_node); private: int64_t _table_id {0}; int64_t _index_id {0}; std::string _task_id; bool _is_global_index = false; bool _is_rollup_index = false; + + std::unique_ptr _lock_primary = nullptr; + std::vector> _lock_secondaries; }; } // namespace baikaldbame diff --git a/include/exec/join_node.h b/include/exec/join_node.h index f4217594d..690d43a45 100644 --- a/include/exec/join_node.h +++ b/include/exec/join_node.h @@ -109,15 +109,39 @@ class JoinNode : public Joiner { std::vector& tuple_order, std::vector& conditions); + void handle_join_equal_key_cast_type(ExprNode* outer_expr, + ExprNode* inner_expr, + std::string outer_tmp_col_name, + std::string inner_tmp_col_name, + std::vector& outer_keys, + std::vector& inner_keys, + std::unordered_map& outer_projection_temp_col, + std::unordered_map& inner_projection_temp_col); + int handle_join_filter_key_expressions(RuntimeState* state, + int idx, + ExprNode* condition, + bool left_child_is_outer, + std::vector& outer_keys, + std::vector& inner_keys, + std::vector& sub_exprs, + std::unordered_map& outer_projection_temp_col, + std::unordered_map& inner_projection_temp_col); + int try_transfer_filter_to_column_first(RuntimeState* state, + int idx, + ExprNode* condition, + std::vector& outer_keys, + std::vector& inner_keys, + std::vector& sub_exprs, + std::unordered_map& outer_projection_temp_col, + std::unordered_map& inner_projection_temp_col); int build_table_arrow_declaration(RuntimeState* state, arrow::acero::Declaration& dec, ExecNode* node, std::unordered_set& tuple_ids, std::vector& mem_rows, - const std::unordered_map>& cast_string_slot_ids, + std::unordered_map& projection_temp_col, bool need_add_index_colletor_node, - bool remove_useless_sort, - bool need_add_join_key); + bool remove_useless_sort); void get_need_add_index_collector_cond_nodes(ExecNode* node, std::set& need_add_nodes); virtual int build_arrow_declaration(RuntimeState* state); @@ -126,26 +150,16 @@ class JoinNode : public Joiner { virtual int set_partition_property_and_schema(QueryContext* ctx); + bool suitable_for_broadcast_join(QueryContext* ctx); + + void make_broadcast_join_property(std::shared_ptr* small_table_property, + ExecNode* small_table_node, + std::shared_ptr* other_property, + NodePartitionProperty* other_child_property); + void get_hash_partitions(NodePartitionProperty& outer_property, NodePartitionProperty& inner_property, - const std::unordered_set& cast_string_hash_columns) { - outer_property.type = _partition_property.type; - inner_property.type = _partition_property.type; - if (_partition_property.type == pb::HashPartitionType) { - outer_property.hash_partition_propertys.emplace_back(_partition_property.hash_partition_propertys[0]); - inner_property.hash_partition_propertys.emplace_back(_partition_property.hash_partition_propertys[1]); - } - outer_property.need_cast_string_columns = _partition_property.need_cast_string_columns; - inner_property.need_cast_string_columns = _partition_property.need_cast_string_columns; - for (auto& name : cast_string_hash_columns) { - outer_property.need_cast_string_columns.insert(name); - inner_property.need_cast_string_columns.insert(name); - if (_on_condition_column_map.count(name) > 0) { - outer_property.need_cast_string_columns.insert(_on_condition_column_map[name]); - inner_property.need_cast_string_columns.insert(_on_condition_column_map[name]); - } - } - } + const std::unordered_set& cast_string_hash_columns); }; } diff --git a/include/exec/joiner.h b/include/exec/joiner.h index 639f2f46c..40830b38e 100644 --- a/include/exec/joiner.h +++ b/include/exec/joiner.h @@ -115,7 +115,7 @@ class Joiner : public ExecNode { _use_index_join = use; } void get_join_on_condition_filter(RuntimeState* state, ExprNode** condition_filter); - static int do_plan_router(RuntimeState* state, const std::vector& scan_nodes, bool& index_has_null, bool is_explain); + static int do_plan_router(RuntimeState* state, const std::vector& scan_nodes, bool& index_has_null, bool is_explain, bool in_acero); int vectorize_index_collector(RuntimeState* state, std::shared_ptr batch); void clear_in_conditions() { _outer_join_values.clear(); @@ -183,6 +183,8 @@ class Joiner : public ExecNode { std::vector _outer_equal_slot; std::vector _inner_equal_slot; + std::vector _equal_slot_cast_types; + //从左边取到的等值条件的value ExprValueSet _outer_join_values; diff --git a/include/exec/mpp_property.h b/include/exec/mpp_property.h index 20340f96e..30eedf78a 100644 --- a/include/exec/mpp_property.h +++ b/include/exec/mpp_property.h @@ -29,6 +29,7 @@ struct HashPartitionColumns { } } } + pb::PartitionPropertyType type = pb::HashPartitionType; // // arrow_name(tupleid_slotid / tmp_1(agg产生的临时列) / hash_1) -> ColumnInfo // std::unordered_map hash_columns; // 需要额外加projection产生hash列, hash_1 -> expr(a+b) @@ -113,7 +114,7 @@ struct NodePartitionProperty { std::string ret = "type: " + pb::PartitionPropertyType_Name(type) + ", hash_partition_propertys size: " + std::to_string(hash_partition_propertys.size()); for (auto m : hash_partition_propertys) { - ret += "\n hash_cols: ["; + ret += "\n type: " + pb::PartitionPropertyType_Name(m->type) + " hash_cols: ["; for (auto c : m->hash_columns) { ret += c.first + ":" + pb::ExprNodeType_Name(c.second->node_type()); } diff --git a/include/exec/packet_node.h b/include/exec/packet_node.h index 6208b55ca..d3c9215fd 100644 --- a/include/exec/packet_node.h +++ b/include/exec/packet_node.h @@ -101,7 +101,8 @@ class PacketNode : public ExecNode { int pack_binary_row(MemRow* row, RuntimeState* state, std::vector>* columns = nullptr, int row_idx = 0); int pack_eof(); int fatch_expr_subquery_results(RuntimeState* state); - + static std::map sample_analyze_regions( + const std::map& all_regions, int64_t sample_cnt); private: bool _binary_protocol = false; pb::OpType _op_type; diff --git a/include/exec/parquet_scan_node.h b/include/exec/parquet_scan_node.h index 6355a9a52..9fe711bf6 100644 --- a/include/exec/parquet_scan_node.h +++ b/include/exec/parquet_scan_node.h @@ -18,12 +18,6 @@ #include "file_manager.h" namespace baikaldb { - -struct ReaderInfo{ - std::shared_ptr read_contents; - std::shared_ptr<::arrow::RecordBatchReader> reader; -}; - class ParquetScanNode : public ScanNode { public: ParquetScanNode() {} @@ -41,11 +35,8 @@ class ParquetScanNode : public ScanNode { std::unordered_map* get_field_id2info_map() { return &_field_id2info_map; } - std::unordered_map* get_parquet_file2reader_map() { - return &_parquet_file2reader_map; - } - std::unordered_map>* get_parquet_file_not_exist_column_map() { - return &_parquet_file_not_exist_column_map; + std::vector>* get_parquet_file_readers() { + return &_parquet_file_readers; } private: @@ -53,23 +44,21 @@ class ParquetScanNode : public ScanNode { // 获取每个parquet文件符合条件数据的RecordBatchReader // 增删列场景下，每个parquet文件的schema可能不相同，需要记录查询需要但是parquet文件中不存在的列，对这些列需要填充null或默认值 - int get_qualified_record_batch_readers( - std::unordered_map& parquet_file2reader_map, - std::unordered_map>& parquet_file_not_exist_column_map); + int get_qualified_record_batch_readers(std::vector>& parquet_file_readers); + void get_qualified_parquet_file_readers(std::vector>& parquet_file_readers, + const std::vector>& parquet_files, std::shared_ptr schema); private: ParquetFileManager* _file_manager = nullptr; SchemaFactory* _factory = nullptr; - SmartTable _table_info; + SmartTable _table_info = nullptr; + SmartIndex _pri_info = nullptr; int64_t _region_id = -1; - std::vector _key_ranges; - std::vector> _parquet_files; + pb::PossibleIndex _possible_index; // key: field_id, value: field_info std::unordered_map _field_id2info_map; - // key: parquet_file_name, value: parquet RecordBatchReader - std::unordered_map _parquet_file2reader_map; - // key: parquet_file_name, value: 本次需要获取的但parquet文件中不存在的列 - std::unordered_map> _parquet_file_not_exist_column_map; + std::unordered_map _field_name2info_map; + std::vector> _parquet_file_readers; }; class ParquetVectorizedReader : public arrow::RecordBatchReader { @@ -87,13 +76,15 @@ class ParquetVectorizedReader : public arrow::RecordBatchReader { RuntimeState* _state = nullptr; ParquetScanNode* _parquet_scan_node = nullptr; std::unordered_map* _field_id2info_map = nullptr; - std::unordered_map* _parquet_file2reader_map = nullptr; - std::unordered_map>* _parquet_file_not_exist_column_map = nullptr; + std::unordered_map* _field_name2info_map = nullptr; + std::vector>* _parquet_file_readers = nullptr; std::shared_ptr<::arrow::Schema> _arrow_schema; // key: baikaldb column name, value: parquet column name std::unordered_map _column_name_map; - std::unordered_map::iterator _reader_iter; + // key: baikaldb column name, value: baikaldb column type + std::unordered_map _column_type_map; + std::vector>::iterator _reader_iter; std::shared_ptr<::arrow::RecordBatch> _record_batch; int64_t _row_idx_in_record_batch = 0; diff --git a/include/exec/property.h b/include/exec/property.h index 2f113a5f6..041de174b 100755 --- a/include/exec/property.h +++ b/include/exec/property.h @@ -23,6 +23,7 @@ struct Property { int64_t expected_cnt = -1; // 向量索引使用 int32_t efsearch = -1; + int32_t nprobe = -1; Property() { } Property(const std::vector& slot_order_exprs_, diff --git a/include/exec/rocksdb_scan_node.h b/include/exec/rocksdb_scan_node.h index ea5c7ae71..af2438c39 100644 --- a/include/exec/rocksdb_scan_node.h +++ b/include/exec/rocksdb_scan_node.h @@ -154,6 +154,14 @@ class RocksdbScanNode : public ScanNode { int get_key_points(RuntimeState* state, RowBatch* batch, bool* eos); + bool is_ddl_update() const { + return _ddl_work_type == pb::DDL_COLUMN && _is_ddl_update; + } + + bool is_ddl_delete() const { + return _ddl_work_type == pb::DDL_COLUMN && !_is_ddl_update; + }; + private: int get_next_by_table_get(RuntimeState* state, RowBatch* batch, bool* eos); int get_next_by_table_seek(RuntimeState* state, RowBatch* batch, bool* eos); @@ -199,6 +207,7 @@ class RocksdbScanNode : public ScanNode { int64_t _region_id; bool _use_get = false; bool _is_ddl_work = false; + bool _is_ddl_update = false; bool _is_global_index = false; bool _has__weight = false; pb::DDLType _ddl_work_type = pb::DDL_NONE; @@ -243,6 +252,7 @@ class RocksdbScanNode : public ScanNode { std::string _vector_word; int _topk = 10; int _efsearch = 16; + int _nprobe = 5; uint64_t _separate_value = 0; SmartTable _table_info; diff --git a/include/exec/scan_node.h b/include/exec/scan_node.h index 5896d1648..9ce2359eb 100644 --- a/include/exec/scan_node.h +++ b/include/exec/scan_node.h @@ -256,7 +256,7 @@ class ScanNode : public ExecNode { return _old_region_infos; } void set_old_region_infos(google::protobuf::RepeatedPtrField&& region_infos) { - _old_region_infos = region_infos; + _old_region_infos = std::move(region_infos); } pb::Engine engine() { return _engine; @@ -421,6 +421,10 @@ class ScanNode : public ExecNode { return _is_mysql_scan_node; } + bool is_file_scan_node() const { + return _is_file_scan_node; + } + void calc_index_range() { _main_path.path(_select_idx)->calc_index_range(_partition_field_id, _expr_partition_map); if (!_main_path.path(_select_idx)->index_info_ptr->is_global && _select_idx != _table_id) { @@ -445,7 +449,7 @@ class ScanNode : public ExecNode { } bool can_use_no_index_join() { - if (_select_idx == _select_index_for_join && _select_idx != _table_id) { + if (_select_idx != -1 && _select_idx == _select_index_for_join && _select_idx != _table_id) { // 当join on条件推不推都选择同一个索引, 且不是主键的时候, 可以使用no index join return true; } @@ -538,6 +542,7 @@ class ScanNode : public ExecNode { bool _has_index = false; bool _is_rocksdb_scan_node = false; bool _is_mysql_scan_node = false; + bool _is_file_scan_node = false; pb::LockCmdType _lock = pb::LOCK_NO; RouterPolicy _router_policy = RouterPolicy::RP_RANGE; google::protobuf::RepeatedPtrField _old_region_infos; diff --git a/include/exec/select_manager_node.h b/include/exec/select_manager_node.h index 211a86030..8e4bfdbc0 100755 --- a/include/exec/select_manager_node.h +++ b/include/exec/select_manager_node.h @@ -58,7 +58,8 @@ class SelectManagerNode : public ExecNode { virtual int get_next(RuntimeState* state, RowBatch* batch, bool* eos); virtual void close(RuntimeState* state) { ExecNode::close(state); - if (_pb_node.derive_node().select_manager_node().slot_order_exprs_size() > 0) { + if (_pb_node.derive_node().select_manager_node().slot_order_exprs_size() > 0 + || _steal_slot_order_exprs) { for (auto& expr : _slot_order_exprs) { ExprNode::destroy_tree(expr); } @@ -77,6 +78,10 @@ class SelectManagerNode : public ExecNode { _is_null_first = sort_node->is_null_first(); return 0; } + void steal_slot_order_exprs(SortNode* sort_node) { + sort_node->clear_slot_order_exprs(); + _steal_slot_order_exprs = true; + } int single_fetcher_store_open(FetcherInfo* fetcher, RuntimeState* state, ExecNode* exec_node); void multi_fetcher_store_open(FetcherInfo* self_fetcher, FetcherInfo* other_fetcher, @@ -107,6 +112,8 @@ class SelectManagerNode : public ExecNode { SmartIndex pri_info, LimitNode* limit); + int dblink_scan_run(RuntimeState* state, ExecNode* exec_node); + int subquery_open(RuntimeState* state); int delay_fetcher_store(RuntimeState* state); @@ -141,14 +148,33 @@ class SelectManagerNode : public ExecNode { ExecNode* dual_scan = get_node(pb::DUAL_SCAN_NODE); return (dual_scan != nullptr); } - bool is_mysql_scan() { + bool is_dblink_scan() { ExecNode* scan = get_node(pb::SCAN_NODE); if (scan == nullptr) { return false; } - return static_cast(scan)->is_mysql_scan_node(); + return static_cast(scan)->is_file_scan_node() || + static_cast(scan)->is_mysql_scan_node(); + } + void set_db_request_map(const std::map& db_request_map) { + _db_request_map = db_request_map; } + void add_conditions(std::vector& conditions) { + _conditions.insert(_conditions.end(), conditions.begin(), conditions.end()); + } + std::vector& is_asc() { + return _is_asc; + } + std::vector& is_null_first() { + return _is_null_first; + } + std::vector& slot_order_exprs() { + return _slot_order_exprs; + } + std::vector& conditions() { + return _conditions; + } private: //允许fetcher回来后排序 std::vector _slot_order_exprs; @@ -158,6 +184,8 @@ class SelectManagerNode : public ExecNode { std::shared_ptr _sorter; SchemaFactory* _factory = nullptr; int32_t _scan_tuple_id = 0; + std::vector _conditions; + bool _steal_slot_order_exprs = false; // vectorized std::vector> _arrow_responses; @@ -166,9 +194,13 @@ class SelectManagerNode : public ExecNode { std::vector _region_batches; std::shared_ptr _arrow_io_executor; std::shared_ptr _index_collector_cond; + arrow::compute::Expression _vectorize_conditions; // mpp bool _has_er_child = false; + + // + std::map _db_request_map; }; class FetcherStoreVectorizedReader : public arrow::RecordBatchReader { @@ -182,6 +214,8 @@ class FetcherStoreVectorizedReader : public arrow::RecordBatchReader { } int init(SelectManagerNode* select_node, RuntimeState* state); + + void init_tuple(SelectManagerNode* select_node, RuntimeState* state); arrow::Status ReadNext(std::shared_ptr* out) override; diff --git a/include/exec/sort_node.h b/include/exec/sort_node.h index 5e1c6f7d1..ab5080f42 100644 --- a/include/exec/sort_node.h +++ b/include/exec/sort_node.h @@ -93,7 +93,9 @@ class SortNode : public ExecNode { std::vector& slot_order_exprs() { return _slot_order_exprs; } - + void clear_slot_order_exprs() { + _slot_order_exprs.clear(); + } std::vector* mutable_order_exprs() { return &_order_exprs; } @@ -104,7 +106,7 @@ class SortNode : public ExecNode { bool is_monotonic() { return _monotonic; } - + bool need_projection(); virtual int show_explain(QueryContext* ctx, std::vector>& output, int& next_id, int display_id) { int return_id = ExecNode::show_explain(ctx, output, next_id, display_id); if (output.empty()) { diff --git a/include/exec/vectorize_helpper.h b/include/exec/vectorize_helpper.h index 44906a433..0383f14cb 100644 --- a/include/exec/vectorize_helpper.h +++ b/include/exec/vectorize_helpper.h @@ -19,6 +19,11 @@ #include #include #include +#include +#include +#include +#include "arrow_function.h" +#include "arrow_io_excutor.h" namespace baikaldb { class VectorizeHelpper { public: @@ -84,6 +89,26 @@ class VectorizeHelpper { return std::make_shared(fields); } + static std::shared_ptr get_arrow_schema(const std::vector& tuples) { + std::vector> fields; + fields.reserve(5); + for (auto tuple : tuples) { + if (tuple == nullptr) { + continue; + } + for (auto& slot : tuple->slots()) { + std::string name = std::to_string(slot.tuple_id()) + "_" + std::to_string(slot.slot_id()); + std::shared_ptr field = construct_arrow_field(name, slot.slot_type()); + if (field == nullptr) { + DB_WARNING("Fail to construct arrow field, name: %s", name.c_str()); + return nullptr; + } + fields.emplace_back(field); + } + } + return std::make_shared(fields); + } + static ExprValue get_vectorized_value(const arrow::ChunkedArray* chunked_array, int row_idx, int32_t float_precision_len = -1) { ExprValue ret; ret.type = pb::NULL_TYPE; @@ -262,8 +287,13 @@ class VectorizeHelpper { // in: store 返回的recordbatch static int change_arrow_record_batch_schema(std::shared_ptr schema, std::shared_ptr in, std::shared_ptr* out, - bool need_cast = false) { + bool need_check_memory_limit = false) { std::vector> columns; + arrow::MemoryPool* pool = arrow::default_memory_pool(); + if (need_check_memory_limit) { + pool = GetMemoryPoolForRead(); + } + arrow::compute::ExecContext exec_ctx(pool); for (auto& f : schema->fields()) { std::shared_ptr array = in->GetColumnByName(f->name()); if (array == nullptr) { @@ -276,7 +306,7 @@ class VectorizeHelpper { auto get_res = f->metadata()->Get("default_value"); if (get_res.status().ok()) { const std::string& default_value = *get_res; - default_array = make_array_from_str(f->type(), default_value, in->num_rows()); + default_array = make_array_from_str(f->type(), default_value, in->num_rows(), &exec_ctx); if (default_array == nullptr) { DB_WARNING("Fail to make_array_from_str"); return -1; @@ -294,11 +324,9 @@ class VectorizeHelpper { array = default_array; columns.emplace_back(array); } else { - if (need_cast && array->type()->id() != f->type()->id()) { + if (array->type()->id() != f->type()->id()) { // 整数大类型转小类型，比如int32转int8，需要避免溢出返回错误 - ::arrow::compute::CastOptions cast_options; - cast_options.allow_int_overflow = true; - auto cast_array_ret = ::arrow::compute::Cast(*array, f->type(), cast_options); + auto cast_array_ret = ::arrow::compute::Cast(*array, f->type(), arrow::compute::CastOptions::Unsafe(f->type()), &exec_ctx); if (!cast_array_ret.ok()) { DB_WARNING("cast array fail, %s", cast_array_ret.status().ToString().c_str()); return -1; @@ -314,16 +342,22 @@ class VectorizeHelpper { return 0; } - static int change_arrow_record_batch_schema(const std::unordered_map& column_name_map, - std::shared_ptr schema, - std::shared_ptr in, - std::shared_ptr* out) { + static int change_arrow_record_batch_schema( + const std::unordered_map& column_name_map, + const std::unordered_map& column_type_map, + std::shared_ptr schema, + std::shared_ptr in, + std::shared_ptr* out) { std::vector> columns; for (const auto& f : schema->fields()) { std::shared_ptr array; if (column_name_map.find(f->name()) != column_name_map.end()) { array = in->GetColumnByName(column_name_map.at(f->name())); } + pb::PrimitiveType to_baikaldb_type = pb::INVALID_TYPE; + if (column_type_map.find(f->name()) != column_type_map.end()) { + to_baikaldb_type = column_type_map.at(f->name()); + } if (array == nullptr) { // 兼容加列场景db/store心跳不一致 auto null_array = arrow::MakeArrayOfNull(f->type(), in->num_rows()); @@ -336,10 +370,7 @@ class VectorizeHelpper { } else { // 新schema的列类型和旧schema的列类型可能需要转换，比如列类型发生变更 if (array->type()->id() != f->type()->id()) { - // 整数大类型转小类型，比如int32转int8，需要避免溢出返回错误 - ::arrow::compute::CastOptions cast_options; - cast_options.allow_int_overflow = true; - auto cast_array_ret = ::arrow::compute::Cast(*array, f->type(), cast_options); + auto cast_array_ret = Cast(array, array->type(), f->type(), to_baikaldb_type); if (!cast_array_ret.ok()) { DB_WARNING("cast array fail, %s", cast_array_ret.status().ToString().c_str()); return -1; @@ -357,7 +388,8 @@ class VectorizeHelpper { static int vectorize_filter(std::shared_ptr record_batch, arrow::Expression* conjuncts, - std::shared_ptr* out) { + std::shared_ptr* out, + bool* recordbatch_memory_cross = nullptr) { if (conjuncts == nullptr) { return -1; } @@ -371,20 +403,39 @@ class VectorizeHelpper { const auto& mask_scalar = filter_mask->scalar_as(); if (mask_scalar.is_valid && mask_scalar.value == true) { *out = record_batch; + if (recordbatch_memory_cross != nullptr) { + *recordbatch_memory_cross = true; + } } else { - *out = record_batch->Slice(0, 0); + //*out = record_batch->Slice(0, 0); + auto empty_batch = arrow::RecordBatch::MakeEmpty(record_batch->schema(), GetMemoryPoolForRead()); + if (!empty_batch.ok()) { + DB_FATAL("arrow make empty batch fail, %s", empty_batch.status().ToString().c_str()); + return -1; + } + *out = *empty_batch; + if (recordbatch_memory_cross != nullptr) { + *recordbatch_memory_cross = false; + } } return 0; } auto mask = filter_mask->array_as(); arrow::Datum record_batch_datum(record_batch); - arrow::Result filter_data = arrow::compute::Filter(record_batch_datum, mask->data()); + arrow::compute::ExecContext exec_ctx(GetMemoryPoolForRead()); + arrow::Result filter_data = arrow::compute::Filter(record_batch_datum, + mask->data(), + arrow::compute::FilterOptions::Defaults(), + &exec_ctx); if (!filter_data.ok()) { // TODO 类似 where a这种会filter会报错: Filter argment不是boolean类型, 加上is true DB_FATAL("arrow filter fail, %s", filter_data.status().ToString().c_str()); return -1; } *out = filter_data->record_batch(); + if (recordbatch_memory_cross != nullptr) { + *recordbatch_memory_cross = false; + } return 0; } @@ -394,7 +445,7 @@ class VectorizeHelpper { std::shared_ptr& data_buffer, arrow::Compression::type compression_type = arrow::Compression::UNCOMPRESSED) { arrow::Result> schema_ret = - arrow::ipc::SerializeSchema(*(concatenate_record_batch->schema()), arrow::default_memory_pool()); + arrow::ipc::SerializeSchema(*(concatenate_record_batch->schema()), GetMemoryPoolForRead()); if (!schema_ret.ok()) { DB_FATAL("Fail to SerializeSchema"); return -1; @@ -427,18 +478,23 @@ class VectorizeHelpper { static int concatenate_record_batches( std::shared_ptr schema, std::vector>& batches, - std::shared_ptr& out) { + std::shared_ptr& out, + bool need_check_memory_limit = false) { if (schema == nullptr) { DB_WARNING("schema is nullptr"); return -1; } + arrow::MemoryPool* pool = arrow::default_memory_pool(); + if (need_check_memory_limit) { + pool = GetMemoryPoolForRead(); + } arrow::Result> build_table = arrow::Table::FromRecordBatches(schema, batches); if (!build_table.ok()) { DB_FATAL("FromRecordBatches fail: %s", build_table.status().ToString().c_str()); return -1; } std::shared_ptr table = *build_table; - arrow::Result> record_batch_result = table->CombineChunksToBatch(); + arrow::Result> record_batch_result = table->CombineChunksToBatch(pool); if (!record_batch_result.ok()) { DB_FATAL("CombineChunksToBatch fail: %s", record_batch_result.status().ToString().c_str()); return -1; @@ -496,9 +552,7 @@ class VectorizeHelpper { return nullptr; } std::shared_ptr array = *array_ret; - ::arrow::compute::CastOptions cast_options; - cast_options.allow_int_overflow = true; - auto cast_array_ret = ::arrow::compute::Cast(*array, field->type(), cast_options); + auto cast_array_ret = ::arrow::compute::Cast(*array, field->type(), arrow::compute::CastOptions::Unsafe(field->type())); if (!cast_array_ret.ok()) { DB_WARNING("arrow cast array fail, %s", cast_array_ret.status().message().c_str()); return nullptr; @@ -541,22 +595,358 @@ class VectorizeHelpper { } static std::shared_ptr make_array_from_str( - std::shared_ptr type, const std::string& str, const int32_t length) { + std::shared_ptr type, const std::string& str, const int32_t length, arrow::compute::ExecContext* exec_ctx) { auto array_ret = arrow::MakeArrayFromScalar(arrow::LargeBinaryScalar(arrow::Buffer::FromString(str)), length); if (!array_ret.ok()) { DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); return nullptr; } std::shared_ptr array = *array_ret; - ::arrow::compute::CastOptions cast_options; - cast_options.allow_int_overflow = true; - auto cast_array_ret = ::arrow::compute::Cast(*array, type, cast_options); + auto cast_array_ret = ::arrow::compute::Cast(*array, type, arrow::compute::CastOptions::Unsafe(type), exec_ctx); if (!cast_array_ret.ok()) { DB_WARNING("arrow cast array fail, %s", cast_array_ret.status().message().c_str()); return nullptr; } return *cast_array_ret; } + static std::shared_ptr make_schema(const pb::TupleDescriptor* tuple_desc) { + if (tuple_desc == nullptr) { + DB_WARNING("tuple_desc is nullptr"); + return nullptr; + } + std::vector> arrow_fields; + arrow_fields.reserve(tuple_desc->slots().size()); + for (const auto& slot : tuple_desc->slots()) { + const std::string& field_name = std::to_string(slot.tuple_id()) + "_" + std::to_string(slot.slot_id()); + std::shared_ptr arrow_field = + make_field(field_name, arrow::Type::type(primitive_to_arrow_type(slot.slot_type()))); + if (arrow_field == nullptr) { + DB_WARNING("Fail to make arrow field"); + return nullptr; + } + arrow_fields.emplace_back(arrow_field); + } + return arrow::schema(arrow_fields); + } + static std::shared_ptr make_field(const std::string& name, arrow::Type::type type) { + switch (type) { + case arrow::Type::type::BOOL: + return std::make_shared(name, arrow::boolean()); + case arrow::Type::type::INT8: + return std::make_shared(name, arrow::int8()); + case arrow::Type::type::UINT8: + return std::make_shared(name, arrow::uint8()); + case arrow::Type::type::INT16: + return std::make_shared(name, arrow::int16()); + case arrow::Type::type::UINT16: + return std::make_shared(name, arrow::uint16()); + case arrow::Type::type::INT32: + return std::make_shared(name, arrow::int32()); + case arrow::Type::type::UINT32: + return std::make_shared(name, arrow::uint32()); + case arrow::Type::type::INT64: + return std::make_shared(name, arrow::int64()); + case arrow::Type::type::UINT64: + return std::make_shared(name, arrow::uint64()); + case arrow::Type::type::FLOAT: + return std::make_shared(name, arrow::float32()); + case arrow::Type::type::DOUBLE: + return std::make_shared(name, arrow::float64()); + case arrow::Type::type::LARGE_BINARY: + return std::make_shared(name, arrow::large_binary()); + default: + return nullptr; + } + } + static std::shared_ptr make_array_from_exprvalue( + const pb::PrimitiveType type, const ExprValue& expr_value, const int length) { + bool is_null = expr_value.is_null(); + switch (type) { + case pb::BOOL: { + arrow::BooleanScalar scalar; + if (!is_null) { + scalar = arrow::BooleanScalar(expr_value.get_numberic()); + } + auto array_ret = arrow::MakeArrayFromScalar(scalar, length); + if (!array_ret.ok()) { + DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); + return nullptr; + } else { + return *array_ret; + } + break; + } + case pb::INT8: { + arrow::Int8Scalar scalar; + if (!is_null) { + scalar = arrow::Int8Scalar(expr_value.get_numberic()); + } + auto array_ret = arrow::MakeArrayFromScalar(scalar, length); + if (!array_ret.ok()) { + DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); + return nullptr; + } else { + return *array_ret; + } + break; + } + case pb::UINT8: { + arrow::UInt8Scalar scalar; + if (!is_null) { + scalar = arrow::UInt8Scalar(expr_value.get_numberic()); + } + auto array_ret = arrow::MakeArrayFromScalar(scalar, length); + if (!array_ret.ok()) { + DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); + return nullptr; + } else { + return *array_ret; + } + break; + } + case pb::INT16: { + arrow::Int16Scalar scalar; + if (!is_null) { + scalar = arrow::Int16Scalar(expr_value.get_numberic()); + } + auto array_ret = arrow::MakeArrayFromScalar(scalar, length); + if (!array_ret.ok()) { + DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); + return nullptr; + } else { + return *array_ret; + } + break; + } + case pb::UINT16: { + arrow::UInt16Scalar scalar; + if (!is_null) { + scalar = arrow::UInt16Scalar(expr_value.get_numberic()); + } + auto array_ret = arrow::MakeArrayFromScalar(scalar, length); + if (!array_ret.ok()) { + DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); + return nullptr; + } else { + return *array_ret; + } + break; + } + case pb::INT32: + case pb::TIME: { + arrow::Int32Scalar scalar; + if (!is_null) { + scalar = arrow::Int32Scalar(expr_value.get_numberic()); + } + auto array_ret = arrow::MakeArrayFromScalar(scalar, length); + if (!array_ret.ok()) { + DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); + return nullptr; + } else { + return *array_ret; + } + break; + } + case pb::UINT32: + case pb::DATE: + case pb::TIMESTAMP: { + arrow::UInt32Scalar scalar; + if (!is_null) { + scalar = arrow::UInt32Scalar(expr_value.get_numberic()); + } + auto array_ret = arrow::MakeArrayFromScalar(scalar, length); + if (!array_ret.ok()) { + DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); + return nullptr; + } else { + return *array_ret; + } + break; + } + case pb::INT64: { + arrow::Int64Scalar scalar; + if (!is_null) { + scalar = arrow::Int64Scalar(expr_value.get_numberic()); + } + auto array_ret = arrow::MakeArrayFromScalar(scalar, length); + if (!array_ret.ok()) { + DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); + return nullptr; + } else { + return *array_ret; + } + break; + } + case pb::UINT64: + case pb::DATETIME: { + arrow::UInt64Scalar scalar; + if (!is_null) { + scalar = arrow::UInt64Scalar(expr_value.get_numberic()); + } + auto array_ret = arrow::MakeArrayFromScalar(scalar, length); + if (!array_ret.ok()) { + DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); + return nullptr; + } else { + return *array_ret; + } + break; + } + case pb::FLOAT: { + arrow::FloatScalar scalar; + if (!is_null) { + scalar = arrow::FloatScalar(expr_value.get_numberic()); + } + auto array_ret = arrow::MakeArrayFromScalar(scalar, length); + if (!array_ret.ok()) { + DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); + return nullptr; + } else { + return *array_ret; + } + break; + } + case pb::DOUBLE: { + arrow::DoubleScalar scalar; + if (!is_null) { + scalar = arrow::DoubleScalar(expr_value.get_numberic()); + } + auto array_ret = arrow::MakeArrayFromScalar(scalar, length); + if (!array_ret.ok()) { + DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); + return nullptr; + } else { + return *array_ret; + } + break; + } + case pb::STRING: + case pb::HLL: + case pb::BITMAP: + case pb::TDIGEST: { + auto array_ret = arrow::MakeArrayFromScalar( + is_null ? arrow::LargeBinaryScalar(): + arrow::LargeBinaryScalar(arrow::Buffer::FromString(expr_value.get_string())), + length); + if (!array_ret.ok()) { + DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); + return nullptr; + } else { + return *array_ret; + } + break; + } + default: + DB_WARNING("Invalid type: %d", type); + return nullptr; + } + return nullptr; + } + static arrow::Result> Cast( + const std::shared_ptr& array, + const std::shared_ptr& type, + const std::shared_ptr& to_type, + const pb::PrimitiveType to_baikaldb_type = pb::INVALID_TYPE) { + bool is_string_type = false; + switch (type->id()) { + case arrow::Type::type::BINARY: + case arrow::Type::type::LARGE_BINARY: + case arrow::Type::type::STRING: + case arrow::Type::type::LARGE_STRING: + is_string_type = true; + break; + default: + break; + } + std::string cast_func_name; + if (is_string_type) { + switch (to_baikaldb_type) { + case pb::DATE: + cast_func_name = "expr_value_to_date"; + break; + case pb::DATETIME: + cast_func_name = "expr_value_to_datetime"; + break; + case pb::TIME: + cast_func_name = "expr_value_to_time"; + break; + case pb::TIMESTAMP: + cast_func_name = "expr_value_to_timestamp"; + break; + default: + break; + } + } + arrow::compute::ExecContext exec_ctx(GetMemoryPoolForRead()); + if (!cast_func_name.empty()) { + ExprValueCastFunctionOptions options(pb::STRING); + auto cast_array_ret = ::arrow::compute::CallFunction(cast_func_name, {array}, &options, &exec_ctx); + if (!cast_array_ret.ok()) { + DB_WARNING("cast array fail, %s", cast_array_ret.status().ToString().c_str()); + return arrow::Status::Invalid("cast array fail");; + } + return (*cast_array_ret).make_array(); + } else { + // 整数大类型转小类型，比如int32转int8，需要避免溢出返回错误 + return ::arrow::compute::Cast(*array, to_type, arrow::compute::CastOptions::Unsafe(to_type), &exec_ctx); + } + } + + static int init_conditions(RuntimeState* state, + std::vector& conditions, + arrow::compute::Expression& vectorize_conditions, + int32_t scan_tuple_id, + std::shared_ptr schema_ptr, + pb::TraceNode* trace_node = nullptr) { + if (conditions.empty()) { + return 0; + } + for (auto condition : conditions) { + if (condition == nullptr) { + DB_FATAL_STATE(state, "condition is nullptr"); + return -1; + } + int ret = condition->open(); + if (ret < 0) { + DB_FATAL_STATE(state, "expr open fail, ret:%d", ret); + return ret; + } + } + if (state->execute_type == pb::EXEC_ARROW_ACERO) { + std::vector sub_exprs; + sub_exprs.reserve(conditions.size()); + for (auto condition : conditions) { + if (0 != condition->transfer_to_arrow_expression()) { + DB_FATAL_STATE(state, "transfer to arrow expr fail"); + return -1; + } + sub_exprs.emplace_back(condition->arrow_expr()); + } + vectorize_conditions = arrow::compute::and_(sub_exprs); + std::shared_ptr schema; + if (schema_ptr != nullptr) { + schema = schema_ptr; // exchange_receiver + } else { + auto tuple = state->get_tuple_desc(scan_tuple_id); // select_manager_node + if (tuple == nullptr) { + DB_FATAL_STATE(state, "no such scan_tuple id: %d", scan_tuple_id); + return -1; + } + schema = VectorizeHelpper::get_arrow_schema({tuple}); + } + if (schema == nullptr) { + DB_FATAL_STATE(state, "get arrow schema fail for filter condition"); + return -1; + } + arrow::Result bind_expr = vectorize_conditions.Bind(*schema); + if (!bind_expr.ok()) { + DB_FATAL_STATE(state, "bind expr fail:%s", bind_expr.status().ToString().c_str()); + return -1; + } + vectorize_conditions = *bind_expr; + } + return 0; + } }; // [ARROW TODO, 复用FetcherStoreVectorizedReader] diff --git a/include/expr/agg_fn_call.h b/include/expr/agg_fn_call.h index 35f548522..4532a67b1 100644 --- a/include/expr/agg_fn_call.h +++ b/include/expr/agg_fn_call.h @@ -39,6 +39,7 @@ struct ExprValueComparator { using ExprValueUniqSet = std::unordered_set; struct AvgIntermediate { + // 只能携带原始类型，否则可能出错 double sum; int64_t count; AvgIntermediate() : sum(0.0), count(0) { @@ -222,6 +223,10 @@ class AggFnCall : public ExprNode { bool is_merge, std::vector& generate_projection_exprs, std::vector& generate_projection_exprs_names); + int transfer_to_arrow_group_concat(std::vector& aggs, + bool is_merge, + std::vector& generate_projection_exprs, + std::vector& generate_projection_exprs_names); int transfer_to_arrow_agg_function(std::vector& aggs, bool is_merge, std::vector& generate_projection_exprs, diff --git a/include/expr/arrow_function.h b/include/expr/arrow_function.h index e0169d9e6..d34974cb2 100644 --- a/include/expr/arrow_function.h +++ b/include/expr/arrow_function.h @@ -97,6 +97,8 @@ int arrow_substr(std::vector& children, pb::Function* fn, const pb::P int arrow_upper(std::vector& children, pb::Function* fn, const pb::PrimitiveType& return_type, arrow::compute::Expression& out); int arrow_lower(std::vector& children, pb::Function* fn, const pb::PrimitiveType& return_type, arrow::compute::Expression& out); int arrow_repeat(std::vector& children, pb::Function* fn, const pb::PrimitiveType& return_type, arrow::compute::Expression& out); +int arrow_substring_index(std::vector& children, pb::Function* fn, const pb::PrimitiveType& return_type, arrow::compute::Expression& out); +int arrow_replace(std::vector& children, pb::Function* fn, const pb::PrimitiveType& return_type, arrow::compute::Expression& out); // 类型转换 int arrow_cast_to_string(std::vector& children, pb::Function* fn, const pb::PrimitiveType& return_type, arrow::compute::Expression& out); int arrow_cast_to_date(std::vector& children, pb::Function* fn, const pb::PrimitiveType& return_type, arrow::compute::Expression& out); @@ -132,7 +134,7 @@ int arrow_unix_timestamp(std::vector& children, pb::Function* fn, con int arrow_week(std::vector& children, pb::Function* fn, const pb::PrimitiveType& return_type, arrow::compute::Expression& out); int arrow_yearweek(std::vector& children, pb::Function* fn, const pb::PrimitiveType& return_type, arrow::compute::Expression& out); int arrow_timestampdiff(std::vector& children, pb::Function* fn, const pb::PrimitiveType& return_type, arrow::compute::Expression& out); - +int arrow_datediff(std::vector& children, pb::Function* fn, const pb::PrimitiveType& return_type, arrow::compute::Expression& out); /* * 注册使用的FunctionOptions */ @@ -147,13 +149,34 @@ struct ExprValueCastState : public arrow::compute::KernelState { explicit ExprValueCastState(pb::PrimitiveType type) : type(type) {} }; -class CommonTimeFunctionOptions : public arrow::compute::FunctionOptions { +class CommonFunctionOptions : public arrow::compute::FunctionOptions { public: - CommonTimeFunctionOptions(const std::string& str_value, int64_t int_value = 0); + CommonFunctionOptions(const std::string& str_value, int64_t int_value = 0); + std::string str_value; + int64_t int_value = 0; +}; +struct CommonState : public arrow::compute::KernelState { std::string str_value; int64_t int_value = 0; + CommonState(const std::string& conf, int64_t value) : str_value(conf), int_value(value) {} +}; + +class GroupConcatOptions : public arrow::compute::FunctionOptions { + public: + explicit GroupConcatOptions(const std::string& sep = "", const std::vector& asc = {}); + static constexpr char const kTypeName[] = "GroupConcatOptions"; + static GroupConcatOptions Defaults() { return GroupConcatOptions{}; } + + std::string separator_; + // TODO + // for group concat orderby, can use make_struct args like {string, orderby feild1, orderby feild2, ...} + // or sort indicates first, and then use make_struct args like {string, orderby indicate id} + std::vector asc_; }; +arrow::Result> InitCommonState(arrow::compute::KernelContext*, + const arrow::compute::KernelInitArgs& args); + class ArrowFunctionManager : public ObjectManager< ArrowExprBuildFun, ArrowFunctionManager> { diff --git a/include/expr/predicate.h b/include/expr/predicate.h index a2a17ec2a..bc57ea6f4 100644 --- a/include/expr/predicate.h +++ b/include/expr/predicate.h @@ -245,7 +245,12 @@ class IsTruePredicate : public ScalarFnCall { if (_children[0]->transfer_to_arrow_expression() != 0) { return -1; } - _arrow_expr = arrow::compute::call("cast", {_children[0]->arrow_expr()}, arrow::compute::CastOptions::Unsafe(arrow::boolean())); + auto bool_expr = arrow_cast(_children[0]->arrow_expr(), _children[0]->col_type(), pb::BOOL); + _arrow_expr = arrow::compute::call("if_else", { + arrow::compute::call("is_null", {bool_expr}), + arrow::compute::literal(false), + bool_expr + }); return 0; } virtual bool can_use_arrow_vector() { @@ -466,9 +471,7 @@ class NotPredicate : public ScalarFnCall { if (_children[0]->transfer_to_arrow_expression() != 0) { return -1; } - _arrow_expr = arrow::compute::not_(arrow::compute::call("cast", - {_children[0]->arrow_expr()}, - arrow::compute::CastOptions::Unsafe(arrow::boolean()))); + _arrow_expr = arrow::compute::not_(arrow_cast(_children[0]->arrow_expr(), _children[0]->col_type(), pb::BOOL)); return 0; } virtual bool can_use_arrow_vector() { diff --git a/include/expr/row_expr.h b/include/expr/row_expr.h index aef7ec39c..75a41063a 100755 --- a/include/expr/row_expr.h +++ b/include/expr/row_expr.h @@ -70,7 +70,7 @@ class RowExpr : public ExprNode { virtual bool can_use_arrow_vector() { return false; } - virtual bool can_use_arrow_vector_for_compare_sclar_exrpr() { + virtual bool children_can_use_arrow_vector() { for (auto& c : _children) { if (!c->can_use_arrow_vector()) { return false; diff --git a/include/expr/scalar_fn_call.h b/include/expr/scalar_fn_call.h index 7789657a8..31922c5af 100644 --- a/include/expr/scalar_fn_call.h +++ b/include/expr/scalar_fn_call.h @@ -44,6 +44,19 @@ class ScalarFnCall : public ExprNode { // vectorized virtual int transfer_to_arrow_expression(); virtual bool can_use_arrow_vector(); + bool is_compare_op() { + switch (_fn.fn_op()) { + case parser::FT_EQ: + case parser::FT_NE: + case parser::FT_GE: + case parser::FT_GT: + case parser::FT_LE: + case parser::FT_LT: + return true; + default: + return false; + } + } private: ExprValue multi_eq_value(MemRow* row) { for (size_t i = 0; i < children(0)->children_size(); i++) { @@ -124,7 +137,6 @@ class ScalarFnCall : public ExprNode { protected: pb::Function _fn; - std::string _origin_fn_name; bool _is_row_expr = false; std::function&)> _fn_call; // FT_COMMON构建arrow expression diff --git a/include/expr/window_fn_call.h b/include/expr/window_fn_call.h index 2763c07f1..4a82e4a9d 100644 --- a/include/expr/window_fn_call.h +++ b/include/expr/window_fn_call.h @@ -67,6 +67,17 @@ struct WindowValueIntermediate { ExprValue value; }; +typedef std::variant< + WindowCountIntermediate, + WindowSumIntermediate, + WindowAvgIntermediate, + WindowMinMaxIntermediate, + WindowRowNumberIntermediate, + WindowRankIntermediate, + WindowNtileIntermediate, + WindowLeadLagIntermediate, + WindowValueIntermediate> Intermediate; + class WindowFnCall : public ExprNode { public: enum WindowType { @@ -127,7 +138,7 @@ class WindowFnCall : public ExprNode { private: WindowType _window_type; - std::string _window_intermediate; // 窗口函数中间结果 + Intermediate _window_intermediate; // 窗口函数中间结果 std::shared_ptr _mem_row_compare; }; diff --git a/include/logical_plan/ddl_planner.h b/include/logical_plan/ddl_planner.h index 0a243ab8d..3e6d83e05 100644 --- a/include/logical_plan/ddl_planner.h +++ b/include/logical_plan/ddl_planner.h @@ -92,6 +92,7 @@ class DDLPlanner : public LogicalPlanner { // DBLINK int check_dblink_table_valid(const pb::SchemaInfo& table); + int check_alter_dblink_table_valid(const TableInfo& table, const parser::AlterTableStmt& stmt); std::map _column_can_null; diff --git a/include/logical_plan/query_context.h b/include/logical_plan/query_context.h index d2c085ea3..8844f9136 100644 --- a/include/logical_plan/query_context.h +++ b/include/logical_plan/query_context.h @@ -421,6 +421,7 @@ class QueryContext { std::set current_table_tuple_ids; bool open_binlog = false; bool no_binlog = false; // 用于控制DM导入是否写binlog + bool disable_on_update = false; // 为true时 ON UPDATE CURRENT_TIMESTAMP 不会自动更新 SignExecType sql_exec_type_defined = SignExecType::SIGN_EXEC_NOT_SET; bool use_mpp = false; @@ -442,6 +443,7 @@ class QueryContext { bool is_cancelled = false; bool execute_global_flow = false; std::shared_ptr kill_ctx; + bool kill_without_raft = false; std::vector> sub_query_plans; std::unordered_map long_data_vars; std::vector param_type; @@ -500,6 +502,7 @@ class QueryContext { // 向量索引使用 int32_t efsearch = -1; + int32_t nprobe = -1; bool dumped_slow_sql = false; @@ -509,6 +512,12 @@ class QueryContext { // 查询是否包含DBLink Mysql表 bool has_dblink_mysql = false; + // 查询是否包含DBLink File表 + bool has_dblink_file = false; + + // 查询是否必须使用向量化执行 + bool must_vectorize = false; + private: std::vector _tuple_descs; bthread::Mutex _kill_lock; diff --git a/include/meta_server/region_manager.h b/include/meta_server/region_manager.h index 6df47df3d..220bd5903 100644 --- a/include/meta_server/region_manager.h +++ b/include/meta_server/region_manager.h @@ -65,7 +65,7 @@ class RegionManager { void update_region(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void restore_region(const pb::MetaManagerRequest& request, pb::MetaManagerResponse* response); void drop_region(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); - void split_region(const pb::MetaManagerRequest& request, braft::Closure* done); + void split_region(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void send_remove_region_request(const std::vector& drop_region_ids); // MIGRATE/DEAD可以任意使用add_peer_for_store/delete_all_region_for_store @@ -107,15 +107,25 @@ class RegionManager { const pb::LeaderHeartBeat& leader_region, const SmartRegionInfo& master_region_info, pb::StoreHeartBeatResponse* response); - bool peer_is_equal(const pb::RegionInfo& left, const pb::RegionInfo& right) { + bool peer_is_equal(const pb::RegionInfo& left, const pb::RegionInfo& right, bool check_learner) { size_t hash_left = 0; size_t hash_right = 0; for (auto& state : left.peers()) { hash_left += std::hash{}(state); } + if (check_learner) { + for (auto& state : left.learners()) { + hash_left += std::hash{}(state); + } + } for (auto& state : right.peers()) { hash_right += std::hash{}(state); } + if (check_learner) { + for (auto& state : right.learners()) { + hash_right += std::hash{}(state); + } + } return hash_left == hash_right; } //是否有超过replica_num数量的region, 这种region需要删掉多余的peer diff --git a/include/meta_server/table_manager.h b/include/meta_server/table_manager.h index 65647aee6..b19f6f93b 100644 --- a/include/meta_server/table_manager.h +++ b/include/meta_server/table_manager.h @@ -204,6 +204,8 @@ class TableManager { void update_byte_size(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void update_split_lines(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void update_charset(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); + int pre_process_for_add_partition(const pb::MetaManagerRequest* request, pb::MetaManagerResponse* response, + uint64_t log_id, google::protobuf::Closure* done); void add_partition(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void drop_partition(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void modify_partition(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); @@ -213,15 +215,18 @@ class TableManager { void update_statistics(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void update_dists(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void update_ttl_duration(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); + void update_ttl_info(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void update_resource_tag(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void update_table_comment(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void update_dynamic_partition_attr(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void drop_partition_ts(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void specify_split_keys(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); - + int pre_process_for_add_index(const pb::MetaManagerRequest* request, pb::MetaManagerResponse* response, + uint64_t log_id, google::protobuf::Closure* done); void add_field(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void add_index(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); - int do_add_index(const pb::SchemaInfo& mem_schema_pb, const int64_t apply_index, braft::Closure* done, + int do_add_index(const pb::MetaManagerRequest& request, const pb::SchemaInfo& mem_schema_pb, + const int64_t apply_index, braft::Closure* done, const int64_t table_id, pb::IndexInfo& index_info); void drop_index(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); void drop_field(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done); @@ -306,11 +311,23 @@ class TableManager { void on_leader_start(); void on_leader_stop(); - + int get_ttl_field_name(int64_t table_id, std::string& ttl_field_name); bool is_create_table_support_engine(pb::Engine engine) { return (engine == pb::ROCKSDB || engine == pb::ROCKSDB_CSTORE || engine == pb::BINLOG); } - + static bool is_ttl_support_index_type(const pb::IndexInfo& index_info) { + if (index_info.index_type() == pb::I_PRIMARY || index_info.index_type() == pb::I_KEY) { + return true; + } + if (index_info.index_type() == pb::I_FULLTEXT || index_info.index_type() == pb::I_VECTOR) { + return true; + } + if (index_info.index_type() == pb::I_UNIQ) { + return true; + } + return false; + } + public: void set_max_table_id(int64_t max_table_id) { BAIDU_SCOPED_LOCK(_table_mutex); @@ -1688,10 +1705,12 @@ class TableManager { _table_timer.init(3600 * 1000); // 1h _dblink_mysql_table_timer.init(FLAGS_baikal_heartbeat_interval_us / 1000); } - int write_schema_for_not_level(TableMem& table_mem, + int write_schema_for_not_level(const pb::MetaManagerRequest& request, + TableMem& table_mem, braft::Closure* done, int64_t max_table_id_tmp, - bool has_auto_increment); + bool has_auto_increment, + const int64_t apply_index); int send_auto_increment_request(const pb::MetaManagerRequest& request); int update_schema_for_rocksdb(int64_t table_id, @@ -1719,9 +1738,10 @@ class TableManager { int check_index(const pb::IndexInfo& index_info_to_check, const pb::SchemaInfo& schema_info, int64_t& index_id); - + static int check_gflag(const IdcInfo& idc, const std::string& gflag_name, const std::string& expected_value); int alloc_field_id(pb::SchemaInfo& table_info, bool& has_auto_increment, TableMem& table_mem); - int alloc_index_id(pb::SchemaInfo& table_info, TableMem& table_mem, int64_t& max_table_id_tmp); + int alloc_index_id(pb::SchemaInfo& table_info, TableMem& table_mem, int64_t& max_table_id_tmp, + std::string resource_tag); void construct_common_region(pb::RegionInfo* region_info, int32_t replica_num) { region_info->set_version(1); region_info->set_conf_version(1); @@ -1760,7 +1780,8 @@ class TableManager { (index_info.index_type() == pb::I_UNIQ || index_info.index_type() == pb::I_KEY); } - int init_global_index_region(const pb::SchemaInfo& schema_info, int64_t table_id, braft::Closure* done, pb::IndexInfo& index_info); + int init_global_index_region(const pb::MetaManagerRequest& request, const pb::SchemaInfo& schema_info, + int64_t table_id, braft::Closure* done, pb::IndexInfo& index_info, const int64_t apply_index); bool partition_check_region_when_update(int64_t table_id, std::string min_start_key, @@ -1820,6 +1841,10 @@ class TableManager { bool check_vector_index(const pb::SchemaInfo& mem_schema_pb, const pb::IndexInfo& index_info); + static bool if_pk_contains_ttl_field(const pb::SchemaInfo& table_info); + + static int check_ttl_info(const pb::SchemaInfo& table_schema, std::string& err_msg); + private: bthread_mutex_t _table_mutex; bthread_mutex_t _load_virtual_to_memory_mutex; @@ -1835,6 +1860,9 @@ class TableManager { DBLinkMysqlTableTimer _dblink_mysql_table_timer; DoubleBufferedTableSchedulingInfo _table_scheduling_infos; + + static const std::set _need_check_wordrank_types; + static const std::set _need_check_wordweight_types; }; //class }//namespace diff --git a/include/physical_plan/fragment.h b/include/physical_plan/fragment.h index 8482e5038..445f929c7 100644 --- a/include/physical_plan/fragment.h +++ b/include/physical_plan/fragment.h @@ -64,6 +64,8 @@ struct FragmentInfo { int64_t open_cost = 0; int64_t exec_cost = 0; + pb::TraceNode trace_node; + arrow::Result> wait() { if (executor) { executor->wait(); diff --git a/include/physical_plan/index_selector.h b/include/physical_plan/index_selector.h index 7b03b15fd..340136603 100644 --- a/include/physical_plan/index_selector.h +++ b/include/physical_plan/index_selector.h @@ -55,13 +55,14 @@ class IndexSelector { std::map& field_range_type, const std::string& sample_sql, const IndexSelectorOptions& options); + static void hit_field_range(ExprNode* expr, std::map& field_range_map, bool* index_predicate_is_null, + int64_t table_id, FulltextInfoNode* fulltext_index_node); + private: - void hit_row_field_range(ExprNode* expr, std::map& field_range_map, bool* index_predicate_is_null); - void hit_match_against_field_range(ExprNode* expr, + static void hit_row_field_range(ExprNode* expr, std::map& field_range_map, bool* index_predicate_is_null); + static void hit_match_against_field_range(ExprNode* expr, std::map& field_range_map, FulltextInfoNode* fulltext_index_node, int64_t table_id); - void hit_field_range(ExprNode* expr, std::map& field_range_map, bool* index_predicate_is_null, - int64_t table_id, FulltextInfoNode* fulltext_index_node); - void hit_field_or_like_range(ExprNode* expr, std::map& field_range_map, + static void hit_field_or_like_range(ExprNode* expr, std::map& field_range_map, int64_t table_id, FulltextInfoNode* fulltext_index_node); bool check_rollup_index_valid(SmartTable& table_info, const IndexInfo& index_info, @@ -77,14 +78,15 @@ class IndexSelector { FilterNode* filter_node, std::map& field_range_map); - bool is_field_has_reverse_index(int64_t table_id, int64_t field_id, int64_t* index_id_ptr) { - auto table_ptr = _factory->get_table_info_ptr(table_id); + static bool is_field_has_reverse_index(int64_t table_id, int64_t field_id, int64_t* index_id_ptr) { + auto factory = SchemaFactory::get_instance(); + auto table_ptr = factory->get_table_info_ptr(table_id); if (table_ptr != nullptr) { // 优先选择arrow格式 auto iter = table_ptr->arrow_reverse_fields.find(field_id); if (iter != table_ptr->arrow_reverse_fields.end()) { *index_id_ptr = iter->second; - auto index_ptr = _factory->get_index_info_ptr(*index_id_ptr); + auto index_ptr = factory->get_index_info_ptr(*index_id_ptr); if (index_ptr != nullptr && index_ptr->state == pb::IS_PUBLIC && index_ptr->index_hint_status == pb::IHS_NORMAL) { @@ -94,7 +96,7 @@ class IndexSelector { iter = table_ptr->reverse_fields.find(field_id); if (iter != table_ptr->reverse_fields.end()) { *index_id_ptr = iter->second; - auto index_ptr = _factory->get_index_info_ptr(*index_id_ptr); + auto index_ptr = factory->get_index_info_ptr(*index_id_ptr); if (index_ptr != nullptr && index_ptr->state == pb::IS_PUBLIC && index_ptr->index_hint_status == pb::IHS_NORMAL) { @@ -106,6 +108,11 @@ class IndexSelector { } int select_partition(SmartTable& table_info, ScanNode* scan_node, std::map& field_range_map); + // 获取需要查询的离线分区文件集合 + int select_partition_files(SmartTable& table_info, + ScanNode* scan_node, + FilterNode* filter_node, + ExprNode* join_on_conditions); int64_t index_merge_selector(const std::vector& tuple_descs, ScanNode* scan_node, diff --git a/include/physical_plan/mpp_analyzer.h b/include/physical_plan/mpp_analyzer.h index 1e84044c3..3b939b682 100644 --- a/include/physical_plan/mpp_analyzer.h +++ b/include/physical_plan/mpp_analyzer.h @@ -54,6 +54,9 @@ class MppAnalyzer { int build_fragment(QueryContext* ctx, SmartFragment& fragment); + // 构建包含ScanNode的Fragment + int build_scan_fragment(QueryContext* ctx, SmartFragment& fragment); + int do_separate(QueryContext* ctx, ExecNode* node, std::vector& child_fragment_exchange_senders); diff --git a/include/physical_plan/physical_planner.h b/include/physical_plan/physical_planner.h index 7454ea1c9..4d5c312e2 100644 --- a/include/physical_plan/physical_planner.h +++ b/include/physical_plan/physical_planner.h @@ -29,6 +29,7 @@ #include "fragment.h" #include "mpp_analyzer.h" #include "join_type_analyzer.h" +#include "condition_optimizer.h" namespace baikaldb { class PhysicalPlanner { diff --git a/include/physical_plan/plan_router.h b/include/physical_plan/plan_router.h index 6420dd86d..ec0be75e1 100644 --- a/include/physical_plan/plan_router.h +++ b/include/physical_plan/plan_router.h @@ -29,13 +29,13 @@ class PlanRouter { /* 通过主键索引获取所在的regions */ int analyze(QueryContext* ctx); - int scan_plan_router(RocksdbScanNode* scan_node, + int scan_plan_router(ScanNode* scan_node, const std::function& get_slot_id, const std::function& get_tuple_desc, bool has_join, const std::set& escape_get_region_infos); - int scan_node_analyze(RocksdbScanNode* scan_node, + int scan_node_analyze(ScanNode* scan_node, QueryContext* ctx, bool has_join, const std::set& escape_get_region_infos); private: diff --git a/include/protocol/db_service.h b/include/protocol/db_service.h index a29feba6d..7fbbd93d1 100644 --- a/include/protocol/db_service.h +++ b/include/protocol/db_service.h @@ -27,6 +27,48 @@ class DBInteract { return &_instance; } + int construct_mpp_dag_request(QueryContext* ctx, + const std::set& fragment_ids, + pb::DAGFragmentRequest& request) { + if (ctx == nullptr) { + DB_WARNING("ctx is nullptr"); + return -1; + } + SmartState state = ctx->get_runtime_state(); + if (state == nullptr) { + DB_WARNING("state is nullptr"); + return -1; + } + request.set_op(pb::OP_FRAGMENT_START); + request.set_log_id(state->log_id()); + request.set_sql_sign(state->sign); + for (const auto& fragment_id : fragment_ids) { + SmartFragment fragment_info; + if (ctx->fragments.find(fragment_id) != ctx->fragments.end()) { + fragment_info = ctx->fragments[fragment_id]; + } + if (fragment_info == nullptr) { + DB_WARNING("fragment_info is nullpt, fragment_id: %d", fragment_id); + return -1; + } + pb::Plan fragment_plan; + pb::FragmentInfo* fragment_info_pb = request.add_fragments(); + fragment_info_pb->set_fragment_id(fragment_id); + pb::RuntimeState* pb_rs = fragment_info_pb->mutable_runtime_state(); + if (fragment_info->runtime_state == nullptr) { + DB_WARNING("fragment_info->runtime_state is nullptr"); + return -1; + } + fragment_info->runtime_state->to_proto(pb_rs); + ExecNode::create_pb_plan(0, &fragment_plan, fragment_info->root); + fragment_info_pb->mutable_plan()->CopyFrom(fragment_plan); + } + if (ctx->user_info != nullptr) { + request.set_username(ctx->user_info->username); + } + return 0; + } + int handle_mpp_dag_fragment(const pb::DAGFragmentRequest& request, const std::string& db_address) { brpc::ChannelOptions channel_opt; diff --git a/include/protocol/show_helper.h b/include/protocol/show_helper.h index 290fcd4e9..e07fd25c5 100644 --- a/include/protocol/show_helper.h +++ b/include/protocol/show_helper.h @@ -27,6 +27,7 @@ const std::string SQL_SHOW_CREATE_TABLE = "create_table"; // s const std::string SQL_SHOW_CREATE_VIEW = "create_view"; // show create view test; const std::string SQL_SHOW_COLLATION = "collation"; // show collation; const std::string SQL_SHOW_DATABASES = "databases"; // show databases; +const std::string SQL_SHOW_SCHEMAS = "schemas"; // show schemas; const std::string SQL_SHOW_NAMESPACE = "namespace"; // show namespace const std::string SQL_SHOW_META = "meta"; // show meta; const std::string SQL_SHOW_TABLE_STATUS = "table"; // show table status; diff --git a/include/reverse/boolean_engine/boolean_executor.h b/include/reverse/boolean_engine/boolean_executor.h index ddb4e6567..33dd1b2aa 100644 --- a/include/reverse/boolean_engine/boolean_executor.h +++ b/include/reverse/boolean_engine/boolean_executor.h @@ -31,15 +31,27 @@ typedef std::string PrimaryIdT; typedef int (*MergeFuncT)(PostingNodeT&, const PostingNodeT&, BoolArg*); +template +class RindexNodeParser; + +template +struct ReverseDelayInitContext { + std::vector*> children; + std::vector reverse_rocksdb_keys; + std::vector terms; +}; + // 获取并遍历倒排链表的接口 template class RindexNodeParser { public: + using ReverseListSptr = typename Schema::ReverseListSptr; + RindexNodeParser(Schema* schema) : _schema(schema) { } virtual ~RindexNodeParser() { } - virtual int init(const std::string& term) = 0; + virtual int init(const std::string& term, ReverseDelayInitContext* delay_init_context = nullptr) = 0; //第一次调用返回第一个元素 virtual const PostingNodeT* current_node() = 0; virtual const PrimaryIdT* current_id() = 0; @@ -47,7 +59,8 @@ class RindexNodeParser { virtual const PostingNodeT* next() = 0; //如果倒排链表是有序数组，用二分查找优化 //大于等于target_id的第一个元素（包括当前元素） - virtual const PostingNodeT* advance(const PrimaryIdT& target_id) = 0; + virtual const PostingNodeT* advance(const PrimaryIdT& target_id) = 0; + virtual int delay_init(const std::string& term, ReverseListSptr new_list_ptr, ReverseListSptr old_list_ptr) = 0; protected: Schema* _schema; }; diff --git a/include/reverse/boolean_engine/logical_query.h b/include/reverse/boolean_engine/logical_query.h index 319732cb1..333ea6739 100644 --- a/include/reverse/boolean_engine/logical_query.h +++ b/include/reverse/boolean_engine/logical_query.h @@ -45,9 +45,15 @@ template class LogicalQuery { public: typedef typename Schema::Parser Parser; - LogicalQuery(Schema *schema) : _schema(schema) {} + using ReverseListSptr = typename ReverseIndex::ReverseListSptr; + using ReverseList = typename ReverseIndex::ReverseList; + + LogicalQuery(Schema *schema, myrocksdb::Transaction* txn, bool is_fast = false) + : _schema(schema), _txn(txn), _is_fast(is_fast) { + _rocksdb = RocksWrapper::get_instance(); + } ~LogicalQuery(){} - BooleanExecutor* create_executor(); + BooleanExecutor* create_executor(); ExecutorNode _root; private: BooleanExecutor* parse_executor_node(const ExecutorNode& node); @@ -55,7 +61,18 @@ class LogicalQuery { BooleanExecutor* parse_op_node(const ExecutorNode& node); void and_or_add_subnode(const ExecutorNode&, OperatorBooleanExecutor*); void weight_add_subnode(const ExecutorNode&, OperatorBooleanExecutor*); + int _multi_get_reverse_list_two( + std::vector& list_new_ptrs, + std::vector& list_old_ptrs); + + // input and ouput are both stored in _delay_init_context + int _multi_get_level_reverse_list(uint8_t level, std::vector& list_ptrs); + Schema *_schema; + bool _is_fast; + RocksWrapper* _rocksdb; + myrocksdb::Transaction* _txn; + std::shared_ptr> _delay_init_context = nullptr; }; } // namespace logical_query diff --git a/include/reverse/boolean_engine/logical_query.hpp b/include/reverse/boolean_engine/logical_query.hpp index 0a38a8630..e3a8296b4 100644 --- a/include/reverse/boolean_engine/logical_query.hpp +++ b/include/reverse/boolean_engine/logical_query.hpp @@ -27,8 +27,34 @@ BooleanExecutor* LogicalQuery::parse_executor_node( return parse_term_node(executor_node); case AND : case OR : - case WEIGHT : - return parse_op_node(executor_node); + case WEIGHT : { + bool need_delay_init = false; + if (_is_fast && _delay_init_context == nullptr) { + // 只在fast模式下进行延迟初始化 + _delay_init_context = std::make_shared>(); + need_delay_init = true; + } + BooleanExecutor* ret = parse_op_node(executor_node); + if (need_delay_init) { + std::vector list_new_ptrs(_delay_init_context->children.size(), nullptr); + std::vector list_old_ptrs(_delay_init_context->children.size(), nullptr); + if (_multi_get_reverse_list_two(list_new_ptrs, list_old_ptrs) != 0) { + DB_WARNING("multi get reverse list failed"); + return nullptr; + } + for (int i = 0; i < _delay_init_context->children.size(); i++) { + int r = _delay_init_context->children[i]->delay_init( + _delay_init_context->terms[i], + list_new_ptrs[i], + list_old_ptrs[i]); + if (r != 0) { + DB_FATAL("init logicalQuery failed."); + return nullptr; + } + } + } + return ret; + } default : DB_WARNING("boolean executor type (%d) is invalid", executor_node._type); return NULL; @@ -36,10 +62,9 @@ BooleanExecutor* LogicalQuery::parse_executor_node( } template -BooleanExecutor* LogicalQuery::parse_term_node( - const ExecutorNode& node) { +BooleanExecutor* LogicalQuery::parse_term_node(const ExecutorNode& node) { Parser* parser = new Parser(_schema); - parser->init(node._term); + parser->init(node._term, _delay_init_context.get()); return new TermBooleanExecutor(parser, node._term, _schema->executor_type, node._arg); } @@ -91,6 +116,85 @@ void LogicalQuery::and_or_add_subnode( } } } +template +int LogicalQuery::_multi_get_reverse_list_two( + std::vector& list_new_ptrs, + std::vector& list_old_ptrs) { + if (_delay_init_context == nullptr) { + return 0; + } + rocksdb::ReadOptions roptions; + roptions.prefix_same_as_start = true; + roptions.fill_cache = false; + auto data_cf = _rocksdb->get_data_handle(); + if (data_cf == nullptr) { + DB_WARNING("get rocksdb data column family failed"); + return -1; + } + + if (_is_fast) { + _multi_get_level_reverse_list(2, list_new_ptrs); + } else { + // should not reach here + } + _multi_get_level_reverse_list(3, list_old_ptrs); + return 0; +} + +template +int LogicalQuery::_multi_get_level_reverse_list(uint8_t level, std::vector& list_ptrs) { + std::vector complete_key_idx(_delay_init_context->reverse_rocksdb_keys.size(), 0); + std::vector rocksdb_keys; + rocksdb_keys.reserve(complete_key_idx.size()); + std::vector rocks_key_string; + rocks_key_string.reserve(complete_key_idx.size()); + std::unordered_map key_index_map; + + for (size_t i = 0; i < _delay_init_context->reverse_rocksdb_keys.size(); ++i) { + std::string key = _delay_init_context->reverse_rocksdb_keys[i]; + key.append((char*)&level, sizeof(uint8_t)); + key.append(_delay_init_context->terms[i]); + if (key_index_map.find(key) == key_index_map.end()) { + rocks_key_string.emplace_back(key); + rocksdb_keys.emplace_back(rocks_key_string.back()); + key_index_map[key] = rocksdb_keys.size() - 1; + } + complete_key_idx[i] = key_index_map[key]; + } + + rocksdb::ReadOptions roptions; + auto data_cf = _rocksdb->get_data_handle(); + if (data_cf == nullptr) { + DB_WARNING("get rocksdb data column family failed"); + return -1; + } + + std::vector statuses(rocksdb_keys.size()); + std::vector values(rocksdb_keys.size()); + _txn->MultiGet(roptions, data_cf, rocksdb_keys, values, statuses, false); + std::vector reverse_list_ptrs(statuses.size(), nullptr); + for (int i = 0; i < statuses.size(); ++i) { + if (statuses[i].ok()) { + ReverseListSptr tmp_ptr(new ReverseList()); + if (!tmp_ptr->ParseFromString(values[i].ToString())) { + DB_FATAL("parse second level list from pb/arrow failed"); + return -1; + } + reverse_list_ptrs[i] = tmp_ptr; + } else if (statuses[i].IsNotFound()) { + + } else { + DB_WARNING("rocksdb get error: code=%d, msg=%s", + statuses[i].code(), statuses[i].ToString().c_str()); + return -1; + } + } + list_ptrs.resize(complete_key_idx.size()); + for (int i = 0; i < complete_key_idx.size(); ++i) { + list_ptrs[i] = reverse_list_ptrs[complete_key_idx[i]]; + } + return 0; +} template void LogicalQuery::weight_add_subnode( diff --git a/include/reverse/reverse_arrow.h b/include/reverse/reverse_arrow.h index eee867f28..129e070f9 100755 --- a/include/reverse/reverse_arrow.h +++ b/include/reverse/reverse_arrow.h @@ -49,7 +49,7 @@ inline std::shared_ptr get_arrow_schema () { std::vector> { arrow::field("key", arrow::binary()), arrow::field("flag", arrow::uint8()), - arrow::field("weight", arrow::float32()) + arrow::field("weight", arrow::float64()) } ); return schema; @@ -125,7 +125,6 @@ class ArrowReverseList { _key_builder.Append(key); _flag_builder.Append(flag); _weight_builder.Append(weight); - _current_node_index++; ++_rows; } @@ -133,7 +132,6 @@ class ArrowReverseList { _key_builder.Append(node.key()); _flag_builder.Append(node.flag()); _weight_builder.Append(node.weight()); - _current_node_index++; ++_rows; } @@ -170,20 +168,20 @@ class ArrowReverseList { void set_internal_info() { _rows = _result->num_rows(); _keys_ptr = - static_cast(_result->column(0).get()); + static_cast(_result->column(0).get()); _flags_ptr = - static_cast(_result->column(1).get()); + static_cast(_result->column(1).get()); _weights_ptr = static_cast(_result->column(2).get()); } private: std::shared_ptr _result {nullptr}; - arrow::StringArray* _keys_ptr = nullptr; - arrow::Int8Array* _flags_ptr = nullptr; + arrow::BinaryArray* _keys_ptr = nullptr; + arrow::UInt8Array* _flags_ptr = nullptr; arrow::DoubleArray* _weights_ptr = nullptr; arrow::MemoryPool* pool = arrow::default_memory_pool(); - arrow::StringBuilder _key_builder {pool}; - arrow::Int8Builder _flag_builder {pool}; + arrow::BinaryBuilder _key_builder {pool}; + arrow::UInt8Builder _flag_builder {pool}; arrow::DoubleBuilder _weight_builder {pool}; int64_t _rows = 0; int64_t _current_node_index = -1; diff --git a/include/reverse/reverse_common.h b/include/reverse/reverse_common.h index 121079fb2..76cd550ec 100644 --- a/include/reverse/reverse_common.h +++ b/include/reverse/reverse_common.h @@ -24,7 +24,7 @@ #include "key_encoder.h" #include "lru_cache.h" #include "boolean_executor.h" -#ifdef BAIDU_INTERNAL +#if defined(BAIDU_INTERNAL) && !defined(__aarch64__) #include #include #include @@ -41,7 +41,7 @@ typedef std::pair KeyRange; extern std::atomic_long g_statistic_insert_key_num; extern std::atomic_long g_statistic_delete_key_num; -#ifdef BAIDU_INTERNAL +#if defined(BAIDU_INTERNAL) && !defined(__aarch64__) extern drpc::NLPCClient* wordrank_client; extern drpc::NLPCClient* wordseg_client; extern drpc::NLPCClient* wordweight_client; @@ -64,9 +64,9 @@ class Tokenizer { }; int init(); -#ifdef BAIDU_INTERNAL - template - int nlpc_seg(drpc::NLPCClient& client, const std::string& word, OUT& s_output, IN& s_input); +#if defined(BAIDU_INTERNAL) && !defined(__aarch64__) + template + int nlpc_seg(drpc::NLPCClient& client, const std::string& word, OUTPUT& s_output, INPUT& s_input); int wordrank(std::string word, std::map& term_map, const pb::Charset& charset); int wordweight(std::string word, std::map& term_map, const pb::Charset& charset, bool is_filter = false, bool is_same_weight = false); diff --git a/include/reverse/reverse_common.hpp b/include/reverse/reverse_common.hpp index ef5d64a21..ed2f6510e 100644 --- a/include/reverse/reverse_common.hpp +++ b/include/reverse/reverse_common.hpp @@ -13,12 +13,12 @@ // limitations under the License. namespace baikaldb { -#ifdef BAIDU_INTERNAL -template +#if defined(BAIDU_INTERNAL) && !defined(__aarch64__) +template int Tokenizer::nlpc_seg(drpc::NLPCClient& client, const std::string& word, - OUT& s_output, - IN& s_input) + OUTPUT& s_output, + INPUT& s_input) { if (word.empty()) { return 0; diff --git a/include/reverse/reverse_index.h b/include/reverse/reverse_index.h index 0950e44b9..ea1912f6e 100644 --- a/include/reverse/reverse_index.h +++ b/include/reverse/reverse_index.h @@ -156,7 +156,7 @@ class SchemaBase { void set_table_info(SmartTable& table_info) { _table_info = table_info; } - virtual int create_executor(const std::string& search_data, + virtual int create_executor(myrocksdb::Transaction* txn, const std::string& search_data, pb::MatchMode mode, pb::SegmentType segment_type, const pb::Charset& charset) = 0; virtual int next(SmartRecord record) = 0; @@ -311,6 +311,14 @@ class ReverseIndex : public ReverseIndexBase { return schema_info->schema->get_query_words(); } + + int64_t get_region_id() const { + return _region_id; + } + + int64_t get_index_id() const { + return _index_id; + } private: struct BthreadLocal { Schema* schema = nullptr; diff --git a/include/reverse/reverse_index.hpp b/include/reverse/reverse_index.hpp index 61e0874a8..c2d38f5a5 100644 --- a/include/reverse/reverse_index.hpp +++ b/include/reverse/reverse_index.hpp @@ -282,7 +282,7 @@ int ReverseIndex::create_executor( schema_info->schema->set_index_info(index_info); schema_info->schema->set_table_info(table_info); schema_info->schema->set_index_search(this); - int ret = schema_info->schema->create_executor(search_data, mode, _segment_type, _charset); + int ret = schema_info->schema->create_executor(txn, search_data, mode, _segment_type, _charset); schema_info->schema->statistic().bool_engine_time += timer.get_time(); if (ret < 0) { DB_WARNING("create_executor fail, region:%ld, index:%ld", _region_id, _index_id); @@ -588,7 +588,7 @@ int ReverseIndex::_get_level_reverse_list( } } std::string value; - auto get_res = txn->Get(roptions, data_cf, key, &value); + auto get_res = txn->Get(roptions, data_cf, key, &value); time.reset(); if (get_res.ok()) { //deserialize diff --git a/include/reverse/reverse_interface.h b/include/reverse/reverse_interface.h index 9245973bb..7df1e383d 100644 --- a/include/reverse/reverse_interface.h +++ b/include/reverse/reverse_interface.h @@ -38,19 +38,30 @@ class CommRindexNodeParser : public RindexNodeParser { } ~CommRindexNodeParser() { } - int init(const std::string& term); + int init(const std::string& term, ReverseDelayInitContext* delay_init_context); //return nullptr 代表遍历结束 const ReverseNode* current_node(); const PrimaryIdT* current_id(); //只进不退 const ReverseNode* next(); const ReverseNode* advance(const PrimaryIdT& target_id); + int delay_init(const std::string& term, ReverseListSptr new_list_ptr, ReverseListSptr old_list_ptr) override { + auto* exist_parser = this->_schema->get_term(term); + if (exist_parser != NULL) { + *this = *exist_parser; + return 0; + } + return _init(term, new_list_ptr, old_list_ptr); + } private: //二分查找，大于或等于 uint32_t binary_search(uint32_t first, uint32_t last, const PrimaryIdT& target_id, ReverseList* list); + int _create_reverse_key_prefix(std::string& key); + int _init(const std::string& term, ReverseListSptr new_list_ptr, ReverseListSptr old_list_ptr); + ReverseListSptr _new_list_ptr; ReverseListSptr _old_list_ptr; ReverseList* _new_list; @@ -93,8 +104,8 @@ class NewSchema : public SchemaBase { //search_data 字符串格式 //"hello world" - int create_executor( - const std::string& search_data, pb::MatchMode mode, pb::SegmentType segment_type, const pb::Charset& charset); + int create_executor(myrocksdb::Transaction* txn, const std::string& search_data, + pb::MatchMode mode, pb::SegmentType segment_type, const pb::Charset& charset); int next(SmartRecord record); bool_executor_type executor_type = ReverseTrait::executor_type; void set_term(const std::string& term, Parser* parse) { @@ -121,6 +132,14 @@ class NewSchema : public SchemaBase { return _query_words; } + int64_t get_region_id() const { + return _index_ptr->get_region_id(); + } + + int64_t get_index_id() const { + return _index_ptr->get_index_id(); + } + private: FieldInfo* _weight_field = nullptr; std::string _query_words; diff --git a/include/reverse/reverse_interface.hpp b/include/reverse/reverse_interface.hpp index b26701bbb..a8300b270 100644 --- a/include/reverse/reverse_interface.hpp +++ b/include/reverse/reverse_interface.hpp @@ -19,16 +19,22 @@ namespace baikaldb { -template -int CommRindexNodeParser::init(const std::string& term) { - auto* exist_parser = this->_schema->get_term(term); - if (exist_parser != NULL) { - *this = *exist_parser; - return 0; - } - this->_schema->get_reverse_list(term, _new_list_ptr, _old_list_ptr); - _new_list = (ReverseList*)_new_list_ptr.get(); - _old_list = (ReverseList*)_old_list_ptr.get(); +template +int CommRindexNodeParser::_create_reverse_key_prefix(std::string& key) { + uint64_t region_encode = KeyEncoder::to_endian_u64(KeyEncoder::encode_i64(this->_schema->get_region_id())); + key.append((char*)®ion_encode, sizeof(uint64_t)); + uint64_t table_encode = KeyEncoder::to_endian_u64(KeyEncoder::encode_i64(this->_schema->get_index_id())); + key.append((char*)&table_encode, sizeof(uint64_t)); + return 0; +} + +template +int CommRindexNodeParser::_init(const std::string& term, + ReverseListSptr new_list_ptr, ReverseListSptr old_list_ptr) { + _new_list_ptr = new_list_ptr; + _old_list_ptr = old_list_ptr; + _new_list = static_cast(new_list_ptr.get()); + _old_list = static_cast(old_list_ptr.get()); _curr_node = nullptr; if (_new_list != nullptr && _new_list->reverse_nodes_size() > 0) { _list_size_new = _new_list->reverse_nodes_size(); @@ -48,7 +54,7 @@ int CommRindexNodeParser::init(const std::string& term) { _cmp_res = _curr_id_new->compare(*_curr_id_old); if (_cmp_res > 0) { _curr_node = _old_list->mutable_reverse_nodes(0); - } + } } else { _curr_node = _old_list->mutable_reverse_nodes(0); } @@ -60,6 +66,28 @@ int CommRindexNodeParser::init(const std::string& term) { } this->_schema->set_term(term, this); return 0; +} + +template +int CommRindexNodeParser::init(const std::string& term, ReverseDelayInitContext* delay_init_context) { + if (delay_init_context != nullptr) { + // 延迟初始化，调用multiget减小RocksDB读取压力 + delay_init_context->children.emplace_back(this); + delay_init_context->terms.emplace_back(term); + std::string key; + _create_reverse_key_prefix(key); + delay_init_context->reverse_rocksdb_keys.emplace_back(std::move(key)); + } else { + auto* exist_parser = this->_schema->get_term(term); + if (exist_parser != NULL) { + *this = *exist_parser; + return 0; + } + ReverseListSptr new_list_ptr, old_list_ptr; + this->_schema->get_reverse_list(term, new_list_ptr, old_list_ptr); + _init(term, new_list_ptr, old_list_ptr); + } + return 0; } template @@ -205,7 +233,7 @@ const ReverseNode* } template -int NewSchema::create_executor(const std::string& search_data, +int NewSchema::create_executor(myrocksdb::Transaction* txn, const std::string& search_data, pb::MatchMode mode, pb::SegmentType segment_type, const pb::Charset& charset) { _weight_field = get_field_info_by_name(_table_info->fields, "__weight"); _query_words_field = get_field_info_by_name(_table_info->fields, "__querywords"); @@ -225,7 +253,7 @@ int NewSchema::create_executor(const std::string& search_data, // 报告需求，like语法用|表示'或' Tokenizer::get_instance()->split_str(search_data, or_search, '|', charset); } - LogicalQuery logical_query(this); + LogicalQuery logical_query(this, txn, _is_fast); ExecutorNode* parent = nullptr; ExecutorNode* root = &logical_query._root; if (or_search.size() == 0) { @@ -263,7 +291,7 @@ int NewSchema::create_executor(const std::string& search_data, case pb::S_ES_STANDARD: ret = Tokenizer::get_instance()->es_standard(or_item, term_map, charset); break; -#ifdef BAIDU_INTERNAL +#if defined(BAIDU_INTERNAL) && !defined(__aarch64__) case pb::S_WORDRANK: ret = Tokenizer::get_instance()->wordrank(or_item, term_map, charset); break; diff --git a/include/runtime/arrow_io_excutor.h b/include/runtime/arrow_io_excutor.h index 87c5b306c..ddc137b1f 100644 --- a/include/runtime/arrow_io_excutor.h +++ b/include/runtime/arrow_io_excutor.h @@ -16,6 +16,7 @@ #include "arrow/util/thread_pool.h" #include "arrow/util/functional.h" #include "arrow/util/cancel.h" +#include "arrow/memory_pool.h" #include "common.h" #include "fragment.h" #ifdef BAIDU_INTERNAL @@ -97,5 +98,8 @@ class GlobalArrowExecutor { static int init(); static void execute(RuntimeState* state, arrow::Result>* result); static void execute_fragment(RuntimeState* state, FragmentInfo* fragment); + static void shutdown(); }; + +arrow::MemoryPool* GetMemoryPoolForRead(); }; \ No newline at end of file diff --git a/include/runtime/row_batch.h b/include/runtime/row_batch.h index fc08c6695..f7ece4998 100644 --- a/include/runtime/row_batch.h +++ b/include/runtime/row_batch.h @@ -49,6 +49,9 @@ class RowBatch { bool is_full() { return size() >= _capacity; } + bool empty() { + return size() == 0; + } bool chunk_exceed_max_size() { if (_use_memrow || _chunk == nullptr) { return false; @@ -154,7 +157,9 @@ class RowBatch { DB_FATAL("chunk init failed"); return 0; } - *schema = _chunk->get_arrow_schema(); + if (schema != nullptr) { + *schema = _chunk->get_arrow_schema(); + } return 0; } int set_chunk_tmp_row_value(int tuple_id, int slot_id, const ExprValue& value) { diff --git a/include/runtime/runtime_state.h b/include/runtime/runtime_state.h index 7fab2020f..ffeee294b 100644 --- a/include/runtime/runtime_state.h +++ b/include/runtime/runtime_state.h @@ -522,8 +522,10 @@ class RuntimeState { std::shared_ptr> statistics_types; // 表总行数 int64_t table_rows = 0; - // 采样行数 + // 单个region采样行数 int64_t sample_rows = 0; + // 采样region数 + int64_t max_sample_regions = 0; // -------- 此部分皆为收集统计信息相关 ---------------- int64_t last_insert_id = INT64_MIN; //存储baikalStore last_insert_id(expr)更新的字段 pb::StoreRes* response = nullptr; diff --git a/include/session/user_info.h b/include/session/user_info.h index b758cbd3c..5fdba888c 100644 --- a/include/session/user_info.h +++ b/include/session/user_info.h @@ -195,7 +195,7 @@ struct UserInfo { } bool allow_op(pb::OpType op_type, int64_t db, int64_t tbl, const std::string& table_name) { - if (acl_v2_is_valid()) { // 优先用v2版本权限控制 + if (acl_v2_is_valid() && (acl_database.count(db) > 1 || acl_table.count(tbl) > 0)) { //防止用新方法加库后，权限全丢 return allow_op_v2(op_type, db, tbl); } else { // 兼容v1版本权限控制 return allow_op_v1(op_type, db, tbl, table_name); diff --git a/include/sqlparser/gen_source.sh b/include/sqlparser/gen_source.sh index 500798271..5c9a02077 100644 --- a/include/sqlparser/gen_source.sh +++ b/include/sqlparser/gen_source.sh @@ -1,6 +1,7 @@ #!/bin/bash -prefix='/opt/compiler/gcc-4.8.2/bin/' +# prefix='/opt/compiler/gcc-4.8.2/bin/' +prefix=$1 cur_dir='.' if [[ $2 == "opensource" ]]; then diff --git a/include/sqlparser/sql_parse.y b/include/sqlparser/sql_parse.y index 3d22d129d..f7873020a 100644 --- a/include/sqlparser/sql_parse.y +++ b/include/sqlparser/sql_parse.y @@ -6047,6 +6047,13 @@ AlterSpec: spec->where = $5; $$ = spec; } + | MODIFY COLUMN DELETE WhereClauseOptional + { + AlterTableSpec* spec = new_node(AlterTableSpec); + spec->spec_type = ALTER_SPEC_MODIFY_COLUMN; + spec->where = $4; + $$ = spec; + } | MODIFY ColumnKwdOpt ColumnDef ColumnPosOpt { AlterTableSpec* spec = new_node(AlterTableSpec); diff --git a/include/store/meta_writer.h b/include/store/meta_writer.h index 3f3989918..893dd5684 100644 --- a/include/store/meta_writer.h +++ b/include/store/meta_writer.h @@ -49,6 +49,7 @@ class MetaWriter { static const std::string COLUMN_HOT_FILE_INDENTIFY; static const std::string COLUMN_COLD_FILE_INDENTIFY; static const std::string BINLOG_DATA_CF_OLDEST_IDENTIFY; + static const std::string COLUMN_TXN_LOG_INDEX_IDENTIFY; virtual ~MetaWriter() {} @@ -71,7 +72,10 @@ class MetaWriter { int write_batch(rocksdb::WriteBatch* updates, int64_t region_id); int write_meta_after_commit(int64_t region_id, int64_t num_table_lines, int64_t applied_index, int64_t data_index, uint64_t txn_id, bool need_write_rollback); - int write_meta_begin_index(int64_t region_id, int64_t log_index, int64_t data_index, uint64_t txn_id); + int delete_column_txn_log_index(int64_t region_id, std::vector txn_ids); + int64_t read_meta_begin_index(int64_t region_id, uint64_t txn_id); + int write_meta_begin_index(int64_t region_id, int64_t log_index, int64_t data_index, + uint64_t txn_id, bool has_column_engine = false); int write_meta_index_and_num_table_lines(int64_t region_id, int64_t log_index, int64_t data_index, int64_t num_table_lines, SmartTransaction txn); int ingest_meta_sst(const std::string& meta_sst_file, int64_t region_id); @@ -138,6 +142,7 @@ class MetaWriter { std::string num_table_lines_key(int64_t region_id) const; std::string transcation_log_index_key(int64_t region_id, uint64_t txn_id) const; std::string log_index_key_prefix(int64_t region_id) const; + std::string column_txn_log_index_key(int64_t region_id, uint64_t txn_id) const; std::string transcation_pb_key(int64_t region_id, uint64_t txn_id, int64_t log_index) const; std::string transcation_pb_key_prefix(int64_t region_id) const; std::string pre_commit_key_prefix(int64_t region_id) const; diff --git a/include/store/region.h b/include/store/region.h index 531021f90..f1e297138 100644 --- a/include/store/region.h +++ b/include/store/region.h @@ -60,7 +60,7 @@ #include "concurrency.h" #include "backup.h" #include "file_manager.h" -#include "region_column.h" +#include "dml_node.h" #ifdef BAIDU_INTERNAL #else @@ -130,6 +130,41 @@ struct BinlogDesc { TimeCost time; }; +struct BinlogReadFields { + FieldInfo* ts_field = nullptr; + FieldInfo* binlog_type_field = nullptr; + FieldInfo* start_ts_field = nullptr; + FieldInfo* binlog_row_cnt_field = nullptr; + int64_t get_ts(SmartRecord& record) { + if (ts_field != nullptr) { + auto f = record->get_field_by_idx(ts_field->pb_idx); + return record->get_value(f).get_numberic(); + } + return 0; + } + BinlogType get_binlog_type(SmartRecord& record) { + if (binlog_type_field != nullptr) { + auto f = record->get_field_by_idx(binlog_type_field->pb_idx); + return static_cast(record->get_value(f).get_numberic()); + } + return FAKE_BINLOG; + } + int64_t get_start_ts(SmartRecord& record) { + if (start_ts_field != nullptr) { + auto f = record->get_field_by_idx(start_ts_field->pb_idx); + return record->get_value(f).get_numberic(); + } + return 0; + } + int64_t get_binlog_row_cnt(SmartRecord& record) { + if (binlog_row_cnt_field != nullptr) { + auto f = record->get_field_by_idx(binlog_row_cnt_field->pb_idx); + return record->get_value(f).get_numberic(); + } + return 0; + } +}; + struct ApproximateInfo { int64_t table_lines = 0; uint64_t region_size = 0; @@ -686,7 +721,7 @@ friend class Backup; void reverse_merge_doing_ddl(); // other thread void ttl_remove_expired_data(); - void vector_schema_change(); + void vector_schema_change(std::unordered_map& vector_indexs); int restore_faiss(const std::string& path, const std::vector& files); @@ -884,10 +919,10 @@ friend class Backup; std::lock_guard lock(_region_lock); return _region_info.used_size(); } - int64_t get_table_id() { + int64_t get_table_id() const { return _table_id; } - int64_t get_global_index_id() { + int64_t get_global_index_id() const { return _global_index_id; } bool is_leader() { @@ -1070,24 +1105,18 @@ friend class Backup; _txn_pool.clear_transactions(this); _multi_thread_cond.decrease_signal(); } - void update_ttl_info() { - if (_shutdown || !_init_success || get_version() <= 0) { - return; - } - TTLInfo ttl_info = _factory->get_ttl_duration(get_table_id()); - if (ttl_info.ttl_duration_s > 0 && ttl_info.online_ttl_expire_time_us > 0) { - // online TTL - if (ttl_info.online_ttl_expire_time_us != _online_ttl_base_expire_time_us) { - _online_ttl_base_expire_time_us = ttl_info.online_ttl_expire_time_us; - _use_ttl = true; - _txn_pool.update_ttl_info(_use_ttl, _online_ttl_base_expire_time_us); - DB_WARNING("table_id: %ld, region_id: %ld, ttl_duration_s: %ld, online_ttl_expire_time_us: %ld, %s", - get_table_id(), _region_id, ttl_info.ttl_duration_s, - ttl_info.online_ttl_expire_time_us, timestamp_to_str(ttl_info.online_ttl_expire_time_us/1000000).c_str()); - } - } - } + // no side effect + int get_ttl_timestamp(IndexInfo &index_info, + const rocksdb::Slice &rocksdb_key, + const rocksdb::Slice &rocksdb_value, + const SmartRecord &record_template, + uint64_t &ttl_timestamp) const; + + int update_ttl_info(); + + SmartState create_ttl_delete_runtime_status(); + //blacklist中新增的sign全部cancel void cancel_all_blacklist_sign() { if (_shutdown || !_init_success || get_version() <= 0) { @@ -1190,7 +1219,7 @@ friend class Backup; int add_reverse_index(int64_t table_id, const std::set& index_ids); - void process_download_sst(brpc::Controller* controller, + void process_download_sst(brpc::Controller* controller, std::vector& req_vec, SstBackupType type); void process_upload_sst(brpc::Controller* controller, bool is_ingest); @@ -1330,7 +1359,7 @@ friend class Backup; void column_on_snapshot_save(braft::SnapshotWriter* writer, braft::Closure* done); void column_snapshot_save(); int column_snapshot_save(const std::string& snapshot_path, std::vector& files); - int get_column_files(const std::vector& key_ranges, std::vector>& files); + int get_column_files(const pb::PossibleIndex& pos_index, std::vector>& files); bool use_column_storage(const pb::Plan& plan, SmartTable table); bool use_userid_statis(bool is_eq, const google::protobuf::RepeatedPtrField& key_ranges); std::vector> column_link_files(const std::vector& file_infos, int64_t min_version, int64_t max_version); @@ -1450,7 +1479,6 @@ friend class Backup; int write_binlog_record(SmartRecord record); int write_binlog_value(const std::map& field_value_map); int64_t binlog_get_int64_val(const std::string& name, const std::map& field_value_map); - int64_t read_data_cf_oldest_ts(); bool flash_back_need_read(const pb::StoreReq* request, const std::map& field_value_map, const std::set& req_db_tables, @@ -1458,8 +1486,8 @@ friend class Backup; std::string binlog_get_str_val(const std::string& name, const std::map& field_value_map); - void binlog_get_scan_fields(std::map& field_ids, std::vector& field_slot, - SmartTable& binlog_table, SmartIndex& binlog_pri); + void binlog_get_scan_fields(std::map& field_ids, BinlogReadFields& binlog_fields, std::vector& field_slot, + SmartTable& binlog_table, SmartIndex& binlog_pri, bool read_all = true); void binlog_get_field_values(std::map& field_value_map, SmartRecord& record, SmartTable& binlog_table); int binlog_reset_on_snapshot_load_restart(); @@ -1512,7 +1540,7 @@ friend class Backup; } else { SplitCompactionFilter::get_instance()->set_filter_region_info( _region_id, region_info.end_key(), - _use_ttl, _online_ttl_base_expire_time_us); + use_normal_ttl(), _online_ttl_base_expire_time_us); } DB_WARNING("region_id: %ld, start_key: %s, end_key: %s", _region_id, rocksdb::Slice(region_info.start_key()).ToString(true).c_str(), @@ -1609,7 +1637,11 @@ friend class Backup; const pb::ExchangeSenderNode& pb_exchange_sender_node, const bool is_merge, const int64_t region_id, const ::google::protobuf::RepeatedPtrField& region_infos); + bool need_decode_ttl_field() const { return _ttl_field_type >= TTLFieldType::T_FIELD_KEY; } + private: + static constexpr int KEY_PREFIX_LENGTH = sizeof(int64_t) * 2; + //Singleton RocksWrapper* _rocksdb; SchemaFactory* _factory; @@ -1654,8 +1686,6 @@ friend class Backup; bool _restart = false; //计算存储分离开关，在store定时任务中更新，避免每次dml都访问schema factory bool _storage_compute_separate = false; - bool _use_ttl = false; // online TTL会更新，只会false 变为true - int64_t _online_ttl_base_expire_time_us = 0; // 存量数据过期时间，仅online TTL的表使用 std::atomic _need_vector_compact{false}; //split的数据，把vector compact一次 std::atomic _reverse_remove_range{false}; //split的数据，把拉链过滤一遍, safe reverse index合并 std::atomic _reverse_unsafe_remove_range{false};//unsafe reverse index合并 @@ -1722,6 +1752,7 @@ friend class Backup; bthread::Mutex _commit_ts_map_lock; bthread::Mutex _binlog_param_mutex; BinlogParam _binlog_param; + Cache _binlog_cache; // offline binlog, only for binlog_backup_days>0 binlog tables bthread::Mutex _offline_binlog_param_mutex; OfflineBinlogParam _offline_binlog_param; @@ -1751,6 +1782,8 @@ friend class Backup; bool _ready_for_follower_read = true; // 解决零星写时主从延迟高,有写入时每100ms发一条NO OP, 停写5min后不再发NO OP NoOpTimer _no_op_timer; + bool _raft_status_error = false; + // olap std::atomic _olap_state {pb::OLAP_ACTIVE}; @@ -1758,6 +1791,36 @@ friend class Backup; std::mutex _snapshot_closure_mutex; std::unique_ptr _snapshot_closure = nullptr; // 使用 + // ttl 相关 + enum class TTLFieldType { + // 请勿更改顺序，部分判断依赖于大小比较 + T_NON_TTL = 0, // region 没有ttl + T_NORMAL_TTL = 1, // 常规ttl，时间戳在value前缀里，为过期时间 + T_FIELD_KEY = 2, // 指定字段ttl, 时间戳在索引的field内 + T_FIELD_PK_VALUE = 3, // 指定字段ttl, region是主键索引，时间戳在value内 + T_FIELD_GI_KEY = 4, // 指定字段ttl, region是全局索引，时间戳在pk_fields内，在rocksdb的key内 + T_FIELD_GUI_VALUE = 5, // 指定字段ttl, region是全局唯一索引，时间戳在pk_fields内，在value内 + }; + + bool is_field_ttl() const { + return _ttl_field_type > TTLFieldType::T_NORMAL_TTL; + } + + bool use_normal_ttl() const { + return _ttl_field_type == TTLFieldType::T_NORMAL_TTL; + } + + bool use_ttl() const { + return _ttl_field_type > TTLFieldType::T_NON_TTL; + } + + TTLFieldType _ttl_field_type = TTLFieldType::T_NON_TTL; + int64_t _online_ttl_base_expire_time_us = 0; // 存量数据过期时间，仅online TTL的表使用 + int64_t _ttl_duration_s; + std::map _decode_ttl_field_map; // field ttl时只有一个ttl字段，其他情况下为空 + std::shared_ptr _ttl_field = nullptr; + bthread::Mutex _ttl_mutex; + //NOT_LEADER分类报警 struct NotLeaderAlarm { enum AlarmType { diff --git a/include/store/store.h b/include/store/store.h index b88c963f7..ec2c68730 100644 --- a/include/store/store.h +++ b/include/store/store.h @@ -348,6 +348,8 @@ class Store : public pb::StoreService { DB_WARNING("_multi_thread_cond wait finish"); _rocksdb->close(); DB_WARNING("rockdb close, quit success"); + ParquetFileManager::get_instance()->close(); + DB_WARNING("parquet file manager close"); } MetaServerInteract& get_meta_server_interact() { return _meta_server_interact; @@ -376,7 +378,7 @@ class Store : public pb::StoreService { bool now_in_interval_period() { struct tm ptm; time_t timep = time(NULL); - localtime_r(&timep, &ptm); + localtime_fixed_r(&timep, &ptm); int now = ptm.tm_hour; // 跨夜 if (_end_hour < _start_hour) { @@ -400,7 +402,7 @@ class Store : public pb::StoreService { void update_schema_info(const pb::SchemaInfo& table, std::map>* reverse_index_map, - std::unordered_set* vector_table_set = nullptr); + std::unordered_map>* vector_index_map = nullptr); //判断分裂在3600S内是否完成，不完成，则自动删除该region void check_region_legal_complete(int64_t region_id); diff --git a/include/vector_index/vector_index.h b/include/vector_index/vector_index.h index 273e42367..0dcbb34f6 100755 --- a/include/vector_index/vector_index.h +++ b/include/vector_index/vector_index.h @@ -34,6 +34,11 @@ const std::string FAISS_SCALAR_DATA_PREFIX = "/faissindex_scalardata_"; const std::string FAISS_DEL_BITMAP_PREFIX = "/faissindex_delbitmap_"; const std::string FAISS_NOT_CACHE_FIELDS_PREFIX = "/faissindex_notcachefields_"; +struct VectorSearchParam { + int32_t nprobe = -1; + int32_t efsearch = -1; +}; + class VectorIndex { struct FaissIndex; using SmartFaissIndex = std::shared_ptr; @@ -171,7 +176,8 @@ class VectorIndex { SmartTransaction& txn, const std::string& word, const std::string& pk, - SmartRecord record); + SmartRecord record, + pb::IndexState index_status); int search_vector( myrocksdb::Transaction* txn, const uint64_t separate_value, @@ -179,7 +185,7 @@ class VectorIndex { SmartTable& table_info, const std::string& search_data, int64_t topk, - int32_t efsearch, + VectorSearchParam& search_param, std::vector& records, std::vector& vector_filter_exprs, std::vector& scan_filter_exprs, @@ -188,7 +194,8 @@ class VectorIndex { SmartFaissIndex faiss_index, const std::string& word, int64_t cache_idx, - SmartRecord record); + SmartRecord record, + bool enable_train); int search( myrocksdb::Transaction* txn, SmartFaissIndex faiss_index, @@ -196,7 +203,7 @@ class VectorIndex { SmartTable& table_info, const std::string& search_data, int64_t topk, - int32_t efsearch, + VectorSearchParam& search_param, std::vector& records, std::vector& vector_filter_exprs, std::vector& scan_filter_exprs, @@ -211,7 +218,8 @@ class VectorIndex { myrocksdb::Transaction* txn, SmartFaissIndex faiss_index, const std::string& pk, - VFlag flag); + VFlag flag, + pb::IndexState index_status); int construct_records( myrocksdb::Transaction* txn, const uint64_t separate_value, @@ -230,6 +238,19 @@ class VectorIndex { // 根据主表行数和faiss索引行数判断是否需要compact bool need_compact(const int64_t table_lines); + bool need_train() { + return _is_ivf || _is_pq || _is_pqfs; + } + + void cal_new_train_vec_size(int64_t valid_vec_size); + + bool is_removed() { + return _removed; + } + + void set_removed(bool r) { + _removed = r; + } private: int init_faiss_index(SmartFaissIndex faiss_index); int reset_faiss_index(SmartFaissIndex faiss_index); @@ -365,6 +386,8 @@ class VectorIndex { } bool _is_separate = false; // 向量索引是否进行隔离 + int32_t _separate_field_id = -1; + SmartFaissIndex _faiss_index = nullptr; // 正常向量索引 std::shared_ptr> _separate_faiss_index_map; // 向量隔离索引 @@ -377,13 +400,32 @@ class VectorIndex { int64_t _table_id = 0; int32_t _dimension = 0; // 维度 faiss::MetricType _metrix_type = faiss::METRIC_L2; + // -- IVF int32_t _nprobe = 5; + int32_t _ivf_nlist = 50; + // -- HNSW int32_t _efsearch = 16; int32_t _efconstruction = 40; + + // metric_type + bool _is_l2norm = false; + + // vector index type bool _is_flat = false; bool _is_hnsw = false; - bool _is_l2norm = false; - int32_t _separate_field_id = -1; + bool _is_ivf = false; + bool _is_pq = false; + bool _is_pqfs = false; + + // for train + int64_t _train_vec_size = -1; + int64_t _trigger_retrain_vector_num = -1; + + // for slowdown + bool _begin_train = false; + TimeCost _train_begin_time; + + bool _removed = false; }; } // namespace baikaldb diff --git a/proto/common.proto b/proto/common.proto index 0f9af597b..fca6f2307 100755 --- a/proto/common.proto +++ b/proto/common.proto @@ -124,6 +124,7 @@ message SchemaConf { optional int32 olap_pre_split_cnt = 23; // OLAP预分裂的个数 optional bool cold_use_column_only = 24; // 是否冷数据只使用列存 optional bool force_column_storage = 25; // 是否强制使用列存 + optional bool column_only_read_base = 26; // 列存只读base文件，可以避免merge on read }; enum Engine { diff --git a/proto/compaction.interface.proto b/proto/compaction.interface.proto index d1a2b6db9..416a0d2a2 100644 --- a/proto/compaction.interface.proto +++ b/proto/compaction.interface.proto @@ -7,7 +7,7 @@ message RocksdbGFLAGS { optional bool rocks_use_partitioned_index_filters = 1; optional bool rocks_use_ribbon_filter = 2; optional bool olap_table_only = 3; - optional bool olap_import_mode = 4; + //optional bool olap_import_mode = 4; optional bool rocks_use_sst_partitioner_fixed_prefix = 5; optional int32 key_point_collector_interval = 6; optional int64 rocks_block_cache_size_mb = 7; diff --git a/proto/dms.proto b/proto/dms.proto index 64b628808..7e6ddc9dd 100644 --- a/proto/dms.proto +++ b/proto/dms.proto @@ -36,12 +36,17 @@ enum DMSOpType { OP_ADD_TMP_TABLE_TASK = 14; // 添加临时表任务 OP_QUERY_TMP_TABLE_TASK = 15; // 查询临时表任务 OP_QUERY_BNS = 16; // 排查工具使用，查询baikaldb bns、baikalstore bns + OP_CREATE_FLASH_BACK_TASK = 17; // 创建闪回任务 + OP_CREATE_BINLOG_TASK = 18; // 创建binlog信息 + OP_QUERY_BINLOG_TASK = 19; // 查询binlog信息 }; enum SqlStatus { SQL_SUCCESS = 0; SQL_FAILED = 1; SQL_DOING = 2; + SQL_UPDATE = 3; // SQL闪回时更新SQL, 需要重新执行 + SQL_WAITING = 4; // 爱速搭不支持把''当枚举值，无法表示待执行的状态 }; enum TaskStatus { @@ -55,6 +60,7 @@ enum TaskStatus { TS_IDLE = 7; // 任务待执行（导入回溯任务查询使用） TS_LORENTZ = 8; // 任务状态为lorentz（凤阁场景使用） TS_INITED = 9; // 导入任务创建校验成功，未启动状态 + TS_UPDATE = 10; // SQL闪回时更新SQL, 需要重新执行 TS_UNKNOWN = 1001; // 未知状态 }; @@ -70,6 +76,7 @@ message SqlDetail { optional bytes errmsg = 3; // SQL错误信息 optional int64 affected_row = 4; // SQL影响行数 optional bytes sql = 5; // SQL文本 + optional int64 old_id = 6; // 旧SQL序号 }; // SQL预估和SQL执行返回结果复用该结构 @@ -80,6 +87,7 @@ message TaskDetail { optional int64 total_sql_num = 4; // 任务SQL总数量 optional int64 done_sql_num = 5; // 任务SQL已完成数量 repeated SqlDetail sql_details = 6; // SQL处理信息 + optional bytes afs_path = 7; // 任务afs路径 }; // 导入任务信息 @@ -106,6 +114,7 @@ message ImporterTaskInfo { optional bytes baikaldb_username = 20; // 用于获取baikalmeta、namespace（凤阁场景使用） optional bytes baikaldb_password = 21; // 用于获取baikalmeta、namespace（凤阁场景使用） optional bytes lorentz_done_json = 22; // 用于检查的done_json（凤阁场景使用） + optional bytes local_done_json = 23; // _SUCCESS模式下不和done_file共用 }; // 回溯任务信息 @@ -135,6 +144,17 @@ message TroubleShootingResponse { repeated bytes baikalstore_bns = 2; }; +message BinlogRequest { + optional bytes baikal_meta = 1; + optional bytes baikal_namespace = 2; + optional bytes db_name = 3; + optional bytes table_name = 4; + optional bytes start_datetime = 5; + optional bytes end_datetime = 6; + optional int64 filter_partition_key = 7; + optional bytes find_str = 8; +} + message TaskConfig { optional bool need_backup = 1; // 是否需要备份 optional bool need_parallel = 2; // 是否需要并行 @@ -150,6 +170,8 @@ message DMSRequest { optional ImporterTaskInfo importer_task_info = 7; // 导入任务信息 optional BacktrackImporterTaskInfo backtrack_importer_task_info = 8; // 回溯任务信息 optional TroubleShootingRequest trouble_shooting_request = 9; // 排查工具请求 + optional BinlogRequest binlog_request = 10; // 查找binlog请求 + optional uint32 limit_tps = 11; // 限制sql issue执行速度 }; message DMSResponse { diff --git a/proto/expr.proto b/proto/expr.proto index c884afe8f..7ec23aa58 100755 --- a/proto/expr.proto +++ b/proto/expr.proto @@ -43,6 +43,7 @@ message Function { repeated PrimitiveType arg_types = 3; optional PrimitiveType return_type = 4; optional bool has_var_args = 5; + optional string origin_name = 6; // 原始函数名 }; message DeriveExprNode { diff --git a/proto/meta.interface.proto b/proto/meta.interface.proto index 8eccab357..005fd578d 100755 --- a/proto/meta.interface.proto +++ b/proto/meta.interface.proto @@ -173,6 +173,22 @@ message PartitionInfo { enum DBLinkType { LT_BAIKALDB = 1; LT_MYSQL = 2; + LT_FILE = 3; +}; + +enum FileFormat { + CSV = 1; + PARQUET = 2; +}; + +message FileInfo { + optional bytes cluster = 1; + optional bytes path = 2; + optional bytes username = 3; + optional bytes password = 4; + repeated bytes partition_fields = 5; + optional FileFormat format = 6; + optional bytes delimiter = 7; }; message DBLinkInfo { @@ -183,6 +199,7 @@ message DBLinkInfo { optional string table_name = 5; optional int64 meta_id = 6; optional MysqlInfo mysql_info = 7; + optional FileInfo file_info = 8; }; // 增加新字段时，考虑是否对该字段增删meta标识，对应dblink表场景!!! @@ -239,6 +256,7 @@ message SchemaInfo { optional DBLinkInfo dblink_info = 50; // DBLink表信息，namespace_name心跳时不需要删除或增加meta标识 optional bool or_replace = 51; optional string view_select_stmt = 52; + optional FieldInfo ttl_field = 53; }; message PartitionRegion { @@ -576,6 +594,16 @@ message VirtualIndexInfluence{ optional string influenced_sql = 3; }; +message IndexInstance { + optional string index_name = 1; + repeated string instances = 2; +} + +message PartitionIndexInstance { + optional int64 partition_id = 1; + repeated IndexInstance index_instances = 2; +} + message MetaManagerRequest { required OpType op_type = 1; //一级操作符 optional SchemaInfo table_info = 2; //建/删表 @@ -606,6 +634,7 @@ message MetaManagerRequest { optional int64 learner_replica_num = 27; optional bool is_dynamic_change = 28; // 标识是否为动态分区变更 optional bool is_force_setting = 29; // 标识是否为强制设置，用于删除默认配置 + repeated PartitionIndexInstance partition_index_instance = 30; // 建表 加索引 加分区状态及外选实例 }; message CreateTableResponse { diff --git a/proto/optype.proto b/proto/optype.proto index a3c6fb6df..80963fcdb 100755 --- a/proto/optype.proto +++ b/proto/optype.proto @@ -86,7 +86,7 @@ enum OpType { OP_ADD_FIELD = 133; //加列 OP_ADD_INDEX = 134; //加索引 OP_DROP_INDEX = 135; //删索引 - OP_DROP_FIELD = 136; //删索引 + OP_DROP_FIELD = 136; //删列 OP_RENAME_FIELD = 137; //改字段名称 OP_RENAME_INDEX = 138; //改索引名称 OP_MODIFY_FIELD = 139; //改字段定义，包括改字段类型, 暂不支持 diff --git a/proto/plan.proto b/proto/plan.proto index c29c0894c..3e16af155 100755 --- a/proto/plan.proto +++ b/proto/plan.proto @@ -58,6 +58,7 @@ message SelectManagerNode { repeated bool is_asc = 3; repeated bool is_null_first = 4; optional bool is_return_empty = 5; + repeated Expr filter_conditions = 6; } enum TxnCmdType { @@ -105,6 +106,7 @@ message PossibleIndex { optional int32 topk = 16; optional uint64 separate_value = 17; // 向量索引使用 optional int32 efsearch = 18; // 向量索引使用 + optional int32 nprobe = 19; // 向量索引使用 }; message SortIndex { @@ -149,6 +151,11 @@ enum DDLType { DDL_COLUMN = 3; }; +message PartitionFile { + repeated bytes partition_vals = 1; // 文件所在分区 + optional bytes file_path = 2; // 文件路径 +}; + message ScanNode { required int32 tuple_id = 1; //tuple中记录有读取列信息与table信息 required int64 table_id = 2; @@ -167,6 +174,7 @@ message ScanNode { optional uint64 watt_stats_version = 15; // watt stats写入version幂等, 用于ROLLUP基准 optional int32 union_tuple_id = 16; // union子查询对应的union_tuple_id optional bool use_column_storage = 17; // 是否使用列存，读取时使用 + repeated PartitionFile files = 18; // 离线查询文件集合，用于FileScanNode场景 }; message LimitNode { @@ -371,7 +379,7 @@ enum PartitionPropertyType { SinglePartitionType = 2; HashPartitionType = 3; BroadcastPartitionType = 4; - RandomPartitionType = 5; // TODO + RandomPartitionType = 5; }; message PartitionProperty { @@ -413,6 +421,7 @@ message ExchangeReceiverNode { repeated Expr slot_order_exprs = 10; repeated bool is_asc = 11; repeated bool is_null_first = 12; + repeated Expr filter_conditions = 13; }; enum FrameType { diff --git a/proto/store.interface.proto b/proto/store.interface.proto index c660ed133..a2f0e5220 100755 --- a/proto/store.interface.proto +++ b/proto/store.interface.proto @@ -148,8 +148,9 @@ message ExtraReq { optional int64 sign_latency = 2; optional OlapRegionInfo olap_info = 3; optional RegionOfflineBinlogInfo offline_binlog_info = 4; - optional uint64 watt_stats_version = 5; + optional uint64 watt_stats_version = 5; optional ColumnRegionInfo column_info = 6; + optional bool no_raft_log = 7; }; message ExtraRes { diff --git a/src/column/column_record.cpp b/src/column/column_record.cpp index 9276aed06..e7992b65d 100644 --- a/src/column/column_record.cpp +++ b/src/column/column_record.cpp @@ -1,5 +1,6 @@ #include "mut_table_key.h" #include "column_record.h" +#include "vectorize_helpper.h" namespace baikaldb { void ColumnRecord::TEST_print_record_batch(const std::shared_ptr& record_batch) { @@ -74,6 +75,96 @@ std::string ColumnRecord::encode_row_key(std::shared_ptr rec return key.data(); } +std::shared_ptr ColumnRecord::make_column_schema(int64_t tableid, SmartTable table_info, SmartIndex pri_info, + const std::unordered_map& field_id2info_map) { + if (table_info == nullptr) { + table_info = SchemaFactory::get_instance()->get_table_info_ptr(tableid); + } + if (pri_info == nullptr) { + pri_info = SchemaFactory::get_instance()->get_index_info_ptr(tableid); + } + if (table_info == nullptr || pri_info == nullptr) { + DB_FATAL("table or index is null, tableid:%ld", tableid); + return nullptr; + } + auto schema_ptr = std::make_shared(); + schema_ptr->index_info = pri_info; + schema_ptr->table_info = table_info; + schema_ptr->key_fields = pri_info->fields; + + std::set key_field_ids; + for (const auto& field : pri_info->fields) { + key_field_ids.insert(field.id); + } + + std::set field_ids_need_sum; + for (const auto& field : table_info->fields_need_sum) { + field_ids_need_sum.insert(field.id); + } + + int value_idx = schema_ptr->key_fields.size(); + schema_ptr->value_fields.reserve(table_info->fields.size() - pri_info->fields.size()); + for (const auto& field : table_info->fields) { + if (key_field_ids.count(field.id) > 0 || field.deleted) { + continue; + } + + if (!field_id2info_map.empty()) { + if (field_id2info_map.count(field.id) == 0) { + continue; + } + } + + schema_ptr->value_fields.emplace_back(field); + if (field_ids_need_sum.count(field.id) > 0) { + schema_ptr->need_sum_idx.insert(value_idx); + } + value_idx++; + } + + schema_ptr->uniq_size = schema_ptr->key_fields.size(); + schema_ptr->keytype_idx = schema_ptr->key_fields.size() + schema_ptr->value_fields.size(); + schema_ptr->raft_index_idx = schema_ptr->keytype_idx + 1; + schema_ptr->batch_pos_idx = schema_ptr->raft_index_idx + 1; + + std::vector> arrow_fields; + arrow_fields.reserve(schema_ptr->batch_pos_idx + 1); + for (const auto& field : schema_ptr->key_fields) { + auto arrow_type = primitive_to_arrow_type(field.type); + if (arrow_type < 0) { + DB_COLUMN_FATAL("field: %s primitive type:%d to arrow type failed", field.lower_short_name.c_str(), field.type); + return nullptr; + } + auto arrow_field = VectorizeHelpper::make_field(field.lower_short_name, arrow::Type::type(arrow_type)); + if (arrow_field == nullptr) { + DB_COLUMN_FATAL("field: %s make arrow schema failed", field.lower_short_name.c_str()); + return nullptr; + } + arrow_fields.emplace_back(arrow_field); + } + + for (const auto& field : schema_ptr->value_fields) { + auto arrow_type = primitive_to_arrow_type(field.type); + if (arrow_type < 0) { + DB_COLUMN_FATAL("field: %s primitive type:%d to arrow type failed", field.lower_short_name.c_str(), field.type); + return nullptr; + } + auto arrow_field = VectorizeHelpper::make_field(field.lower_short_name, arrow::Type::type(arrow_type)); + if (arrow_field == nullptr) { + DB_COLUMN_FATAL("field: %s make arrow schema failed", field.lower_short_name.c_str()); + return nullptr; + } + arrow_fields.emplace_back(arrow_field); + } + + arrow_fields.emplace_back(VectorizeHelpper::make_field(KEY_TYPE_NAME, arrow::Type::type::INT32)); + schema_ptr->schema = std::make_shared(arrow_fields); + arrow_fields.emplace_back(VectorizeHelpper::make_field(RAFT_INDEX_NAME, arrow::Type::type::INT64)); + arrow_fields.emplace_back(VectorizeHelpper::make_field(BATCH_POS_NAME, arrow::Type::type::INT32)); + schema_ptr->schema_with_order_info = std::make_shared(arrow_fields); + return schema_ptr; +} + ExprValue ColumnRecord::get_vectorized_value(const std::shared_ptr& array, int row_idx) { ExprValue ret; ret.type = pb::NULL_TYPE; @@ -174,222 +265,6 @@ ExprValue ColumnRecord::get_vectorized_value(const std::shared_ptr return ret; } -std::shared_ptr ColumnRecord::make_schema(const std::string& name, arrow::Type::type type) { - switch (type) { - case arrow::Type::type::BOOL: - return std::make_shared(name, arrow::boolean()); - case arrow::Type::type::INT8: - return std::make_shared(name, arrow::int8()); - case arrow::Type::type::UINT8: - return std::make_shared(name, arrow::uint8()); - case arrow::Type::type::INT16: - return std::make_shared(name, arrow::int16()); - case arrow::Type::type::UINT16: - return std::make_shared(name, arrow::uint16()); - case arrow::Type::type::INT32: - return std::make_shared(name, arrow::int32()); - case arrow::Type::type::UINT32: - return std::make_shared(name, arrow::uint32()); - case arrow::Type::type::INT64: - return std::make_shared(name, arrow::int64()); - case arrow::Type::type::UINT64: - return std::make_shared(name, arrow::uint64()); - case arrow::Type::type::FLOAT: - return std::make_shared(name, arrow::float32()); - case arrow::Type::type::DOUBLE: - return std::make_shared(name, arrow::float64()); - case arrow::Type::type::LARGE_BINARY: - return std::make_shared(name, arrow::large_binary()); - default: - return nullptr; - } -} - -std::shared_ptr ColumnRecord::make_array_from_exprvalue( - const pb::PrimitiveType type, const ExprValue& expr_value, const int length) { - bool is_null = expr_value.is_null(); - switch (type) { - case pb::BOOL: { - arrow::BooleanScalar scalar; - if (!is_null) { - scalar = arrow::BooleanScalar(expr_value.get_numberic()); - } - auto array_ret = arrow::MakeArrayFromScalar(scalar, length); - if (!array_ret.ok()) { - DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); - return nullptr; - } else { - return *array_ret; - } - break; - } - case pb::INT8: { - arrow::Int8Scalar scalar; - if (!is_null) { - scalar = arrow::Int8Scalar(expr_value.get_numberic()); - } - auto array_ret = arrow::MakeArrayFromScalar(scalar, length); - if (!array_ret.ok()) { - DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); - return nullptr; - } else { - return *array_ret; - } - break; - } - case pb::UINT8: { - arrow::UInt8Scalar scalar; - if (!is_null) { - scalar = arrow::UInt8Scalar(expr_value.get_numberic()); - } - auto array_ret = arrow::MakeArrayFromScalar(scalar, length); - if (!array_ret.ok()) { - DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); - return nullptr; - } else { - return *array_ret; - } - break; - } - case pb::INT16: { - arrow::Int16Scalar scalar; - if (!is_null) { - scalar = arrow::Int16Scalar(expr_value.get_numberic()); - } - auto array_ret = arrow::MakeArrayFromScalar(scalar, length); - if (!array_ret.ok()) { - DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); - return nullptr; - } else { - return *array_ret; - } - break; - } - case pb::UINT16: { - arrow::UInt16Scalar scalar; - if (!is_null) { - scalar = arrow::UInt16Scalar(expr_value.get_numberic()); - } - auto array_ret = arrow::MakeArrayFromScalar(scalar, length); - if (!array_ret.ok()) { - DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); - return nullptr; - } else { - return *array_ret; - } - break; - } - case pb::INT32: - case pb::TIME: { - arrow::Int32Scalar scalar; - if (!is_null) { - scalar = arrow::Int32Scalar(expr_value.get_numberic()); - } - auto array_ret = arrow::MakeArrayFromScalar(scalar, length); - if (!array_ret.ok()) { - DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); - return nullptr; - } else { - return *array_ret; - } - break; - } - case pb::UINT32: - case pb::DATE: - case pb::TIMESTAMP: { - arrow::UInt32Scalar scalar; - if (!is_null) { - scalar = arrow::UInt32Scalar(expr_value.get_numberic()); - } - auto array_ret = arrow::MakeArrayFromScalar(scalar, length); - if (!array_ret.ok()) { - DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); - return nullptr; - } else { - return *array_ret; - } - break; - } - case pb::INT64: { - arrow::Int64Scalar scalar; - if (!is_null) { - scalar = arrow::Int64Scalar(expr_value.get_numberic()); - } - auto array_ret = arrow::MakeArrayFromScalar(scalar, length); - if (!array_ret.ok()) { - DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); - return nullptr; - } else { - return *array_ret; - } - break; - } - case pb::UINT64: - case pb::DATETIME: { - arrow::UInt64Scalar scalar; - if (!is_null) { - scalar = arrow::UInt64Scalar(expr_value.get_numberic()); - } - auto array_ret = arrow::MakeArrayFromScalar(scalar, length); - if (!array_ret.ok()) { - DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); - return nullptr; - } else { - return *array_ret; - } - break; - } - case pb::FLOAT: { - arrow::FloatScalar scalar; - if (!is_null) { - scalar = arrow::FloatScalar(expr_value.get_numberic()); - } - auto array_ret = arrow::MakeArrayFromScalar(scalar, length); - if (!array_ret.ok()) { - DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); - return nullptr; - } else { - return *array_ret; - } - break; - } - case pb::DOUBLE: { - arrow::DoubleScalar scalar; - if (!is_null) { - scalar = arrow::DoubleScalar(expr_value.get_numberic()); - } - auto array_ret = arrow::MakeArrayFromScalar(scalar, length); - if (!array_ret.ok()) { - DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); - return nullptr; - } else { - return *array_ret; - } - break; - } - case pb::STRING: - case pb::HLL: - case pb::BITMAP: - case pb::TDIGEST: { - auto array_ret = arrow::MakeArrayFromScalar( - is_null ? arrow::LargeBinaryScalar(): - arrow::LargeBinaryScalar(arrow::Buffer::FromString(expr_value.get_string())), - length); - if (!array_ret.ok()) { - DB_WARNING("arrow make array from scalar fail, %s", array_ret.status().ToString().c_str()); - return nullptr; - } else { - return *array_ret; - } - break; - } - default: - DB_WARNING("Invalid type: %d", type); - return nullptr; - } - return nullptr; -} - int ColumnRecord::init() { _field_num = _schema->num_fields(); _builders.reserve(_field_num); diff --git a/src/column/file_manager.cpp b/src/column/file_manager.cpp index 27ea0b7c6..5fbb0abae 100644 --- a/src/column/file_manager.cpp +++ b/src/column/file_manager.cpp @@ -1,13 +1,14 @@ #include "file_manager.h" #include "meta_writer.h" #include "rocksdb_filesystem.h" +#include "vectorize_helpper.h" namespace baikaldb { DEFINE_bool(use_row_ranges, false, "GetRecordBatchReader use row_ranges"); DEFINE_bool(parquet_read_pre_buffer, false, "parquet read use pre_buffer"); DEFINE_bool(parquet_read_use_threads, false, "parquet read use use_threads"); -DEFINE_bool(parquet_file_enable_lru, true, "parquet_file_enable_lru"); +DEFINE_bool(parquet_file_enable_lru, true, "Enable LRU for parquet files, default: true"); DEFINE_int32(parquet_file_lru_cache_capacity, 1024, "parquet_file_lru_cache_capacity, default(1024)"); DECLARE_int32(chunk_size); @@ -19,6 +20,7 @@ DEFINE_int64(column_minor_compaction_interval_s, 60, "column_minor_compaction_in DEFINE_int64(column_minor_compaction_raft_interval, 1000, "column_minor_compaction_raft_interval, default(1000)"); DEFINE_int64(column_major_compaction_minor_interval, 10, "column_major_compaction_minor_interval, default(10)"); DEFINE_int64(column_base_compaction_interval_s, 2 * 24 * 3600, "column_major_compaction_interval_s, default(2day)"); +DEFINE_int64(column_fast_base_compaction_interval_s, 1800, "column_fast_base_compaction_interval_s, default(1800s)"); DEFINE_int64(column_base_compaction_threshold_mb, 256, "column_base_compaction_threshold_mb, default(256MB)"); DEFINE_int64(column_base_compaction_cumulative_threshold_mb, 128, "column_base_compaction_cumulative_threshold_mb, default(128MB)"); DEFINE_int64(column_cumulative_file_max_size_mb, 16, "column_cumulative_file_max_size_mb, default(16MB)"); @@ -57,22 +59,25 @@ int ParquetFile::open() { } _file_path = path; std::shared_ptr<::arrow::io::RandomAccessFile> infile = std::move(res).ValueOrDie(); - + arrow::MemoryPool* pool = GetMemoryPoolForRead(); ::parquet::arrow::FileReaderBuilder builder; // Set GetRecordBatchReader batch size, default: 65536 + ::arrow::io::IOContext io_context(pool); ::parquet::ArrowReaderProperties arrow_properties = ::parquet::default_arrow_reader_properties(); arrow_properties.set_batch_size(FLAGS_chunk_size); arrow_properties.set_pre_buffer(FLAGS_parquet_read_pre_buffer); arrow_properties.set_use_threads(FLAGS_parquet_read_use_threads); + arrow_properties.set_io_context(io_context); ::arrow::io::CacheOptions cache_options = ::arrow::io::CacheOptions::LazyDefaults(); cache_options.hole_size_limit = FLAGS_parquet_read_hole_size_limit; cache_options.range_size_limit = FLAGS_parquet_read_range_size_limit; arrow_properties.set_cache_options(cache_options); builder.properties(arrow_properties); + builder.memory_pool(pool); // Avoid reading whole file data into memory at once - ::parquet::ReaderProperties read_properties = ::parquet::default_reader_properties(); + ::parquet::ReaderProperties read_properties(pool); read_properties.enable_buffered_stream(); read_properties.set_buffer_size(FLAGS_file_buffer_size * 1024 * 1024ULL); auto status = builder.Open(infile, read_properties); @@ -159,7 +164,7 @@ int ParquetFile::init_kv_metadata(const std::vector& keys, const st return -1; } -bool ParquetFile::check_interval_overlapped(const pb::PossibleIndex::Range& index_range, +bool ParquetFile::check_interval_overlapped(const pb::PossibleIndex::Range& index_range, bool is_eq, bool is_left_open, bool is_right_open, const std::string& file_start_key, const std::string& file_end_key) { // 文件区间理论上不为空，BUG if (file_start_key.empty() || file_end_key.empty()) { @@ -167,16 +172,16 @@ bool ParquetFile::check_interval_overlapped(const pb::PossibleIndex::Range& inde } // 索引右区间为空 - if (index_range.right_key().empty()) { - if (index_range.left_open()) { + if (index_range.right_key().empty() && !is_eq) { + if (is_left_open) { return file_end_key > index_range.left_key(); } else { return file_end_key >= index_range.left_key(); } } - MutTableKey end_key(index_range.right_key()); - if (!index_range.right_open()) { + MutTableKey end_key(is_eq ? index_range.left_key() : index_range.right_key()); + if (!is_right_open) { end_key.append_u64(UINT64_MAX); end_key.append_u64(UINT64_MAX); end_key.append_u64(UINT64_MAX); @@ -189,10 +194,10 @@ bool ParquetFile::check_interval_overlapped(const pb::PossibleIndex::Range& inde if (max_start_key < min_end_key) { return true; } else if (max_start_key == min_end_key) { - if (index_end_key == file_start_key && index_range.right_open()) { + if (index_end_key == file_start_key && is_right_open) { return false; } - if (index_start_key == file_end_key && index_range.left_open()) { + if (index_start_key == file_end_key && is_left_open) { return false; } return true; @@ -220,7 +225,7 @@ void print_rg_and_rrs(std::vector& rowgroup_indices, } int ParquetFile::get_qualified_rowgroup_and_rowranges( - const std::vector& key_ranges, + const pb::PossibleIndex& possible_index, std::vector& rowgroup_indices, std::vector>>& rowranges) { if (_sparse_index_map.empty()) { @@ -232,14 +237,15 @@ int ParquetFile::get_qualified_rowgroup_and_rowranges( const std::string& file_end_key = _sparse_index_map.rbegin()->first; // 获取rowgroup和rowranges std::map>> rowgroup_rowranges_map; - for (const auto& key_range : key_ranges) { + for (const auto& key_range : possible_index.ranges()) { // 跳过和该parquet文件未重叠的key_range - if (!check_interval_overlapped(key_range, file_start_key, file_end_key)) { + if (!check_interval_overlapped(key_range, possible_index.is_eq(), possible_index.left_open(), + possible_index.right_open(), file_start_key, file_end_key)) { continue; } - MutTableKey range_right_key(key_range.right_key()); - if (!key_range.right_open()) { + MutTableKey range_right_key(possible_index.is_eq() ? key_range.left_key() : key_range.right_key()); + if (!possible_index.right_open()) { range_right_key.append_u64(UINT64_MAX); range_right_key.append_u64(UINT64_MAX); range_right_key.append_u64(UINT64_MAX); @@ -252,7 +258,7 @@ int ParquetFile::get_qualified_rowgroup_and_rowranges( } // 找到大于等于right_key的最小元素的位置 auto right_iter = _sparse_index_map.lower_bound(range_right_key.data()); - if (key_range.right_key().empty()) { + if (key_range.right_key().empty() && !possible_index.is_eq()) { right_iter = _sparse_index_map.end(); } if (right_iter == _sparse_index_map.end()) { @@ -348,15 +354,79 @@ int ParquetFileReader::init() { if (_init) { return 0; } - int ret = _parquet_file->open(); - if (ret < 0) { - DB_COLUMN_FATAL("Fail to open parquet file"); - return -1; + + const std::string& file_short_name = _parquet_file->get_file_short_name(); + // 获取column_indices，以及在parquet文件里不存在的列 + // 在parquet文件中不存在的列需要补充默认值或NULL + const std::unordered_map& column_name2index_map = _parquet_file->get_column_name2index_map(); + std::vector exist_column_indices; + exist_column_indices.reserve(column_name2index_map.size()); + for (const auto& f : _options.schema->fields()) { + const std::string& field_name = f->name(); + auto iter = column_name2index_map.find(field_name); + if (iter != column_name2index_map.end()) { + exist_column_indices.emplace_back(iter->second); + // DB_WARNING("file: %s, field_name: %s, index: %d", file_short_name.c_str(), field_name.c_str(), iter->second); + } else { + DB_WARNING("file: %s, field_name: %s, not exist", file_short_name.c_str(), field_name.c_str()); + } } - auto s = _parquet_file->GetRecordBatchReader(&_reader); + std::vector row_group_indices; + bool fill_cache = false; + if (_options.pos_index != nullptr && _options.pos_index->ranges_size() > 0) { + // 获取符合条件的row_group_indices和row_ranges + std::vector>> row_ranges; + if (_parquet_file->get_qualified_rowgroup_and_rowranges(*_options.pos_index, row_group_indices, row_ranges) != 0) { + DB_WARNING("%s Fail to get_qualified_rowgroup_and_rowranges", file_short_name.c_str()); + return -1; + } + if (row_group_indices.empty()) { + DB_WARNING("parquet file: %s has no qualified data", file_short_name.c_str()); + return -1; + } + fill_cache = true; + } else { + // 获取所有rowgroup的数据 + std::shared_ptr<::parquet::FileMetaData> file_metadata = _parquet_file->get_file_metadata(); + if (file_metadata == nullptr) { + DB_WARNING("file_metadata:%s is nullptr", file_short_name.c_str()); + return -1; + } + if (file_metadata->num_row_groups() == 0) { + DB_WARNING("parquet file: %s has no row group", file_short_name.c_str()); + return -1; + } + + row_group_indices.reserve(file_metadata->num_row_groups()); + for (int i = 0; i < file_metadata->num_row_groups(); ++i) { + row_group_indices.emplace_back(i); + } + } + + std::shared_ptr file_info = _parquet_file->get_file_info(); + _raftindex = file_info->end_version; + if (fill_cache) { + std::vector read_ranges; + read_ranges.reserve(row_group_indices.size() * exist_column_indices.size()); + _parquet_file->parser_position(row_group_indices, exist_column_indices, read_ranges); + _read_contents = std::make_shared(); + _read_contents->fill_cache = true; + _read_contents->file_short_name = file_short_name; + _read_contents->ranges.swap(read_ranges); + _read_contents->region_id = file_info->region_id; + _read_contents->start_version = file_info->start_version; + _read_contents->end_version = file_info->end_version; + _read_contents->file_idx = ColumnFileInfo::get_file_idx(file_short_name); + if (_read_contents->file_idx < 0) { + DB_COLUMN_FATAL("Fail to get file idx : %s", file_short_name.c_str()); + return -1; + } + } + + ::arrow::Status s = _parquet_file->GetRecordBatchReader(row_group_indices, exist_column_indices, &_reader); if (!s.ok()) { - DB_COLUMN_FATAL("Fail to get_record_batch_reader"); + DB_COLUMN_FATAL("%s Fail to get_record_batch_reader", file_short_name.c_str()); return -1; } _init = true; @@ -364,6 +434,10 @@ int ParquetFileReader::init() { } ::arrow::Status ParquetFileReader::ReadNext(std::shared_ptr<::arrow::RecordBatch>* batch) { + ON_SCOPE_EXIT([]() { + ParquetCache::get_instance()->set_bthread_local(nullptr); + }); + ParquetCache::get_instance()->set_bthread_local(_read_contents.get()); int ret = init(); if (ret < 0) { return ::arrow::Status::IOError("Fail to init"); @@ -379,9 +453,9 @@ ::arrow::Status ParquetFileReader::ReadNext(std::shared_ptr<::arrow::RecordBatch batch->reset(); return ::arrow::Status::OK(); } - const auto& key_fields = _options.schema_info->key_fields; - const auto& value_fields = _options.schema_info->value_fields; - const auto& schema = _options.need_order_info ? _options.schema_info->schema_with_order_info : _options.schema_info->schema; + + const auto& fields = _options.lower_short_name_fields; + const auto& schema = _options.schema; std::vector> columns; for (int i = 0; i < schema->num_fields(); ++i) { const auto& f = schema->field(i); @@ -394,23 +468,23 @@ ::arrow::Status ParquetFileReader::ReadNext(std::shared_ptr<::arrow::RecordBatch if (f->name() == RAFT_INDEX_NAME) { ExprValue raft_index; raft_index.type = pb::INT64; - raft_index._u.int64_val = _options.raftindex; - array = ColumnRecord::make_array_from_exprvalue( + raft_index._u.int64_val = _raftindex; + array = VectorizeHelpper::make_array_from_exprvalue( pb::INT64, raft_index, tmp_batch->num_rows()); } else if (f->name() == BATCH_POS_NAME) { ExprValue batch_pos; batch_pos.type = pb::INT32; batch_pos._u.int32_val = 0; - array = ColumnRecord::make_array_from_exprvalue( + array = VectorizeHelpper::make_array_from_exprvalue( pb::INT32, batch_pos, tmp_batch->num_rows()); } else { - if (i < key_fields.size()) { - DB_COLUMN_FATAL("%s Fail to find key field", _parquet_file->get_file_path().c_str()); - return arrow::Status::IOError("Fail to find key field"); + auto it = fields.find(f->name()); + if (it == fields.end()) { + DB_COLUMN_FATAL("%s Fail to find field: %s", _parquet_file->get_file_path().c_str(), f->name().c_str()); + return arrow::Status::IOError("Fail to find field"); } - const auto& field = value_fields[i - key_fields.size()]; - array = ColumnRecord::make_array_from_exprvalue( - field.type, field.default_expr_value, tmp_batch->num_rows()); + array = VectorizeHelpper::make_array_from_exprvalue( + it->second.type, it->second.default_expr_value, tmp_batch->num_rows()); } if (array == nullptr) { DB_COLUMN_FATAL("%s Fail to make array from expr value", _parquet_file->get_file_path().c_str()); @@ -696,7 +770,7 @@ int ColumnFileManager::pick_major_compact_file(std::vector>& file_infos) { +int ColumnFileManager::pick_base_compact_file(std::vector>& file_infos, bool only_read_base) { std::unique_lock l(_mutex); if (_column_status != pb::CS_NORMAL) { DB_WARNING("region_id: %ld, column status invalid", _region_id); @@ -713,6 +787,11 @@ int ColumnFileManager::pick_base_compact_file(std::vector FLAGS_column_fast_base_compaction_interval_s * 1000 * 1000LL)) { + need_check_cumulatives_size = false; + } + bool need_base_compact = false; if (need_check_cumulatives_size) { if (_base_file_size > FLAGS_column_base_compaction_threshold_mb * 1024 * 1024LL && diff --git a/src/column/parquet_cache.cpp b/src/column/parquet_cache.cpp index 56e72d67c..2c9ae688f 100644 --- a/src/column/parquet_cache.cpp +++ b/src/column/parquet_cache.cpp @@ -92,6 +92,10 @@ ::arrow::Result ParquetArrowReadableFile::ReadAt(int64_t position, int6 } ::arrow::Result ParquetArrowReadableFile::read_from_device(int64_t position, int64_t nbytes, void* buffer) { + if (nbytes == 0) { + DB_WARNING("read from device, filename: %s, offset: %ld, len: 0", _file_reader->file_name().c_str(), position); + return 0; + } int64_t r = -1; int64_t left = nbytes; char* ptr = static_cast(buffer); diff --git a/src/column/region_column.cpp b/src/column/region_column.cpp index d9a3110d5..3170035ba 100644 --- a/src/column/region_column.cpp +++ b/src/column/region_column.cpp @@ -21,6 +21,7 @@ #include "row2column.h" #include "rocks_wrapper.h" #include "rocksdb_filesystem.h" +#include "vectorize_helpper.h" namespace baikaldb { DEFINE_int32(parquet_userid_statis_batch_count, 5, "parquet_userid_statis_batch_count(5)"); @@ -29,13 +30,14 @@ DEFINE_int32(parquet_rowgroup_max_length, 1000000, "parquet_rowgroup_max_length( DEFINE_int32(parquet_file_max_length, 20000000, "parquet_file_max_length(2000w)"); DEFINE_int32(raftlog_read_batch_size, 10000, "raftlog_read_batch_size"); DEFINE_int32(column_snapshot_timeout_s, 180, "column_snapshot_timeout_s"); -DEFINE_bool(column_minor_compaction_use_acero, false, "column_minor_compaction_use_acero"); -DEFINE_bool(column_major_compaction_use_acero, false, "column_major_compaction_use_acero"); +DEFINE_bool(column_minor_compaction_use_acero, false, "Use Acero for column minor compaction, default: false"); +DEFINE_bool(column_major_compaction_use_acero, false, "Use Acero for column major compaction, default: false"); DEFINE_int64(column_major_compaction_use_acero_max_rows, 2000000, "column_major_compaction_use_acero_max_rows, default(200w)"); DEFINE_int64(column_automatic_judgment_max_in_count, 100, "column_automatic_judgment_max_in_count, default(100)"); -DEFINE_bool(parquet_read_use_userid_statis, true, "parquet_read_use_userid_statis"); +DEFINE_bool(parquet_read_use_userid_statis, true, "Use userid statistics in parquet reading, default: true"); DEFINE_int64(column_row2column_flush_delay_h, 6, "column_row2column_flush_delay_h, default(6)"); -DEFINE_bool(column_cold_parquet_clear, false, "column_cold_parquet_clear"); +DEFINE_bool(column_cold_parquet_clear, false, "Enable cold parquet clearing, default: false"); +DEFINE_bool(enable_column_engine, false, "Enable column engine, default: false"); #define IF_DONE_SET_RESPONSE(done, errcode, err_message) \ do {\ if (done != nullptr && ((ColumnOPClosure*)done)->response != nullptr) {\ @@ -60,85 +62,6 @@ struct SnapshotManager { const rocksdb::Snapshot* cold_snapshot = nullptr; }; -std::shared_ptr make_column_schema(int64_t tableid) { - auto table = SchemaFactory::get_instance()->get_table_info_ptr(tableid); - auto index = SchemaFactory::get_instance()->get_index_info_ptr(tableid); - if (table == nullptr || index == nullptr) { - DB_FATAL("table or index is null, tableid:%ld", tableid); - return nullptr; - } - auto schema_ptr = std::make_shared(); - schema_ptr->index_info = index; - schema_ptr->table_info = table; - schema_ptr->key_fields = index->fields; - - std::set key_field_ids; - for (const auto& field : index->fields) { - key_field_ids.insert(field.id); - } - - std::set field_ids_need_sum; - for (const auto& field : table->fields_need_sum) { - field_ids_need_sum.insert(field.id); - } - - int value_idx = schema_ptr->key_fields.size(); - schema_ptr->value_fields.reserve(table->fields.size() - index->fields.size()); - for (const auto& field : table->fields) { - if (key_field_ids.count(field.id) > 0 || field.deleted) { - continue; - } - - schema_ptr->value_fields.emplace_back(field); - if (field_ids_need_sum.count(field.id) > 0) { - schema_ptr->need_sum_idx.insert(value_idx); - } - value_idx++; - } - - schema_ptr->uniq_size = schema_ptr->key_fields.size(); - schema_ptr->keytype_idx = schema_ptr->key_fields.size() + schema_ptr->value_fields.size(); - schema_ptr->raft_index_idx = schema_ptr->keytype_idx + 1; - schema_ptr->batch_pos_idx = schema_ptr->raft_index_idx + 1; - - std::vector> arrow_fields; - arrow_fields.reserve(schema_ptr->batch_pos_idx + 1); - for (const auto& field : schema_ptr->key_fields) { - auto arrow_type = primitive_to_arrow_type(field.type); - if (arrow_type < 0) { - DB_COLUMN_FATAL("field: %s primitive type:%d to arrow type failed", field.lower_short_name.c_str(), field.type); - return nullptr; - } - auto arrow_field = ColumnRecord::make_schema(field.lower_short_name, arrow::Type::type(arrow_type)); - if (arrow_field == nullptr) { - DB_COLUMN_FATAL("field: %s make arrow schema failed", field.lower_short_name.c_str()); - return nullptr; - } - arrow_fields.emplace_back(arrow_field); - } - - for (const auto& field : schema_ptr->value_fields) { - auto arrow_type = primitive_to_arrow_type(field.type); - if (arrow_type < 0) { - DB_COLUMN_FATAL("field: %s primitive type:%d to arrow type failed", field.lower_short_name.c_str(), field.type); - return nullptr; - } - auto arrow_field = ColumnRecord::make_schema(field.lower_short_name, arrow::Type::type(arrow_type)); - if (arrow_field == nullptr) { - DB_COLUMN_FATAL("field: %s make arrow schema failed", field.lower_short_name.c_str()); - return nullptr; - } - arrow_fields.emplace_back(arrow_field); - } - - arrow_fields.emplace_back(ColumnRecord::make_schema(KEY_TYPE_NAME, arrow::Type::type::INT32)); - schema_ptr->schema = std::make_shared(arrow_fields); - arrow_fields.emplace_back(ColumnRecord::make_schema(RAFT_INDEX_NAME, arrow::Type::type::INT64)); - arrow_fields.emplace_back(ColumnRecord::make_schema(BATCH_POS_NAME, arrow::Type::type::INT32)); - schema_ptr->schema_with_order_info = std::make_shared(arrow_fields); - return schema_ptr; -} - std::string files_name(const std::vector>& files) { std::ostringstream os; os << "["; @@ -350,14 +273,15 @@ int Region::column_snapshot_save(const std::string& snapshot_path, std::vector& key_ranges, std::vector>& files) { +int Region::get_column_files(const pb::PossibleIndex& pos_index, std::vector>& files) { std::shared_ptr column_file_set = _column_mgr.get_column_fileset(); DB_DEBUG("region_id: %ld, get column files, key_ranges size: %ld", _region_id, key_ranges.size()); for (const auto& [_, info] : column_file_set->column_files) { - if (!key_ranges.empty()) { + if (pos_index.ranges_size() > 0) { bool is_overlap = false; - for (const auto& range : key_ranges) { - if (ParquetFile::check_interval_overlapped(range, info->start_key, info->end_key)) { + for (const auto& range : pos_index.ranges()) { + if (ParquetFile::check_interval_overlapped(range, pos_index.is_eq(), pos_index.left_open(), pos_index.right_open(), + info->start_key, info->end_key)) { is_overlap = true; break; } @@ -535,7 +459,8 @@ bool Region::can_do_column_compact() { // 定时执行, 快速将raft log 刷成parquet文件 int Region::column_minor_compact() { TimeCost cost; - std::shared_ptr schema_info = make_column_schema(get_table_id()); + std::unordered_map field_id2info_map; + std::shared_ptr schema_info = ColumnRecord::make_column_schema(get_table_id(), nullptr, nullptr, field_id2info_map); if (schema_info == nullptr) { DB_FATAL("region_id: %ld, make column schema failed", _region_id); return -1; @@ -638,6 +563,10 @@ int Region::column_minor_compact() { column_delete_files(new_files); return -1; } + + // compaction成功，删除本次涉及的column_txn_log_index_key + raftlog_reader->delete_column_txn_log_index(); + auto new_files_name = files_name(new_files); DB_NOTICE("region_id: %ld, minor compact success, read lines[%ld, %ld, %ld], write lines: %ld, use_acero_flag: %d, cost: %ld, new_files: %s", _region_id, raftlog_reader->row_count(), raftlog_reader->put_count(), raftlog_reader->delete_count(), writer->row_count(), @@ -646,6 +575,12 @@ int Region::column_minor_compact() { } void Region::column_flush() { + if (!FLAGS_enable_column_engine) { + if (_column_mgr.column_status() != pb::CS_INVALID) { + _column_mgr.remove_column_data(pb::CS_INVALID, 0); + } + return; + } auto table = _factory->get_table_info_ptr(get_table_id()); if (table == nullptr) { return; @@ -710,7 +645,8 @@ void Region::column_flush() { void Region::column_major_compact(bool is_base) { TimeCost cost; - std::shared_ptr schema_info = make_column_schema(get_table_id()); + std::unordered_map field_id2info_map; + std::shared_ptr schema_info = ColumnRecord::make_column_schema(get_table_id(), nullptr, nullptr, field_id2info_map); if (schema_info == nullptr) { DB_FATAL("region_id: %ld, get schema info failed", _region_id); return; @@ -718,7 +654,12 @@ void Region::column_major_compact(bool is_base) { std::vector> old_files; int ret = 0; if (is_base) { - ret = _column_mgr.pick_base_compact_file(old_files); + bool only_read_base = false; + auto table = _factory->get_table_info_ptr(get_table_id()); + if (table != nullptr && table->schema_conf.column_only_read_base()) { + only_read_base = true; + } + ret = _column_mgr.pick_base_compact_file(old_files, only_read_base); } else { ret = _column_mgr.pick_major_compact_file(old_files); } @@ -739,10 +680,19 @@ void Region::column_major_compact(bool is_base) { min_version = is_base ? 0 : std::min(file->start_version, min_version); max_version = std::max(file->end_version, max_version); ParquetFileReaderOptions options; - options.raftindex = file->end_version; - options.schema_info = schema_info; - options.file_info = file; - auto parquet_reader = std::make_shared(options); + for (const auto& f : schema_info->key_fields) { + options.lower_short_name_fields[f.lower_short_name] = f; + } + for (const auto& f : schema_info->value_fields) { + options.lower_short_name_fields[f.lower_short_name] = f; + } + options.schema = schema_info->schema_with_order_info; + auto parquet_file = ParquetFileManager::get_instance()->get_parquet_file(file); + if (parquet_file == nullptr) { + DB_WARNING("open file:%s failed", file->full_path().c_str()); + return; + } + auto parquet_reader = std::make_shared(options, parquet_file); parquet_readers.emplace_back(parquet_reader); } @@ -815,7 +765,8 @@ void Region::column_base_row2column() { } } TimeCost cost; - std::shared_ptr schema_info = make_column_schema(get_table_id()); + std::unordered_map field_id2info_map; + std::shared_ptr schema_info = ColumnRecord::make_column_schema(get_table_id(), nullptr, nullptr, field_id2info_map); if (schema_info == nullptr) { DB_FATAL("make column schema failed, region_id: %ld", _region_id); return; diff --git a/src/column/row2column.cpp b/src/column/row2column.cpp index b538fe18f..77cf87c27 100644 --- a/src/column/row2column.cpp +++ b/src/column/row2column.cpp @@ -131,8 +131,10 @@ arrow::Status RocksdbBaseReader::ReadNext(std::shared_ptr* o _read_times++; if (_is_finish) { out->reset(); - DB_NOTICE("read base rocksdb is finish, region_id: %ld, need_index[%ld, %ld], read times: %d, total rows: %d, cost: %ld", - _options.region_id, _options.start_index, _options.end_index, _read_times, _total_row_nums, _cost.get_time()); + DB_NOTICE("read base rocksdb is finish, region_id: %ld, need_index[%ld, %ld], " + "read times: %d, total rows: %d, cost: %ld", + _options.region_id, _options.start_index, _options.end_index, + _read_times, _total_row_nums, _cost.get_time()); return arrow::Status::OK(); } std::string prefix; @@ -168,8 +170,10 @@ arrow::Status RocksdbBaseReader::ReadNext(std::shared_ptr* o if (_column_record->size() == 0) { out->reset(); - DB_NOTICE("read base rocksdb is finish, region_id: %ld, need_index[%ld, %ld], read times: %d, total rows: %d, cost: %ld", - _options.region_id, _options.start_index, _options.end_index, _read_times, _total_row_nums, _cost.get_time()); + DB_NOTICE("read base rocksdb is finish, region_id: %ld, need_index[%ld, %ld], " + "read times: %d, total rows: %d, cost: %ld", + _options.region_id, _options.start_index, _options.end_index, + _read_times, _total_row_nums, _cost.get_time()); return arrow::Status::OK(); } @@ -191,14 +195,12 @@ int RaftLogReader::init() { return -1; } TimeCost cost; - std::string log_entry; MutTableKey log_data_key; log_data_key.append_i64(_options.region_id).append_u8(MyRaftLogStorage::LOG_DATA_IDENTIFY).append_i64(_options.start_index); MutTableKey prefix; MutTableKey end; prefix.append_i64(_options.region_id).append_u8(MyRaftLogStorage::LOG_DATA_IDENTIFY); end.append_i64(_options.region_id).append_u8(MyRaftLogStorage::LOG_DATA_IDENTIFY).append_i64(_options.end_index + 1); - std::string log_value; rocksdb::ReadOptions options; rocksdb::Slice upper_bound_slice = end.data(); options.iterate_upper_bound = &upper_bound_slice; @@ -211,14 +213,17 @@ int RaftLogReader::init() { } std::unique_ptr iter(iter_ptr); iter->Seek(log_data_key.data()); + RocksdbVars::get_instance()->raft_log_scan_times_count << 1; for (; iter->Valid(); iter->Next()) { if (!iter->key().starts_with(prefix.data())) { - DB_WARNING("read end info, region_id: %ld, key:%s", _options.region_id, iter->key().ToString(true).c_str()); + DB_WARNING("read end info, region_id: %ld, key:%s", + _options.region_id, iter->key().ToString(true).c_str()); return -1; } int64_t log_index = TableKey(iter->key()).extract_i64(sizeof(int64_t) + 1); if (log_index > _options.end_index) { - DB_WARNING("region_id:%ld, log_index:%ld, end_log_index:%ld", _options.region_id, log_index, _options.end_index); + DB_WARNING("region_id:%ld, log_index:%ld, end_log_index:%ld", + _options.region_id, log_index, _options.end_index); break; } @@ -235,7 +240,8 @@ int RaftLogReader::init() { value_slice.remove_prefix(MyRaftLogStorage::LOG_HEAD_SIZE); if (head.type != braft::ENTRY_TYPE_DATA) { ++_skip_count; - DB_WARNING("log entry is not data, region_id: %ld head.type: %d, raft index: %ld", _options.region_id, head.type, log_index); + DB_WARNING("log entry is not data, region_id: %ld head.type: %d, raft index: %ld", + _options.region_id, head.type, log_index); continue; } @@ -245,6 +251,12 @@ int RaftLogReader::init() { return -1; } + int64_t txn_id = 0; + bool optimize_1pc = false; + if (request.txn_infos_size() > 0) { + txn_id = request.txn_infos(0).txn_id(); + optimize_1pc = request.txn_infos(0).optimize_1pc(); + } if (request.op_type() == pb::OP_KV_BATCH) { auto s = MetaWriter::get_instance()->get_skip_watt_stats_version(_options.region_id, log_index); if (s.ok()) { @@ -253,8 +265,43 @@ int RaftLogReader::init() { // 找到说明已经这个点被跳过 continue; } + if (txn_id != 0) { + insert(txn_id, log_index, request); + continue; + } else { + insert(txn_id, log_index, request); + commit(txn_id, log_index); + } + } else if (request.op_type() == pb::OP_PREPARE) { + if (optimize_1pc) { + commit(txn_id, log_index); + } else { + continue; + } + } else if (request.op_type() == pb::OP_COMMIT) { + commit(txn_id, log_index); + } else if (request.op_type() == pb::OP_ROLLBACK) { + rollback(txn_id, log_index); + continue; + } else if (request.op_type() == pb::OP_PARTIAL_ROLLBACK) { + // TODO: 暂时不处理partial rollback + DB_COLUMN_FATAL("column read raft log is partial rollback, region_id: %ld", _options.region_id); + continue; + } else { + DB_WARNING("column read raft log is not kv batch, region_id: %ld, op_type: %s", + _options.region_id, pb::OpType_Name(request.op_type()).c_str()); + continue; + } + + std::map log_index_req_map; + ret = get(txn_id, log_index_req_map); + if (ret < 0) { + DB_COLUMN_FATAL("get raft log fail, region_id: %ld", _options.region_id); + return -1; + } + for (const auto& iter : log_index_req_map) { int idx = 0; - for (auto& kv_op : request.kv_ops()) { + for (auto& kv_op : iter.second.kv_ops()) { int ret = 0; pb::OpType op_type = kv_op.op_type(); int prefix_len = 2 * sizeof(int64_t); @@ -268,13 +315,13 @@ int RaftLogReader::init() { key.remove_prefix(prefix_len); if (op_type == pb::OP_PUT_KV) { ++_put_count; - ret = row2col(key, kv_op.value(), COLUMN_KEY_PUT, log_index, idx++); + ret = row2col(key, kv_op.value(), COLUMN_KEY_PUT, iter.first, idx++); } else if (op_type == pb::OP_MERGE_KV) { ++_merge_count; - ret = row2col(key, kv_op.value(), COLUMN_KEY_MERGE, log_index, idx++); + ret = row2col(key, kv_op.value(), COLUMN_KEY_MERGE, iter.first, idx++); } else { ++_delete_count; - ret = row2col(key, kv_op.value(), COLUMN_KEY_DELETE, log_index, idx++); + ret = row2col(key, kv_op.value(), COLUMN_KEY_DELETE, iter.first, idx++); } if (ret < 0) { DB_FATAL("row2col fail, region_id: %ld", _options.region_id); @@ -291,7 +338,7 @@ int RaftLogReader::init() { return -1; } _batchs.push_back(out); - if ( _total_row_nums >= FLAGS_column_minor_compact_read_raft_rows) { + if (_total_row_nums >= FLAGS_column_minor_compact_read_raft_rows) { break; } } @@ -309,16 +356,78 @@ int RaftLogReader::init() { if (_first_index != _options.start_index && _total_row_nums > 0) { // 可能丢数据，报警 - DB_COLUMN_FATAL("column read raft log is not start index, region_id: %ld, start_index: %ld, first_index: %ld", _options.region_id, _options.start_index, _first_index); + DB_COLUMN_FATAL("column read raft log is not start index, region_id: %ld, start_index: %ld, first_index: %ld", + _options.region_id, _options.start_index, _first_index); } - DB_NOTICE("read raft log is finish, region_id: %ld, need_index[%ld, %ld], read_index[%ld, %ld], total rows: %d, cost: %ld, skip_count: %ld", - _options.region_id, _options.start_index, _options.end_index, _first_index, _last_index, _total_row_nums, cost.get_time(), _skip_count); + DB_NOTICE("read raft log is finish, region_id: %ld, need_index[%ld, %ld], read_index[%ld, %ld], " + "total rows: %d, cost: %ld, skip_count: %ld", + _options.region_id, _options.start_index, _options.end_index, + _first_index, _last_index, _total_row_nums, cost.get_time(), _skip_count); _init = true; return 0; } +int RaftLogReader::get_raft_log(int64_t start_index, int64_t end_index, uint64_t txn_id, std::map& pre_reqs) { + TimeCost cost; + MutTableKey log_data_key; + log_data_key.append_i64(_options.region_id).append_u8(MyRaftLogStorage::LOG_DATA_IDENTIFY).append_i64(start_index); + MutTableKey prefix; + MutTableKey end; + prefix.append_i64(_options.region_id).append_u8(MyRaftLogStorage::LOG_DATA_IDENTIFY); + end.append_i64(_options.region_id).append_u8(MyRaftLogStorage::LOG_DATA_IDENTIFY).append_i64(end_index + 1); + rocksdb::ReadOptions options; + rocksdb::Slice upper_bound_slice = end.data(); + options.iterate_upper_bound = &upper_bound_slice; + options.prefix_same_as_start = true; + options.total_order_seek = false; + options.fill_cache = false; + auto iter_ptr = RocksWrapper::get_instance()->new_iterator(options, RocksWrapper::RAFT_LOG_CF); + if (iter_ptr == nullptr) { + return -1; + } + std::unique_ptr iter(iter_ptr); + iter->Seek(log_data_key.data()); + RocksdbVars::get_instance()->raft_log_scan_times_count << 1; + for (; iter->Valid(); iter->Next()) { + if (!iter->key().starts_with(prefix.data())) { + DB_WARNING("read end info, region_id: %ld, key:%s", + _options.region_id, iter->key().ToString(true).c_str()); + return -1; + } + int64_t log_index = TableKey(iter->key()).extract_i64(sizeof(int64_t) + 1); + if (log_index > end_index) { + DB_WARNING("region_id:%ld, log_index:%ld, end_log_index:%ld", + _options.region_id, log_index, end_index); + break; + } + + rocksdb::Slice value_slice(iter->value()); + LogHead head(value_slice); + value_slice.remove_prefix(MyRaftLogStorage::LOG_HEAD_SIZE); + if (head.type != braft::ENTRY_TYPE_DATA) { + DB_WARNING("log entry is not data, region_id: %ld head.type: %d, raft index: %ld", + _options.region_id, head.type, log_index); + continue; + } + + pb::StoreReq request; + if (!request.ParseFromArray(value_slice.data(), value_slice.size())) { + DB_FATAL("Fail to parse request fail, region_id: %ld", _options.region_id); + return -1; + } + + if (request.txn_infos_size() > 0 && txn_id == request.txn_infos(0).txn_id() && request.op_type() == pb::OP_KV_BATCH) { + pre_reqs[log_index].Swap(&request); + DB_NOTICE("find pre request, region_id: %ld, log_index: %ld, txn_id: %lu", + _options.region_id, log_index, txn_id); + } + } + + return 0; +} + arrow::Status RaftLogReader::ReadNext(std::shared_ptr* out) { int ret = init(); if (ret < 0) { @@ -336,4 +445,109 @@ arrow::Status RaftLogReader::ReadNext(std::shared_ptr* out) return arrow::Status::OK(); } +void RaftLogReader::commit(int64_t txn_id, int64_t raft_index) { + if (txn_id != 0) { + _txn_ids.emplace_back(txn_id); + } + + auto iter = _raft_log_cache->txn_id_raft_log_map.find(txn_id); + if (iter != _raft_log_cache->txn_id_raft_log_map.end()) { + iter->second->commit_index = raft_index; + } else { + DB_COLUMN_FATAL("txn_id not exist, region_id: %ld, txn_id: %ld, raft_index: %ld", + _region_id, txn_id, raft_index); + auto logiter = std::make_shared(); + logiter->commit_index = raft_index; + _raft_log_cache->txn_id_raft_log_map[txn_id] = logiter; + } + _commited_txn_id = txn_id; +} + +void RaftLogReader::rollback(int64_t txn_id, int64_t raft_index) { + if (txn_id != 0) { + _txn_ids.emplace_back(txn_id); + } + + auto iter = _raft_log_cache->txn_id_raft_log_map.find(txn_id); + if (iter != _raft_log_cache->txn_id_raft_log_map.end()) { + _raft_log_cache->txn_id_raft_log_map.erase(iter); + DB_WARNING("txn rollbacked, region_id: %ld, txn_id: %ld, raft_index: %ld", + _region_id, txn_id, raft_index); + } else { + DB_COLUMN_FATAL("txn_id not exist, region_id: %ld, txn_id: %ld, raft_index: %ld", + _region_id, txn_id, raft_index); + } +} + +void RaftLogReader::insert(int64_t txn_id, int64_t raft_index, pb::StoreReq& request) { + auto iter = _raft_log_cache->txn_id_raft_log_map.find(txn_id); + if (iter != _raft_log_cache->txn_id_raft_log_map.end()) { + iter->second->log_index_req_map[raft_index].Swap(&request); + return; + } + + auto logiter = std::make_shared(); + logiter->log_index_req_map[raft_index].Swap(&request); + _raft_log_cache->txn_id_raft_log_map[txn_id] = logiter; +} + +int RaftLogReader::get(int64_t txn_id, std::map& log_index_req_map) { + if (txn_id != _commited_txn_id) { + DB_COLUMN_FATAL("txn_id not commited, region_id: %ld, txn_id: %ld, commited_txn_id: %ld", + _region_id, txn_id, _commited_txn_id); + return -1; + } + + auto iter = _raft_log_cache->txn_id_raft_log_map.find(txn_id); + if (iter == _raft_log_cache->txn_id_raft_log_map.end()) { + DB_COLUMN_FATAL("txn_id not exist, region_id: %ld, txn_id: %ld", _region_id, txn_id); + return -1; + } + + if (txn_id == 0) { + log_index_req_map.swap(iter->second->log_index_req_map); + _raft_log_cache->txn_id_raft_log_map.erase(iter); + return 0; + } + + int64_t txn_begin_index = -1; + if (iter->second->log_index_req_map.empty()) { + txn_begin_index = iter->second->commit_index; + } else { + txn_begin_index = iter->second->log_index_req_map.begin()->first; + } + + int64_t begin_index = MetaWriter::get_instance()->read_meta_begin_index(_region_id, txn_id); + if (begin_index < 0) { + DB_COLUMN_FATAL("read begin index failed, region_id: %ld, txn_id: %ld, begin_index: %ld", + _region_id, txn_id, begin_index); + return -1; + } + + if (begin_index < txn_begin_index) { + DB_WARNING("begin index not match, region_id: %ld, txn_id: %ld, begin_index: %ld, txn_begin_index: %ld", + _region_id, txn_id, begin_index, txn_begin_index); + std::map pre_reqs; + int ret = get_raft_log(begin_index, txn_begin_index, txn_id, pre_reqs); + if (ret < 0) { + DB_COLUMN_FATAL("get raft log failed, region_id: %ld, txn_id: %ld", _region_id, txn_id); + return -1; + } + for (auto& req : pre_reqs) { + iter->second->log_index_req_map[req.first].Swap(&req.second); + } + } + + log_index_req_map.swap(iter->second->log_index_req_map); + _raft_log_cache->txn_id_raft_log_map.erase(iter); + + std::ostringstream os; + for (auto& req : log_index_req_map) { + os << req.first << ","; + } + + DB_WARNING("region_id: %ld, txn_id: %ld, get success req_size: %ld, raft_index: %s", + _region_id, txn_id, log_index_req_map.size(), os.str().c_str()); + return 0; +} } // namespace baikaldb \ No newline at end of file diff --git a/src/common/baikal_heartbeat.cpp b/src/common/baikal_heartbeat.cpp index 9edbcddb7..2a269cddc 100755 --- a/src/common/baikal_heartbeat.cpp +++ b/src/common/baikal_heartbeat.cpp @@ -20,7 +20,7 @@ namespace baikaldb { DEFINE_bool(enable_dblink, false, "enable dblink"); -DEFINE_bool(can_do_ddlwork, true, "can_do_ddlwork"); +DEFINE_bool(can_do_ddlwork, true, "Enable DDL work capability, default: true"); DECLARE_int32(baikal_heartbeat_interval_us); DEFINE_string(baikal_resource_tag, "", "resource tag"); DECLARE_int32(baikal_port); @@ -366,15 +366,22 @@ int BaseBaikalHeartBeat::heartbeat(bool is_sync) { BaikalHeartBeat::construct_heart_beat_request(request); request.set_can_do_ddlwork(false); request.set_need_heartbeat_table(true); - for (const auto& full_table_table : _table_names) { + for (const auto& full_table : _table_names) { auto* heartbeat_table = request.add_heartbeat_tables(); if (heartbeat_table == nullptr) { DB_WARNING("baikal_heartbeat_table is nullptr"); return -1; } - heartbeat_table->set_namespace_name(full_table_table.namespace_name); - heartbeat_table->set_database(full_table_table.database); - heartbeat_table->set_table_name(full_table_table.table_name); + heartbeat_table->set_namespace_name(full_table.namespace_name); + heartbeat_table->set_database(full_table.database); + heartbeat_table->set_table_name(full_table.table_name); + for (const auto& partition_id : full_table.partition_ids) { + if (partition_id < 0) { + // 兼容-1000场景 + continue; + } + heartbeat_table->add_partition_ids(partition_id); + } } int64_t construct_req_cost = cost.get_time(); @@ -656,7 +663,7 @@ int BinlogNetworkServer::update_table_infos() { if (id < 0) { DB_FATAL("table[%s] cant find field[%s]", table->name.c_str(), field_name.c_str()); find_all_fields = false; - break; + return -1; } table_ids.fields.insert(id); } @@ -666,7 +673,7 @@ int BinlogNetworkServer::update_table_infos() { if (id < 0) { DB_FATAL("table[%s] cant find field[%s]", table->name.c_str(), field_name.c_str()); find_all_fields = false; - break; + return -1; } table_ids.monitor_fields.insert(id); } diff --git a/src/common/common.cpp b/src/common/common.cpp index 9b8f6f108..efb709b0d 100644 --- a/src/common/common.cpp +++ b/src/common/common.cpp @@ -88,6 +88,8 @@ DEFINE_int32(store_port, 8110, "Server port"); DEFINE_string(secondary_db_path, "./rocks_db_secondary", "rocks db path"); DEFINE_string(compaction_db_path, "./rocks_db_compaction", "rocks db path for compaction server"); DEFINE_int32(remote_compaction_server_concurrency, 10, "remote_compaction_server_concurrency"); +DEFINE_int64(print_time_us, 10000, "print log when time_cost > print_time_us(us)"); +BRPC_VALIDATE_GFLAG(print_time_us, brpc::NonNegativeInteger); int64_t timestamp_diff(timeval _start, timeval _end) { return (_end.tv_sec - _start.tv_sec) * 1000000 @@ -672,10 +674,12 @@ bool same_with_container_id_and_address(const std::string& container_id, const s std::vector instances; int ret = get_instance_from_bns(&bns_ret, container_id, instances, false); if (ret != 0) { - return true; + DB_WARNING("diff with container_id:%s bns error and address:%s", container_id.c_str(), address.c_str()); + return false; } if (instances.size() != 1) { - return true; + DB_WARNING("diff with container_id:%s(size:%lu) and address:%s", container_id.c_str(), instances.size(), address.c_str()); + return false; } if (instances[0] == address) { return true; diff --git a/src/common/datetime.cpp b/src/common/datetime.cpp index 3ce6d0b86..638975198 100755 --- a/src/common/datetime.cpp +++ b/src/common/datetime.cpp @@ -18,6 +18,86 @@ #include "expr_value.h" namespace baikaldb { +inline bool is_leap(int y) { + return y % 4 == 0 && (y % 100 != 0 || y % 400 == 0); +} + +static inline int64_t div_floor(int64_t a, int64_t b) { + return (a >= 0 ? a : a + 1 - b) / b; +} +// https://github.com/HowardHinnant/date +// https://howardhinnant.github.io/date_algorithms.html#days_from_civil +// 核心：将年月日转换为自 1970-01-01 起的天数（Howard Hinnant 算法） +static inline int64_t days_from_civil(int64_t y, unsigned m, unsigned d) noexcept { + y -= m <= 2; + const int64_t era = div_floor(y, 400); + const unsigned yoe = static_cast(y - era * 400); + const unsigned doy = (153 * (m + (m > 2 ? -3 : 9)) + 2) / 5 + d - 1; + const unsigned doe = yoe * 365 + yoe / 4 - yoe / 100 + doy; + return era * 146097 + static_cast(doe) - 719468; +} + +// 将 struct tm（本地时间，不考虑 DST） 转为 time_t（UTC timestamp） +time_t mktime_fixed_r(const struct tm* tm, int tz_offset_hours) { + int year = tm->tm_year + 1900; + unsigned mon = tm->tm_mon + 1; + unsigned mday = tm->tm_mday; + int64_t days = days_from_civil(year, mon, mday); + int64_t secs = days * 86400LL + + tm->tm_hour * 3600LL + + tm->tm_min * 60LL + + tm->tm_sec; + // 北京时间表示的时刻减去 tz_offset_hours 得到 UTC timestamp + return static_cast(secs - tz_offset_hours * 3600LL); +} + +struct tm* localtime_fixed_r(const time_t *timep, + struct tm *result, + int tz_offset_hours) { + int64_t t = (int64_t)*timep + tz_offset_hours * 3600LL; + int64_t days = t / 86400LL; + int64_t rem = t % 86400LL; + if (rem < 0) { + rem += 86400LL; + --days; + } + + // 反向 civil from days => year, month, day + // Use days + 719468 to get proleptic civil day count since 0000‑03‑01 + int64_t z = days + 719468; + const int64_t era = div_floor(z, 146097); + const unsigned doe = (unsigned)(z - era * 146097); + const unsigned yoe = (doe - doe/1460 + doe/36524 - doe/146096) / 365; + int year = (int)(yoe + era * 400); + const unsigned doy = doe - (365*yoe + yoe/4 - yoe/100); + const unsigned mp = (5*doy +2)/153; + unsigned day = doy - (153*mp+2)/5 +1; + unsigned month = mp + (mp < 10 ? 3 : -9); + year += (month <= 2); + + result->tm_year = year - 1900; + result->tm_mon = month -1; + result->tm_mday = day; + // 预计算各月累积天数（平年/闰年） + static constexpr int kMonthDays[2][12] = { + {0,31,59,90,120,151,181,212,243,273,304,334}, // 平年 + {0,31,60,91,121,152,182,213,244,274,305,335} // 闰年 + }; + const int* month_days = kMonthDays[is_leap(year)]; + result->tm_yday = month_days[result->tm_mon] + day - 1; + //result->tm_yday = days ? (doy + begin_days) - year_days : (doy + begin_days); + + result->tm_wday = (int)(days >= -4 ? (days + 4) % 7 : (days + 5) % 7 + 6); + + result->tm_hour = (int)(rem / 3600); + rem %= 3600; + result->tm_min = (int)(rem / 60); + result->tm_sec = (int)(rem % 60); + result->tm_isdst = 0; + + return result; +} + std::string timestamp_to_str(time_t timestamp, bool is_utc) { // 内部存储采用了uint32，因此小于0的都不合法 if (timestamp <= 0) { @@ -25,14 +105,9 @@ std::string timestamp_to_str(time_t timestamp, bool is_utc) { } struct tm tm; if (is_utc) { - gmtime_r(×tamp, &tm); + localtime_fixed_r(×tamp, &tm, 0); } else { - localtime_r(×tamp, &tm); - // 夏令时影响 - if (tm.tm_isdst == 1) { - timestamp = timestamp - 3600; - localtime_r(×tamp, &tm); - } + localtime_fixed_r(×tamp, &tm); } char str_time[21] = {0}; strftime(str_time, sizeof(str_time), "%Y-%m-%d %H:%M:%S", &tm); @@ -251,7 +326,7 @@ time_t datetime_to_timestamp(uint64_t datetime) { tm.tm_year -= 1900; tm.tm_mon--; - time_t t = mktime(&tm); + time_t t = mktime_fixed_r(&tm); return t <= 0 ? 0 : t; } @@ -270,7 +345,7 @@ time_t snapshot_to_timestamp(uint64_t snapshot) { tm.tm_sec = 0; tm.tm_isdst = 0; - time_t ts = mktime(&tm); + time_t ts = mktime_fixed_r(&tm); return ts; } @@ -281,7 +356,7 @@ uint64_t timestamp_to_datetime(time_t timestamp) { uint64_t datetime = 0; struct tm tm; - localtime_r(×tamp, &tm); + localtime_fixed_r(×tamp, &tm); tm.tm_year += 1900; tm.tm_mon++; uint64_t year_month = tm.tm_year * 13 + tm.tm_mon; @@ -312,12 +387,7 @@ void datetime_to_time_struct(uint64_t datetime, DateTime& time_struct, uint8_t t } else if (type == MYSQL_TYPE_TIMESTAMP) { struct tm tm; time_t timestamp = (time_t)datetime; - localtime_r(×tamp, &tm); - // 夏令时影响 - if (tm.tm_isdst == 1) { - timestamp = timestamp - 3600; - localtime_r(×tamp, &tm); - } + localtime_fixed_r(×tamp, &tm); time_struct.year = tm.tm_year + 1900; time_struct.month = tm.tm_mon + 1; time_struct.day = tm.tm_mday; @@ -536,14 +606,14 @@ bool tz_to_second(const char* time_zone, int32_t& result) { struct tm tm_local; time(&time_utc); - localtime_r(&time_utc, &tm_local); + localtime_fixed_r(&time_utc, &tm_local); time_t time_local; struct tm tm_gmt; - time_local = mktime(&tm_local); + time_local = mktime_fixed_r(&tm_local); - gmtime_r(&time_utc, &tm_gmt); + localtime_fixed_r(&time_utc, &tm_gmt, 0); int hour = tm_local.tm_hour - tm_gmt.tm_hour; if (hour < -12) { hour += 24; @@ -584,16 +654,13 @@ size_t date_format_internal(char* s, size_t maxsize, const char* format, const s if (tp == nullptr || format == nullptr) { return 0; } + size_t i = 0; std::string f = ""; - char tmp[20] = {0}; - int hour12 = tp->tm_hour % 12; - if (hour12 == 0) { - hour12 = 12; - } - while (format[i] != '\0') { + bool need_convert = false; + while (format[i] != '\0' && !need_convert) { if (format[i] != '%' ) { - f += format[i++]; + i++; continue; } i++; @@ -602,104 +669,138 @@ size_t date_format_internal(char* s, size_t maxsize, const char* format, const s } switch (format[i]) { case 'c': - f += std::to_string(tp->tm_mon + 1); - break; case 'D': - f += std::to_string(tp->tm_mday); - if (tp->tm_mday == 1 || tp->tm_mday == 21 || tp->tm_mday == 31) { - f += "st"; - } - else if (tp->tm_mday == 2 || tp->tm_mday == 22) { - f += "nd"; - } - else if (tp->tm_mday == 3 || tp->tm_mday == 23) { - f += "rd"; - } else { - f += "th"; - } - break; case 'e': - //mysql为月的天，strftime中<10时会带个空格 - f += std::to_string(tp->tm_mday); - break; case 'f': - //微妙数 - f += "000000"; - break; case 'h': case 'I': - if (hour12 < 10) { - f += "0"; - } - f += std::to_string(hour12); - break; case 'i': - f += "%M"; - break; case 'l': - f += std::to_string(hour12); - break; case 'M': - f += "%B"; - break; case 'p': - if (tp->tm_hour % 24 >= 12){ - f += "PM"; - } else { - f += "AM"; - } - break; case 'r': - //02:12:00 AM， - memset(tmp, 0, sizeof(tmp)); - snprintf(tmp, sizeof(tmp), "%02d:%02d:%02d ", hour12, tp->tm_min, tp->tm_sec); - f += tmp; - if (tp->tm_hour % 24 >= 12) { - f += "PM"; - } else { - f += "AM"; - } - break; case 'v': - f += "%V"; - break; case 'W': - f += "%A"; - break; case 'X': case 'x': - f += "%Y"; - break; case 's': - f += "%S"; - break; case 'k': - f += std::to_string(tp->tm_hour); - break; case 'u': - f += "%W"; + need_convert = true; break; - case 'Y': - case 'y': - case 'j': - case 'm': - case 'H': - default: - f += "%"; - f += format[i]; } i++; } - return strftime(s, maxsize, f.c_str(), tp); + if (need_convert) { + i = 0; + char tmp[20] = {0}; + int hour12 = tp->tm_hour % 12; + if (hour12 == 0) { + hour12 = 12; + } + while (format[i] != '\0') { + if (format[i] != '%' ) { + f += format[i++]; + continue; + } + i++; + if (format[i] == '\0') { + break; + } + switch (format[i]) { + case 'c': + f += std::to_string(tp->tm_mon + 1); + break; + case 'D': + f += std::to_string(tp->tm_mday); + if (tp->tm_mday == 1 || tp->tm_mday == 21 || tp->tm_mday == 31) { + f += "st"; + } + else if (tp->tm_mday == 2 || tp->tm_mday == 22) { + f += "nd"; + } + else if (tp->tm_mday == 3 || tp->tm_mday == 23) { + f += "rd"; + } else { + f += "th"; + } + break; + case 'e': + //mysql为月的天，strftime中<10时会带个空格 + f += std::to_string(tp->tm_mday); + break; + case 'f': + //微妙数 + f += "000000"; + break; + case 'h': + case 'I': + if (hour12 < 10) { + f += "0"; + } + f += std::to_string(hour12); + break; + case 'i': + f += "%M"; + break; + case 'l': + f += std::to_string(hour12); + break; + case 'M': + f += "%B"; + break; + case 'p': + if (tp->tm_hour % 24 >= 12){ + f += "PM"; + } else { + f += "AM"; + } + break; + case 'r': + //02:12:00 AM， + memset(tmp, 0, sizeof(tmp)); + snprintf(tmp, sizeof(tmp), "%02d:%02d:%02d ", hour12, tp->tm_min, tp->tm_sec); + f += tmp; + if (tp->tm_hour % 24 >= 12) { + f += "PM"; + } else { + f += "AM"; + } + break; + case 'v': + f += "%V"; + break; + case 'W': + f += "%A"; + break; + case 'X': + case 'x': + f += "%Y"; + break; + case 's': + f += "%S"; + break; + case 'k': + f += std::to_string(tp->tm_hour); + break; + case 'u': + f += "%W"; + break; + default: + f += "%"; + f += format[i]; + } + i++; + } + format = f.c_str(); + } + return strftime(s, maxsize, format, tp); } int64_t timestamp_to_ts(uint32_t timestamp) { return (((int64_t)timestamp) * 1000 - tso::base_timestamp_ms) << 18; } -bool is_leap_year(int year) { - return (year % 4 == 0 && year % 100 != 0) || year % 400 == 0; -} // 判断字符串是否为YYYYMMDD(不包含hour)/YYYYMMDDHH(包含hour)格式 bool is_valid_date(const std::string& date_str, const bool has_hour) { static const int32_t days_in_month[12] = {31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31}; @@ -724,7 +825,7 @@ bool is_valid_date(const std::string& date_str, const bool has_hour) { return false; } int mday_max = days_in_month[mon - 1]; - if (mon == 2 && is_leap_year(year)) { + if (mon == 2 && is_leap(year)) { mday_max++; } if (mday < 1 || mday > mday_max) { @@ -751,7 +852,7 @@ int date_add_interval(time_t& ts, const int64_t interval, const TimeUnit time_un case TimeUnit::YEAR: case TimeUnit::MONTH: { struct tm tm; - localtime_r(&ts, &tm); + localtime_fixed_r(&ts, &tm); int year = tm.tm_year + 1900; int mon = tm.tm_mon; int day = tm.tm_mday; @@ -761,7 +862,7 @@ int date_add_interval(time_t& ts, const int64_t interval, const TimeUnit time_un // 2024-02-29 -> 2023-02-28 // 2024-02-20 -> 2023-02-20 year += interval; - if (mon + 1 == 2 && day == 29 && !is_leap_year(year)) { + if (mon + 1 == 2 && day == 29 && !is_leap(year)) { day = 28; } } else { @@ -775,7 +876,7 @@ int date_add_interval(time_t& ts, const int64_t interval, const TimeUnit time_un mon = (mon % 12 + 12) % 12; if (day > days_in_month[mon]) { day = days_in_month[mon]; - if (mon + 1 == 2 && is_leap_year(year)) { + if (mon + 1 == 2 && is_leap(year)) { day++; } } @@ -783,7 +884,7 @@ int date_add_interval(time_t& ts, const int64_t interval, const TimeUnit time_un tm.tm_year = year - 1900; tm.tm_mon = mon; tm.tm_mday = day; - ts = mktime(&tm); + ts = mktime_fixed_r(&tm); break; } case TimeUnit::DAY: @@ -813,27 +914,27 @@ int get_current_day_timestamp(time_t& current_day_ts, time_t current_ts) { if (current_ts == -1) { current_ts = ::time(NULL); } - localtime_r(¤t_ts, &tm); + localtime_fixed_r(¤t_ts, &tm); tm.tm_hour = 0; tm.tm_min = 0; tm.tm_sec = 0; - current_day_ts = mktime(&tm); + current_day_ts = mktime_fixed_r(&tm); return 0; } int get_current_month_timestamp(const int start_day_of_month, time_t& current_month_ts) { struct tm tm; time_t current_ts = ::time(NULL); - localtime_r(¤t_ts, &tm); + localtime_fixed_r(¤t_ts, &tm); tm.tm_mday = start_day_of_month; tm.tm_hour = 0; tm.tm_min = 0; tm.tm_sec = 0; - current_month_ts = mktime(&tm); + current_month_ts = mktime_fixed_r(&tm); return 0; } @@ -842,7 +943,7 @@ int timestamp_to_format_str(const time_t ts, const char* format, std::string& st return -1; } struct tm tm; - localtime_r(&ts, &tm); + localtime_fixed_r(&ts, &tm); char str_time[21] = {0}; strftime(str_time, sizeof(str_time), format, &tm); diff --git a/src/common/default_room_define.cpp b/src/common/default_room_define.cpp index 9b965ad19..b69722646 100644 --- a/src/common/default_room_define.cpp +++ b/src/common/default_room_define.cpp @@ -16,7 +16,7 @@ #include namespace baikaldb{ -DEFINE_string(default_logical_room, "default", "default_logical_room"); -DEFINE_string(default_physical_room, "default", "default_physical_room"); +DEFINE_string(default_logical_room, "default", "Default logical room name, default: default"); +DEFINE_string(default_physical_room, "default", "Default physical room name, default: default"); } /* vim: set expandtab ts=4 sw=4 sts=4 tw=100: */ diff --git a/src/common/expr_value.cpp b/src/common/expr_value.cpp index 6a5f4cac3..4b2f2a9d0 100644 --- a/src/common/expr_value.cpp +++ b/src/common/expr_value.cpp @@ -16,7 +16,7 @@ #include "hll_common.h" namespace baikaldb { -DEFINE_bool(use_double_conversion, true, "use_double_conversion"); +DEFINE_bool(use_double_conversion, true, "Enable double conversion, default: true"); DEFINE_bool(double_use_all_precision, false, "Double precision output compatibility with MySQL, eg: 1.003*100=100.29999999999998"); SerializeStatus ExprValue::serialize_to_mysql_text_packet(char* buf, size_t size, size_t& len) const { if (size < 1) { diff --git a/src/common/file_system.cpp b/src/common/file_system.cpp index 65decd18d..840f891f6 100644 --- a/src/common/file_system.cpp +++ b/src/common/file_system.cpp @@ -3,8 +3,11 @@ namespace baikaldb { DEFINE_int32(file_buffer_size, 1, "read file buf size (MBytes)"); +DEFINE_int32(file_block_size, 100, "split file to block to handle(MBytes)"); -#ifdef BAIDU_INTERNAL +const char* AFS_CLIENT_CONF_PATH = "./conf/client.conf"; + +#if defined(BAIDU_INTERNAL) // AfsFileWriter int64_t AfsFileWriter::append(const char* buf, uint32_t count) { if (_writer == nullptr) { @@ -59,7 +62,7 @@ int64_t PosixFileReader::read(size_t pos, char* buf, size_t buf_size) { return size; } -#ifdef BAIDU_INTERNAL +#if defined(BAIDU_INTERNAL) // AfsFileReader int64_t AfsFileReader::seek(int64_t position) { if (_reader == nullptr) { @@ -158,6 +161,26 @@ int PosixFileSystem::read_dir(const std::string& path, std::vector& return 1; } +int PosixFileSystem::read_dir(const std::string& path, std::vector& direntrys) { + dir_iter iter(path); + dir_iter end; + for (; iter != end; ++iter) { + std::string child_path = iter->path().c_str(); + FileInfo file_info; + if (get_file_info(child_path, file_info, nullptr) != 0) { + return -1; + } + std::vector split_vec; + boost::split(split_vec, child_path, boost::is_any_of("/")); + std::string out_path = split_vec.back(); + DirEntry dir_entry; + dir_entry.mode = file_info.mode; + dir_entry.path = out_path; + direntrys.emplace_back(std::move(dir_entry)); + } + return 1; +} + int PosixFileSystem::get_file_info(const std::string& path, FileInfo& file_info, std::string* err_msg) { if (boost::filesystem::is_directory(path)) { file_info.mode = FileMode::I_DIR; @@ -173,7 +196,8 @@ int PosixFileSystem::get_file_info(const std::string& path, FileInfo& file_info, return 0; } -#ifdef BAIDU_INTERNAL + +#if defined(BAIDU_INTERNAL) // AfsFileSystem int AfsFileSystem::init() { // 创建一个AfsFileSystem实例 @@ -311,6 +335,35 @@ int AfsFileSystem::read_dir(const std::string& path, std::vector& d return 1; } +int AfsFileSystem::read_dir(const std::string& path, std::vector& direntrys) { + std::vector afs_entrys; + int afs_res = _afs->Readdir(path.c_str(), &afs_entrys); + if (afs_res < 0){ + DB_WARNING("fail to readdir %s, errno:%d, errmsg:%s\n", path.c_str(), afs_res, afs::Rc2Str(afs_res)); + return -1; + } + for (size_t i = 0; i < afs_entrys.size(); i++) { + FileMode mode; + afs::FileType file_type = afs::GetFileTypeFromInode(afs_entrys[i].inode); + if (file_type == afs::FT_DIRECTORY) { + mode = FileMode::I_DIR; + } else if (file_type == afs::FT_SYMLINK) { + mode = FileMode::I_LINK; + } else if (file_type == afs::FT_REGULAR) { + mode = FileMode::I_FILE; + } else { + DB_WARNING("Invalid file_type: %d", file_type); + return -1; + } + DirEntry dir_entry; + dir_entry.mode = mode; + dir_entry.path = afs_entrys[i].name; + direntrys.emplace_back(dir_entry); + } + return 1; +} + + int AfsFileSystem::get_file_info(const std::string& path, FileInfo& file_info, std::string* err_msg) { int afs_res = _afs->Exist(path.c_str()); if (afs_res < 0) { @@ -569,13 +622,54 @@ int ReadDirImpl::next_entry(std::string& entry){ return 0; } +int ReadDirImpl::get_all_files(FileSystem* fs, const std::string& path, std::vector& files) { + if (fs == nullptr) { + DB_WARNING("fs is nullptr"); + return -1; + } + std::vector direntrys; + int ret = fs->read_dir(path, direntrys); + if (ret < 0) { + DB_WARNING("Fail to read_dir, path: %s", path.c_str()); + return -1; + } + for (const auto& direntry : direntrys) { + if (direntry.mode == FileMode::I_FILE) { + files.emplace_back(direntry.path); + } + } + return 0; +} + +int ReadDirImpl::get_all_dirs(FileSystem* fs, const std::string& path, std::vector& dirs) { + if (fs == nullptr) { + DB_WARNING("fs is nullptr"); + return -1; + } + std::vector direntrys; + int ret = fs->read_dir(path, direntrys); + if (ret < 0) { + DB_WARNING("Fail to read_dir, path: %s", path.c_str()); + return -1; + } + for (const auto& direntry : direntrys) { + if (direntry.mode == FileMode::I_DIR) { + dirs.emplace_back(direntry.path); + } + } + return 0; +} + std::shared_ptr create_filesystem(const std::string& cluster_name, const std::string& user_name, const std::string& password, const std::string& conf_file) { std::shared_ptr fs; if (cluster_name.find("afs") != std::string::npos) { -#ifdef BAIDU_INTERNAL +#if !defined(BAIDU_INTERNAL) + DB_FATAL("doesn't support AfsFileSystem!"); + return nullptr; +#else fs.reset(new (std::nothrow) AfsFileSystem(cluster_name, user_name, password, conf_file)); #endif } else { diff --git a/src/common/information_schema.cpp b/src/common/information_schema.cpp index 0fa6b8eb9..b6a5c89d6 100644 --- a/src/common/information_schema.cpp +++ b/src/common/information_schema.cpp @@ -24,6 +24,7 @@ #include "external_filesystem.h" namespace baikaldb { +DECLARE_bool(sign_with_resouce_tag); int InformationSchema::init() { init_partition_split_info(); init_region_status(); @@ -323,7 +324,7 @@ void InformationSchema::init_region_status() { record->set_string(record->get_field_by_name("peers"), boost::join(region.peers(), ",")); time_t t = region.timestamp(); struct tm t_result; - localtime_r(&t, &t_result); + localtime_fixed_r(&t, &t_result); char s[100]; strftime(s, sizeof(s), "%F %T", &t_result); record->set_string(record->get_field_by_name("create_time"), s); @@ -1165,7 +1166,7 @@ void InformationSchema::init_afs_partitions() { } std::map>> table_id_name_partitions; -#ifdef BAIDU_INTERNAL +#if defined(BAIDU_INTERNAL) int ret = ExtFileSystemGC::get_all_partitions_from_store(table_id_name_partitions); if (ret < 0) { DB_WARNING("get all partitions from store failed"); @@ -1786,6 +1787,8 @@ void InformationSchema::init_sign_list() { {"database_name",pb::STRING}, {"table_name",pb::STRING}, {"sign",pb::STRING}, + {"sign_without_resource_tag",pb::STRING}, + {"sql",pb::STRING} }; int64_t blacklist_table_id = construct_table("SIGN_BLACKLIST", fields); @@ -1799,7 +1802,9 @@ void InformationSchema::init_sign_list() { std::vector records; records.reserve(10); auto blacklist_table = SchemaFactory::get_instance()->get_table_info_ptr(blacklist_table_id); - auto func = [&records, &blacklist_table](const SmartTable& table) -> bool { + std::unordered_map special_signs; + SchemaFactory::get_instance()->get_special_signs(special_signs); + auto func = [&records, &blacklist_table, &special_signs](const SmartTable& table) -> bool { for (auto sign : table->sign_blacklist) { auto record = SchemaFactory::get_instance()->new_record(*blacklist_table); record->set_string(record->get_field_by_name("namespace"), table->namespace_); @@ -1812,6 +1817,15 @@ void InformationSchema::init_sign_list() { record->set_string(record->get_field_by_name("database_name"), db_name); record->set_string(record->get_field_by_name("table_name"),table->short_name); record->set_string(record->get_field_by_name("sign"), std::to_string(sign)); + if (special_signs.count(sign) > 0 && special_signs[sign] != "") { + uint64_t out[2]; + butil::MurmurHash3_x64_128(special_signs[sign].c_str(), special_signs[sign].size(), 0x1234, out); + record->set_string(record->get_field_by_name("sign_without_resource_tag"), std::to_string(out[0])); + record->set_string(record->get_field_by_name("sql"), special_signs[sign]); + } else { + record->set_string(record->get_field_by_name("sign_without_resource_tag"), "NOT FOUND"); + record->set_string(record->get_field_by_name("sql"), "NOT FOUND"); + } records.emplace_back(record); } return false; @@ -1826,7 +1840,9 @@ void InformationSchema::init_sign_list() { std::vector records; records.reserve(10); auto forcelearner_table = SchemaFactory::get_instance()->get_table_info_ptr(forcelearner_table_id); - auto func = [&records, &forcelearner_table](const SmartTable& table) -> bool { + std::unordered_map special_signs; + SchemaFactory::get_instance()->get_special_signs(special_signs); + auto func = [&records, &forcelearner_table, &special_signs](const SmartTable& table) -> bool { for (auto sign : table->sign_forcelearner) { auto record = SchemaFactory::get_instance()->new_record(*forcelearner_table); record->set_string(record->get_field_by_name("namespace"), table->namespace_); @@ -1839,6 +1855,15 @@ void InformationSchema::init_sign_list() { record->set_string(record->get_field_by_name("database_name"), db_name); record->set_string(record->get_field_by_name("table_name"),table->short_name); record->set_string(record->get_field_by_name("sign"), std::to_string(sign)); + if (special_signs.count(sign) > 0 && special_signs[sign] != "") { + uint64_t out[2]; + butil::MurmurHash3_x64_128(special_signs[sign].c_str(), special_signs[sign].size(), 0x1234, out); + record->set_string(record->get_field_by_name("sign_without_resource_tag"), std::to_string(out[0])); + record->set_string(record->get_field_by_name("sql"), special_signs[sign]); + } else { + record->set_string(record->get_field_by_name("sign_without_resource_tag"), "NOT FOUND"); + record->set_string(record->get_field_by_name("sql"), "NOT FOUND"); + } records.emplace_back(record); } return false; @@ -1853,7 +1878,9 @@ void InformationSchema::init_sign_list() { std::vector records; records.reserve(10); auto rolling_table = SchemaFactory::get_instance()->get_table_info_ptr(rolling_table_id); - auto func = [&records, &rolling_table](const SmartTable& table) -> bool { + std::unordered_map special_signs; + SchemaFactory::get_instance()->get_special_signs(special_signs); + auto func = [&records, &rolling_table, &special_signs](const SmartTable& table) -> bool { for (auto sign : table->sign_rolling) { auto record = SchemaFactory::get_instance()->new_record(*rolling_table); record->set_string(record->get_field_by_name("namespace"), table->namespace_); @@ -1866,33 +1893,15 @@ void InformationSchema::init_sign_list() { record->set_string(record->get_field_by_name("database_name"), db_name); record->set_string(record->get_field_by_name("table_name"),table->short_name); record->set_string(record->get_field_by_name("sign"), std::to_string(sign)); - records.emplace_back(record); - } - return false; - }; - std::vector database_table; - SchemaFactory::get_instance()->get_table_by_filter(database_table, func); - return records; - }; - - _calls[forceindex_table_id] = [forceindex_table_id](RuntimeState* state,std::vector& conditions) -> - std::vector { - std::vector records; - records.reserve(10); - auto forceindex_table = SchemaFactory::get_instance()->get_table_info_ptr(forceindex_table_id); - auto func = [&records, &forceindex_table](const SmartTable& table) -> bool { - for (auto sign_index : table->sign_forceindex) { - auto record = SchemaFactory::get_instance()->new_record(*forceindex_table); - record->set_string(record->get_field_by_name("namespace"), table->namespace_); - std::string db_name; - std::vector vec; - boost::split(vec, table->name, boost::is_any_of(".")); - if (!vec.empty()) { - db_name = vec[0]; + if (special_signs.count(sign) > 0 && special_signs[sign] != "") { + uint64_t out[2]; + butil::MurmurHash3_x64_128(special_signs[sign].c_str(), special_signs[sign].size(), 0x1234, out); + record->set_string(record->get_field_by_name("sign_without_resource_tag"), std::to_string(out[0])); + record->set_string(record->get_field_by_name("sql"), special_signs[sign]); + } else { + record->set_string(record->get_field_by_name("sign_without_resource_tag"), "NOT FOUND"); + record->set_string(record->get_field_by_name("sql"), "NOT FOUND"); } - record->set_string(record->get_field_by_name("database_name"), db_name); - record->set_string(record->get_field_by_name("table_name"),table->short_name); - record->set_string(record->get_field_by_name("sign"), sign_index); records.emplace_back(record); } return false; @@ -1907,7 +1916,9 @@ void InformationSchema::init_sign_list() { std::vector records; records.reserve(10); auto forceindex_table = SchemaFactory::get_instance()->get_table_info_ptr(forceindex_table_id); - auto func = [&records, &forceindex_table](const SmartTable& table) -> bool { + std::unordered_map special_signs; + SchemaFactory::get_instance()->get_special_signs(special_signs); + auto func = [&records, &forceindex_table, &special_signs](const SmartTable& table) -> bool { for (auto sign_index : table->sign_forceindex) { auto record = SchemaFactory::get_instance()->new_record(*forceindex_table); record->set_string(record->get_field_by_name("namespace"), table->namespace_); @@ -1920,6 +1931,21 @@ void InformationSchema::init_sign_list() { record->set_string(record->get_field_by_name("database_name"), db_name); record->set_string(record->get_field_by_name("table_name"),table->short_name); record->set_string(record->get_field_by_name("sign"), sign_index); + vec.clear(); + boost::split(vec, sign_index, boost::is_any_of(":")); + uint64_t sign = 0; + if (vec.size() == 2) { + sign = strtoull(vec[0].c_str(), NULL, 10); + } + if (special_signs.count(sign) > 0 && special_signs[sign] != "") { + uint64_t out[2]; + butil::MurmurHash3_x64_128(special_signs[sign].c_str(), special_signs[sign].size(), 0x1234, out); + record->set_string(record->get_field_by_name("sign_without_resource_tag"), std::to_string(out[0])); + record->set_string(record->get_field_by_name("sql"), special_signs[sign]); + } else { + record->set_string(record->get_field_by_name("sign_without_resource_tag"), "NOT FOUND"); + record->set_string(record->get_field_by_name("sql"), "NOT FOUND"); + } records.emplace_back(record); } return false; @@ -1935,7 +1961,9 @@ void InformationSchema::init_sign_list() { std::vector records; records.reserve(10); auto exec_type_table = SchemaFactory::get_instance()->get_table_info_ptr(exec_type_table_id); - auto func = [&records, &exec_type_table](const SmartTable& table) -> bool { + std::unordered_map special_signs; + SchemaFactory::get_instance()->get_special_signs(special_signs); + auto func = [&records, &exec_type_table, &special_signs](const SmartTable& table) -> bool { for (auto sign_index : table->sign_exec_type) { auto record = SchemaFactory::get_instance()->new_record(*exec_type_table); record->set_string(record->get_field_by_name("namespace"), table->namespace_); @@ -1949,11 +1977,22 @@ void InformationSchema::init_sign_list() { record->set_string(record->get_field_by_name("table_name"),table->short_name); vec.clear(); boost::split(vec, sign_index, boost::is_any_of(":")); + uint64_t exec_type = 0; + uint64_t sign = 0; if (vec.size() == 2) { - uint64_t exec_type = strtoll(vec[1].c_str(), NULL, 10); - sign_index = vec[0] + ": " + explain_type_to_str(exec_type); + exec_type = strtoull(vec[1].c_str(), NULL, 10); + sign = strtoull(vec[0].c_str(), NULL, 10); + } + record->set_string(record->get_field_by_name("sign"), std::to_string(sign) + ": " + explain_type_to_str(exec_type)); + if (special_signs.count(sign) > 0 && special_signs[sign] != "") { + uint64_t out[2]; + butil::MurmurHash3_x64_128(special_signs[sign].c_str(), special_signs[sign].size(), 0x1234, out); + record->set_string(record->get_field_by_name("sign_without_resource_tag"), std::to_string(out[0])); + record->set_string(record->get_field_by_name("sql"), special_signs[sign]); + } else { + record->set_string(record->get_field_by_name("sign_without_resource_tag"), "NOT FOUND"); + record->set_string(record->get_field_by_name("sql"), "NOT FOUND"); } - record->set_string(record->get_field_by_name("sign"), sign_index); records.emplace_back(record); } return false; diff --git a/src/common/memory_profile.cpp b/src/common/memory_profile.cpp index a0bd4f5c6..04ec2e47a 100755 --- a/src/common/memory_profile.cpp +++ b/src/common/memory_profile.cpp @@ -13,6 +13,7 @@ // limitations under the License. #include "memory_profile.h" +#include "arrow/memory_pool.h" #ifdef BAIKAL_TCMALLOC #include #endif @@ -70,6 +71,10 @@ void MemoryGCHandler::memory_gc_thread() { DB_WARNING("tcmalloc release memory about size: %ld cast: %ld", total_bytes_to_gc, cost.get_time()); } } + + SQL_TRACE("arrow DefaultMemoryPool allocate: %ld, LimitMemoryPool allocate: %ld", + arrow::default_memory_pool()->bytes_allocated(), + arrow::system_memory_pool_with_limit()->bytes_allocated()); bthread_usleep_fast_shutdown(FLAGS_memory_gc_interval_s * 1000 * 1000LL, _shutdown); } #endif diff --git a/src/common/meta_server_interact.cpp b/src/common/meta_server_interact.cpp index 9dda3a567..848d8dddc 100644 --- a/src/common/meta_server_interact.cpp +++ b/src/common/meta_server_interact.cpp @@ -22,7 +22,7 @@ DEFINE_int32(meta_request_timeout, 30000, DEFINE_int32(meta_connect_timeout, 5000, "meta as server connect timeout, default:5000ms"); DEFINE_string(meta_server_bns, "group.opera-qa-baikalMeta-000-yz.FENGCHAO.all", "meta server bns"); -DEFINE_string(backup_meta_server_bns, "", "backup_meta_server_bns"); +DEFINE_string(backup_meta_server_bns, "", "Backup meta server BNS name, default: empty"); DEFINE_int64(time_between_meta_connect_error_ms, 0, "time_between_meta_connect_error_ms. default(0ms)"); DEFINE_bool(auto_update_meta_list, false, "auto_update_meta_list, default false"); diff --git a/src/common/schema_factory.cpp b/src/common/schema_factory.cpp index a61362726..956658f21 100644 --- a/src/common/schema_factory.cpp +++ b/src/common/schema_factory.cpp @@ -30,13 +30,16 @@ using google::protobuf::FileDescriptor; namespace baikaldb { -DEFINE_bool(need_health_check, true, "need_health_check"); +DEFINE_bool(need_health_check, true, "Enable health checking, default: true"); DECLARE_int32(baikal_faulty_interval_times); DECLARE_int32(baikal_heartbeat_interval_us); DECLARE_string(meta_server_bns); DECLARE_int32(baikal_port); DECLARE_int64(mpp_min_statistics_rows); DECLARE_int64(mpp_min_statistics_bytes); +DEFINE_int64(broadcast_table_max_rows, 500000, "broadcast_table_max_rows, -1 mean no use broadcast join"); +DEFINE_int64(broadcast_table_max_mb, 4, "broadcast_table_max_mb, -1 mean no use broadcast join"); + BthreadLocal SchemaFactory::use_backup; int SchemaFactory::init(bool is_db, bool is_backup) { if (_is_inited) { @@ -475,6 +478,7 @@ int SchemaFactory::update_table_internal(SchemaMapping& background, const pb::Sc tbl_info.arrow_reverse_fields.clear(); tbl_info.has_global_not_none = false; tbl_info.has_rollup_index = false; + tbl_info.rollup_indexs.clear(); tbl_info.has_index_write_only_or_write_local = false; tbl_info.sign_blacklist.clear(); tbl_info.sign_forcelearner.clear(); @@ -616,18 +620,6 @@ int SchemaFactory::update_table_internal(SchemaMapping& background, const pb::Sc tbl_info.region_num = table.region_num(); } - if (table.has_ttl_duration()) { - tbl_info.ttl_info.ttl_duration_s = table.ttl_duration(); - if (table.has_online_ttl_expire_time_us()) { - tbl_info.ttl_info.online_ttl_expire_time_us = table.online_ttl_expire_time_us(); - } - - DB_WARNING("table:%s ttl_duration:%ld, online_ttl_expire_time_us:%ld, %s", - tbl_info.name.c_str(), tbl_info.ttl_info.ttl_duration_s, - tbl_info.ttl_info.online_ttl_expire_time_us, - timestamp_to_str(tbl_info.ttl_info.online_ttl_expire_time_us / 1000000).c_str()); - } - tbl_info.learner_resource_tags.clear(); for (auto& learner_resource : table.learner_resource_tags()) { tbl_info.learner_resource_tags.emplace_back(learner_resource); @@ -741,6 +733,35 @@ int SchemaFactory::update_table_internal(SchemaMapping& background, const pb::Sc tbl_info.fields.push_back(field_info); //DB_WARNING("field_name:%s, field_id:%d", field_info.name.c_str(), field_info.id); } + + if (table.has_ttl_duration()) { + // 只有第一次设置ttl info允许修改ttl field + if (table.has_ttl_field() && !table.ttl_field().field_name().empty() + && tbl_info.ttl_info.ttl_duration_s <= 0 && tbl_info.ttl_info.ttl_field == nullptr) { + std::string ttl_field_name = table.ttl_field().field_name(); + auto iter = std::find_if(tbl_info.fields.begin(), tbl_info.fields.end(), + [ttl_field_name](const FieldInfo& field_info) {return field_info.short_name == ttl_field_name;}); + if (iter == tbl_info.fields.end()) { + // should never happen + DB_FATAL("ttl field not exist! table_id: %ld, field_name: %s", table_id, ttl_field_name.c_str()); + return -1; + } + tbl_info.ttl_info.ttl_field = std::make_shared(*iter); + } + + tbl_info.ttl_info.ttl_duration_s = table.ttl_duration(); + if (table.has_online_ttl_expire_time_us()) { + tbl_info.ttl_info.online_ttl_expire_time_us = table.online_ttl_expire_time_us(); + } + + int ttl_field_id = tbl_info.ttl_info.ttl_field == nullptr ? -1 : tbl_info.ttl_info.ttl_field->id; + + DB_WARNING("table:%s ttl_duration:%ld, ttl_field_id: %d, online_ttl_expire_time_us:%ld, %s", + tbl_info.name.c_str(), tbl_info.ttl_info.ttl_duration_s, ttl_field_id, + tbl_info.ttl_info.online_ttl_expire_time_us, + timestamp_to_str(tbl_info.ttl_info.online_ttl_expire_time_us / 1000000).c_str()); + } + tbl_info.link_field_map.clear(); tbl_info.binlog_target_ids.clear(); tbl_info.is_linked = false; @@ -921,6 +942,7 @@ int SchemaFactory::update_table_internal(SchemaMapping& background, const pb::Sc tbl_info.has_fulltext = true; } else if (cur.index_type() == pb::I_ROLLUP) { tbl_info.has_rollup_index = true; + tbl_info.rollup_indexs.insert(index_id); } else if (cur.index_type() == pb::I_VECTOR) { tbl_info.has_vector_index = true; } @@ -1721,9 +1743,15 @@ void SchemaFactory::update_statistics(const StatisticsVec& statistics) { int SchemaFactory::update_statistics_internal(SchemaMapping& background, const std::map& mapping) { auto& table_statistics_mapping = background.table_statistics_mapping; + auto& tables = background.table_info_mapping; for (auto iter = mapping.begin(); iter != mapping.end(); iter++) { auto origin_iter = table_statistics_mapping.find(iter->first); + auto table_ptr_pair = tables.find(iter->first); + if (table_ptr_pair == tables.end()) { + DB_WARNING("update table statistics but table not found, table_id: %ld", iter->first); + continue; + } if (origin_iter != table_statistics_mapping.end()) { if (iter->second->version() > origin_iter->second->version()) { table_statistics_mapping[iter->first] = iter->second; @@ -1731,6 +1759,7 @@ int SchemaFactory::update_statistics_internal(SchemaMapping& background, const s } else { table_statistics_mapping[iter->first] = iter->second; } + table_ptr_pair->second->have_statistics = table_statistics_mapping[iter->first]->is_valid(); } return 1; @@ -1824,7 +1853,7 @@ double SchemaFactory::get_eq_field_ratio(int64_t table_id, int field_id, const E auto& table_statistics_mapping = table_ptr->table_statistics_mapping; auto iter = table_statistics_mapping.find(table_id); if (iter != table_statistics_mapping.end()) { - if (iter->second->is_cms_exist()) { + if (iter->second->is_cms_exist(field_id)) { return iter->second->get_cmsketch_ratio(field_id, value); } else { return iter->second->get_hll_ratio(field_id); @@ -1881,6 +1910,9 @@ void SchemaFactory::table_with_statistics_info(std::vector& databas if (meta_id != 0) { continue; } + if (!st.second->is_valid()) { + continue; + } auto table = table_info_mapping.find(st.first); if (table != table_info_mapping.end()) { database_table.push_back(table->second->namespace_ + "." + table->second->name); @@ -2598,6 +2630,9 @@ void SchemaFactory::get_table_by_filter(std::vector& database_table dblink_info = table.second->dblink_info.mysql_info().addr() + "/" + table.second->dblink_info.mysql_info().database_name() + "." + table.second->dblink_info.mysql_info().table_name(); + } else if (table.second->dblink_info.type() == pb::LT_FILE) { + dblink_info = table.second->dblink_info.file_info().cluster() + "/" + + table.second->dblink_info.file_info().path(); } else { DB_WARNING("Invalid dblink type: %d, table: %s", table.second->dblink_info.type(), table.second->name.c_str()); @@ -2937,6 +2972,27 @@ int SchemaFactory::get_partition_regions(int64_t table_id, const std::vector& partition_ids) { + DoubleBufferedTableRegionInfo::ScopedPtr table_region_mapping_ptr; + if (_table_region_mapping.Read(&table_region_mapping_ptr) != 0) { + DB_WARNING("DoubleBufferedTableRegion read scoped ptr error."); + return -1; + } + auto it = table_region_mapping_ptr->find(table_id); + if (it == table_region_mapping_ptr->end()) { + DB_WARNING("index id[%ld] not in table_region_mapping", table_id); + return -1; + } + auto frontground = it->second; + auto& key_region_mapping = frontground->key_region_mapping; + for (const auto& [partition_id, _] : key_region_mapping) { + partition_ids.emplace_back(partition_id); + } + return 0; +} + // 检测table下region范围是否连续没有空洞 int SchemaFactory::check_region_ranges_consecutive(int64_t table_id) { DoubleBufferedTableRegionInfo::ScopedPtr table_region_mapping_ptr; @@ -3049,9 +3105,9 @@ int SchemaFactory::get_region_by_key(int64_t main_table_id, int range_size = primary->ranges_size(); for (int i = 0; i < range_size; ++i) { const auto& range = primary->ranges(i); - bool like_prefix = template_primary.has_like_prefix() ? template_primary.like_prefix() : range.like_prefix(); - bool left_open = template_primary.has_left_open() ? template_primary.left_open() : range.left_open(); - bool right_open = template_primary.has_right_open() ? template_primary.right_open() : range.right_open(); + bool like_prefix = template_primary.has_like_prefix() ? template_primary.like_prefix() : range.like_prefix(); + bool left_open = template_primary.has_left_open() ? template_primary.left_open() : range.left_open(); + bool right_open = template_primary.has_right_open() ? template_primary.right_open() : range.right_open(); MutTableKey start; MutTableKey end; if (!range.left_key().empty()) { @@ -3826,6 +3882,37 @@ int SchemaFactory::get_binlog_regions(int64_t binlog_id, int64_t partition_index return -1; } } + +bool SchemaFactory::table_suitable_for_broadcast_join(int64_t table_id) { + if (FLAGS_broadcast_table_max_rows <= 0 && FLAGS_broadcast_table_max_mb <= 0) { + return false; + } + DoubleBufferedTableRegionInfo::ScopedPtr table_region_mapping_ptr; + if (_table_region_mapping.Read(&table_region_mapping_ptr) != 0) { + DB_WARNING("DoubleBufferedTableRegion read scoped ptr error."); + return false; + } + auto it = table_region_mapping_ptr->find(table_id); + if (it == table_region_mapping_ptr->end()) { + DB_WARNING("table id[%ld] not in table_region_mapping.", table_id); + return false; + } + auto& region_map = it->second->region_info_mapping; + if (region_map.size() > 1) { + return false; + } + // TODO, 下面不是实时同步, 启动/version变更才会同步, 后期可以考虑如何同步小表的统计信息 + int64_t total_rows = 0; + int64_t total_bytes = 0; + for (auto& [_, region] : region_map) { + total_rows += region.region_info.num_table_lines(); + total_bytes += region.region_info.used_size(); + } + if (total_rows <= FLAGS_broadcast_table_max_rows || total_bytes <= FLAGS_broadcast_table_max_mb * 1024 * 1024ULL) { + return true; + } + return false; +} }//namespace /* vim: set expandtab ts=4 sw=4 sts=4 tw=100: */ diff --git a/src/engine/external_filesystem.cpp b/src/engine/external_filesystem.cpp index f0da7839c..5bc352427 100644 --- a/src/engine/external_filesystem.cpp +++ b/src/engine/external_filesystem.cpp @@ -23,13 +23,13 @@ namespace baikaldb { DECLARE_string(meta_server_bns); DECLARE_string(cold_rocksdb_afs_infos); DEFINE_int64(afs_double_read_interval_us, 1000 * 1000LL, "afs_double_read_interval_us"); -DEFINE_bool(afs_open_reader_async_switch, true, "afs_open_reader_async_switch"); +DEFINE_bool(afs_open_reader_async_switch, true, "Enable async AFS reader, default: true"); DEFINE_int64(afs_gc_interval_s, 24 * 3600LL, "default 1 day"); DEFINE_int64(afs_gc_count, 10, "afs_gc_count"); DEFINE_int64(afs_gc_delay_days, 30, "afs_gc_delay_days"); DEFINE_int64(afs_gc_allow_dead_store_count, 3, "afs_gc_allow_dead_store_count"); -DEFINE_bool(afs_gc_enable, false, "afs_gc_enable"); -DEFINE_bool(need_ext_fs_gc, false, "need_ext_fs_gc"); +DEFINE_bool(afs_gc_enable, false, "Enable AFS garbage collection, default: false"); +DEFINE_bool(need_ext_fs_gc, false, "Need external filesystem garbage collection, default: false"); DEFINE_int64(compaction_sst_cache_max_block, 8192, "compaction_sst_cache_max_block"); int get_size_by_external_file_name(uint64_t* size, uint64_t* lines, const std::string& external_file) { @@ -89,8 +89,7 @@ int get_size_by_external_file_name(uint64_t* size, uint64_t* lines, const std::s } return 0; } -#ifdef BAIDU_INTERNAL - +#if defined(BAIDU_INTERNAL) // uri,user,password,conf_file,root_path 多组afs ugi使用英文分号分割用户名密码等信息使用英文逗号分割 int get_afs_infos(std::vector& ugi_infos) { ugi_infos.clear(); @@ -347,7 +346,6 @@ bool AfsExtFileWriter::close() { return all_succ ? true : false; } - #endif int64_t CompactionExtFileReader::read(char* buf, uint32_t count, uint32_t offset, bool* eof) { @@ -490,8 +488,7 @@ bool CompactionExtFileWriter::close() { return true; } -#ifdef BAIDU_INTERNAL - +#if defined(BAIDU_INTERNAL) AfsExtFileSystem::~AfsExtFileSystem() { for (auto& info : _ugi_infos) { if (info.afs != nullptr) { @@ -924,7 +921,6 @@ int AfsExtFileSystem::readdir(const std::string& full_name, std::setSetBackgroundThreads(2, rocksdb::Env::HIGH); db_options.listeners.emplace_back(my_listener); - if (FLAGS_olap_import_mode) { - db_options.max_background_flushes = 4; - db_options.env->SetBackgroundThreads(4, rocksdb::Env::HIGH); - db_options.allow_concurrent_memtable_write = false; - db_options.memtable_factory.reset(new rocksdb::VectorRepFactory(1024*1024)); - } if (FLAGS_enable_remote_compaction) { db_options.compaction_service.reset(new MyCompactionService(path, FLAGS_compaction_server_bns)); } @@ -259,8 +252,9 @@ int32_t RocksWrapper::init(const std::string& path) { if (FLAGS_enable_bottommost_compression) { _data_cf_option.bottommost_compression_opts.enabled = true; _data_cf_option.bottommost_compression = rocksdb::kZSTD; - _data_cf_option.bottommost_compression_opts.max_dict_bytes = 1 << 14; // 16KB - _data_cf_option.bottommost_compression_opts.zstd_max_train_bytes = 1 << 18; // 256KB + _data_cf_option.bottommost_compression_opts.max_dict_bytes = FLAGS_max_dict_bytes; + _data_cf_option.bottommost_compression_opts.zstd_max_train_bytes = FLAGS_zstd_max_train_bytes; + _data_cf_option.bottommost_compression_opts.max_dict_buffer_bytes = 1 << 22; // 4MB 防止SstFileWriter内存太大 } // 等待心跳成功后再开启rocksdb compaction @@ -278,60 +272,6 @@ int32_t RocksWrapper::init(const std::string& path) { _meta_info_option.compaction_pri = rocksdb::kOldestSmallestSeqFirst; _meta_info_option.level_compaction_dynamic_level_bytes = FLAGS_rocks_data_dynamic_level_bytes; _meta_info_option.max_write_buffer_number_to_maintain = _meta_info_option.max_write_buffer_number; - if (FLAGS_olap_import_mode) { - _log_cf_option.max_write_buffer_number_to_maintain = 0; - _log_cf_option.min_write_buffer_number_to_merge = 1; - _binlog_cf_option.max_write_buffer_number_to_maintain = 0; - _binlog_cf_option.min_write_buffer_number_to_merge = 1; - _data_cf_option.max_write_buffer_number_to_maintain = 0; - _data_cf_option.min_write_buffer_number_to_merge = 1; - _meta_info_option.max_write_buffer_number_to_maintain = 0; - _meta_info_option.min_write_buffer_number_to_merge = 1; - _data_cf_option.compression_per_level = {rocksdb::CompressionType::kLZ4Compression, - rocksdb::CompressionType::kLZ4Compression, - rocksdb::CompressionType::kLZ4Compression, - rocksdb::CompressionType::kLZ4Compression, - rocksdb::CompressionType::kLZ4Compression, - rocksdb::CompressionType::kLZ4Compression, - rocksdb::CompressionType::kLZ4Compression}; - - // 参考rocksdb PrepareForBulkLoad - // never slowdown ingest. - _data_cf_option.level0_file_num_compaction_trigger = (1<<30); - _data_cf_option.level0_slowdown_writes_trigger = (1<<30); - _data_cf_option.level0_stop_writes_trigger = (1<<30); - _data_cf_option.soft_pending_compaction_bytes_limit = 0; - _data_cf_option.hard_pending_compaction_bytes_limit = 0; - - // no auto compactions please. The application should issue a - // manual compaction after all data is loaded into L0. - _data_cf_option.disable_auto_compactions = true; - // A manual compaction run should pick all files in L0 in - // a single compaction run. - _data_cf_option.max_compaction_bytes = (static_cast(1) << 60); - - // It is better to have only 2 levels, otherwise a manual - // compaction would compact at every possible level, thereby - // increasing the total time needed for compactions. - // num_levels = 2; - - // Need to allow more write buffers to allow more parallism - // of flushes. - _data_cf_option.max_write_buffer_number = 6; - _data_cf_option.min_write_buffer_number_to_merge = 1; - - // When compaction is disabled, more parallel flush threads can - // help with write throughput. - // _data_cf_option.max_background_flushes = 4; - - // Prevent a memtable flush to automatically promote files - // to L1. This is helpful so that all files that are - // input to the manual compaction are all at L0. - // _data_cf_option.max_background_compactions = 2; - - // The compaction would create large files in L1. - _data_cf_option.target_file_size_base = 256 * 1024 * 1024; - } _db_path = path; // List Column Family @@ -507,7 +447,7 @@ int32_t RocksWrapper::init_cold_rocksdb(const std::string& path) { _cold_option.compression = rocksdb::kZSTD; _cold_option.compression_opts.max_dict_bytes = FLAGS_max_dict_bytes; _cold_option.compression_opts.zstd_max_train_bytes = FLAGS_zstd_max_train_bytes; - _cold_option.compression_opts.max_dict_buffer_bytes = 128 * 1024 * 1024ULL; + _cold_option.compression_opts.max_dict_buffer_bytes = 1 << 22; // 4MB 防止SstFileWriter内存太大 } else { _cold_option.compression = rocksdb::kLZ4Compression; } @@ -631,9 +571,6 @@ void RocksWrapper::set_table_options(rocksdb::BlockBasedTableOptions& table_opti // olap集群关闭bloomfilter table_options.filter_policy = nullptr; } - if (rocksdb_gflags.olap_import_mode()) { - table_options.filter_policy = nullptr; - } return; } diff --git a/src/engine/rocksdb_compaction_service.cpp b/src/engine/rocksdb_compaction_service.cpp index f6cf23e06..babe3721e 100644 --- a/src/engine/rocksdb_compaction_service.cpp +++ b/src/engine/rocksdb_compaction_service.cpp @@ -26,7 +26,6 @@ void set_rocksdb_flags(pb::RocksdbGFLAGS* rocksdb_gflags) { rocksdb_gflags->set_rocks_use_partitioned_index_filters(FLAGS_rocks_use_partitioned_index_filters); rocksdb_gflags->set_rocks_use_ribbon_filter(FLAGS_rocks_use_ribbon_filter); rocksdb_gflags->set_olap_table_only(FLAGS_olap_table_only); - rocksdb_gflags->set_olap_import_mode(FLAGS_olap_import_mode); rocksdb_gflags->set_rocks_use_sst_partitioner_fixed_prefix(FLAGS_rocks_use_sst_partitioner_fixed_prefix); rocksdb_gflags->set_key_point_collector_interval(FLAGS_key_point_collector_interval); rocksdb_gflags->set_rocks_block_cache_size_mb(FLAGS_rocks_block_cache_size_mb); @@ -108,6 +107,7 @@ rocksdb::CompactionServiceScheduleResponse MyCompactionService::Schedule(const r _time_cost_map[remote_compaction_id_str] = TimeCost(); _doing_map[remote_compaction_id_str] = closure; } + DB_WARNING("start remote_compaction: %s", remote_compaction_id_str.c_str()); // 异步调用 RPC，并传入 OnRPCDone 回调 stub.do_compaction(&closure->cntl, &request, &closure->_response, closure.get()); #endif @@ -165,12 +165,14 @@ rocksdb::CompactionServiceJobStatus MyCompactionService::Wait(const std::string& } }); - brpc::CallId call_id; + brpc::CallId call_id = {0}; pb::RemoteCompactionResponse response; { BAIDU_SCOPED_LOCK(_mutex); if (_doing_map.find(remote_compaction_id_str) != _doing_map.end()) { call_id = _doing_map[remote_compaction_id_str]->cntl.call_id(); + } else { + DB_FATAL("not find remote_compaction_id: %s", remote_compaction_id_str.c_str()); } } diff --git a/src/engine/rocksdb_merge_operator.cpp b/src/engine/rocksdb_merge_operator.cpp index 133f1a0ce..2e67d3cc0 100644 --- a/src/engine/rocksdb_merge_operator.cpp +++ b/src/engine/rocksdb_merge_operator.cpp @@ -14,6 +14,7 @@ #include "rocksdb_merge_operator.h" namespace baikaldb { +DEFINE_bool(olap_enable_partial_merge, false, "Enable partial merge for OLAP, default: false"); bool OLAPMergeOperator::FullMergeV2(const rocksdb::MergeOperator::MergeOperationInput& merge_in, rocksdb::MergeOperator::MergeOperationOutput* merge_out) const { // 扫表或compaction时同一个线程大量处理同一个table的merge，可以thread local将table信息缓存 OLAPTODO @@ -30,7 +31,7 @@ bool OLAPMergeOperator::FullMergeV2(const rocksdb::MergeOperator::MergeOperation static thread_local TimeCost table_cache_time; static TimeCost print_log_time; static bool first_print_log = true; //确保启动时第一次能打印日志 - if (table_info == nullptr || table_info->id != table_id || table_cache_time.get_time() > 60 * 1000 * 1000LL) { + if (table_info == nullptr || (table_info->id != table_id && table_info->rollup_indexs.count(table_id) <= 0) || table_cache_time.get_time() > 60 * 1000 * 1000LL) { auto info = factory->get_table_info_ptr(table_id); if (info == nullptr) { info = factory->get_table_info_ptr_by_index(table_id); @@ -48,7 +49,12 @@ bool OLAPMergeOperator::FullMergeV2(const rocksdb::MergeOperator::MergeOperation } bool is_rollup_key = false; if (table_info->id != table_id) { - is_rollup_key = true; + if (table_info->rollup_indexs.count(table_id) > 0) { + is_rollup_key = true; + } else { + DB_FATAL("table_id: %ld, region_id: %ld, get table failed", table_id, region_id); + return true; + } } int begin_idx = 0; @@ -153,19 +159,103 @@ bool OLAPMergeOperator::FullMergeV2(const rocksdb::MergeOperator::MergeOperation return true; } -// bool OLAPMergeOperator::PartialMerge(const rocksdb::Slice& key, const rocksdb::Slice& left_operand, -// const rocksdb::Slice& right_operand, std::string* new_value, -// rocksdb::Logger* /*logger*/) const { -// TableKey table_key(key); -// int64_t region_id = table_key.extract_i64(0); -// int64_t table_id = table_key.extract_i64(sizeof(int64_t)); -// SmartRecord left_record = SchemaFactory::get_instance()->new_record(table_id); -// SmartRecord right_record = SchemaFactory::get_instance()->new_record(table_id); -// left_record->decode(left_operand.data(), left_operand.size()); -// right_record->decode(right_operand.data(), right_operand.size()); -// DB_WARNING("table_id: %ld, region_id: %ld, left_record: %s, right_operand: %s", table_id, region_id, -// left_record->debug_string().c_str(), right_record->debug_string().c_str()); -// return false; -// } +bool OLAPMergeOperator::PartialMerge(const rocksdb::Slice& key, const rocksdb::Slice& left_operand, + const rocksdb::Slice& right_operand, std::string* new_value, + rocksdb::Logger* /*logger*/) const { + if (!FLAGS_olap_enable_partial_merge) { + return false; + } + + auto factory = SchemaFactory::get_instance(); + TableKey table_key(key); + int64_t region_id = table_key.extract_i64(0); + int64_t table_id = table_key.extract_i64(sizeof(int64_t)); + + // table_info会在thread_local短暂cache，如果加列有概率会造成新列短暂无法识别 + static thread_local SmartTable table_info = nullptr; + static thread_local TimeCost table_cache_time; + if (table_info == nullptr || (table_info->id != table_id && table_info->rollup_indexs.count(table_id) <= 0) || table_cache_time.get_time() > 60 * 1000 * 1000LL) { + auto info = factory->get_table_info_ptr(table_id); + if (info == nullptr) { + info = factory->get_table_info_ptr_by_index(table_id); + if (info == nullptr) { + return false; + } + } + + table_cache_time.reset(); + table_info = info; + } + + bool is_rollup_key = false; + if (table_info->id != table_id) { + if (table_info->rollup_indexs.count(table_id) > 0) { + is_rollup_key = true; + } else { + DB_FATAL("table_id: %ld, region_id: %ld, not found", table_id, region_id); + return false; + } + } + + SmartRecord left_record = factory->new_record(*table_info); + SmartRecord right_record = factory->new_record(*table_info); + left_record->decode(left_operand.data(), left_operand.size()); + right_record->decode(right_operand.data(), right_operand.size()); + DB_DEBUG("table_id: %ld, region_id: %ld, left_record: %s, right_operandright_record: %s", table_id, region_id, + left_record->debug_string().c_str(), right_record->debug_string().c_str()); + if (table_info->has_version && !is_rollup_key) { + auto left_version_desc = left_record->get_field_by_idx(table_info->version_field.pb_idx); + if (left_version_desc == nullptr) { + DB_FATAL("table_id: %ld, region_id: %ld, get left version desc failed", table_id, region_id); + return false; + } + + ExprValue left_version = left_record->get_value(left_version_desc); + if (left_version.is_null()) { + DB_FATAL("table_id: %ld, region_id: %ld, left version is null", table_id, region_id); + return false; + } + + auto right_version_desc = right_record->get_field_by_idx(table_info->version_field.pb_idx); + if (right_version_desc == nullptr) { + DB_FATAL("table_id: %ld, region_id: %ld, get right version desc failed", table_id, region_id); + return false; + } + + ExprValue right_version = right_record->get_value(right_version_desc); + if (right_version.is_null()) { + DB_FATAL("table_id: %ld, region_id: %ld, right version is null", table_id, region_id); + return false; + } + + if (left_version.get_numberic() != 0 || right_version.get_numberic() != 0) { + return false; + } + } + + for (const FieldInfo& f : table_info->fields_need_sum) { + auto left_field = left_record->get_field_by_idx(f.pb_idx); + if (left_field == nullptr) { + return false; + } + + ExprValue left_value = left_record->get_value(left_field); + if (left_value.is_null()) { + return false; + } + + auto right_field = right_record->get_field_by_idx(f.pb_idx); + if (right_field == nullptr) { + return false; + } + + right_record->add_value(right_field, left_value); + } + + DB_DEBUG("table_id: %ld, region_id: %ld, left_record: %s, right_operand: %s", table_id, region_id, + left_record->debug_string().c_str(), right_record->debug_string().c_str()); + right_record->encode(*new_value); + return true; +} } // namespace baikaldb \ No newline at end of file diff --git a/src/engine/table_iterator.cpp b/src/engine/table_iterator.cpp index 1e34ef754..2b80d9e20 100644 --- a/src/engine/table_iterator.cpp +++ b/src/engine/table_iterator.cpp @@ -17,7 +17,7 @@ #include "tuple_record.h" namespace baikaldb { -DEFINE_bool(cstore_scan_fill_cache, true, "cstore_scan_fill_cache"); +DEFINE_bool(cstore_scan_fill_cache, true, "Fill cache during column store scan, default: true"); DEFINE_bool(scan_fill_cache, true, "iterator_prefix_same_as_start"); TableIterator* Iterator::scan_binlog_primary( @@ -96,7 +96,7 @@ int Iterator::open(const IndexRange& range, std::map& field _fields = fields; _field_slot = field_slot; if (txn != nullptr) { - _use_ttl = txn->use_ttl(); + _use_normal_ttl = txn->use_normal_ttl(); _read_ttl_timestamp_us = txn->read_ttl_timestamp_us(); _online_ttl_base_expire_time_us = txn->online_ttl_base_expire_time_us(); _txn = txn->get_txn(); @@ -453,10 +453,10 @@ int TableIterator::get_next_internal(SmartRecord* record, int32_t tuple_id, std: return -1; } rocksdb::Slice value_slice; - if (_use_ttl || _mode != KEY_ONLY) { + if (_use_normal_ttl || _mode != KEY_ONLY) { value_slice = _iter->value(); } - if (_use_ttl) { + if (_use_normal_ttl) { int64_t row_ttl_timestamp_us = ttl_decode(value_slice, _index_info, _online_ttl_base_expire_time_us); if (_read_ttl_timestamp_us > row_ttl_timestamp_us) { //expired @@ -475,7 +475,7 @@ int TableIterator::get_next_internal(SmartRecord* record, int32_t tuple_id, std: TupleRecord tuple_record(value_slice); // only decode the required field (field_ids stored in fields) if (0 != tuple_record.decode_fields(_fields, &_field_slot, record, tuple_id, mem_row, chunk)) { - DB_WARNING("decode value failed: %ld, _use_ttl:%d", _index_info->id, _use_ttl); + DB_WARNING("decode value failed: %ld, _use_normal_ttl:%d", _index_info->id, _use_normal_ttl); _valid = false; return -1; } @@ -606,10 +606,10 @@ int IndexIterator::get_next_internal(SmartRecord* record, int32_t tuple_id, std: return -1; } rocksdb::Slice iter_value; - if (_idx_type == pb::I_UNIQ || _use_ttl) { + if (_idx_type == pb::I_UNIQ || _use_normal_ttl) { iter_value = _iter->value(); } - if (_use_ttl) { + if (_use_normal_ttl) { int64_t row_ttl_timestamp_us = ttl_decode(iter_value, _index_info, _online_ttl_base_expire_time_us); if (_read_ttl_timestamp_us > row_ttl_timestamp_us) { //expired diff --git a/src/engine/transaction.cpp b/src/engine/transaction.cpp index d6c0f373d..dd1a97d15 100644 --- a/src/engine/transaction.cpp +++ b/src/engine/transaction.cpp @@ -24,7 +24,7 @@ DEFINE_bool(disable_wal, false, "disable rocksdb interanal WAL log, only use raf DECLARE_int32(rocks_transaction_lock_timeout_ms); DEFINE_int64(exec_1pc_out_fsm_timeout_ms, 5 * 1000, "exec 1pc out of fsm, timeout"); DEFINE_int64(exec_1pc_in_fsm_timeout_ms, 100, "exec 1pc in fsm, timeout"); -DEFINE_bool(leader_merge_in_raft, false, "leader_merge_in_raft"); +DEFINE_bool(leader_merge_in_raft, false, "Enable leader merge in raft, default: false"); /** * @brief 解码TTL @@ -186,13 +186,13 @@ int Transaction::begin(const rocksdb::TransactionOptions& txn_opt) { _txn = new myrocksdb::Transaction(txn, _use_cold_db, _cold_data_cf); if (_pool != nullptr) { - _use_ttl = _pool->use_ttl(); + _use_normal_ttl = _pool->use_normal_ttl(); _online_ttl_base_expire_time_us = _pool->online_ttl_base_expire_time_us(); DB_DEBUG(); } last_active_time = butil::gettimeofday_us(); begin_time = last_active_time; - if (_use_ttl) { + if (_use_normal_ttl) { _read_ttl_timestamp_us = last_active_time; } _in_process = true; @@ -589,8 +589,8 @@ rocksdb::Status Transaction::put_kv_without_lock(const std::string& key, const s value_slices[0].size_ = sizeof(uint64_t); value_slices[1].data_ = value.data(); value_slices[1].size_ = value.size(); - DB_DEBUG("key:%s, v:%s, ,use_ttl:%d ttl_timestamp_us:%ld", str_to_hex(key).c_str(), str_to_hex(value).c_str(), _use_ttl, ttl_timestamp_us); - if (_use_ttl && ttl_timestamp_us > 0) { + DB_DEBUG("key:%s, v:%s, ,use_normal_ttl:%d ttl_timestamp_us:%ld", str_to_hex(key).c_str(), str_to_hex(value).c_str(), _use_normal_ttl, ttl_timestamp_us); + if (_use_normal_ttl && ttl_timestamp_us > 0) { value_slice_parts.parts = value_slices; value_slice_parts.num_parts = 2; } else { @@ -748,7 +748,7 @@ int Transaction::get_update_primary( DB_DEBUG("lock ok and key exist"); if (mode == GET_ONLY || mode == GET_LOCK) { rocksdb::Slice value_slice(pin_slice); - if (_use_ttl && _read_ttl_timestamp_us > 0) { + if (_use_normal_ttl && _read_ttl_timestamp_us > 0) { int64_t row_ttl_timestamp_us = ttl_decode(value_slice, &pk_index, _online_ttl_base_expire_time_us); if (_read_ttl_timestamp_us > row_ttl_timestamp_us) { DB_DEBUG("expired _read_ttl_timestamp_us:%ld row_ttl_timestamp_us:%ld", @@ -848,7 +848,7 @@ int Transaction::multiget_primary( if (statuses[i].ok()) { rocksdb::Slice value_slice(values[i]); read_disk_size += rocksdb_keys[i].size() + values[i].size(); - if (_use_ttl && _read_ttl_timestamp_us > 0) { + if (_use_normal_ttl && _read_ttl_timestamp_us > 0) { int64_t row_ttl_timestamp_us = ttl_decode(value_slice, &pk_index, _online_ttl_base_expire_time_us); if (_read_ttl_timestamp_us > row_ttl_timestamp_us) { DB_DEBUG("expired _read_ttl_timestamp_us:%ld row_ttl_timestamp_us:%ld", @@ -867,7 +867,7 @@ int Transaction::multiget_primary( if (0 != tuple_record.decode_fields(fields, &field_slot, nullptr, tuple_id, use_memrow ? &mem_row : nullptr, row_batch->get_chunk())) { - DB_WARNING("decode value failed: %ld, _use_ttl:%d", pk_index.id, _use_ttl); + DB_WARNING("decode value failed: %ld, _use_normal_ttl:%d", pk_index.id, _use_normal_ttl); continue; } } else { @@ -992,9 +992,8 @@ int Transaction::multiget_primary( DB_FATAL("Fail to append_index, reg:%ld, tab:%ld", region, pk_index.id); return -1; } - MutTableKey _key; + MutTableKey& _key = raw_read_keys.emplace_back(); _key.append_i64(region).append_i64(pk_index.id).append_index(pk_key); - raw_read_keys.emplace_back(_key); rocksdb_keys.emplace_back(_key.data()); ++num_keys; } @@ -1050,9 +1049,8 @@ int Transaction::multiget_primary( continue; } } - MutTableKey _key; + MutTableKey& _key = raw_read_keys.emplace_back(); _key.append_i64(region).append_i64(pk_index.id).append_index(left_key); - raw_read_keys.emplace_back(_key); rocksdb_keys.emplace_back(_key.data()); ++num_keys; } @@ -1158,9 +1156,8 @@ int Transaction::multiget_secondary( int64_t num_keys = 0; for (auto key : read_keys) { auto left_key = key->left_key(); - MutTableKey _key; + MutTableKey& _key = raw_read_keys.emplace_back(); _key.append_i64(region).append_i64(index.id).append_index(left_key); - raw_read_keys.emplace_back(_key); rocksdb_keys.emplace_back(_key.data()); ++num_keys; } @@ -1179,7 +1176,7 @@ int Transaction::multiget_secondary( if (statuses[i].ok()) { rocksdb::Slice value_slice(values[i]); read_disk_size += rocksdb_keys[i].size() + values[i].size(); - if (_use_ttl && _read_ttl_timestamp_us > 0) { + if (_use_normal_ttl && _read_ttl_timestamp_us > 0) { int64_t row_ttl_timestamp_us = ttl_decode(value_slice, &index, _online_ttl_base_expire_time_us); if (_read_ttl_timestamp_us > row_ttl_timestamp_us) { //expired @@ -1325,7 +1322,7 @@ int Transaction::get_update_secondary( rocksdb::Slice value(pin_slice); read_disk_size = _key.size() + value.size(); - if (_use_ttl && _read_ttl_timestamp_us > 0) { + if (_use_normal_ttl && _read_ttl_timestamp_us > 0) { int64_t row_ttl_timestamp_us = ttl_decode(value, &index, _online_ttl_base_expire_time_us); if (_read_ttl_timestamp_us > row_ttl_timestamp_us) { //expired diff --git a/src/engine/transaction_pool.cpp b/src/engine/transaction_pool.cpp index a2a1e0e9f..7b496de3d 100644 --- a/src/engine/transaction_pool.cpp +++ b/src/engine/transaction_pool.cpp @@ -37,9 +37,9 @@ DEFINE_int64(1pc_out_fsm_interval_us, 20 * 1000 * 1000LL, DEFINE_int32(transaction_query_primary_region_interval_ms, 15 * 1000, "interval duration send request to primary region"); -int TransactionPool::init(int64_t region_id, bool use_ttl, int64_t online_ttl_base_expire_time_us) { +int TransactionPool::init(int64_t region_id, bool use_normal_ttl, int64_t online_ttl_base_expire_time_us) { _region_id = region_id; - _use_ttl = use_ttl; + _use_normal_ttl = use_normal_ttl; _online_ttl_base_expire_time_us = online_ttl_base_expire_time_us; _meta_writer = MetaWriter::get_instance(); return 0; diff --git a/src/exec/access_path.cpp b/src/exec/access_path.cpp index 418876fe3..62109e2b5 100755 --- a/src/exec/access_path.cpp +++ b/src/exec/access_path.cpp @@ -22,6 +22,7 @@ namespace baikaldb { using namespace range; DEFINE_uint64(max_in_records_num, 10000, "max_in_records_num"); +DEFINE_uint64(cut_huge_in_filter_size, 500000, "cut_huge_in_filter_size, maybe 50w is suitable"); DEFINE_int64(index_use_for_learner_delay_s, 3600, "1h"); DEFINE_bool(date_range_to_in, false, "date range to in"); @@ -71,10 +72,14 @@ void AccessPath::calc_row_expr_range(std::vector& range_fields, ExprNod return; } size_t row_idx = 0; - for (; row_idx < range_fields.size() && - field_idx < index_info_ptr->fields.size(); row_idx++, field_idx++) { - if (index_info_ptr->fields[field_idx].id == range_fields[row_idx]) { - key.append_value(values[row_idx].cast_to(index_info_ptr->fields[field_idx].type)); + auto avaliable_key_fields = index_info_ptr->fields; + if (index_info_ptr->type == pb::I_KEY) { + avaliable_key_fields.insert(avaliable_key_fields.end(), index_info_ptr->pk_fields.begin(), index_info_ptr->pk_fields.end()); + } + for (; row_idx < range_fields.size() && + field_idx < avaliable_key_fields.size(); row_idx++, field_idx++) { + if (avaliable_key_fields[field_idx].id == range_fields[row_idx]) { + key.append_value(values[row_idx].cast_to(avaliable_key_fields[field_idx].type)); hit_index_field_ids.insert(range_fields[row_idx]); } else { break; @@ -91,7 +96,7 @@ bool AccessPath::check_sort_use_index(Property& sort_property) { std::vector& order_exprs = sort_property.slot_order_exprs; SlotRef* slot_ref = static_cast(order_exprs[0]); size_t idx = 0; - std::vectorfields(index_info_ptr->fields.begin(), index_info_ptr->fields.end()); + std::vector fields(index_info_ptr->fields.begin(), index_info_ptr->fields.end()); fields.insert(fields.end(),index_info_ptr->pk_fields.begin(), index_info_ptr->pk_fields.end()); for (; idx < fields.size(); ++idx) { if (tuple_id == slot_ref->tuple_id() && fields[idx].id == slot_ref->field_id()) { @@ -131,7 +136,16 @@ void AccessPath::calc_normal(Property& sort_property) { // hit_fields_cnt: in_row_expr谓词匹配的字段个数,用于判断是否可以剪切 std::map< ExprNode*, std::pair> in_row_expr_map; // size_t in_records_size = 1; - for (auto& field : index_info_ptr->fields) { + // 主键a b, 索引c a 且包含变长字段， 则avaliable_key_fields为c a a b；遇到第二个a会直接结束匹配 + // 如果发生主键压缩，主键a b, 索引c b, 则avaliable_key_fields 为 c b a；会一直匹配到底 + std::vector avaliable_key_fields = index_info_ptr->fields; + if (index_info_ptr->type == pb::I_KEY) { + avaliable_key_fields.insert(avaliable_key_fields.end(), index_info_ptr->pk_fields.begin(), index_info_ptr->pk_fields.end()); + } + for (auto& field : avaliable_key_fields) { + if (0 != hit_index_field_ids.count(field.id)) { + break; + } bool field_break = false; auto iter = field_range_map->find(field.id); if (iter == field_range_map->end()) { @@ -160,8 +174,8 @@ void AccessPath::calc_normal(Property& sort_property) { need_cut_index_range_condition.insert(range.left_expr); } else if (range.left.size() > 1) { size_t row_idx = 0; - while (row_idx < range.left_row_field_ids.size() && field_idx < index_info_ptr->fields.size()) { - if (index_info_ptr->fields[field_idx].id == range.left_row_field_ids[row_idx]) { + while (row_idx < range.left_row_field_ids.size() && field_idx < avaliable_key_fields.size()) { + if (avaliable_key_fields[field_idx].id == range.left_row_field_ids[row_idx]) { hit_index_field_ids.insert(range.left_row_field_ids[row_idx]); } else { break; @@ -183,8 +197,8 @@ void AccessPath::calc_normal(Property& sort_property) { need_cut_index_range_condition.insert(range.right_expr); } else if (range.right.size() > 1) { size_t row_idx = 0; - while (row_idx < range.right_row_field_ids.size() && field_idx < index_info_ptr->fields.size()) { - if (index_info_ptr->fields[field_idx].id == range.right_row_field_ids[row_idx]) { + while (row_idx < range.right_row_field_ids.size() && field_idx < avaliable_key_fields.size()) { + if (avaliable_key_fields[field_idx].id == range.right_row_field_ids[row_idx]) { hit_index_field_ids.insert(range.right_row_field_ids[row_idx]); } else { break; @@ -247,6 +261,12 @@ void AccessPath::calc_normal(Property& sort_property) { in_row_expr_map[*range.conditions.begin()].second++; } else { // 第一个in不限制FLAGS_max_in_records_num + if (FLAGS_cut_huge_in_filter_size > 0 + && (index_type == pb::I_KEY && !index_info_ptr->is_global) + && range.eq_in_values.size() > FLAGS_cut_huge_in_filter_size) { + field_break = true; + break; + } if (in_records_size > 1 && in_records_size * range.eq_in_values.size() > FLAGS_max_in_records_num) { field_break = true; break; @@ -344,7 +364,14 @@ void AccessPath::calc_index_range( // hit_fields_cnt: in_row_expr谓词匹配的字段个数,用于判断是否可以剪切 std::map< ExprNode*, std::pair> in_row_expr_map; // int field_cnt = 0; - for (auto& field : index_info_ptr->fields) { + // 索引或者主键包含变长字段时可能包含重复字段，如 + // 主键a b, 索引c a 且包含变长字段， 则avaliable_key_fields为c a a b + // 如果发生主键压缩，主键a b, 索引c b, 则avaliable_key_fields 为 c b a + std::vector avaliable_key_fields = index_info_ptr->fields; + if (index_info_ptr->type == pb::I_KEY) { + avaliable_key_fields.insert(avaliable_key_fields.end(), index_info_ptr->pk_fields.begin(), index_info_ptr->pk_fields.end()); + } + for (auto& field : avaliable_key_fields) { auto iter = field_range_map->find(field.id); field_cnt++; if (field_cnt > hit_index_field_ids.size()) { @@ -520,7 +547,7 @@ void AccessPath::calc_index_range( } filter.insert(rg.left_key.data()); auto range = pos_index.add_ranges(); - if (_left_field_cnt == index_info_ptr->fields.size() + if (_left_field_cnt == avaliable_key_fields.size() && (index_type == pb::I_PRIMARY || index_type == pb::I_UNIQ) && !_like_prefix) { rg.left_key.set_full(true); @@ -533,7 +560,7 @@ void AccessPath::calc_index_range( range->set_left_key(rg.left_key.data()); range->set_left_full(rg.left_key.get_full()); if (!_is_eq_or_in) { - if (_right_field_cnt == index_info_ptr->fields.size() + if (_right_field_cnt == avaliable_key_fields.size() && (index_type == pb::I_PRIMARY || index_type == pb::I_UNIQ) && !_like_prefix) { rg.right_key.set_full(true); @@ -545,12 +572,12 @@ void AccessPath::calc_index_range( } else { is_possible = true; auto range = pos_index.add_ranges(); - if (_left_field_cnt == index_info_ptr->fields.size() + if (_left_field_cnt == avaliable_key_fields.size() && (index_type == pb::I_PRIMARY || index_type == pb::I_UNIQ) && !_like_prefix) { left_key.set_full(true); } - if (_right_field_cnt == index_info_ptr->fields.size() + if (_right_field_cnt == avaliable_key_fields.size() && (index_type == pb::I_PRIMARY || index_type == pb::I_UNIQ) && !_like_prefix) { right_key.set_full(true); @@ -664,6 +691,7 @@ void AccessPath::calc_fulltext(Property& sort_property) { } range->set_topk(sort_property.expected_cnt); range->set_efsearch(sort_property.efsearch); + range->set_nprobe(sort_property.nprobe); if (first_field_expr_values != nullptr && !first_field_expr_values->empty()) { ExprValue& expr_value = *(first_field_expr_values->begin()); uint64_t separate_value = expr_value.cast_to(pb::UINT64).get_numberic(); diff --git a/src/exec/apply_node.cpp b/src/exec/apply_node.cpp index c57c50031..49cd8e3bb 100644 --- a/src/exec/apply_node.cpp +++ b/src/exec/apply_node.cpp @@ -226,7 +226,7 @@ int ApplyNode::check_unique_key(RuntimeState* state, const std::map 0 ? _row_ttl_duration : _table_info->ttl_info.ttl_duration_s; - if (ttl_duration > 0) { + if (ttl_duration > 0 && _table_info->ttl_info.ttl_field == nullptr) { _ttl_timestamp_us = butil::gettimeofday_us() + ttl_duration * 1000 * 1000LL; } DB_DEBUG("table_id: %ld, region_id: %ld, _row_ttl_duration: %ld, table ttl duration: %ld", @@ -634,7 +634,8 @@ int DMLNode::remove_row(RuntimeState* state, SmartRecord record, return ret; } ret = vector_index_map[info.id]->delete_vector(_txn, - word, pk_str, record); + word, pk_str, record, + info.state); if (ret < 0) { DB_WARNING_STATE(state, "vector_index fail delete, index_id: %ld", info.id); return ret; @@ -759,7 +760,7 @@ int DMLNode::update_row(RuntimeState* state, SmartRecord record, MemRow* row) { if (field == nullptr) { state->error_code = ER_BAD_FIELD_ERROR; state->error_msg << "Unknown column id " << slot.field_id() << " in 'field list'"; - DB_WARNING_STATE(state, "Unknown column, table_id: %d, field_id: %d", slot.table_id(), slot.field_id()); + DB_WARNING_STATE(state, "Unknown column, table_id: %ld, field_id: %d", slot.table_id(), slot.field_id()); return -1; } if (field->type == pb::FLOAT || field->type == pb::DOUBLE || field->type == pb::DATETIME) { diff --git a/src/exec/exchange_receiver_node.cpp b/src/exec/exchange_receiver_node.cpp index 60f0b6302..e782ad747 100644 --- a/src/exec/exchange_receiver_node.cpp +++ b/src/exec/exchange_receiver_node.cpp @@ -257,6 +257,13 @@ int DataStreamReceiver::add_record_batch(const pb::TransmitDataParam& param, DB_WARNING_ER("get process key failed"); return -1; } + // filter first + if (_condition != nullptr && record_batch != nullptr) { + if (0 != VectorizeHelpper::vectorize_filter(record_batch, _condition, &record_batch)) { + DB_WARNING_ER("filter record batch failed"); + return -1; + } + } BAIDU_SCOPED_LOCK(_mtx); // 重复数据以第一次为准，merge场景可能重复发送同一个region的数据 if (_done_set.find(process_key) != _done_set.end()) { @@ -332,12 +339,26 @@ int DataStreamReceiver::add_record_batch(const pb::TransmitDataParam& param, int DataStreamReceiver::add_local_pass_through_chunk(uint64_t sender_fragment_instance_id, const std::vector>& record_batchs) { + std::vector> filtered_record_batchs; + if (_condition != nullptr) { + filtered_record_batchs.reserve(record_batchs.size()); + for (auto& record_batch : record_batchs) { + std::shared_ptr out; + if (0 != VectorizeHelpper::vectorize_filter(record_batch, _condition, &out)) { + DB_WARNING_ER("filter record batch failed"); + return -1; + } + filtered_record_batchs.emplace_back(out); + } + } else { + filtered_record_batchs = record_batchs; + } BAIDU_SCOPED_LOCK(_mtx); _local_pass_through_record_batchs[sender_fragment_instance_id]; if (record_batchs.empty()) { return 0; } - for (auto& record_batch : record_batchs) { + for (auto& record_batch : filtered_record_batchs) { if (record_batch != nullptr) { DB_DEBUG("local pass through record batch num rows: %ld, _fragment_instance_id: %lu, sender_fragment_instance_id: %lu", record_batch->num_rows(), _fragment_instance_id, sender_fragment_instance_id); @@ -426,6 +447,26 @@ int ExchangeReceiverNode::init(const pb::PlanNode& node) { for (bool is_null_first : er_node.is_null_first()) { _is_null_first.push_back(is_null_first); } + _conditions.reserve(er_node.filter_conditions_size()); + for (auto& expr : er_node.filter_conditions()) { + ExprNode* expr_node = nullptr; + ret = ExprNode::create_tree(expr, &expr_node); + if (ret < 0) { + DB_FATAL("create expr fail"); + ExprNode::destroy_tree(expr_node); + return ret; + } + _conditions.emplace_back(expr_node); + } + return 0; +} + +int ExchangeReceiverNode::init_condition_and_sort_info(ExecNode* select_manager) { + SelectManagerNode* select_manager_node = static_cast(select_manager); + _slot_order_exprs = select_manager_node->slot_order_exprs(); + _is_asc = select_manager_node->is_asc(); + _is_null_first = select_manager_node->is_null_first(); + _conditions = select_manager_node->conditions(); return 0; } @@ -435,6 +476,21 @@ int ExchangeReceiverNode::open(RuntimeState* state) { DB_FATAL_ER("ExecNode::open fail, ret:%d", ret); return ret; } + + // 主db的schema和region在构建物理计划的时候生成, pb里没有 + // 非主db的schema和region在init(pb)的时候生成 + if (_exchange_sender_node != nullptr && _exchange_sender_node->children_size() > 0) { + auto data_schema = _exchange_sender_node->children(0)->data_schema(); + _arrow_schema = VectorizeHelpper::get_arrow_schema(data_schema); + } + if (_arrow_schema == nullptr) { + DB_FATAL_ER("transfer receiver arrow schema fail"); + return -1; + } + if (0 != VectorizeHelpper::init_conditions(state, _conditions, _vectorize_conditions, -1, _arrow_schema)) { + DB_FATAL_ER("init condition fail"); + return -1; + } _data_stream_receiver = DataStreamManager::get_instance()->create_receiver(_log_id, _fragment_instance_id, @@ -446,16 +502,8 @@ int ExchangeReceiverNode::open(RuntimeState* state) { return -1; } _data_stream_receiver->set_fragment_id(_fragment_id); - - // 主db的schema和region在构建物理计划的时候生成, pb里没有 - // 非主db的schema和region在init(pb)的时候生成 - if (_exchange_sender_node != nullptr && _exchange_sender_node->children_size() > 0) { - auto data_schema = _exchange_sender_node->children(0)->data_schema(); - _arrow_schema = VectorizeHelpper::get_arrow_schema(data_schema); - } - if (_arrow_schema == nullptr) { - DB_FATAL_ER("transfer receiver arrow schema fail"); - return -1; + if (_conditions.size() > 0) { + _data_stream_receiver->set_conditions(&_vectorize_conditions); } set_node_exec_type(pb::EXEC_ARROW_ACERO); return 0; @@ -488,6 +536,11 @@ void ExchangeReceiverNode::close(RuntimeState* state) { ExprNode::destroy_tree(expr); } } + if (_pb_node.derive_node().exchange_receiver_node().filter_conditions_size() > 0) { + for (auto& expr : _conditions) { + ExprNode::destroy_tree(expr); + } + } if (_arrow_io_executor != nullptr) { _arrow_io_executor.reset(); } @@ -508,14 +561,16 @@ void ExchangeReceiverNode::transfer_pb(int64_t region_id, pb::PlanNode* pb_node) } if (_relate_select_manager_node != nullptr) { - auto regions = _relate_select_manager_node->region_infos(); - if (regions.empty()) { - DB_FATAL_ER("relate select manager node region is empty"); - return; - } - for (auto& region : regions) { - // TODO 如果region多, 改成[region_id, version]? 确认实际执行用了region info哪些信息 - er_node->add_regions()->CopyFrom(region.second); + if (!static_cast(_relate_select_manager_node)->is_dblink_scan()) { + auto regions = _relate_select_manager_node->region_infos(); + if (regions.empty()) { + DB_FATAL_ER("relate select manager node region is empty"); + return; + } + for (auto& region : regions) { + // TODO 如果region多, 改成[region_id, version]? 确认实际执行用了region info哪些信息 + er_node->add_regions()->CopyFrom(region.second); + } } } // 序列化schema @@ -529,7 +584,7 @@ void ExchangeReceiverNode::transfer_pb(int64_t region_id, pb::PlanNode* pb_node) DB_FATAL_ER("transfer receiver arrow schema fail"); return; } - arrow::Result> schema_ret = arrow::ipc::SerializeSchema(*_arrow_schema, arrow::default_memory_pool()); + arrow::Result> schema_ret = arrow::ipc::SerializeSchema(*_arrow_schema, GetMemoryPoolForRead()); if (!schema_ret.ok()) { DB_FATAL_ER("arrow serialize schema fail, status: %s", schema_ret.status().ToString().c_str()); return; @@ -546,6 +601,10 @@ void ExchangeReceiverNode::transfer_pb(int64_t region_id, pb::PlanNode* pb_node) for (bool null_first : _is_null_first) { er_node->add_is_null_first(null_first); } + // filter + for (ExprNode* expr : _conditions) { + ExprNode::create_pb_expr(er_node->add_filter_conditions(), expr); + } return; } @@ -567,6 +626,9 @@ int ExchangeReceiverNode::build_arrow_declaration(RuntimeState* state) { arrow::acero::Declaration dec{"record_batch_source", arrow::acero::RecordBatchSourceNodeOptions{reader->schema(), std::move(iter_maker), _arrow_io_executor.get()}}; LOCAL_TRACE_ARROW_PLAN_WITH_SCHEMA(dec, reader->schema(), nullptr); + if (_conditions.size() > 0) { + LOCAL_TRACE_ARROW_FILTER(&_vectorize_conditions, _limit); + } state->append_acero_declaration(dec); if (!_slot_order_exprs.empty()) { @@ -656,7 +718,7 @@ arrow::Status ExchangeReceiverVectorizedReader::ReadNext(std::shared_ptrschema(), need_send_record_batches, - concatenate_record_batch)) { + concatenate_record_batch, + /*need_check_memory_limit=*/true)) { DB_FATAL_CHANNEL("Fail to concatenate record batches"); return -1; } @@ -453,7 +454,7 @@ void ExchangeSenderNode::transfer_pb(int64_t region_id, pb::PlanNode* pb_node) { DB_FATAL_ES("transfer receiver arrow schema fail"); return; } - arrow::Result> schema_ret = arrow::ipc::SerializeSchema(*_arrow_schema, arrow::default_memory_pool()); + arrow::Result> schema_ret = arrow::ipc::SerializeSchema(*_arrow_schema, GetMemoryPoolForRead()); if (!schema_ret.ok()) { DB_FATAL_ES("arrow serialize schema fail, status: %s", schema_ret.status().ToString().c_str()); return; @@ -501,8 +502,7 @@ int ExchangeSenderNode::open(RuntimeState* state) { // 具体执行的时候需要使用对应上游receiver的partition_property _partition_property = *(_exchange_receiver_node->partition_property()); } - - _is_db_fragment = (get_node_pass_subquery(pb::EXCHANGE_RECEIVER_NODE) != nullptr); + _is_db_fragment = is_db_fragment(); _region_id = state->region_id(); for (int i = 0; i < _receiver_destinations.size(); ++i) { bool is_local_pass_through = false; @@ -530,14 +530,7 @@ int ExchangeSenderNode::open(RuntimeState* state) { } int ExchangeSenderNode::get_next(RuntimeState* state, RowBatch* batch, bool* eos) { - // 应该不需要往下get_next, 直接转向量化执行, 注意RocksdbScanNode的reader - // 应该不会调用get_next - // int ret = 0; - // ret = _children[0]->get_next(state, batch, eos); - // if (ret < 0) { - // DB_WARNING("_children get_next fail"); - // return ret; - // } + // 不需要往下get_next, 直接转向量化执行, 注意RocksdbScanNode的reader set_node_exec_type(pb::EXEC_ARROW_ACERO); *eos = true; return 0; @@ -631,7 +624,8 @@ int ExchangeSenderNode::send_record_batch(RuntimeState* state, DB_FATAL_ES("Fail to send record batch"); return -1; } - } else if (_partition_property.type == pb::BroadcastPartitionType) { + } else if (_partition_property.type == pb::BroadcastPartitionType + || _partition_property.type == pb::RandomPartitionType) { std::vector> need_send_record_batches; bool need_send_rpc = false; if (0 != _broadcast_record_batch_keepper.add_record_batch(state, record_batch, _is_db_fragment, need_send_record_batches, exchange_state, need_send_rpc)) { @@ -653,7 +647,8 @@ int ExchangeSenderNode::send_record_batch(RuntimeState* state, if (0 != VectorizeHelpper::concatenate_record_batches( record_batch->schema(), need_send_record_batches, - concatenate_record_batch)) { + concatenate_record_batch, + /*need_check_memory_limit=*/true)) { DB_FATAL_ES("Fail to concatenate record batches"); return -1; } @@ -672,12 +667,18 @@ int ExchangeSenderNode::send_record_batch(RuntimeState* state, serialize_cost = time_cost.get_time(); time_cost.reset(); attachment_size = data_buffer->size(); + int channel_idx = 0; // 发送全部上游 for (auto& channel : _channels) { + bool is_not_target_channel_for_random_type = (_partition_property.type == pb::RandomPartitionType && _target_channel_idx_for_random_type != channel_idx); + if (is_not_target_channel_for_random_type && exchange_state == pb::ExchangeState::ES_DOING) { + // 不是目标channel, 最后发一次eof即可 + continue; + } if (0 != channel->send_record_batch_data(state, need_send_record_batches, schema_buffer, - data_buffer, + is_not_target_channel_for_random_type ? _empty_recordbatch_buffer : data_buffer, exchange_state, is_send_query_stat, send_rows)) { @@ -687,6 +688,7 @@ int ExchangeSenderNode::send_record_batch(RuntimeState* state, if (is_send_query_stat) { is_send_query_stat = false; } + channel_idx++; } send_cost = time_cost.get_time(); if (combine_cost + serialize_cost + send_cost > FLAGS_print_time_us) { @@ -767,7 +769,31 @@ int ExchangeSenderNode::init_repartition_param() { // repartition使用 // 获取分区列，获取需要转化成字符串的分区列 _hash_bucket_num = _receiver_destinations.size(); - + if (_hash_bucket_num == 0) { + DB_FATAL_ES("hash_bucket_num is 0"); + return -1; + } + if (_partition_property.type == pb::RandomPartitionType) { + if (_is_db_fragment) { + _target_channel_idx_for_random_type = _fragment_instance_id % _hash_bucket_num; + } else { + _target_channel_idx_for_random_type = _region_id % _hash_bucket_num; + } + auto empty_record_batch = arrow::RecordBatch::MakeEmpty(_arrow_schema); + if (!empty_record_batch.ok()) { + DB_FATAL_ES("Fail to MakeEmpty recordbatch for random partition type"); + return -1; + } + std::shared_ptr record_batch = *empty_record_batch; + arrow::Result> empty_recordbatch_buffer = + arrow::ipc::SerializeRecordBatch(*record_batch, arrow::ipc::IpcWriteOptions::Defaults()); + if (!empty_recordbatch_buffer.ok()) { + DB_FATAL("Fail to Serialize empty RecordBatch"); + return -1; + } + _empty_recordbatch_buffer = *empty_recordbatch_buffer; + return 0; + } if (_partition_property.type == pb::SinglePartitionType || _partition_property.type == pb::BroadcastPartitionType) { return 0; @@ -857,13 +883,16 @@ arrow::Status ExchangeSenderNode::repartition( // cast string const int64_t batch_length = batch->num_rows(); std::unordered_map> cast_array_map; + arrow::compute::ExecContext exec_ctx(GetMemoryPoolForRead()); + arrow::compute::CastOptions cast_options; + cast_options.allow_int_overflow = true; for (int i = 0; i < hash_indices.size(); ++i) { if (need_cast_indices[i]) { if (hash_indices[i] >= batch->column_data().size()) { return arrow::Status::Invalid("hash_indices[i] >= batch->num_columns()"); } auto array = batch->column(hash_indices[i]); - ARROW_ASSIGN_OR_RAISE(cast_array_map[hash_indices[i]], arrow::compute::Cast(*array, arrow::large_binary())); + ARROW_ASSIGN_OR_RAISE(cast_array_map[hash_indices[i]], arrow::compute::Cast(*array, arrow::large_binary(), cast_options, &exec_ctx)); if (cast_array_map[hash_indices[i]]->length() != batch_length) { return arrow::Status::Invalid("cast_array_map[hash_indices[i]]->length() != batch_length"); } @@ -871,7 +900,7 @@ arrow::Status ExchangeSenderNode::repartition( } // 构造arrow ctx arrow::util::TempVectorStack stack; - stack.Init(::arrow::default_memory_pool(), 8 * max_batch_size * sizeof(uint64_t)); + stack.Init(GetMemoryPoolForRead(), 8 * max_batch_size * sizeof(uint64_t)); arrow::compute::LightContext ctx; ctx.hardware_flags = arrow::internal::CpuInfo::GetInstance()->hardware_flags(); ctx.stack = &stack; diff --git a/src/exec/exec_node.cpp b/src/exec/exec_node.cpp index 7d8005607..716316e70 100644 --- a/src/exec/exec_node.cpp +++ b/src/exec/exec_node.cpp @@ -518,7 +518,9 @@ int ExecNode::push_cmd_to_cache(RuntimeState* state, bool ExecNode::shrink_partition_property(std::shared_ptr my_hash_columns, NodePartitionProperty* child_partition_property) { bool is_same_or_shrinked = false; for (auto& child_hash_partition : child_partition_property->hash_partition_propertys) { - if (my_hash_columns->hash_partition_is_contain(child_hash_partition.get())) { + if (child_partition_property->type == pb::HashPartitionType + && child_hash_partition->type == pb::HashPartitionType + && my_hash_columns->hash_partition_is_contain(child_hash_partition.get())) { is_same_or_shrinked = true; for (auto iter = my_hash_columns->hash_columns.begin(); iter != my_hash_columns->hash_columns.end();) { if (child_hash_partition->hash_columns.count(iter->first) == 0) { @@ -527,18 +529,29 @@ bool ExecNode::shrink_partition_property(std::shared_ptr m ++iter; } } - for (auto order_iter = my_hash_columns->ordered_hash_columns.begin(); order_iter != my_hash_columns->ordered_hash_columns.end();) { - if (my_hash_columns->hash_columns.count(*order_iter) == 0) { - order_iter = my_hash_columns->ordered_hash_columns.erase(order_iter); - } else { - ++order_iter; - } - } + // 按照child hash列顺序调整 + // 比如 my: b, a, child: a, b + // 调整my hash partition分区: a, b. arrow hash值计算和列顺序有关 + my_hash_columns->ordered_hash_columns = child_hash_partition->ordered_hash_columns; } } return is_same_or_shrinked; } +void ExecNode::set_child_node_property_any_type() { + _partition_property.set_any_partition(); + if (_node_type == pb::DUAL_SCAN_NODE) { + DualScanNode* dual_scan_node = static_cast(this); + ExecNode* sub_query_node = dual_scan_node->sub_query_node(); + if (sub_query_node != nullptr) { + sub_query_node->set_child_node_property_any_type(); + } + } + for (auto c : _children) { + c->set_child_node_property_any_type(); + } +} + void ExecNode::encode_exprs_key(std::vector& exprs, MemRow* row, MutTableKey& key) { uint8_t null_flag = 0; key.append_u8(null_flag); diff --git a/src/exec/fetcher_store.cpp b/src/exec/fetcher_store.cpp index f8a9e7cb2..cb85b3c7b 100755 --- a/src/exec/fetcher_store.cpp +++ b/src/exec/fetcher_store.cpp @@ -29,6 +29,7 @@ #include "rocksdb_scan_node.h" #include "arrow/util/byte_size.h" #include "filter_node.h" +#include "vectorize_helpper.h" namespace baikaldb { @@ -36,9 +37,7 @@ DEFINE_int64(retry_interval_us, 500 * 1000, "retry interval "); DEFINE_int32(single_store_concurrency, 20, "max request for one store"); DEFINE_int64(max_select_rows, 10000000, "query will be fail when select too much rows"); DEFINE_int64(max_affected_rows, 10000000, "query will be fail when affect too much rows"); -DEFINE_int64(print_time_us, 10000, "print log when time_cost > print_time_us(us)"); DEFINE_int64(baikaldb_alive_time_s, 10 * 60, "obervation time length in baikaldb, default:10 min"); -BRPC_VALIDATE_GFLAG(print_time_us, brpc::NonNegativeInteger); DEFINE_int32(fetcher_request_timeout, 100000, "store as server request timeout, default:100000ms"); DEFINE_int32(fetcher_connect_timeout, 1000, @@ -57,6 +56,7 @@ BRPC_VALIDATE_GFLAG(sql_exec_timeout, brpc::PassValidate); DEFINE_bool(enable_batch_rpc, false, "enable batch rpc"); DEFINE_int32(enable_batch_rpc_region_num, 5, "enable batch rpc region num"); DEFINE_bool(open_prefer_arrow_data, false, "open prefer arrow data in vectore execution"); +DECLARE_int64(print_time_us); bvar::Adder OnRPCDone::async_rpc_region_count {"async_rpc_region_count"}; bvar::LatencyRecorder OnRPCDone::total_send_request {"total_send_request"}; bvar::LatencyRecorder OnRPCDone::add_backup_send_request {"add_backup_send_request"}; @@ -124,6 +124,7 @@ ErrorType OnRPCDone::fill_single_request(pb::StoreReq& single_req, pb::RegionInf if (_trace_node != nullptr) { single_req.set_is_trace(true); } + _request.mutable_extra_req()->set_no_raft_log(_fetcher_store->broadcast_all_peer_without_raft); if (_state->explain_type == ANALYZE_STATISTICS) { pb::AnalyzeInfo* info = single_req.mutable_analyze_info(); bool need_hist = false; @@ -387,6 +388,9 @@ void OnRPCDone::select_addr(pb::RegionInfo& info, std::string& addr, bool& resource_insulate_read, bool& select_without_leader) { + if (_fetcher_store->broadcast_all_peer_without_raft) { + return; + } addr = info.leader(); resource_insulate_read = false; // 是否读learner，或者指定读从集群，进行资源隔离 if (_state->need_learner_backup() && info.learners_size() == 0) { @@ -775,6 +779,19 @@ ErrorType OnRPCDone::handle_version_old(const pb::RegionInfo& info, pb::StoreRes return E_FATAL; } +bool OnRPCDone::need_copy(MemRow* row) { + if (_fetcher_store->conditions == nullptr) { + return true; + } + for (auto& condition : *(_fetcher_store->conditions)) { + ExprValue value = condition->get_value(row); + if (value.is_null() || value.get_numberic() == false) { + return false; + } + } + return true; +} + ErrorType OnRPCDone::handle_single_response(const std::string& remote_side, int64_t region_id, const std::string& addr, @@ -1003,6 +1020,7 @@ ErrorType OnRPCDone::handle_single_response(const std::string& remote_side, return E_BIG_SQL; } uint64_t returned_row_size = 0; + uint64_t row_size_after_filter = 0; // EXEC_ROW: region只会返回memrow // EXEC_ARROW_ACERO: region可能返回memrow, 可能返回arrow if (single_response.execute_type() == pb::EXEC_ROW) { @@ -1029,6 +1047,10 @@ ErrorType OnRPCDone::handle_single_response(const std::string& remote_side, int32_t tuple_id = single_response.tuple_ids(i); row->from_string(tuple_id, pb_row.tuple_values(i)); } + returned_row_size++; + if (!need_copy(row.get())) { + continue; + } row->set_partition_id(info.partition_id()); int64_t row_size = row->used_size(); used_size += row_size; @@ -1049,7 +1071,7 @@ ErrorType OnRPCDone::handle_single_response(const std::string& remote_side, DB_DEBUG("region_id: %ld, ttl_timestamp: %ld", region_id, time_us); } } - returned_row_size = batch->size(); + row_size_after_filter = batch->size(); _fetcher_store->db_handle_bytes += memrow_total_used_size; _fetcher_store->db_handle_rows += returned_row_size; if (global_ddl_with_ttl) { @@ -1089,6 +1111,7 @@ ErrorType OnRPCDone::handle_single_response(const std::string& remote_side, _fetcher_store->start_key_sort.emplace(info.start_key(), region_id); _fetcher_store->region_batch[region_id].set_row_data(batch); } + _fetcher_store->region_batch[region_id].set_partition_id(info.partition_id()); } else { BAIDU_SCOPED_LOCK(_fetcher_store->region_lock); // merge可能会重复请求相同的region_id @@ -1099,6 +1122,7 @@ ErrorType OnRPCDone::handle_single_response(const std::string& remote_side, _fetcher_store->start_key_sort.emplace(info.start_key(), region_id); _fetcher_store->region_batch[region_id].set_arrow_data(out); } + _fetcher_store->region_batch[region_id].set_partition_id(info.partition_id()); } } else { // 处理列存arrow recordbatch @@ -1106,6 +1130,8 @@ ErrorType OnRPCDone::handle_single_response(const std::string& remote_side, const auto& vector_rows = single_response.mutable_extra_res()->vectorized_rows(); const auto& vector_schema = single_response.mutable_extra_res()->vectorized_schema(); std::shared_ptr batch; + bool need_keep_response = true; + int64_t recordbatch_size = vector_rows.size(); if (vector_rows.size() > 0) { // 解析列存格式 std::shared_ptr schema_buffer = std::make_shared(reinterpret_cast(vector_schema.data()), @@ -1135,10 +1161,27 @@ ErrorType OnRPCDone::handle_single_response(const std::string& remote_side, _state->error_msg.str("parse arrow data fail"); return E_FATAL; } + // filter后按需将response的recordbatch释放掉, 避免内存double + if (_fetcher_store->vectorize_conditions != nullptr) { + std::shared_ptr filtered_batch; + if (0 != VectorizeHelpper::vectorize_filter(batch, _fetcher_store->vectorize_conditions, &filtered_batch, &need_keep_response)) { + DB_WARNING("vectorize filter error"); + _state->error_code = ER_EXEC_PLAN_FAILED; + _state->error_msg.str("vectorize filter fail"); + return E_FATAL; + } + if (filtered_batch == nullptr) { + DB_FATAL("filtered_batch is nullptr"); + return E_FATAL; + } + batch = filtered_batch; + recordbatch_size = arrow::util::TotalBufferSize(*filtered_batch); + } + row_size_after_filter = batch->num_rows(); // [ARROW TODO] 如何限制内存? // 向量化是将store返回的数据都存在内存, 等请求结束这部分内存才释放 if (vector_rows.size() > 1024 * 1024LL) { - if (0 != _state->memory_limit_exceeded(returned_row_size, vector_rows.size())) { + if (0 != _state->memory_limit_exceeded(row_size_after_filter, recordbatch_size)) { BAIDU_SCOPED_LOCK(_fetcher_store->region_lock); _state->error_code = ER_TOO_BIG_SELECT; _state->error_msg.str("select reach memory limit"); @@ -1161,9 +1204,13 @@ ErrorType OnRPCDone::handle_single_response(const std::string& remote_side, _fetcher_store->start_key_sort.emplace(info.start_key(), region_id); _fetcher_store->region_batch[region_id].set_arrow_data(batch); } - // buffer和解析出来的recordbatch是直接使用pb里的vectorized_rows内存, 是zero copy的 - // 所以列存要求response生命周期比recordBatch生命周期长 - set_region_vectorized_response(region_id, single_response_ptr); + if (need_keep_response) { + // buffer和解析出来的recordbatch是直接使用pb里的vectorized_rows内存, 是zero copy的 + // 所以列存要求response生命周期比recordBatch生命周期长 + // batch模式, 但凡一个region结果保留, 整个batchresponse都保留下来了 + set_region_vectorized_response(region_id, single_response_ptr); + } + _fetcher_store->region_batch[region_id].set_partition_id(info.partition_id()); } } @@ -1184,7 +1231,11 @@ ErrorType OnRPCDone::handle_single_response(const std::string& remote_side, } } if (cost.get_time() > FLAGS_print_time_us) { - DB_DONE(WARNING, "parse time:%ld rows:%lu", cost.get_time(), returned_row_size); + if (_fetcher_store->vectorize_conditions == nullptr) { + DB_DONE(WARNING, "parse time:%ld rows:%lu", cost.get_time(), returned_row_size); + } else { + DB_DONE(WARNING, "parse time:%ld rows:%lu,rows_after_filter: %lu", cost.get_time(), returned_row_size, row_size_after_filter); + } } return E_OK; } @@ -1295,7 +1346,11 @@ void OnSingleRPCDone::set_region_vectorized_response(int64_t region_id, std::sha } void OnSingleRPCDone::retry_region_task(const std::string& remote_side) { - _fetcher_store->peer_status.set_cannot_access(_info.region_id(), remote_side); + if (_fetcher_store->broadcast_all_peer_without_raft) { + _rpc_ctrl->task_finish(this); + return; + } + _fetcher_store->peer_status.set_cannot_access(_info.region_id(), remote_side); FetcherStore::other_normal_peer_to_leader(_info, _addr); bthread_usleep(_retry_times * FLAGS_retry_interval_us); _rpc_ctrl->task_retry(this); @@ -1303,6 +1358,10 @@ void OnSingleRPCDone::retry_region_task(const std::string& remote_side) { void OnSingleRPCDone::retry_or_finish_task(const std::string& remote_side) { auto err = handle_response(remote_side); + if (_fetcher_store->broadcast_all_peer_without_raft) { + _rpc_ctrl->task_finish(this); + return; + } if (err == E_RETRY) { _rpc_ctrl->task_retry(this); } else { @@ -1441,12 +1500,23 @@ std::vector OnBatchRPCDone::select_addr() { std::string addr; bool resource_insulate_read = false; bool select_without_leader = false; - OnRPCDone::select_addr(info, addr, resource_insulate_read, select_without_leader); - RegionInfoData region_info_data; - region_info_data.region_info = info_ptr; - region_info_data.select_without_leader = select_without_leader; - region_info_data.resource_insulate_read = resource_insulate_read; - addr_batch_info_map[addr].push_back(region_info_data); + if (!_fetcher_store->broadcast_all_peer_without_raft) { + OnRPCDone::select_addr(info, addr, resource_insulate_read, select_without_leader); + RegionInfoData region_info_data; + region_info_data.region_info = info_ptr; + region_info_data.select_without_leader = select_without_leader; + region_info_data.resource_insulate_read = resource_insulate_read; + addr_batch_info_map[addr].emplace_back(region_info_data); + } else { + for (auto& peer : info.peers()) { + OnRPCDone::select_addr(info, addr, resource_insulate_read, select_without_leader); + RegionInfoData region_info_data; + region_info_data.region_info = info_ptr; + region_info_data.select_without_leader = select_without_leader; + region_info_data.resource_insulate_read = resource_insulate_read; + addr_batch_info_map[addr].emplace_back(region_info_data); + } + } } std::vector batch_rpc_done_list; for (auto& iter : addr_batch_info_map) { @@ -1486,6 +1556,10 @@ void OnBatchRPCDone::set_region_vectorized_response(int64_t region_id, std::shar } void OnBatchRPCDone::retry_region_task(const std::string& remote_side) { + if (_fetcher_store->broadcast_all_peer_without_raft) { + _rpc_ctrl->task_finish(this); + return; + } for (auto& info_ptr : _infos) { _fetcher_store->peer_status.set_cannot_access(info_ptr->region_id(), remote_side); FetcherStore::other_normal_peer_to_leader(*info_ptr, _addr); @@ -1502,6 +1576,10 @@ void OnBatchRPCDone::retry_region_task(const std::string& remote_side) { void OnBatchRPCDone::retry_or_finish_task(const std::string& remote_side) { std::vector need_retry_tasks; auto err = handle_response(remote_side, need_retry_tasks); + if (_fetcher_store->broadcast_all_peer_without_raft) { + _rpc_ctrl->task_finish(this); + return; + } if (need_retry_tasks.size() > 0) { for (auto& task : need_retry_tasks) { _rpc_ctrl->add_new_task(task); @@ -1686,11 +1764,22 @@ void FetcherStore::send_request(RuntimeState* state, } } else { for (auto info : infos) { - auto task = new OnSingleRPCDone(this, state, store_request, info, + if (!broadcast_all_peer_without_raft) { + auto task = new OnSingleRPCDone(this, state, store_request, info, info->region_id(), info->region_id(), start_seq_id, current_seq_id, op_type, need_check_memory); - rpc_ctrl.add_new_task(task); - traces.insert(task->get_trace()); + rpc_ctrl.add_new_task(task); + traces.insert(task->get_trace()); + } else { + for (auto peer : info->peers()) { + auto task = new OnSingleRPCDone(this, state, store_request, info, + info->region_id(), info->region_id(), start_seq_id, current_seq_id, + op_type, need_check_memory); + task->_addr = peer; + rpc_ctrl.add_new_task(task); + traces.insert(task->get_trace()); + } + } } } @@ -1732,10 +1821,14 @@ int FetcherStore::run_not_set_state(RuntimeState* state, received_arrow_data = false; batch_region_vectorized_response.clear(); shared_plan.reset(); + broadcast_all_peer_without_raft = false; if (region_infos.size() == 0) { DB_WARNING("region_infos size == 0, op_type:%s", pb::OpType_Name(op_type).c_str()); return 0; } + if (state->ctx() != nullptr && state->ctx()->kill_without_raft) { + broadcast_all_peer_without_raft = true; + } auto scan_node = store_request->get_node(pb::SCAN_NODE); ExecNode* filter_node = (scan_node == nullptr ? nullptr : scan_node->get_parent()); if (filter_node != nullptr diff --git a/src/exec/filter_node.cpp b/src/exec/filter_node.cpp index 456405250..1e980eb4b 100644 --- a/src/exec/filter_node.cpp +++ b/src/exec/filter_node.cpp @@ -33,7 +33,7 @@ namespace baikaldb { DECLARE_bool(open_nonboolean_sql_forbid); DECLARE_bool(open_nonboolean_sql_statistics); - +DECLARE_uint64(cut_huge_in_filter_size); int FilterNode::init(const pb::PlanNode& node) { int ret = ExecNode::init(node); if (ret < 0) { @@ -707,6 +707,22 @@ void FilterNode::transfer_pb(int64_t region_id, pb::PlanNode* pb_node) { } } +void FilterNode::cut_huge_in_condition(std::vector& huge_in_conditions) { + std::vector left_conditions; + left_conditions.reserve(_pruned_conjuncts.size()); + for (auto conjunct : _pruned_conjuncts) { + if (conjunct->node_type() == pb::IN_PREDICATE && conjunct->children_size() > 100000) { + huge_in_conditions.emplace_back(conjunct); + } else { + left_conditions.emplace_back(conjunct); + } + } + if (!huge_in_conditions.empty()) { + modifiy_pruned_conjuncts_by_index(left_conditions); + } + return; +} + inline bool FilterNode::need_copy(MemRow* row) { for (auto conjunct : _pruned_conjuncts) { ExprValue value = conjunct->get_value(row); @@ -827,6 +843,19 @@ int FilterNode::show_explain(QueryContext* ctx, std::vectorchildren_size(); + } + } + return all_expr_child_sum >= FLAGS_cut_huge_in_filter_size; +} + int FilterNode::arrow_steal_conjuncts(std::vector& conjuncts, int64_t& limit) { if (limit == -1 && get_limit() > 0) { limit = get_limit(); diff --git a/src/exec/full_export_node.cpp b/src/exec/full_export_node.cpp index 5df55e54d..2126ec12b 100755 --- a/src/exec/full_export_node.cpp +++ b/src/exec/full_export_node.cpp @@ -176,7 +176,6 @@ int FullExportNode::calc_last_key(RuntimeState* state, MemRow* mem_row) { pos_index.set_left_open(true); pos_index.mutable_ranges(0)->set_left_key(key.data()); pos_index.mutable_ranges(0)->set_left_full(key.get_full()); - pos_index.mutable_ranges(0)->mutable_partition_ids()->Clear(); pos_index.mutable_ranges(0)->add_partition_ids(_current_partition); pos_index.SerializeToString(&scan_index_info.raw_index); diff --git a/src/exec/index_ddl_manager_node.cpp b/src/exec/index_ddl_manager_node.cpp index a2121dbe7..184c33b13 100755 --- a/src/exec/index_ddl_manager_node.cpp +++ b/src/exec/index_ddl_manager_node.cpp @@ -27,6 +27,15 @@ IndexDDLManagerNode::IndexDDLManagerNode() { IndexDDLManagerNode::~IndexDDLManagerNode() { } +bool IndexDDLManagerNode::is_ddl_delete(const pb::ScanNode& scan_node) { + if (scan_node.ddl_work_type() == pb::DDL_COLUMN + && scan_node.has_column_ddl_info() + && scan_node.column_ddl_info().update_exprs_size() == 0) { + return true; + } + return false; +} + int IndexDDLManagerNode::open(RuntimeState* state) { state->set_is_ddl_work(true); int ret = 0; @@ -54,13 +63,20 @@ int IndexDDLManagerNode::open(RuntimeState* state) { std::vector delete_records; insert_records.reserve(limit); std::map region_infos; + + const pb::ScanNode& scan_node_pb = scan_node->pb_node().derive_node().scan_node(); + bool ddl_delete = is_ddl_delete(scan_node_pb); + if (ddl_delete && init_lock_nodes_if_not_exist() < 0) { + DB_FATAL("ddl delete init lock_nodes failed"); + return ret; + } SmartRecord record_template = SchemaFactory::get_instance()->new_record(main_table_id); auto tuple_id = 0; int32_t ddl_scan_size = 0; std::string max_pk_str; std::string max_record; - bool global_ddl_with_ttl = _fetcher_store.region_id_ttl_timestamp_batch.size() > 0 ? true : false; + bool global_ddl_with_ttl = _fetcher_store.region_id_ttl_timestamp_batch.size() > 0; std::map record_ttl_map; auto pk_info = SchemaFactory::get_instance()->get_index_info(_table_id); @@ -213,6 +229,119 @@ int IndexDDLManagerNode::open(RuntimeState* state) { } } } + if (ddl_delete && !insert_records.empty()) { + delete_records.swap(insert_records); + int64_t affected_rows = send_request(state, _lock_primary.get(), std::vector{}, delete_records); + if (affected_rows < 0) { + DB_WARNING("DDL delete pk failed."); + return -1; + } + delete_records.swap(_fetcher_store.index_records[_table_id]); + for (auto& lock_secondary: _lock_secondaries) { + affected_rows = send_request(state, lock_secondary.get(), std::vector{}, delete_records); + if (affected_rows < 0) { + DB_WARNING("DDL delete global index failed."); + return -1; + } + } + // TODO: 需要处理binlog? + } return ret; } + +int IndexDDLManagerNode::create_lock_primary_node( + int64_t table_id, + std::unique_ptr& primary_node) const { + auto table_info = _factory->get_table_info_ptr(table_id); + if (table_info == nullptr) { + return -1; + } + primary_node.reset(new (std::nothrow) LockPrimaryNode); + if (primary_node == nullptr) { + DB_WARNING("create manager_node failed"); + return -1; + } + + std::vector local_affected_indexs; + local_affected_indexs.reserve(3); + for (auto index_id: table_info->indices) { + auto index_info = _factory->get_index_info_ptr(index_id); + if (index_info == nullptr) { + return -1; + } + if (index_info->index_hint_status == pb::IHS_VIRTUAL) { + DB_NOTICE("index info is virtual, skip."); + continue; + } + if (index_info->index_hint_status == pb::IHS_DISABLE + && index_info->state == pb::IS_DELETE_LOCAL) { + continue; + } + if (index_info->is_global) { + continue; + } + local_affected_indexs.emplace_back(index_id); + } + + pb::PlanNode plan_node; + plan_node.set_node_type(pb::LOCK_PRIMARY_NODE); + plan_node.set_limit(-1); + plan_node.set_num_children(0); + auto lock_primary_node = plan_node.mutable_derive_node()->mutable_lock_primary_node(); + lock_primary_node->set_lock_type(pb::LOCK_GET_DML); + lock_primary_node->set_table_id(table_id); + primary_node->init(plan_node); + primary_node->set_affected_index_ids(local_affected_indexs); + return 0; +} + +int IndexDDLManagerNode::create_lock_secondary_node( + int64_t table_id, + std::vector>& secondary_nodes) const { + auto table_info = _factory->get_table_info_ptr(table_id); + if (table_info == nullptr) { + return -1; + } + std::vector global_affected_indexs; + global_affected_indexs.reserve(3); + for (auto index_id : table_info->indices) { + auto index_info = _factory->get_index_info_ptr(index_id); + if (index_info == nullptr) { + return -1; + } + if (index_info->index_hint_status == pb::IHS_VIRTUAL) { + DB_NOTICE("index info is virtual, skip."); + continue; + } + if (index_info->index_hint_status == pb::IHS_DISABLE + && index_info->state == pb::IS_DELETE_LOCAL) { + continue; + } + if (index_info->is_global) { + if (index_info->state == pb::IS_NONE) { + DB_NOTICE("index info is NONE, skip."); + continue; + } + global_affected_indexs.emplace_back(index_id); + } + } + for (auto index_id: global_affected_indexs) { + std::unique_ptr secondary_node(new (std::nothrow) LockSecondaryNode); + if (secondary_node == nullptr) { + DB_WARNING("create manager_node failed"); + return -1; + } + pb::PlanNode plan_node; + plan_node.set_node_type(pb::LOCK_SECONDARY_NODE); + plan_node.set_limit(-1); + plan_node.set_num_children(0); + auto lock_secondary_node = plan_node.mutable_derive_node()->mutable_lock_secondary_node(); + lock_secondary_node->set_lock_type(pb::LOCK_DML); + lock_secondary_node->set_global_index_id(index_id); + lock_secondary_node->set_table_id(table_id); + secondary_node->init(plan_node); + secondary_nodes.emplace_back(std::move(secondary_node)); + } + return 0; +} } diff --git a/src/exec/join_node.cpp b/src/exec/join_node.cpp index 84ef7c72a..1f9f155df 100644 --- a/src/exec/join_node.cpp +++ b/src/exec/join_node.cpp @@ -29,7 +29,7 @@ namespace baikaldb { DECLARE_int32(arrow_multi_threads); - +DECLARE_bool(join_key_cast_like_mysql); int JoinNode::init(const pb::PlanNode& node) { int ret = 0; ret = Joiner::init(node); @@ -268,6 +268,97 @@ bool JoinNode::can_use_arrow_vector(RuntimeState* state) { return true; } +void JoinNode::get_hash_partitions(NodePartitionProperty& outer_property, + NodePartitionProperty& inner_property, + const std::unordered_set& cast_string_hash_columns) { + outer_property.type = _partition_property.type; + inner_property.type = _partition_property.type; + if (_partition_property.hash_partition_propertys.size() == 2) { + outer_property.hash_partition_propertys.emplace_back(_partition_property.hash_partition_propertys[0]); + inner_property.hash_partition_propertys.emplace_back(_partition_property.hash_partition_propertys[1]); + outer_property.type = _partition_property.hash_partition_propertys[0]->type; + inner_property.type = _partition_property.hash_partition_propertys[1]->type; + } + outer_property.need_cast_string_columns = _partition_property.need_cast_string_columns; + inner_property.need_cast_string_columns = _partition_property.need_cast_string_columns; + for (auto& name : cast_string_hash_columns) { + outer_property.need_cast_string_columns.insert(name); + inner_property.need_cast_string_columns.insert(name); + if (_on_condition_column_map.count(name) > 0) { + outer_property.need_cast_string_columns.insert(_on_condition_column_map[name]); + inner_property.need_cast_string_columns.insert(_on_condition_column_map[name]); + } + } + } + +void JoinNode::make_broadcast_join_property(std::shared_ptr* small_table_property, + ExecNode* small_table_node, + std::shared_ptr* other_property, + NodePartitionProperty* other_child_property) { + // small table使用broadcast + (*small_table_property)->type = pb::BroadcastPartitionType; + // 去掉子树的hash属性, 如agg + small_table_node->set_child_node_property_any_type(); + + // other table继承other_child_property的属性 + if (other_child_property->type == pb::AnyType) { + (*other_property)->type = pb::RandomPartitionType; + _partition_property.type = pb::RandomPartitionType; + } else { + (*other_property)->type = other_child_property->type; + _partition_property.type = other_child_property->type; + for (auto& property : other_child_property->hash_partition_propertys) { + if (property->type == pb::HashPartitionType) { + *other_property = property; + break; + } + } + } + return; +} + +bool JoinNode::suitable_for_broadcast_join(QueryContext* ctx) { + auto factory = SchemaFactory::get_instance(); + if (factory == nullptr) { + return false; + } + std::vector outer_scan_nodes; + std::vector inner_scan_nodes; + _outer_node->get_node_pass_subquery(pb::SCAN_NODE, outer_scan_nodes); + _inner_node->get_node_pass_subquery(pb::SCAN_NODE, inner_scan_nodes); + NodePartitionProperty* outer_partition_property = _outer_node->partition_property(); + NodePartitionProperty* inner_partition_property = _inner_node->partition_property(); + bool outer_is_small = false; + bool inner_is_small = false; + if (outer_scan_nodes.size() == 1 + && _join_type == pb::INNER_JOIN + && (!_use_index_join || ctx->runtime_state->sign_exec_type == SignExecType::SIGN_EXEC_ARROW_FORCE_NO_INDEX_JOIN)) { + // index join 驱动表不支持broadcast join, 因为非驱动表不走exchange) + ScanNode* outer_scan_node = static_cast(outer_scan_nodes[0]); + outer_is_small = outer_scan_node->is_rocksdb_scan_node() + && factory->table_suitable_for_broadcast_join(outer_scan_node->table_id()); + } + if (inner_scan_nodes.size() == 1) { + ScanNode* inner_scan_node = static_cast(inner_scan_nodes[0]); + inner_is_small = inner_scan_node->is_rocksdb_scan_node() + && factory->table_suitable_for_broadcast_join(inner_scan_node->table_id()); + } + if (outer_is_small || inner_is_small) { + std::shared_ptr outer = std::make_shared(); + std::shared_ptr inner = std::make_shared(); + + if (outer_is_small) { + make_broadcast_join_property(&outer, _outer_node, &inner, inner_partition_property); + } else { + make_broadcast_join_property(&inner, _inner_node, &outer, outer_partition_property); + } + _partition_property.hash_partition_propertys.emplace_back(outer); // outer对应的hash属性 + _partition_property.hash_partition_propertys.emplace_back(inner); // inner对应的hash属性 + return true; + } + return false; +} + int JoinNode::set_partition_property_and_schema(QueryContext* ctx) { for (auto& c : _children) { if (0 != c->set_partition_property_and_schema(ctx)) { @@ -275,6 +366,13 @@ int JoinNode::set_partition_property_and_schema(QueryContext* ctx) { } _data_schema.insert(c->data_schema().begin(), c->data_schema().end()); } + // 有可能separate会加入selectmanagernode, 这里重新设置 + _outer_node = _children[0]; + _inner_node = _children[1]; + if (_join_type == pb::RIGHT_JOIN) { + _outer_node = _children[1]; + _inner_node = _children[0]; + } if (_outer_node->partition_property()->has_no_input_data && _join_type != pb::FULL_JOIN) { _partition_property.has_no_input_data = true; @@ -283,6 +381,9 @@ int JoinNode::set_partition_property_and_schema(QueryContext* ctx) { _partition_property.set_single_partition(); return 0; } + if (suitable_for_broadcast_join(ctx)) { + return 0; + } // 当前join向量化, key只能是slotref, 但是类型不一致会都做一次类型转换 // 类型不一样, 如int 1和string 1, es产生的hash值不一样, 需要指定cast类型发到es std::shared_ptr outer = std::make_shared(); @@ -308,25 +409,24 @@ int JoinNode::set_partition_property_and_schema(QueryContext* ctx) { NodePartitionProperty* inner_partition_property = _inner_node->partition_property(); bool is_same_or_shrinked = shrink_partition_property(outer, outer_partition_property); if (is_same_or_shrinked) { - if (outer->hash_columns.size() != inner->hash_columns.size()) { - // outer hash partition收缩了, 需要联动更改inner的hash partition - std::unordered_map new_hash_columns; - std::vector new_ordered_col_names; - for (const auto& outer_col : outer->ordered_hash_columns) { - std::string inner_col_name = _on_condition_column_map[outer_col]; - if (inner_col_name.empty()) { - DB_FATAL("column map has no match col: %s", outer_col.c_str()); - } - new_hash_columns[inner_col_name] = inner->hash_columns[inner_col_name]; - new_ordered_col_names.emplace_back(inner_col_name); + // outer hash partition调整了, 需要联动更改inner的hash partition + std::unordered_map new_hash_columns; + std::vector new_ordered_col_names; + new_ordered_col_names.reserve(outer->ordered_hash_columns.size()); + for (const auto& outer_col : outer->ordered_hash_columns) { + std::string inner_col_name = _on_condition_column_map[outer_col]; + if (inner_col_name.empty()) { + DB_FATAL("column map has no match col: %s", outer_col.c_str()); } - inner->hash_columns = new_hash_columns; - inner->ordered_hash_columns = new_ordered_col_names; + new_hash_columns[inner_col_name] = inner->hash_columns[inner_col_name]; + new_ordered_col_names.emplace_back(inner_col_name); } + inner->hash_columns = new_hash_columns; + inner->ordered_hash_columns = new_ordered_col_names; } else { - shrink_partition_property(inner, inner_partition_property); - if (outer->hash_columns.size() != inner->hash_columns.size()) { - // inner hash partition收缩了, 需要联动更改outer的hash partition + is_same_or_shrinked = shrink_partition_property(inner, inner_partition_property); + if (is_same_or_shrinked) { + // inner hash partition调整了, 需要联动更改outer的hash partition std::unordered_map new_hash_columns; std::vector new_ordered_col_names; for (const auto& inner_col : inner->ordered_hash_columns) { @@ -364,10 +464,9 @@ int JoinNode::build_table_arrow_declaration(RuntimeState* state, ExecNode* node, std::unordered_set& tuple_ids, std::vector& mem_rows, - const std::unordered_map>& cast_string_slot_ids, + std::unordered_map& projection_temp_col, bool need_add_index_colletor_node, - bool remove_useless_sort, - bool need_add_join_key) { + bool remove_useless_sort) { START_LOCAL_TRACE(get_trace(), state->get_trace_cost(), OPEN_TRACE, nullptr); if (node->node_exec_type() == pb::EXEC_ROW) { // 返回行, 行转source node @@ -401,7 +500,7 @@ int JoinNode::build_table_arrow_declaration(RuntimeState* state, state->acero_declarations.pop_back(); } } - if (cast_string_slot_ids.size() > 0 || need_add_join_key) { + if (projection_temp_col.size() > 0) { std::vector exprs; std::vector names; for (auto tuple_id : tuple_ids) { @@ -414,18 +513,11 @@ int JoinNode::build_table_arrow_declaration(RuntimeState* state, std::string name = std::to_string(tuple_id) + "_" + std::to_string(slot.slot_id()); exprs.emplace_back(arrow::compute::field_ref(name)); names.emplace_back(name); - auto iter = cast_string_slot_ids.find(tuple_id); - if (iter != cast_string_slot_ids.end() && iter->second.find(slot.slot_id()) != iter->second.end()) { - // 额外加cast string列 - exprs.emplace_back(arrow::compute::call("cast", {arrow::compute::field_ref(name)}, - arrow::compute::CastOptions::Unsafe(arrow::large_binary()))); - names.emplace_back(name + "_cast"); - } } } - if (need_add_join_key) { - exprs.emplace_back(arrow::compute::literal(true)); - names.emplace_back("__fake_join_key"); + for (auto& col : projection_temp_col) { + exprs.emplace_back(col.second); + names.emplace_back(col.first); } arrow::acero::Declaration dec{"project", arrow::acero::ProjectNodeOptions{exprs, names}}; LOCAL_TRACE_ARROW_PLAN(dec); @@ -466,15 +558,273 @@ void JoinNode::get_need_add_index_collector_cond_nodes(ExecNode* node, std::set< return; } +void JoinNode::handle_join_equal_key_cast_type(ExprNode* outer_expr, + ExprNode* inner_expr, + std::string outer_tmp_col_name, + std::string inner_tmp_col_name, + std::vector& outer_keys, + std::vector& inner_keys, + std::unordered_map& outer_projection_temp_col, + std::unordered_map& inner_projection_temp_col) { + auto outer_type = outer_expr->col_type(); + auto inner_type = inner_expr->col_type(); + auto cast_type = pb::STRING; + if (is_uint(outer_type) && is_uint(inner_type)) { + // all cast uint64_t + cast_type = pb::UINT64; + } else if (is_signed(outer_type) && is_signed(inner_type)) { + // all cast int64_t + cast_type = pb::INT64; + } else if (FLAGS_join_key_cast_like_mysql && is_string(outer_type) && !is_string(inner_type)) { + // 和mysql一样, 将string转换为数值型 + // outer string cast numberic + cast_type = inner_type; + if (is_int(inner_type) || is_double(inner_type)) { + cast_type = pb::DOUBLE; + } + } else if (FLAGS_join_key_cast_like_mysql && is_string(inner_type) && !is_string(outer_type)) { + // inner string cast numberic + cast_type = outer_type; + if (is_int(outer_type) || is_double(outer_type)) { + cast_type = pb::DOUBLE; + } + } + outer_projection_temp_col[outer_tmp_col_name] = arrow_cast(outer_expr->arrow_expr(), outer_type, cast_type); + inner_projection_temp_col[inner_tmp_col_name] = arrow_cast(inner_expr->arrow_expr(), inner_type, cast_type); + outer_keys.emplace_back(arrow::FieldRef(outer_tmp_col_name)); + inner_keys.emplace_back(arrow::FieldRef(inner_tmp_col_name)); + return; +} + +int JoinNode::handle_join_filter_key_expressions(RuntimeState* state, + int idx, + ExprNode* condition, + bool left_child_is_outer, + std::vector& outer_keys, + std::vector& inner_keys, + std::vector& sub_exprs, + std::unordered_map& outer_projection_temp_col, + std::unordered_map& inner_projection_temp_col) { + auto outer_child_expr = condition->children(0); + auto inner_child_expr = condition->children(1); + if (!left_child_is_outer) { + outer_child_expr = condition->children(1); + inner_child_expr = condition->children(0); + } + if (outer_child_expr->transfer_to_arrow_expression() < 0) { + DB_FATAL_STATE(state, "transfer join filter outer child to arrow expression failed"); + return -1; + } + if (inner_child_expr->transfer_to_arrow_expression() < 0) { + DB_FATAL_STATE(state, "transfer join filter inner child to arrow expression failed"); + return -1; + } + auto outer_child_tmp_col_name = "outer_tmp_" + std::to_string(idx); + auto inner_child_tmp_col_name = "inner_tmp_" + std::to_string(idx); + auto outer_type = outer_child_expr->col_type(); + auto inner_type = inner_child_expr->col_type(); + auto fn_op = static_cast(condition)->fn().fn_op(); + if (fn_op == parser::FT_EQ) { + // 放在key里 + outer_keys.emplace_back(arrow::FieldRef(outer_child_tmp_col_name)); + inner_keys.emplace_back(arrow::FieldRef(inner_child_tmp_col_name)); + if (outer_type == inner_type) { + outer_projection_temp_col[outer_child_tmp_col_name] = outer_child_expr->arrow_expr(); + inner_projection_temp_col[inner_child_tmp_col_name] = inner_child_expr->arrow_expr(); + } else { + handle_join_equal_key_cast_type(outer_child_expr, + inner_child_expr, + outer_child_tmp_col_name, + inner_child_tmp_col_name, + outer_keys, + inner_keys, + outer_projection_temp_col, + inner_projection_temp_col); + } + } else { + // 放在filter condition + outer_projection_temp_col[outer_child_tmp_col_name] = outer_child_expr->arrow_expr(); + inner_projection_temp_col[inner_child_tmp_col_name] = inner_child_expr->arrow_expr(); + std::string arrow_func_name; + switch (fn_op) + { + case parser::FT_GE: + arrow_func_name = "greater_equal"; + break; + case parser::FT_GT: + arrow_func_name = "greater"; + break; + case parser::FT_LE: + arrow_func_name = "less_equal"; + break; + case parser::FT_LT: + arrow_func_name = "less"; + break; + case parser::FT_NE: + arrow_func_name = "not_equal"; + break; + default: + DB_FATAL_STATE(state, "unsupported fn_op"); + return -1; + } + if (left_child_is_outer) { + sub_exprs.emplace_back(arrow::compute::call(arrow_func_name, + {arrow::compute::field_ref(outer_child_tmp_col_name), arrow::compute::field_ref(inner_child_tmp_col_name)})); + } else { + sub_exprs.emplace_back(arrow::compute::call(arrow_func_name, + {arrow::compute::field_ref(inner_child_tmp_col_name), arrow::compute::field_ref(outer_child_tmp_col_name)})); + } + } + return 0; +} + +int JoinNode::try_transfer_filter_to_column_first(RuntimeState* state, + int idx, + ExprNode* condition, + std::vector& outer_keys, + std::vector& inner_keys, + std::vector& sub_exprs, + std::unordered_map& outer_projection_temp_col, + std::unordered_map& inner_projection_temp_col) { + if (condition->node_type() != pb::FUNCTION_CALL + || !static_cast(condition)->is_compare_op()) { + return 1; + } + if (condition->children_size() != 2) { + return 1; + } + ExprNode* left_child = condition->children(0); + ExprNode* right_child = condition->children(1); + std::unordered_set left_tuple_ids; + std::unordered_set right_tuple_ids; + left_child->get_all_tuple_ids(left_tuple_ids); + right_child->get_all_tuple_ids(right_tuple_ids); + + bool left_tuple_ids_has_inner = false; + bool left_tuple_ids_has_outer = false; + for (auto tuple_id : left_tuple_ids) { + if (_inner_tuple_ids.count(tuple_id) == 1) { + left_tuple_ids_has_inner = true; + } + if (_outer_tuple_ids.count(tuple_id) == 1) { + left_tuple_ids_has_outer = true; + } + } + if (left_tuple_ids_has_inner && left_tuple_ids_has_outer) { + return 1; + } + + bool right_tuple_ids_has_inner = false; + bool right_tuple_ids_has_outer = false; + for (auto tuple_id : right_tuple_ids) { + if (_inner_tuple_ids.count(tuple_id) == 1) { + right_tuple_ids_has_inner = true; + } + if (_outer_tuple_ids.count(tuple_id) == 1) { + right_tuple_ids_has_outer = true; + } + } + if (right_tuple_ids_has_inner && right_tuple_ids_has_outer) { + return 1; + } + + // 双边filter, 如 funcA(table_a.id) > funcB(table_b.id) + if ((left_tuple_ids_has_inner && right_tuple_ids_has_outer) + || (left_tuple_ids_has_outer && right_tuple_ids_has_inner)) { + // left_child, right_child分别作为两表临时列 + return handle_join_filter_key_expressions(state, idx, condition, left_tuple_ids_has_outer, + outer_keys, inner_keys, sub_exprs, + outer_projection_temp_col, inner_projection_temp_col); + } + + // 单边filter, 如 funcA(table_a.id) > 10 + if (condition->transfer_to_arrow_expression() < 0) { + DB_FATAL_STATE(state, "join filter condition expr transfer arrow fail"); + return -1; + } + auto bool_expr = arrow_cast(condition->arrow_expr(), condition->col_type(), pb::BOOL); + if (left_tuple_ids_has_outer || right_tuple_ids_has_outer) { + // condition is true作为outer table临时列; 临时列直接作为join filter + std::string tmp_col_name = "outer_tmp_" + std::to_string(idx); + outer_projection_temp_col[tmp_col_name] = bool_expr; + sub_exprs.emplace_back(arrow::compute::field_ref(tmp_col_name)); + } else if (left_tuple_ids_has_inner || right_tuple_ids_has_inner) { + // condition is true作为inner table临时列; 临时列直接作为join filter + std::string tmp_col_name = "inner_tmp_" + std::to_string(idx); + inner_projection_temp_col[tmp_col_name] = bool_expr; + sub_exprs.emplace_back(arrow::compute::field_ref(tmp_col_name)); + } + return 0; +} + int JoinNode::build_arrow_declaration(RuntimeState* state) { START_LOCAL_TRACE_WITH_PARTITION_PROPERTY(get_trace(), state->get_trace_cost(), &_partition_property, OPEN_TRACE, nullptr); - std::unordered_map> outer_cast_slot_ids; - std::unordered_map> inner_cast_slot_ids; + std::vector outer_keys; + std::vector inner_keys; + std::vector sub_exprs; + std::unordered_map outer_projection_temp_col; + std::unordered_map inner_projection_temp_col; + outer_keys.reserve(_outer_equal_slot.size()); + inner_keys.reserve(_inner_equal_slot.size()); + sub_exprs.reserve(_conditions.size()); + // 处理slot_ref on condition first for (int i = 0; i < _outer_equal_slot.size(); ++i) { - if (_outer_equal_slot[i]->col_type() != _inner_equal_slot[i]->col_type()) { - outer_cast_slot_ids[_outer_equal_slot[i]->tuple_id()].insert(_outer_equal_slot[i]->slot_id()); - inner_cast_slot_ids[_inner_equal_slot[i]->tuple_id()].insert(_inner_equal_slot[i]->slot_id()); + auto& outer_expr = _outer_equal_slot[i]; + auto& inner_expr = _inner_equal_slot[i]; + if (outer_expr->transfer_to_arrow_expression() < 0) { + DB_FATAL_STATE(state, "_outer_equal_slot[%d] expr transfer arrow fail", i); + return -1; + } + if (inner_expr->transfer_to_arrow_expression() < 0) { + DB_FATAL_STATE(state, "_inner_equal_slot[%d] expr transfer arrow fail", i); + return -1; + } + std::string outer_col_name = std::to_string(outer_expr->tuple_id()) + "_" + std::to_string(outer_expr->slot_id()); + std::string inner_col_name = std::to_string(inner_expr->tuple_id()) + "_" + std::to_string(inner_expr->slot_id()); + auto outer_type = outer_expr->col_type(); + auto inner_type = inner_expr->col_type(); + if (outer_type == inner_type) { + outer_keys.emplace_back(arrow::FieldRef(outer_col_name)); + inner_keys.emplace_back(arrow::FieldRef(inner_col_name)); + } else if (outer_type != inner_type) { + std::string outer_tmp_col_name = outer_col_name + "_cast"; + std::string inner_tmp_col_name = inner_col_name + "_cast"; + handle_join_equal_key_cast_type(outer_expr, + inner_expr, + outer_tmp_col_name, + inner_tmp_col_name, + outer_keys, + inner_keys, + outer_projection_temp_col, + inner_projection_temp_col); + } + } + // 处理剩余的filter condition + for (auto idx = 0; idx < _conditions.size(); ++idx) { + int ret = try_transfer_filter_to_column_first(state, idx, _conditions[idx], + outer_keys, inner_keys, sub_exprs, + outer_projection_temp_col, inner_projection_temp_col); + if (ret < 0) { + return -1; + } + if (ret == 0) { + continue; } + ret = _conditions[idx]->transfer_to_arrow_expression(); + if (ret < 0) { + DB_FATAL_STATE(state, "expr transfer arrow fail, ret:%d", ret); + return ret; + } + sub_exprs.emplace_back(_conditions[idx]->arrow_expr()); + } + + if (outer_keys.empty()) { + outer_keys.emplace_back(arrow::FieldRef("__fake_join_key")); + outer_projection_temp_col["__fake_join_key"] = arrow::compute::literal(1); + } + if (inner_keys.empty()) { + inner_keys.emplace_back(arrow::FieldRef("__fake_join_key")); + inner_projection_temp_col["__fake_join_key"] = arrow::compute::literal(1); } arrow::acero::Declaration outer_dec; arrow::acero::Declaration inner_dec; @@ -483,10 +833,9 @@ int JoinNode::build_arrow_declaration(RuntimeState* state) { _outer_node, _outer_tuple_ids, _outer_tuple_data, - outer_cast_slot_ids, + outer_projection_temp_col, _need_add_index_collector_node, - true, - _outer_equal_slot.empty())) { + true)) { DB_FATAL_STATE(state, "outer join node build arrow declaration failed"); return -1; } @@ -495,60 +844,13 @@ int JoinNode::build_arrow_declaration(RuntimeState* state) { _inner_node, _inner_tuple_ids, _inner_tuple_data, - inner_cast_slot_ids, + inner_projection_temp_col, false, - false, - _inner_equal_slot.empty())) { + false)) { DB_FATAL_STATE(state, "inner join node build arrow declaration failed"); return -1; } arrow::acero::JoinType join_type; - std::vector outer_keys; - std::vector inner_keys; - for (auto slot_ref : _outer_equal_slot) { - int ret = slot_ref->transfer_to_arrow_expression(); - if (ret < 0) { - DB_FATAL_STATE(state, "expr transfer arrow fail, ret:%d", ret); - return ret; - } - std::string name = std::to_string(slot_ref->tuple_id()) + "_" + std::to_string(slot_ref->slot_id()); - auto iter = outer_cast_slot_ids.find(slot_ref->tuple_id()); - if (iter != outer_cast_slot_ids.end() && iter->second.find(slot_ref->slot_id()) != iter->second.end()) { - outer_keys.emplace_back(arrow::FieldRef(name + "_cast")); - } else { - outer_keys.emplace_back(arrow::FieldRef(name)); - } - } - for (auto slot_ref : _inner_equal_slot) { - int ret = slot_ref->transfer_to_arrow_expression(); - if (ret < 0) { - DB_FATAL_STATE(state, "expr transfer arrow fail, ret:%d", ret); - return ret; - } - std::string name = std::to_string(slot_ref->tuple_id()) + "_" + std::to_string(slot_ref->slot_id()); - auto iter = inner_cast_slot_ids.find(slot_ref->tuple_id()); - if (iter != inner_cast_slot_ids.end() && iter->second.find(slot_ref->slot_id()) != iter->second.end()) { - inner_keys.emplace_back(arrow::FieldRef(name + "_cast")); - } else { - inner_keys.emplace_back(arrow::FieldRef(name)); - } - } - if (outer_keys.empty()) { - outer_keys.emplace_back(arrow::FieldRef("__fake_join_key")); - } - if (inner_keys.empty()) { - inner_keys.emplace_back(arrow::FieldRef("__fake_join_key")); - } - std::vector sub_exprs; - for (auto& condition : _conditions) { - int ret = condition->transfer_to_arrow_expression(); - if (ret < 0) { - DB_FATAL_STATE(state, "expr transfer arrow fail, ret:%d", ret); - return ret; - } - sub_exprs.emplace_back(condition->arrow_expr()); - } - switch (_join_type) { case pb::LEFT_JOIN: case pb::RIGHT_JOIN: @@ -623,7 +925,7 @@ int JoinNode::hash_join(RuntimeState* state) { std::vector scan_nodes; _outer_node->get_node(pb::SCAN_NODE, scan_nodes); bool index_has_null = false; - if (do_plan_router(state, scan_nodes, index_has_null, _is_explain) != 0) { + if (do_plan_router(state, scan_nodes, index_has_null, _is_explain, false) != 0) { DB_WARNING("Fail to do_plan_router"); return -1; } @@ -739,7 +1041,7 @@ int JoinNode::nested_loop_join(RuntimeState* state) { std::vector scan_nodes; _outer_node->get_node(pb::SCAN_NODE, scan_nodes); bool index_has_null = false; - if (do_plan_router(state, scan_nodes, index_has_null, _is_explain) != 0) { + if (do_plan_router(state, scan_nodes, index_has_null, _is_explain, false) != 0) { DB_WARNING("Fail to do_plan_router"); return -1; } @@ -811,7 +1113,7 @@ int JoinNode::nested_loop_join(RuntimeState* state) { int JoinNode::open(RuntimeState* state) { if (_children.size() < 2) { - DB_WARNING("join node children size is %d", _children.size()); + DB_WARNING("join node children size is %lu", _children.size()); return -1; } TimeCost join_time_cost; diff --git a/src/exec/joiner.cpp b/src/exec/joiner.cpp index a9f9eb0b6..6a3433136 100644 --- a/src/exec/joiner.cpp +++ b/src/exec/joiner.cpp @@ -23,11 +23,12 @@ #include "logical_planner.h" #include "literal.h" #include "vectorize_helpper.h" +#include "condition_optimizer.h" namespace baikaldb { DECLARE_bool(use_arrow_vector); DECLARE_int64(row_number_to_check_memory); - +DEFINE_bool(join_key_cast_like_mysql, false, "Cast join key like MySQL behavior, default: false"); int Joiner::init(const pb::PlanNode& node) { int ret = 0; ret = ExecNode::init(node); @@ -166,6 +167,7 @@ int Joiner::strip_out_equal_slots() { auto iter = _conditions.begin(); _outer_equal_slot.clear(); _inner_equal_slot.clear(); + _equal_slot_cast_types.clear(); while (iter != _conditions.end()) { auto expr = *iter; if (expr_is_equal_condition_and_build_slot(expr)) { @@ -184,10 +186,39 @@ int Joiner::strip_out_equal_slots() { _use_hash_map = false; DB_WARNING("has not eq"); } + // key cast类型: + // mysql: string和数字, string转double; string和时间, string转时间 + for (auto i = 0; i < _outer_equal_slot.size(); ++i) { + auto outer_type = _outer_equal_slot[i]->col_type(); + auto inner_type = _inner_equal_slot[i]->col_type(); + if (outer_type == inner_type) { + _equal_slot_cast_types.emplace_back(outer_type); + } else { + if (is_signed(outer_type) && is_signed(inner_type)) { + _equal_slot_cast_types.emplace_back(pb::INT64); + } else if (is_uint(outer_type) && is_uint(inner_type)) { + _equal_slot_cast_types.emplace_back(pb::UINT64); + } else if (FLAGS_join_key_cast_like_mysql && is_string(outer_type) && !is_string(inner_type)) { + if (is_int(inner_type) || is_double(inner_type)) { + _equal_slot_cast_types.emplace_back(pb::DOUBLE); + } else { + _equal_slot_cast_types.emplace_back(inner_type); + } + } else if (FLAGS_join_key_cast_like_mysql && !is_string(outer_type) && is_string(inner_type)) { + if (is_int(outer_type) || is_double(outer_type)) { + _equal_slot_cast_types.emplace_back(pb::DOUBLE); + } else { + _equal_slot_cast_types.emplace_back(outer_type); + } + } else { + _equal_slot_cast_types.emplace_back(pb::STRING); + } + } + } return 0; } -int Joiner::do_plan_router(RuntimeState* state, const std::vector& scan_nodes, bool& index_has_null, bool is_explain) { +int Joiner::do_plan_router(RuntimeState* state, const std::vector& scan_nodes, bool& index_has_null, bool is_explain, bool in_acero) { QueryContext* ctx = state->ctx(); if (ctx == nullptr) { DB_FATAL("ctx is nullptr"); @@ -195,10 +226,10 @@ int Joiner::do_plan_router(RuntimeState* state, const std::vector& sc } //重新做路由选择 for (auto& exec_node : scan_nodes) { - if (!static_cast(exec_node)->is_rocksdb_scan_node()) { + ScanNode* scan_node = static_cast(exec_node); + if (!scan_node->is_rocksdb_scan_node() && !scan_node->is_file_scan_node()) { continue; } - RocksdbScanNode* scan_node = static_cast(exec_node); ExecNode* parent_node_ptr = scan_node->get_parent(); FilterNode* filter_node = nullptr; if (parent_node_ptr->node_type() == pb::WHERE_FILTER_NODE @@ -252,6 +283,8 @@ int Joiner::do_plan_router(RuntimeState* state, const std::vector& sc auto region_infos = scan_node->region_infos(); //更改scan_node对应的fethcer_node的region信息 related_manager_node->set_region_infos(region_infos); + //判断是否有in条件太大可能造成db oom, 这种情况调整计划不下推in条件 + ConditionOptimizer(ctx).adjust_huge_in_condition(scan_node, in_acero); } } return 0; @@ -295,7 +328,7 @@ int Joiner::runtime_filter(RuntimeState* state, ExecNode* node, std::vectormemory_limit_exceeded(check_mem_batch, batch_estimate_size)) { DB_WARNING("memory limit exceeded, logid: %lu, in_values size:%lu, estimate_size:%ld", state->log_id(), in_values.size(), total_estimate_size); + state->error_code = ER_TOO_BIG_SELECT; + state->error_msg.str("select reach memory limit"); ExprNode::destroy_tree(conjunct); return -1; } @@ -505,6 +540,8 @@ int Joiner::construct_in_condition(RuntimeState* state, if (0 != state->memory_limit_exceeded(check_mem_batch, batch_estimate_size)) { DB_WARNING("memory limit exceeded, logid: %lu, in_values size:%lu, estimate_size:%ld", state->log_id(), in_values.size(), total_estimate_size); + state->error_code = ER_TOO_BIG_SELECT; + state->error_msg.str("select reach memory limit"); ExprNode::destroy_tree(conjunct); return -1; } @@ -571,11 +608,16 @@ bool Joiner::is_satisfy_filter(MemRow* row) { void Joiner::encode_hash_key(MemRow* row, const std::vector& slot_ref_exprs, MutTableKey& key) { + int idx = 0; for (auto& slot_ref_expr : slot_ref_exprs) { ExprValue value = row->get_value(static_cast(slot_ref_expr)->tuple_id(), static_cast(slot_ref_expr)->slot_id()); - // TODO: 同类型可以不转string - key.append_value(value.cast_to(pb::STRING)); + if (idx < _equal_slot_cast_types.size()) { + key.append_value(value.cast_to(_equal_slot_cast_types[idx])); + } else { + key.append_value(value.cast_to(pb::STRING)); + } + ++idx; } } @@ -681,6 +723,7 @@ void Joiner::close(RuntimeState* state) { _need_add_index_collector_node = false; _index_collector_cond.reset(); _on_condition_column_map.clear(); + _equal_slot_cast_types.clear(); } int Joiner::show_explain(QueryContext* ctx, std::vector>& output, int& next_id, int display_id) { diff --git a/src/exec/mysql_scan_node.cpp b/src/exec/mysql_scan_node.cpp index 8a993a9f3..92740d810 100644 --- a/src/exec/mysql_scan_node.cpp +++ b/src/exec/mysql_scan_node.cpp @@ -5,7 +5,7 @@ namespace baikaldb { int MysqlScanNode::init(const pb::PlanNode& node) { int ret = 0; ret = ScanNode::init(node); - if (ret != 0) { + if (ret < 0) { DB_WARNING("ExecNode::init fail, ret: %d", ret); return ret; } @@ -18,6 +18,11 @@ int MysqlScanNode::open(RuntimeState* state) { return -1; } START_LOCAL_TRACE(get_trace(), state->get_trace_cost(), OPEN_TRACE, nullptr); + int ret = ScanNode::open(state); + if (ret < 0) { + DB_WARNING("ScanNode::open fail, ret: %d", ret); + return -1; + } _table_info = SchemaFactory::get_instance()->get_table_info_ptr(_table_id); if (_table_info == nullptr) { DB_WARNING("table info is null, table id: %ld", _table_id); @@ -30,15 +35,11 @@ int MysqlScanNode::open(RuntimeState* state) { DB_WARNING("tuple desc is nullptr"); return -1; } - if (_related_manager_node == nullptr) { - DB_WARNING("related_manager_node is nullptr"); - return -1; - } - if (_related_manager_node->is_delay_fetcher_store()) { + if (_related_manager_node != nullptr && _related_manager_node->is_delay_fetcher_store()) { set_node_exec_type(pb::EXEC_ARROW_ACERO); return 0; } - int ret = query_sql(state); + ret = query_sql(state); if (ret < 0) { DB_WARNING("Fail to query mysql"); return -1; @@ -202,11 +203,10 @@ int MysqlScanNode::build_arrow_declaration(RuntimeState* state) { arrow::Iterator> batch_it = arrow::MakeIteratorFromReader(_vectorized_reader); return batch_it; }; - if (_related_manager_node == nullptr) { - DB_WARNING("related_manager_node is nullptr"); - return -1; + bool is_delay_fetch = false; + if (_related_manager_node != nullptr) { + is_delay_fetch = _related_manager_node->is_delay_fetcher_store(); } - bool is_delay_fetch = _related_manager_node->is_delay_fetcher_store(); if (state->vectorlized_parallel_execution == false) { arrow::acero::Declaration dec{"record_batch_source", arrow::acero::RecordBatchSourceNodeOptions{_vectorized_reader->schema(), std::move(iter_maker)}}; @@ -350,12 +350,10 @@ int MysqlVectorizedReader::init(MysqlScanNode* scan_node, RuntimeState* state) { } SelectManagerNode* related_manager_node = static_cast(_scan_node->get_related_manager_node()); - if (related_manager_node == nullptr) { - DB_WARNING("_related_manager_node is nullptr"); - return -1; + if (related_manager_node != nullptr) { + _is_delay_fetch = related_manager_node->is_delay_fetcher_store(); + _index_cond = related_manager_node->get_index_collector_cond(); } - _is_delay_fetch = related_manager_node->is_delay_fetcher_store(); - _index_cond = related_manager_node->get_index_collector_cond(); return 0; } diff --git a/src/exec/packet_node.cpp b/src/exec/packet_node.cpp index dc8ec3e80..a7c7255b7 100644 --- a/src/exec/packet_node.cpp +++ b/src/exec/packet_node.cpp @@ -194,9 +194,27 @@ int PacketNode::handle_trace(RuntimeState* state) { } pack_head(); pack_fields(); + std::vector> rows; std::vector row; - row.push_back(_trace->DebugString().c_str()); - pack_vector_row(row); + rows.reserve(1); + if (!state->use_mpp) { + row.emplace_back(_trace->DebugString().c_str()); + rows.emplace_back(row); + } else { + auto ctx = state->ctx(); + if (ctx != nullptr) { + for (auto& [id, fragment] : ctx->fragments) { + if (fragment != nullptr) { + row.clear(); + row.emplace_back(fragment->trace_node.DebugString().c_str()); + rows.emplace_back(row); + } + } + } + } + for (auto& row : rows) { + pack_vector_row(row); + } pack_eof(); return 0; } @@ -659,7 +677,7 @@ int PacketNode::start_vectorized_execution(RuntimeState* state) { return 0; } } else { - DB_FATAL("arrow execute fail: arrow acero run fail, status: %s", final_table.status().ToString().c_str()); + DB_FATAL_STATE(state, "arrow execute fail: arrow acero run fail, status: %s", final_table.status().ToString().c_str()); return -1; } return vectorized_pack_rows(state, table, false); @@ -733,6 +751,29 @@ int PacketNode::open(RuntimeState* state) { if (state->explain_type == ANALYZE_STATISTICS) { state->force_single_rpc = true; + + std::vector scan_nodes; + get_node(pb::SCAN_NODE, scan_nodes); + if (scan_nodes.size() != 1) { + DB_WARNING("analyze packet node should have one and only one scan node, but now %ld scan node(s) are found.", + scan_nodes.size()); + return -1; + } + auto* scan_node = static_cast(scan_nodes[0]); + int64_t table_id = scan_node->table_id(); + + if (state->statistics_types->count(pb::ST_CLEAR) != 0) { + if (SchemaFactory::get_instance()->is_switch_open(table_id, TABLE_SWITCH_COST)) { + DB_WARNING("table %ld has cost switch is on, forbid clear operation", table_id); + return -1; + } + // 清除统计信息固定只发一个region，且store侧会立刻返回，不做任何事 + state->max_sample_regions = 1; + } + if (state->max_sample_regions > 0) { + ScanIndexInfo* main_scan_index = scan_node->main_scan_index(); + main_scan_index->region_infos = sample_analyze_regions(main_scan_index->region_infos, state->max_sample_regions); + } } _send_buf = state->send_buf(); @@ -930,10 +971,6 @@ int PacketNode::open_histogram(RuntimeState* state) { rows.push_back(row); } } - if (rows.size() <= 0) { - return -1; - } - _fields.clear(); for (auto& name : names) { ResultField field; @@ -1024,6 +1061,38 @@ int PacketNode::open_cmsketch(RuntimeState* state) { return 0; } +std::map PacketNode::sample_analyze_regions( + const std::map& all_regions, + int64_t sample_cnt) { + + size_t all_cnt = all_regions.size(); + if (sample_cnt >= all_cnt || sample_cnt <= 0) { + return all_regions; + } + + std::vector all_regions_vec; + all_regions_vec.reserve(all_regions.size()); + for (auto& region : all_regions) { + all_regions_vec.emplace_back(®ion.second); + } + + std::vector sample_regions_vec; + sample_regions_vec.reserve(sample_cnt); + double step = static_cast(all_cnt) / sample_cnt; + for (size_t i = 0; i < sample_cnt; i++) { + size_t idx = static_cast(i * step); + if (idx >= all_cnt) { + idx = all_cnt - 1; + } + sample_regions_vec.emplace_back(all_regions_vec[idx]); + } + std::map sample_regions; + for (auto it: sample_regions_vec) { + sample_regions[it->region_id()] = *it; + } + return sample_regions; +} + int PacketNode::open_analyze(RuntimeState* state) { bool eos = false; int ret = 0; diff --git a/src/exec/parquet_scan_node.cpp b/src/exec/parquet_scan_node.cpp index 08b3c8b0b..3bd9f6464 100644 --- a/src/exec/parquet_scan_node.cpp +++ b/src/exec/parquet_scan_node.cpp @@ -1,6 +1,7 @@ #include "parquet_scan_node.h" #include "store.h" #include "vectorize_helpper.h" +#include "sort_merge.h" namespace baikaldb { int ParquetScanNode::init(const pb::PlanNode& node) { @@ -25,6 +26,11 @@ int ParquetScanNode::init(const pb::PlanNode& node) { DB_WARNING("table_info is nullptr"); return -1; } + _pri_info = _factory->get_index_info_ptr(_table_id); + if (_pri_info == nullptr) { + DB_WARNING("primary index info is nullptr"); + return -1; + } return 0; } @@ -55,7 +61,7 @@ int ParquetScanNode::open(RuntimeState* state) { } _field_id2info_map[field_info->id] = field_info; } - if (get_qualified_record_batch_readers(_parquet_file2reader_map, _parquet_file_not_exist_column_map) != 0) { + if (get_qualified_record_batch_readers(_parquet_file_readers) != 0) { DB_WARNING("get qualified record batch readers fail"); return -1; } @@ -74,6 +80,7 @@ void ParquetScanNode::close(RuntimeState* state) { } int ParquetScanNode::build_arrow_declaration(RuntimeState* state) { + START_LOCAL_TRACE_WITH_PARTITION_PROPERTY(get_trace(), state->get_trace_cost(), &_partition_property, OPEN_TRACE, nullptr); int ret = 0; // add SourceNode std::shared_ptr vectorized_reader = std::make_shared(); @@ -89,6 +96,7 @@ int ParquetScanNode::build_arrow_declaration(RuntimeState* state) { arrow::acero::Declaration dec{"record_batch_source", arrow::acero::RecordBatchSourceNodeOptions{vectorized_reader->schema(), std::move(iter_maker)}}; state->append_acero_declaration(dec); + LOCAL_TRACE_ARROW_PLAN_WITH_SCHEMA(dec, vectorized_reader->schema(), nullptr); return 0; } @@ -98,119 +106,121 @@ int ParquetScanNode::process_index(RuntimeState* state) { DB_FATAL_STATE(state, "no index"); return -1; } - pb::PossibleIndex pos_index; - pos_index.ParseFromString(scan_pb.indexes(0)); - if (pos_index.index_id() != _table_id) { + _possible_index.ParseFromString(scan_pb.indexes(0)); + if (_possible_index.index_id() != _table_id) { DB_WARNING_STATE(state, "ParquetScanNode only support primary key index"); return -1; } - if (pos_index.ranges_size() == 0) { + if (_possible_index.ranges_size() == 0) { DB_WARNING_STATE(state, "PossibleIndex has no range"); return 0; } - DB_WARNING("PossibleIndex: %s", pos_index.ShortDebugString().c_str()); - _key_ranges = std::vector(pos_index.ranges().begin(), pos_index.ranges().end()); + DB_WARNING("PossibleIndex: %s", _possible_index.ShortDebugString().c_str()); return 0; } -int ParquetScanNode::get_qualified_record_batch_readers( - std::unordered_map& parquet_file2reader_map, - std::unordered_map>& parquet_file_not_exist_column_map) { +void ParquetScanNode::get_qualified_parquet_file_readers(std::vector>& parquet_file_readers, + const std::vector>& parquet_files, std::shared_ptr schema) { + for (const auto& parquet_file : parquet_files) { + ParquetFileReaderOptions options; + options.pos_index = &_possible_index; + for (const auto& [_, field_info] : _field_id2info_map) { + options.lower_short_name_fields[field_info->lower_short_name] = *field_info; + } + options.schema = schema; + auto parquet_reader = std::make_shared(options, parquet_file); + parquet_file_readers.emplace_back(parquet_reader); + } +} + +int ParquetScanNode::get_qualified_record_batch_readers(std::vector>& record_batch_readers) { // 获取该region最新版本的所有parquet_file集合 SmartRegion region = Store::get_instance()->get_region(_region_id); if (region == nullptr) { DB_WARNING("region is nullptr"); return -1; } - - if (region->get_column_files(_key_ranges, _parquet_files) != 0) { + std::vector> parquet_files_tmp; + if (region->get_column_files(_possible_index, parquet_files_tmp) != 0) { DB_WARNING("Fail to get_column_files"); return -1; } - for (const auto& parquet_file : _parquet_files) { - const std::string file_path = parquet_file->get_file_path(); - // 获取column_indices，以及在parquet文件里不存在的列 - // 在parquet文件中不存在的列需要补充默认值或NULL - const std::unordered_map& column_name2index_map = parquet_file->get_column_name2index_map(); - std::vector exist_column_indices; - std::vector not_exist_columns; - for (const auto& [_, field_info] : _field_id2info_map) { - const std::string& field_name = field_info->lower_short_name; - if (column_name2index_map.find(field_name) != column_name2index_map.end()) { - exist_column_indices.emplace_back(column_name2index_map.at(field_name)); - // DB_WARNING("field_name: %s, index: %d", field_name.c_str(), column_name2index_map.at(field_name)); - } else { - DB_WARNING("field_name: %s, not exist", field_name.c_str()); - not_exist_columns.emplace_back(field_info); - } + std::vector> parquet_files; + parquet_files.reserve(parquet_files_tmp.size()); + int cumulatives_file_count = 0; + bool column_only_read_base = _table_info->schema_conf.column_only_read_base(); + // 只获取base层的parquet_file，可以不用merge_on_read + for (auto f : parquet_files_tmp) { + // base层文件start_version为0 + if (f->get_file_info()->start_version == 0) { + parquet_files.emplace_back(f); + } else if (!column_only_read_base) { + cumulatives_file_count++; + parquet_files.emplace_back(f); } - if (!not_exist_columns.empty()) { - parquet_file_not_exist_column_map[file_path].swap(not_exist_columns); + } + + bool has_put_or_delete = false; + for (const auto& f : parquet_files) { + auto info = f->get_file_info(); + if (info->start_version != 0 && (info->put_count > 0 || info->delete_count > 0)) { + has_put_or_delete = true; + break; } - std::unique_ptr<::arrow::RecordBatchReader> reader; - std::vector row_group_indices; - bool fill_cache = false; - if (_key_ranges.empty()) { - // 获取所有rowgroup的数据 - std::shared_ptr<::parquet::FileMetaData> file_metadata = parquet_file->get_file_metadata(); - if (file_metadata == nullptr) { - DB_WARNING("file_metadata is nullptr"); - return -1; - } - if (file_metadata->num_row_groups() == 0) { - DB_WARNING("parquet file has no row group"); - continue; - } - - row_group_indices.reserve(file_metadata->num_row_groups()); - for (int i = 0; i < file_metadata->num_row_groups(); ++i) { - row_group_indices.emplace_back(i); - } - auto status = parquet_file->GetRecordBatchReader(row_group_indices, exist_column_indices, &reader); - if (!status.ok()) { - DB_WARNING("Fail to get_record_batch_reader"); - return -1; - } - } else { - // 获取符合条件的row_group_indices和row_ranges - std::vector>> row_ranges; - if (parquet_file->get_qualified_rowgroup_and_rowranges(_key_ranges, row_group_indices, row_ranges) != 0) { - DB_WARNING("Fail to get_qualified_rowgroup_and_rowranges"); - return -1; - } - if (row_group_indices.empty()) { - DB_WARNING("parquet file has no qualified data"); + } + + std::unordered_map field_id2info_map = _field_id2info_map; + if (cumulatives_file_count > 0 && has_put_or_delete) { + // merge on read时，需要补充主键 + for (const auto& f : _pri_info->fields) { + FieldInfo* field_info = _table_info->get_field_ptr(f.id); + if (field_info == nullptr) { + DB_WARNING("field not found region_id: %ld, field_id: %d", _region_id, f.id); continue; } - // 获取record_batch_reader - auto status = parquet_file->GetRecordBatchReader(row_group_indices, exist_column_indices, row_ranges, &reader); - if (!status.ok()) { - DB_WARNING("Fail to get_record_batch_reader"); + field_id2info_map[f.id] = field_info; + } + std::shared_ptr schema_info = ColumnRecord::make_column_schema(_table_info->id, _table_info, _pri_info, field_id2info_map); + if (schema_info == nullptr) { + DB_FATAL("get schema info failed"); + return -1; + } + std::vector> parquet_file_readers; + parquet_file_readers.reserve(parquet_files.size()); + get_qualified_parquet_file_readers(parquet_file_readers, parquet_files, schema_info->schema_with_order_info); + std::vector> single_row_readers; + single_row_readers.reserve(parquet_file_readers.size()); + SortMergeOptions merge_options; + merge_options.batch_size = 1024; + merge_options.is_base_compact = true; + merge_options.schema_info = schema_info; + for (auto& r : parquet_file_readers) { + single_row_readers.emplace_back(std::make_shared(r, schema_info.get())); + } + record_batch_readers.emplace_back(std::make_shared(merge_options, single_row_readers)); + } else { + std::vector> arrow_fields; + arrow_fields.reserve(_field_id2info_map.size()); + for (const auto& [_, field_info] : _field_id2info_map) { + std::shared_ptr arrow_field = VectorizeHelpper::make_field(field_info->lower_short_name, + arrow::Type::type(primitive_to_arrow_type(field_info->type))); + if (arrow_field == nullptr) { + DB_WARNING("Fail to make arrow field"); return -1; } - fill_cache = true; + arrow_fields.emplace_back(arrow_field); } - std::vector read_ranges; - read_ranges.reserve(row_group_indices.size() * exist_column_indices.size()); - parquet_file->parser_position(row_group_indices, exist_column_indices, read_ranges); - auto file_info = parquet_file->get_file_info(); - ReaderInfo reader_info; - reader_info.read_contents = std::make_shared(); - reader_info.read_contents->fill_cache = fill_cache; - reader_info.read_contents->file_short_name = parquet_file->get_file_short_name(); - reader_info.read_contents->ranges.swap(read_ranges); - reader_info.read_contents->region_id = file_info->region_id; - reader_info.read_contents->start_version = file_info->start_version; - reader_info.read_contents->end_version = file_info->end_version; - reader_info.read_contents->file_idx = ColumnFileInfo::get_file_idx(parquet_file->get_file_short_name()); - if (reader_info.read_contents->file_idx < 0) { - DB_COLUMN_FATAL("Fail to get file idx : %s", parquet_file->get_file_short_name().c_str()); - return -1; + std::shared_ptr arrow_schema = arrow::schema(arrow_fields); + + std::vector> parquet_file_readers; + parquet_file_readers.reserve(parquet_files.size()); + get_qualified_parquet_file_readers(parquet_file_readers, parquet_files, arrow_schema); + for (auto& r : parquet_file_readers) { + record_batch_readers.emplace_back(r); } - reader_info.reader = std::move(reader); - parquet_file2reader_map[file_path] = reader_info; } + return 0; } @@ -221,14 +231,13 @@ int ParquetVectorizedReader::init(RuntimeState* state, ParquetScanNode* parquet_ return -1; } if (parquet_scan_node == nullptr) { - DB_WARNING("_parquet_scan_node is nullptr"); + DB_WARNING("parquet_scan_node is nullptr"); return -1; } _state = state; _parquet_scan_node = parquet_scan_node; _field_id2info_map = _parquet_scan_node->get_field_id2info_map(); - _parquet_file2reader_map = _parquet_scan_node->get_parquet_file2reader_map(); - _parquet_file_not_exist_column_map = _parquet_scan_node->get_parquet_file_not_exist_column_map(); + _parquet_file_readers = _parquet_scan_node->get_parquet_file_readers(); // 构造schema pb::TupleDescriptor* tuple_desc = _parquet_scan_node->get_tuple(); @@ -241,7 +250,7 @@ int ParquetVectorizedReader::init(RuntimeState* state, ParquetScanNode* parquet_ arrow_fields.reserve(tuple_desc->slots().size()); for (const auto& slot : tuple_desc->slots()) { const std::string& field_name = std::to_string(slot.tuple_id()) + "_" + std::to_string(slot.slot_id()); - std::shared_ptr arrow_field = ColumnRecord::make_schema(field_name, + std::shared_ptr arrow_field = VectorizeHelpper::make_field(field_name, arrow::Type::type(primitive_to_arrow_type(slot.slot_type()))); if (arrow_field == nullptr) { DB_WARNING("Fail to make arrow field"); @@ -261,11 +270,12 @@ int ParquetVectorizedReader::init(RuntimeState* state, ParquetScanNode* parquet_ return -1; } _column_name_map[field_name] = field_info->lower_short_name; + _column_type_map[field_name] = field_info->type; } _arrow_schema = arrow::schema(arrow_fields); // 设置parquet文件读迭代器 - _reader_iter = _parquet_file2reader_map->begin(); + _reader_iter = _parquet_file_readers->begin(); return 0; } @@ -273,9 +283,6 @@ int ParquetVectorizedReader::init(RuntimeState* state, ParquetScanNode* parquet_ // 每次最多读取FLAGS_chunk_size行 arrow::Status ParquetVectorizedReader::ReadNext(std::shared_ptr* out) { out->reset(); - ON_SCOPE_EXIT([]() { - ParquetCache::get_instance()->set_bthread_local(nullptr); - }); const int64_t limit = _parquet_scan_node->get_limit(); if (limit > 0 && _processed_row_cnt >= limit) { return arrow::Status::OK(); @@ -284,46 +291,20 @@ arrow::Status ParquetVectorizedReader::ReadNext(std::shared_ptr record_batch = nullptr; - while (_reader_iter != _parquet_file2reader_map->end()) { - ParquetCache::get_instance()->set_bthread_local(_reader_iter->second.read_contents.get()); - auto status = _reader_iter->second.reader->ReadNext(&record_batch); + while (_reader_iter != _parquet_file_readers->end()) { + auto status = (*_reader_iter)->ReadNext(&record_batch); if (!status.ok()) { DB_WARNING("Fail to read next record batch, %s", status.message().c_str()); return status; } if (record_batch == nullptr) { - DB_WARNING("file: %s, ReadNext cost: %ld", _reader_iter->first.c_str(), cost.get_time()); + DB_WARNING("file:, ReadNext cost: %ld", /*_reader_iter->first.c_str(),*/ cost.get_time()); ++_reader_iter; continue; } - DB_DEBUG("file: %s, ReadNext cost: %ld, num rows: %ld, num columns: %d", _reader_iter->first.c_str(), cost.get_time(), + DB_DEBUG("file ReadNext cost: %ld, num rows: %ld, num columns: %d", /*_reader_iter->first.c_str(), */cost.get_time(), record_batch->num_rows(), record_batch->num_columns()); - // 添加查询需要，但是在parquet文件中不存在的列 - const std::string& file_name = _reader_iter->first; - if (_parquet_file_not_exist_column_map->find(file_name) != _parquet_file_not_exist_column_map->end()) { - const auto& not_exist_columns = _parquet_file_not_exist_column_map->at(file_name); - for (const auto& field_info : not_exist_columns) { - TimeCost add_column_cost; - std::shared_ptr arrow_field = ColumnRecord::make_schema(field_info->lower_short_name, - arrow::Type::type(primitive_to_arrow_type(field_info->type))); - if (arrow_field == nullptr) { - DB_WARNING("Fail to get arrow type, field_id: %d, field_type: %d", field_info->id, field_info->type); - return arrow::Status::IOError("Fail to get arrow type"); - } - std::shared_ptr arrow_array = ColumnRecord::make_array_from_exprvalue( - field_info->type, field_info->default_expr_value, record_batch->num_rows()); - if (arrow_array == nullptr) { - return arrow::Status::IOError("Fail to make array from expr value"); - } - auto new_record_batch_ret = - record_batch->AddColumn(record_batch->num_columns(), arrow_field, arrow_array); - if (!new_record_batch_ret.ok()) { - return arrow::Status::IOError("Fail to add column"); - } else { - record_batch = *new_record_batch_ret; - } - } - } + break; } // eof @@ -333,7 +314,7 @@ arrow::Status ParquetVectorizedReader::ReadNext(std::shared_ptrvector_index_map()[_index_id]; @@ -205,6 +206,7 @@ int RocksdbScanNode::choose_index(RuntimeState* state) { check_memory = true; } bool has_global_param = false; + // 这个兼容性需要长期保留，比如watt基准等很多地方还会使用旧逻辑 if (pos_index.has_left_field_cnt() || pos_index.has_right_field_cnt()) { has_global_param = true; _is_eq = pos_index.is_eq(); @@ -217,17 +219,13 @@ int RocksdbScanNode::choose_index(RuntimeState* state) { for (auto& range : pos_index.ranges()) { if (!has_global_param) { has_global_param = true; - if (range.left_key() == range.right_key()) { - _is_eq = true; - } _left_field_cnt = range.left_field_cnt(); _right_field_cnt = range.right_field_cnt(); _left_open = range.left_open(); _right_open = range.right_open(); _like_prefix = range.like_prefix(); - if (_is_eq) { - _right_field_cnt = _left_field_cnt; - _right_open = _left_open; + if (range.left_key() == range.right_key() && _left_field_cnt == _right_field_cnt && _left_open == _right_open) { + _is_eq = true; } } if (!_is_eq) { @@ -299,6 +297,7 @@ int RocksdbScanNode::init(const pb::PlanNode& node) { return -1; } } else if (_ddl_work_type == pb::DDL_COLUMN) { + // ddl delete和ddl update的区别是update slot是否为空 for (auto& slot : node.derive_node().scan_node().column_ddl_info().update_slots()) { _update_slots.emplace_back(slot); } @@ -332,6 +331,7 @@ int RocksdbScanNode::init(const pb::PlanNode& node) { return ret; } } + _is_ddl_update = !_update_exprs.empty(); } return 0; } @@ -645,13 +645,16 @@ int RocksdbScanNode::open(RuntimeState* state) { auto reverse_index_map = state->reverse_index_map(); //DB_WARNING_STATE(state, "_is_covering_index:%d", _is_covering_index); if (_vector_index != nullptr) { + VectorSearchParam search_param; + search_param.efsearch = _efsearch; + search_param.nprobe = _nprobe; int ret = _vector_index->search_vector(txn->get_txn(), _separate_value, _pri_info, _table_info, _vector_word, _topk, - _efsearch, + search_param, _left_records, _vector_filter_conjuncts, _scan_conjuncts, @@ -666,7 +669,7 @@ int RocksdbScanNode::open(RuntimeState* state) { return -1; } } - DB_WARNING("vector_index search, index:%ld, table:%ld, size:%lu", _index_info->id, _table_info->id, _left_records.size()); + //DB_WARNING("vector_index search, index:%ld, table:%ld, size:%lu", _index_info->id, _table_info->id, _left_records.size()); } else if (_reverse_infos.size() > 1) { //TODO 为不影响原流程暂时保留，后续删除。 for (auto& info : _reverse_infos) { @@ -688,7 +691,7 @@ int RocksdbScanNode::open(RuntimeState* state) { _m_index.search(txn->get_txn(), _pri_info, _table_info, reverse_index_map, !FLAGS_reverse_seek_first_level, _pb_node.derive_node().scan_node().fulltext_index()); - //DB_FATAL("fulltext: %s", _pb_node.derive_node().scan_node().fulltext_index().DebugString().c_str()); + // DB_FATAL("fulltext: %s", _pb_node.derive_node().scan_node().fulltext_index().DebugString().c_str()); } else { DB_FATAL("fulltext storage type error"); return -1; @@ -1148,11 +1151,12 @@ int RocksdbScanNode::column_ddl_work(RuntimeState* state, MemRow* row) { record->set_value(record->get_field_by_tag(slot.field_id()), expr->get_value(row).cast_to(slot.slot_type())); } - - ret = txn->put_primary(_region_id, *_pri_info, record, nullptr); - if (ret < 0) { - DB_WARNING_STATE(state, "put table:%ld fail:%d", _table_id, ret); - return -1; + if (is_ddl_update()) { + ret = txn->put_primary(_region_id, *_pri_info, record, nullptr); + if (ret < 0) { + DB_WARNING_STATE(state, "put table:%ld fail:%d", _table_id, ret); + return -1; + } } return 0; } @@ -1228,6 +1232,45 @@ int RocksdbScanNode::index_ddl_work(RuntimeState* state, MemRow* row) { } return 0; } + if (_ddl_index_info->type == pb::I_VECTOR) { + int64_t index_id = _ddl_index_info->id; + auto& vector_index_map = state->vector_index_map(); + if (vector_index_map.count(index_id) == 0) { + DB_WARNING_STATE(state, "DDL_LOG vector ddl info not found index_id:%ld.", index_id); + return -1; + } + if (_ddl_index_info->fields.size() != 1 && _ddl_index_info->fields.size() != 2) { + DB_WARNING_STATE(state, "DDL_LOG vector ddl fields size error:%ld.", _ddl_index_info->fields.size()); + return -1; + } + MutTableKey pk_key; + ret = record->encode_key(*_pri_info, pk_key, -1, false, false); + if (ret < 0) { + DB_WARNING_STATE(state, "DDL_LOG record [%s] encode key failed[%d].", record->to_string().c_str(), ret); + return -1; + } + int field_idx = 0; + if (_ddl_index_info->fields.size() > 1) { + field_idx = 1; + } + auto field = record->get_field_by_idx(_ddl_index_info->fields[field_idx].pb_idx); + if (record->is_null(field)) { + return 0; + } + std::string word; + ret = record->get_reverse_word(*_ddl_index_info, word); + if (ret < 0) { + DB_WARNING_STATE(state, "index_info to word fail for index_id: %ld", index_id); + return ret; + } + //DB_NOTICE("word:%s", str_to_hex(word).c_str()); + ret = vector_index_map[index_id]->insert_vector(txn, word, pk_key.data(), record); + if (ret < 0) { + DB_WARNING_STATE(state, "vector_index fail insert, index_id: %ld", index_id); + return ret; + } + return 0; + } if (_ddl_index_info->type != pb::I_ROLLUP) { for (auto& pair: _field_ids) { int32_t field_id = pair.first; @@ -1417,15 +1460,24 @@ int RocksdbScanNode::get_next_by_table_seek(RuntimeState* state, RowBatch* batch } } } + bool need_return = true; if (_lock == pb::LOCK_GET && - (_ddl_work_type == pb::DDL_COLUMN || _ddl_work_type == pb::DDL_LOCAL_INDEX)) { + (is_ddl_update() || _ddl_work_type == pb::DDL_LOCAL_INDEX)) { // local index or column返回最大一条数据 batch->replace_row(std::move(row), 0); + } else if (is_ddl_delete()) { + if (need_copy(row.get(), _scan_conjuncts)) { + batch->move_row(std::move(row)); + } else { + need_return = false; + } } else { batch->move_row(std::move(row)); } - ++_num_rows_returned; - ++_num_rows_returned_by_range; + if (need_return) { + ++_num_rows_returned; + ++_num_rows_returned_by_range; + } } else { ret = batch->add_chunk_row(); if (ret != 0) { @@ -1508,15 +1560,24 @@ int RocksdbScanNode::get_next_by_table_seek(RuntimeState* state, RowBatch* batch } } } + bool need_return = true; if (_lock == pb::LOCK_GET && - (_ddl_work_type == pb::DDL_COLUMN || _ddl_work_type == pb::DDL_LOCAL_INDEX)) { + (is_ddl_update() || _ddl_work_type == pb::DDL_LOCAL_INDEX)) { // local index or column返回最大一条数据 batch->replace_row(std::move(row), 0); + } else if (is_ddl_delete()) { + if (filter == nullptr || filter != nullptr && !filter->test(row_batch.index())) { + batch->move_row(std::move(row)); + } else { + need_return = false; + } } else { batch->move_row(std::move(row)); } - ++_num_rows_returned; - ++_num_rows_returned_by_range; + if (need_return) { + ++_num_rows_returned; + ++_num_rows_returned_by_range; + } } } } diff --git a/src/exec/scan_node.cpp b/src/exec/scan_node.cpp index eaf817242..a3fa97bfe 100644 --- a/src/exec/scan_node.cpp +++ b/src/exec/scan_node.cpp @@ -25,6 +25,7 @@ #include "rocksdb_scan_node.h" #include "information_schema_scan_node.h" #include "redis_scan_node.h" +#include "file_scan_node.h" namespace baikaldb { int64_t AccessPathMgr::select_index_common() { @@ -55,7 +56,8 @@ int64_t AccessPathMgr::select_index_common() { index_id, pb::IndexState_Name(index_state).c_str()); continue; } - int field_count = path->hit_index_field_ids.size(); + // 确保不影响存量索引选择 + int field_count = std::min(path->hit_index_field_ids.size(), info_ptr->fields.size()); if (info.fields.size() == 0) { continue; } @@ -229,8 +231,8 @@ int ScanNode::show_explain(QueryContext* ctx, std::vectorpos_index; if (pos_index.ranges_size() == 1) { - int field_cnt = pos_index.has_left_field_cnt() ? pos_index.left_field_cnt() : pos_index.ranges(0).left_field_cnt(); - if (field_cnt == (int)index_info.fields.size() && pos_index.is_eq()) { + int field_cnt = pos_index.left_field_cnt(); + if (field_cnt >= (int)index_info.fields.size() && pos_index.is_eq()) { explain_info["type"] = "eq_ref"; if (index_info.type == pb::I_UNIQ || index_info.type == pb::I_PRIMARY) { explain_info["type"] = "const"; @@ -759,6 +761,8 @@ ScanNode* ScanNode::create_scan_node(const pb::PlanNode& node, const CreateExecO } if (table_info->dblink_info.type() == pb::LT_MYSQL) { return new MysqlScanNode; + } else if (table_info->dblink_info.type() == pb::LT_FILE) { + return new FileScanNode; } break; } @@ -929,7 +933,7 @@ bool ScanNode::need_index_merge() { _conjuncts_without_or.clear(); for (auto conjunct : _filter_node->conjuncts()) { if (_or_conjunct != conjunct) { - _conjuncts_without_or.push_back(conjunct); + _conjuncts_without_or.emplace_back(conjunct); } } _or_sub_conjuncts.clear(); diff --git a/src/exec/select_manager_node.cpp b/src/exec/select_manager_node.cpp index 32e6e18e6..9440bb2e1 100755 --- a/src/exec/select_manager_node.cpp +++ b/src/exec/select_manager_node.cpp @@ -26,6 +26,8 @@ #include "vectorize_helpper.h" #include "arrow_io_excutor.h" #include "arrow_exec_node_manager.h" +#include "file_scan_node.h" +#include "db_service.h" namespace baikaldb { DEFINE_bool(global_index_read_consistent, true, "double check for global and primary region consistency"); @@ -89,9 +91,8 @@ int SelectManagerNode::open(RuntimeState* state) { // mpp return subquery_open(state); } - if (is_mysql_scan()) { - // open子节点 - return subquery_open(state); + if (is_dblink_scan()) { + return dblink_scan_run(state, scan_nodes[0]); } _scan_tuple_id = static_cast(scan_nodes[0])->tuple_id(); if (_delay_fetcher_store) { @@ -148,6 +149,17 @@ int SelectManagerNode::init(const pb::PlanNode& node){ if (select_node.has_is_return_empty()) { _return_empty = select_node.is_return_empty(); } + _conditions.reserve(select_node.filter_conditions_size()); + for (auto& expr : select_node.filter_conditions()) { + ExprNode* expr_node = nullptr; + ret = ExprNode::create_tree(expr, &expr_node); + if (ret < 0) { + DB_FATAL("create expr fail"); + ExprNode::destroy_tree(expr_node); + return ret; + } + _conditions.emplace_back(expr_node); + } return 0; } @@ -162,7 +174,7 @@ void SelectManagerNode::transfer_pb(int64_t region_id, pb::PlanNode* pb_node) { return; } // 序列化schema - arrow::Result> schema_ret = arrow::ipc::SerializeSchema(*_arrow_schema, arrow::default_memory_pool()); + arrow::Result> schema_ret = arrow::ipc::SerializeSchema(*_arrow_schema, GetMemoryPoolForRead()); if (!schema_ret.ok()) { DB_FATAL("arrow serialize schema fail, status: %s", schema_ret.status().ToString().c_str()); return; @@ -179,6 +191,9 @@ void SelectManagerNode::transfer_pb(int64_t region_id, pb::PlanNode* pb_node) { node->add_is_null_first(null_first); } node->set_is_return_empty(_return_empty); + for (ExprNode* expr : _conditions) { + ExprNode::create_pb_expr(node->add_filter_conditions(), expr); + } return; } @@ -192,7 +207,7 @@ int SelectManagerNode::delay_fetcher_store(RuntimeState* state) { } int SelectManagerNode::build_arrow_declaration(RuntimeState* state) { - if (is_dual_scan() || _has_er_child || is_mysql_scan()) { + if (is_dual_scan() || _has_er_child || is_dblink_scan()) { return _children[0]->build_arrow_declaration(state); } RocksdbScanNode* scan_node = static_cast(get_node(pb::SCAN_NODE)); @@ -329,11 +344,14 @@ int SelectManagerNode::get_next(RuntimeState* state, RowBatch* batch, bool* eos) } int SelectManagerNode::mpp_fetcher_global_index(RuntimeState* state) { - RocksdbScanNode* scan_node = static_cast(get_node(pb::SCAN_NODE)); - if (scan_node == nullptr) { + ScanNode* scan = static_cast(get_node(pb::SCAN_NODE)); + if (scan == nullptr) { return 0; } - + if (!scan->is_rocksdb_scan_node()) { + return 0; + } + RocksdbScanNode* scan_node = static_cast(scan); ScanIndexInfo* scan_index_info = scan_node->main_scan_index(); int64_t router_index_id = scan_index_info->router_index_id; int64_t main_table_id = scan_node->table_id(); @@ -439,6 +457,17 @@ int SelectManagerNode::fetcher_store_run(RuntimeState* state, ExecNode* exec_nod << main_scan_index->region_infos.size(); return -1; }; + if (_conditions.size() > 0) { + START_LOCAL_TRACE(get_trace(), state->get_trace_cost(), OPEN_TRACE, nullptr); + if (0 != VectorizeHelpper::init_conditions(state, _conditions, _vectorize_conditions, _scan_tuple_id, nullptr)) { + DB_FATAL("init condition fail, txn_id: %lu, log_id:%lu, scan tuple_id: %d", + state->txn_id, state->log_id(), _scan_tuple_id); + return -1; + } + if (state->execute_type == pb::EXEC_ARROW_ACERO) { + LOCAL_TRACE_ARROW_FILTER(&_vectorize_conditions, _limit); + } + } if (backup_scan_index != nullptr && dynamic_timeout_ms > 0 && state->txn_id == 0) { // 非事务情况下才进行全局二级索引降级，txn_id != 0 情况下state中会有修改，无法多个请求并发使用state // 可以降级 @@ -482,6 +511,10 @@ int SelectManagerNode::fetcher_store_run(RuntimeState* state, ExecNode* exec_nod // 没有开启动态超时，不满足降级条件，只使用main router main_fetcher.scan_index = main_scan_index; fetcher_store = &main_fetcher.fetcher_store; + if (_conditions.size() > 0) { + fetcher_store->conditions = &_conditions; + fetcher_store->vectorize_conditions = &_vectorize_conditions; + } ret = single_fetcher_store_open(&main_fetcher, state, exec_node); if (ret < 0) { state->error_code = fetcher_store->error_code; @@ -512,7 +545,11 @@ int SelectManagerNode::fetcher_store_run(RuntimeState* state, ExecNode* exec_nod } } // 无sort节点时不会排序，按顺序输出 - _sorter->merge_sort(); + if (_conditions.empty()) { + _sorter->merge_sort(); + } else { + _sorter->sort(); + } } else { // EXEC_ARROW_ACERO(多表查询, 或者单表查询有region返回了arrow格式) for (auto& pair : fetcher_store->start_key_sort) { @@ -738,11 +775,10 @@ int SelectManagerNode::construct_primary_possible_index( return -1; } pos_index.set_is_eq(true); + pos_index.set_left_open(false); + pos_index.set_left_field_cnt(pri_info->fields.size()); range->set_left_key(key.data()); range->set_left_full(key.get_full()); - pos_index.set_left_field_cnt(pri_info->fields.size()); - pos_index.set_left_open(false); - range->add_partition_ids(mem_row->get_partition_id()); limit_cnt --; if(!limit_cnt) { @@ -773,8 +809,10 @@ int SelectManagerNode::construct_primary_possible_index_vectorize( std::vector> batchs; std::shared_ptr chunk; auto tuple = state->get_tuple_desc(tuple_id); + bool is_partition_table = pri_info->is_partitioned; for (auto& pair : fetcher_store.start_key_sort) { auto row_batch = fetcher_store.region_batch[pair.second].row_data; + int64_t partition_id = fetcher_store.region_batch[pair.second].partition_id; if (row_batch != nullptr) { if (chunk == nullptr) { chunk = std::make_shared(); @@ -796,11 +834,31 @@ int SelectManagerNode::construct_primary_possible_index_vectorize( && 0 != VectorizeHelpper::change_arrow_record_batch_schema(fetcher_store.arrow_schema, fetcher_store.region_batch[pair.second].arrow_data, &fetcher_store.region_batch[pair.second].arrow_data, - true)) { + /*need_check_momery_limit=*/true)) { DB_FATAL("change arrow record batch schema chunk fail"); return -1; } - batchs.push_back(fetcher_store.region_batch[pair.second].arrow_data); + std::shared_ptr data_without_partition = fetcher_store.region_batch[pair.second].arrow_data; + if (is_partition_table) { + // 处理全局索引partition id + arrow::Int64Scalar scalar = arrow::Int64Scalar(partition_id); + auto partition_id_array = arrow::MakeArrayFromScalar(scalar, data_without_partition->num_rows()); + if (!partition_id_array.ok()) { + DB_FATAL("make partition id array failed: %s", partition_id_array.status().ToString().c_str()); + return -1; + } + auto data_with_partition = data_without_partition->AddColumn( + data_without_partition->num_columns(), + arrow::field("__partition_id__", arrow::int64()), + *partition_id_array); + if (!data_with_partition.ok()) { + DB_FATAL("add partition column failed: %s", data_with_partition.status().ToString().c_str()); + return -1; + } + batchs.emplace_back(*data_with_partition); + } else { + batchs.emplace_back(data_without_partition); + } } if (batchs.size() == 0){ return 0; @@ -842,6 +900,14 @@ int SelectManagerNode::construct_primary_possible_index_vectorize( limit->add_num_rows_skipped(offset_cnt); } std::unordered_map> field_id_to_arrow_array; + std::shared_ptr partition_id_array = nullptr; + if (is_partition_table) { + partition_id_array = table->GetColumnByName("__partition_id__"); + if (partition_id_array == nullptr) { + DB_FATAL("not find partition id array"); + return -1; + } + } for (auto& pri_field : pri_info->fields) { int32_t field_id = pri_field.id; int32_t slot_id = state->get_slot_id(tuple_id, field_id); @@ -868,12 +934,13 @@ int SelectManagerNode::construct_primary_possible_index_vectorize( return -1; } pos_index.set_is_eq(true); + pos_index.set_left_open(false); + pos_index.set_left_field_cnt(pri_info->fields.size()); range->set_left_key(key.data()); range->set_left_full(key.get_full()); - pos_index.set_left_field_cnt(pri_info->fields.size()); - pos_index.set_left_open(false); - // [ARROW todo] 全局索引反查partition分发 - // range->set_partition_id(mem_row->get_partition_id()); + if (is_partition_table) { + range->add_partition_ids(VectorizeHelpper::get_vectorized_value(partition_id_array.get(), row_idx).get_numberic()); + } } //重新做路由选择 pos_index.SerializeToString(&scan_index_info->raw_index); @@ -881,12 +948,9 @@ int SelectManagerNode::construct_primary_possible_index_vectorize( &scan_index_info->region_primary, scan_node->get_partition()); } - - // arrow vectorize -int FetcherStoreVectorizedReader::init(SelectManagerNode* select_node, RuntimeState* state) { - _schema = select_node->get_arrow_schema(); - _need_fetcher_store = select_node->is_delay_fetcher_store(); +void FetcherStoreVectorizedReader::init_tuple(SelectManagerNode* select_node, RuntimeState* state) { + _tuples.clear(); int32_t scan_tuple_id = select_node->get_scan_tuple_id(); pb::TupleDescriptor* tuple = state->get_tuple_desc(scan_tuple_id); if (state->is_simple_select) { @@ -897,28 +961,50 @@ int FetcherStoreVectorizedReader::init(SelectManagerNode* select_node, RuntimeSt } else { _tuples.emplace_back(tuple); } +} + +int FetcherStoreVectorizedReader::init(SelectManagerNode* select_node, RuntimeState* state) { + int32_t scan_tuple_id = select_node->get_scan_tuple_id(); + _schema = select_node->get_arrow_schema(); + _need_fetcher_store = select_node->is_delay_fetcher_store(); + init_tuple(select_node, state); if (_schema == nullptr) { // 当store返回的数据都是row pb, 但是需要列式执行, 如index join场景; 或是非index join延迟访问store // 1. 初始化根据tuple_id(用scannode的tuple_id)生成对应tuple对应的arrow schema // 2. 列式执行过程中将返回的memrow转arrow // 3. 列式执行返回的record batch, 转换为db构建的schema - _chunk = std::make_shared(); - if (_chunk->init(_tuples) != 0) { - DB_FATAL("build arrow schema fail, scan_tuple_id: %d", scan_tuple_id); - return -1; - } - _schema = _chunk->get_arrow_schema(); - if (_schema == nullptr) { - DB_FATAL("build arrow schema fail"); - return -1; - } + _schema = VectorizeHelpper::get_arrow_schema(_tuples); } + if (_schema == nullptr) { + DB_FATAL("build arrow schema fail"); + return -1; + } state->arrow_input_schemas[scan_tuple_id] = _schema; _select_node = select_node; _state = state; _index_cond = select_node->get_index_collector_cond(); return 0; } + +int SelectManagerNode::dblink_scan_run(RuntimeState* state, ExecNode* exec_node) { + if (state == nullptr) { + DB_WARNING("state is nullptr"); + return -1; + } + if (_db_request_map.empty()) { + // 非mpp场景 + return subquery_open(state); + } + // 放在这里启动mpp请求，保证 db fragment（非dblink scan的fragment）都启动了，再启动dblink scan； + for (const auto& [ith_db, request] : _db_request_map) { + int ret = DBInteract::get_instance()->handle_mpp_dag_fragment(request, ith_db); + if (ret != 0) { + DB_FATAL("logid: %lu, send request to other db fail, ret: %d", state->log_id(), ret); + return -1; + } + } + return 0; +} arrow::Status FetcherStoreVectorizedReader::ReadNext(std::shared_ptr* out) { int ret = 0; @@ -970,7 +1056,11 @@ arrow::Status FetcherStoreVectorizedReader::ReadNext(std::shared_ptrEquals(batch.arrow_data->schema()) - && 0 != VectorizeHelpper::change_arrow_record_batch_schema(_schema, batch.arrow_data, &(batch.arrow_data), true)) { + && 0 != VectorizeHelpper::change_arrow_record_batch_schema( + _schema, + batch.arrow_data, + &(batch.arrow_data), + /*need_check_momery_limit=*/true)) { return arrow::Status::IOError("change arrow record batch schema fail"); } } @@ -987,6 +1077,10 @@ arrow::Status FetcherStoreVectorizedReader::ReadNext(std::shared_ptr arrow recordbatch if (_chunk == nullptr) { _chunk = std::make_shared(); + // 如果是join的非驱动表, 且命中全局索引并且不是covering index, 同时sql里不涉及主键 + // acero计划启动(schema不包含主键) -> 查驱动表 -> 对非驱动表planrouter -> 额外将主键加到tuple里 -> 查非驱动表 -> 返回行数据包含主键 + // 在这里重新init_tuple兼容解析包含主键的行数据 -> change schame(去掉主键) + init_tuple(_select_node, _state); if (_chunk->init(_tuples) != 0) { return arrow::Status::IOError("chunk init fail"); } @@ -996,7 +1090,11 @@ arrow::Status FetcherStoreVectorizedReader::ReadNext(std::shared_ptradd_num_rows_returned((*out)->num_rows()); diff --git a/src/exec/sort_node.cpp b/src/exec/sort_node.cpp index 16031ab97..6e1de6abd 100644 --- a/src/exec/sort_node.cpp +++ b/src/exec/sort_node.cpp @@ -175,6 +175,15 @@ int SortNode::build_arrow_declaration(RuntimeState* state) { return 0; } +bool SortNode::need_projection() { + for (auto& expr : _order_exprs) { + if (expr->node_type() != pb::SLOT_REF && expr->node_type() != pb::AGG_EXPR) { + return true; + } + } + return false; +} + int SortNode::build_sort_arrow_declaration(RuntimeState* state, pb::TraceNode* trace_node) { if (state->acero_declarations.size() > 0 && (state->acero_declarations.back().factory_name == "order_by" diff --git a/src/exec/update_manager_node.cpp b/src/exec/update_manager_node.cpp index e43ba1b12..a5a6a6eae 100644 --- a/src/exec/update_manager_node.cpp +++ b/src/exec/update_manager_node.cpp @@ -57,6 +57,7 @@ int UpdateManagerNode::init_update_info(UpdateNode* update_node) { std::vector g_non_unique_indexs; std::vector g_affected_unique_indexs; std::vector g_affected_non_unique_indexs; + std::shared_ptr ttl_field_ptr = _table_info->get_ttl_field(); for (auto index_id : _table_info->indices) { auto info_ptr = SchemaFactory::get_instance()->get_index_info_ptr(index_id); if (info_ptr == nullptr) { @@ -106,8 +107,8 @@ int UpdateManagerNode::init_update_info(UpdateNode* update_node) { global_affected_indices.insert(global_affected_indices.end(), g_affected_unique_indexs.begin(), g_affected_unique_indexs.end()); global_affected_indices.insert(global_affected_indices.end(), g_affected_non_unique_indexs.begin(), g_affected_non_unique_indexs.end()); - // 如果更新主键或者ttl，那么影响了全部索引 - if (_affect_primary || (_table_info->ttl_info.ttl_duration_s > 0)) { + // 如果更新主键或者常规ttl，那么影响了全部索引 + if (_affect_primary || (_table_info->ttl_info.ttl_duration_s > 0 && ttl_field_ptr == nullptr)) { _global_affected_index_ids.insert(_global_affected_index_ids.end(), g_unique_indexs.begin(), g_unique_indexs.end()); _global_affected_index_ids.insert(_global_affected_index_ids.end(), g_non_unique_indexs.begin(), g_non_unique_indexs.end()); } else { diff --git a/src/exec/window_node.cpp b/src/exec/window_node.cpp index 508b3976d..5d7982b6e 100644 --- a/src/exec/window_node.cpp +++ b/src/exec/window_node.cpp @@ -161,15 +161,17 @@ int WindowNode::get_next(RuntimeState* state, RowBatch* batch, bool* eos) { DB_WARNING_STATE(state, "split into partition fail, ret:%d", ret); return ret; } - if (is_first_partition_belong_to_prev) { + if (is_first_partition_belong_to_prev) { get_next_partition(start, end); if (start < 0 || end > _child_row_batch.size() || start >= end) { DB_WARNING_STATE(state, "get next partition fail, start:%d end:%d", start, end); return -1; - } + } for (int i = start; i < end; ++i) { cur_partition_batch.move_row(std::move(_child_row_batch.get_row(i))); } + } else { + break; } } ret = _window_processor->process_one_partition(&cur_partition_batch); diff --git a/src/expr/agg_fn_call.cpp b/src/expr/agg_fn_call.cpp index 4f3ae9312..1f4c56fa4 100644 --- a/src/expr/agg_fn_call.cpp +++ b/src/expr/agg_fn_call.cpp @@ -16,6 +16,8 @@ #include #include "hll_common.h" #include "slot_ref.h" +#include "row_expr.h" +#include "arrow_function.h" #include namespace baikaldb { @@ -1260,11 +1262,17 @@ bool AggFnCall::can_use_arrow_vector() { case MAX: case MULTI_COUNT_DISTINCT: case AVG: + case GROUP_CONCAT: break; default: return false; } - if (_children.size() != 1) { + if (_agg_type == GROUP_CONCAT) { + if (_children.size() > 2) { + // 暂不支持orderby + return false; + } + } else if (_children.size() != 1) { return false; // count(distinct fieldA, filedB) not supported. } if (_agg_type == AVG && _children[0]->is_literal()) { @@ -1273,7 +1281,12 @@ bool AggFnCall::can_use_arrow_vector() { } } for (auto& c : _children) { - if (!c->can_use_arrow_vector()) { + if (_agg_type == GROUP_CONCAT && c->node_type() == pb::ROW_EXPR) { + RowExpr* expr = static_cast(c); + if (!expr->children_can_use_arrow_vector()) { + return false; + } + } else if (!c->can_use_arrow_vector()) { return false; } } @@ -1338,9 +1351,7 @@ int AggFnCall::build_agg_argument(int i, && (_children[i]->is_slot_ref() || (_children[i]->col_type() == pb::STRING) || (_children[i]->arrow_expr().type() != nullptr && arrow::is_base_binary_like(_children[i]->arrow_expr().type()->id())))) { - generate_projection_exprs.emplace_back(arrow::compute::call("cast", - {_children[i]->arrow_expr()}, - arrow::compute::CastOptions::Unsafe(arrow::float64()))); + generate_projection_exprs.emplace_back(arrow_cast(_children[i]->arrow_expr(), _children[i]->col_type(), pb::DOUBLE)); } else { generate_projection_exprs.emplace_back(_children[i]->arrow_expr()); } @@ -1380,13 +1391,81 @@ int AggFnCall::transfer_to_arrow_avg(std::vector& agg return 0; } -int AggFnCall::transfer_to_arrow_agg_function(std::vector& aggs, +int AggFnCall::transfer_to_arrow_group_concat(std::vector& aggs, + bool is_merge, + std::vector& generate_projection_exprs, + std::vector& generate_projection_exprs_names) { + std::string intermediate_slot = std::to_string(_tuple_id) + "_" + std::to_string(_intermediate_slot_id); + std::string final_slot = std::to_string(_tuple_id) + "_" + std::to_string(_final_slot_id); + std::vector args; + std::shared_ptr opt = std::make_shared(_sep); + if (is_merge && !_is_distinct) { + args.emplace_back(arrow::FieldRef(intermediate_slot)); + aggs.emplace_back("group_concat_intermediate", opt, args, final_slot); + if (intermediate_slot != final_slot) { + // agg下推需要保留intermediate_slot列 + aggs.emplace_back("group_concat_intermediate", opt, args, intermediate_slot); + } + } else { + if (_children.size() == 0) { + DB_FATAL("group_concat must have at least one child"); + return -1; + } + ExprNode* first_row_expr = _children[0]; + std::vector first_row_expr_children = first_row_expr->children(); + for (auto& c : first_row_expr_children) { + if (0 != build_arrow_expr_with_cast(c, pb::STRING)) { + return -1; + } + } + if (first_row_expr_children.size() == 1 + && first_row_expr_children[0]->is_slot_ref()) { + if (first_row_expr_children[0]->col_type() == pb::STRING) { + args.emplace_back(*(first_row_expr_children[0]->arrow_expr().field_ref())); + } else { + // cast string临时列 + generate_projection_exprs.emplace_back(first_row_expr_children[0]->arrow_expr()); + std::string tmp_name = "tmp_" + std::to_string(generate_projection_exprs.size()); + generate_projection_exprs_names.emplace_back(tmp_name); + args.emplace_back(arrow::FieldRef(tmp_name)); + } + } else { + // 先concat映射 + std::vector concat_args; + concat_args.reserve(first_row_expr_children.size()); + for (auto& c : first_row_expr_children) { + concat_args.emplace_back(c->arrow_expr()); + } + // 连接符放在最后 + concat_args.emplace_back(arrow::compute::literal(std::make_shared(""))); + std::shared_ptr option = std::make_shared(); + option->null_handling = arrow::compute::JoinOptions::NullHandlingBehavior::SKIP; + arrow::compute::Expression concat_expr = arrow::compute::call("binary_join_element_wise", concat_args, option); + + generate_projection_exprs.emplace_back(concat_expr); + std::string tmp_name = "tmp_" + std::to_string(generate_projection_exprs.size()); + generate_projection_exprs_names.emplace_back(tmp_name); + args.emplace_back(arrow::FieldRef(tmp_name)); + } + aggs.emplace_back("group_concat_intermediate", opt, args, final_slot); + if (intermediate_slot != final_slot) { + // agg下推需要保留intermediate_slot列 + aggs.emplace_back("group_concat_intermediate", opt, args, intermediate_slot); + } + } + return 0; +} + +int AggFnCall::transfer_to_arrow_agg_function(std::vector& aggs, bool is_merge, std::vector& generate_projection_exprs, std::vector& generate_projection_exprs_names) { if (_agg_type == AVG) { return transfer_to_arrow_avg(aggs, is_merge, generate_projection_exprs, generate_projection_exprs_names); } + if (_agg_type == GROUP_CONCAT) { + return transfer_to_arrow_group_concat(aggs, is_merge, generate_projection_exprs, generate_projection_exprs_names); + } // arrow里有没有group by用的两套算子, 有group by算子需要加前缀hash_ std::string func_name = "hash_"; bool need_check_type = false; // 处理sum(string)需要先加cast diff --git a/src/expr/arrow_agg_function.cpp b/src/expr/arrow_agg_function.cpp index 03bffa505..c8f2cee10 100644 --- a/src/expr/arrow_agg_function.cpp +++ b/src/expr/arrow_agg_function.cpp @@ -21,6 +21,8 @@ #include #include #include +#include +#include #include #include "slot_ref.h" #include "agg_fn_call.h" @@ -431,6 +433,135 @@ struct GroupedAvgFinalizeFactory { arrow::compute::InputType argument_type; }; +/* + * group concat 不带order by + * 1. store一阶段agg -> append string (group by groupby cols) + * 2. db 三阶段agg -> append string (group by groupby cols) + * + * group concat 带order by (TODO) + * 1. store一阶段agg -> append string (group by groupby cols, orderby cols) + * 2. db 二阶段agg -> append string (group by groupby cols, orderby cols) + * 3. db 三阶段agg -> sort -> final string value (group by groupby cols) + * + * GROUP_CONCAT generate intermediate string(agg) + */ +static auto kCountOptionsType = + arrow::compute::internal::GetFunctionOptionsType(arrow::internal::DataMember("separator_", &GroupConcatOptions::separator_)); +GroupConcatOptions::GroupConcatOptions(const std::string& sep, const std::vector& asc) + : FunctionOptions(kCountOptionsType), separator_(sep), asc_(asc) {} +constexpr char GroupConcatOptions::kTypeName[]; + +struct GroupeConcatIntermediateImpl final : public GroupedAggregator { + arrow::Status Init(arrow::compute::ExecContext* ctx, const arrow::compute::KernelInitArgs& args) override { + ctx_ = ctx; + null_bitmap_ = arrow::TypedBufferBuilder(ctx->memory_pool()); + options_ = arrow::internal::checked_cast(*args.options); + return arrow::Status::OK(); + } + + arrow::Status Resize(int64_t new_num_groups) override { + auto added_groups = new_num_groups - num_groups_; + DCHECK_GE(added_groups, 0); + num_groups_ = new_num_groups; + values_.resize(new_num_groups); + RETURN_NOT_OK(null_bitmap_.Append(added_groups, false)); + return arrow::Status::OK(); + } + + arrow::Status Consume(const arrow::compute::ExecSpan& batch) override { + int64_t values_size = values_.size(); + return VisitGroupedValues( + batch, + [&](uint32_t g, std::string_view val) -> arrow::Status { + if (values_size <= g) { + return arrow::Status::Invalid("Group index ", g, " out of range"); + } + if (!values_[g].has_value()) { + values_[g].emplace(val.data(), val.size()); + } else { + values_[g]->append(options_.separator_); + values_[g]->append(val.data(), val.size()); + } + arrow::bit_util::SetBit(null_bitmap_.mutable_data(), g); + return arrow::Status::OK(); + }, + [&](uint32_t g) -> arrow::Status { return arrow::Status::OK(); }); + } + + arrow::Status Merge(GroupedAggregator&& raw_other, const arrow::ArrayData& group_id_mapping) override { + auto other = arrow::internal::checked_cast(&raw_other); + auto g = group_id_mapping.GetValues(1); + for (uint32_t other_g = 0; static_cast(other_g) < group_id_mapping.length; ++other_g, ++g) { + if (!other->values_[other_g].has_value()) { + continue; + } + if (!values_[*g].has_value()) { + values_[*g].emplace(other->values_[other_g]->data(), other->values_[other_g]->size()); + } else { + values_[*g]->append(options_.separator_); + values_[*g]->append(other->values_[other_g]->data(), other->values_[other_g]->size()); + } + arrow::bit_util::SetBit(null_bitmap_.mutable_data(), *g); + } + return arrow::Status::OK(); + } + + arrow::Result Finalize() override { + if (ctx_ == nullptr) { + return arrow::Status::Invalid("ctx is nullptr"); + } + using offset_type = int64_t; + ARROW_ASSIGN_OR_RAISE(auto nulls, null_bitmap_.Finish()); + auto result_out = arrow::ArrayData::Make(out_type(), num_groups_, {std::move(nulls), nullptr}); + ARROW_ASSIGN_OR_RAISE( + auto raw_offsets, + AllocateBuffer((1 + values_.size()) * sizeof(offset_type), ctx_->memory_pool())); + auto* offsets = reinterpret_cast(raw_offsets->mutable_data()); + offsets[0] = 0; + offsets++; + if (result_out->buffers.size() < 2) { + return arrow::Status::Invalid("Result array want 2 buffers but only has ", result_out->buffers.size()); + } + const uint8_t* null_bitmap = result_out->buffers[0]->data(); + offset_type total_length = 0; + for (size_t i = 0; i < values_.size(); i++) { + if (arrow::bit_util::GetBit(null_bitmap, i)) { + const std::optional& value = values_[i]; + DCHECK(value.has_value()); + if (arrow::internal::AddWithOverflow( + total_length, static_cast(value->size()), &total_length)) { + return arrow::Status::Invalid("Result is too large to fit in ", *result_out->type, + " cast to large_ variant of type"); + } + } + offsets[i] = total_length; + } + ARROW_ASSIGN_OR_RAISE(auto data, AllocateBuffer(total_length, ctx_->memory_pool())); + int64_t offset = 0; + for (size_t i = 0; i < values_.size(); i++) { + if (arrow::bit_util::GetBit(null_bitmap, i)) { + const std::optional& value = values_[i]; + DCHECK(value.has_value()); + std::memcpy(data->mutable_data() + offset, value->data(), value->size()); + offset += value->size(); + } + } + result_out->buffers[1] = std::move(raw_offsets); + result_out->buffers.push_back(std::move(data)); + return result_out; + } + + std::shared_ptr out_type() const override { + return arrow::large_binary(); + } + + int64_t num_groups_ = 0; + std::vector> values_; + arrow::TypedBufferBuilder null_bitmap_; + arrow::compute::ExecContext* ctx_; + GroupConcatOptions options_; +}; + arrow::Status ArrowFunctionManager::RegisterAllHashAggFunction() { auto registry = arrow::compute::GetFunctionRegistry(); { @@ -451,6 +582,13 @@ arrow::Status ArrowFunctionManager::RegisterAllHashAggFunction() { DCHECK_OK(AddHashAggKernels(arrow::BaseBinaryTypes(), GroupedAvgFinalizeFactory::Make, func.get())); // 一阶段agg产出中间结果 DCHECK_OK(registry->AddFunction(std::move(func))); } + { + static auto default_group_concat_options = GroupConcatOptions::Defaults(); + auto func = std::make_shared("group_concat_intermediate", arrow::compute::Arity::Binary(), + /*doc=*/arrow::compute::FunctionDoc::Empty(), &default_group_concat_options); + DCHECK_OK(func->AddKernel(MakeKernel(arrow::compute::InputType::Any(), HashAggregateInit))); + DCHECK_OK(registry->AddFunction(std::move(func))); + } return arrow::Status::OK(); } } \ No newline at end of file diff --git a/src/expr/arrow_function.cpp b/src/expr/arrow_function.cpp index bf9d335a9..e4596be3d 100644 --- a/src/expr/arrow_function.cpp +++ b/src/expr/arrow_function.cpp @@ -21,12 +21,51 @@ #include #include #include +#include #include "slot_ref.h" #include "row_expr.h" namespace baikaldb { -DEFINE_bool(enable_arrow_complex_func, false, "enable_arrow_complex_func"); +/* + * 通用config + */ +class CommonOptionsType : public arrow::compute::FunctionOptionsType { +public: + static const arrow::compute::FunctionOptionsType* GetInstance() { + static std::unique_ptr instance(new CommonOptionsType()); + return instance.get(); + } + const char* type_name() const override { return "CommonOptionType"; } + std::string Stringify(const arrow::compute::FunctionOptions& options) const override { + return type_name(); + } + bool Compare(const arrow::compute::FunctionOptions& options, + const arrow::compute::FunctionOptions& other) const override { + const auto& lop = static_cast(options); + const auto& rop = static_cast(other); + return lop.str_value == rop.str_value && lop.int_value == rop.int_value;; + } + std::unique_ptr Copy(const arrow::compute::FunctionOptions& options) const override { + const auto& opts = static_cast(options); + return std::make_unique(opts.str_value, opts.int_value); + } +}; + +arrow::Result> InitCommonState(arrow::compute::KernelContext*, + const arrow::compute::KernelInitArgs& args) { + auto func_options = static_cast(args.options); + if (func_options == nullptr) { + return std::make_unique("", 0); + } + return std::make_unique(func_options->str_value, func_options->int_value); +} +CommonFunctionOptions::CommonFunctionOptions(const std::string& value, int64_t int_value) + : arrow::compute::FunctionOptions(CommonOptionsType::GetInstance()), str_value(value), int_value(int_value) {} + +/* + * cast type map + */ const std::unordered_map> cast_types = { {pb::BOOL, arrow::boolean()}, @@ -81,7 +120,7 @@ bool check_row_expr_is_support(pb::Function& fn, ExprNode* node) { || fn.fn_op() == parser::FT_LE || fn.fn_op() == parser::FT_LT || fn.fn_op() == parser::FT_IN) { - return static_cast(node)->can_use_arrow_vector_for_compare_sclar_exrpr(); + return static_cast(node)->children_can_use_arrow_vector(); } return false; } @@ -104,6 +143,23 @@ bool is_same_type(const pb::PrimitiveType& type1, const pb::PrimitiveType& type2 } return false; } + +bool need_special_cast(const pb::PrimitiveType& from_type, const pb::PrimitiveType& to_type) { + if (from_type == to_type) { + return false; + } + // 时间类型+字符串，之间相互转换 + if ((is_datetime_specic(from_type) || from_type == pb::STRING) + && (is_datetime_specic(to_type) || to_type == pb::STRING)) { + return true; + } + // 字符串转数值型 + if (from_type == pb::STRING && is_numberic(to_type)) { + return true; + } + // 基础类型直接用arrow cast + return false; +} /* * Transfer to arrow inline compute Function * args需要cast成fn对应的类型, @@ -135,9 +191,7 @@ int build_arrow_expr_with_cast(ExprNode* node, const pb::PrimitiveType& col_type return 0; } BUILD_ARROW_EXPR_RET(node); - if (node->col_type() != col_type - && (is_datetime_specic(node->col_type()) || node->col_type() == pb::STRING) - && (is_datetime_specic(col_type) || col_type == pb::STRING)) { + if (need_special_cast(node->col_type(), col_type)) { arrow::Expression cast_expr = arrow_cast(node->arrow_expr(), node->col_type(), col_type); node->set_arrow_expr(cast_expr); return 0; @@ -165,6 +219,9 @@ int build_arrow_expr_with_cast(ExprNode* node, const pb::PrimitiveType& col_type arrow::compute::Expression arrow_cast(const arrow::compute::Expression& expr, const pb::PrimitiveType& type, const pb::PrimitiveType& cast_type) { arrow::compute::Expression cast_expr = expr; + if (type == cast_type) { + return cast_expr; + } if ((is_datetime_specic(type) || type == pb::STRING) && (is_datetime_specic(cast_type) || cast_type == pb::STRING)) { std::string cast_func_name; @@ -187,6 +244,13 @@ arrow::compute::Expression arrow_cast(const arrow::compute::Expression& expr, co } ExprValueCastFunctionOptions option(type); cast_expr = arrow::compute::call(cast_func_name, {expr}, std::move(option)); + } else if (type == pb::STRING && is_numberic(cast_type)){ + auto iter = cast_types.find(cast_type); + if (iter == cast_types.end()) { + return cast_expr; + } + std::string cast_name = "string_cast_" + iter->second->ToString(); + cast_expr = arrow::compute::call(cast_name, {expr}); } else { auto iter = cast_types.find(cast_type); if (iter == cast_types.end()) { @@ -226,11 +290,11 @@ int get_all_arrow_argments_for_add_minus_multiple(std::vector& childr for (auto& c : children) { BUILD_ARROW_EXPR_RET(c); if (has_double(args_types)) { - arguments.emplace_back(arrow::compute::call("cast", {c->arrow_expr()}, arrow::compute::CastOptions::Unsafe(arrow::float64()))); + arguments.emplace_back(arrow_cast(c->arrow_expr(), c->col_type(), pb::DOUBLE)); } else if (has_uint(args_types)) { - arguments.emplace_back(arrow::compute::call("cast", {c->arrow_expr()}, arrow::compute::CastOptions::Unsafe(arrow::uint64()))); + arguments.emplace_back(arrow_cast(c->arrow_expr(), c->col_type(), pb::UINT64)); } else { - arguments.emplace_back(arrow::compute::call("cast", {c->arrow_expr()}, arrow::compute::CastOptions::Unsafe(arrow::int64()))); + arguments.emplace_back(arrow_cast(c->arrow_expr(), c->col_type(), pb::INT64)); } } return 0; @@ -242,7 +306,7 @@ int arrow_add(std::vector& children, pb::Function* fn, const pb::Prim if (0 != get_all_arrow_argments_for_add_minus_multiple(children, arguments)) { return -1; } - out = arrow::compute::call("add_checked", arguments); + out = arrow::compute::call("add", arguments); return 0; } @@ -252,7 +316,7 @@ int arrow_minus(std::vector& children, pb::Function* fn, const pb::Pr if (0 != get_all_arrow_argments_for_add_minus_multiple(children, arguments)) { return -1; } - out = arrow::compute::call("subtract_checked", arguments); + out = arrow::compute::call("subtract", arguments); return 0; } @@ -262,25 +326,25 @@ int arrow_multiplies(std::vector& children, pb::Function* fn, const p if (0 != get_all_arrow_argments_for_add_minus_multiple(children, arguments)) { return -1; } - out = arrow::compute::call("multiply_checked", arguments); + out = arrow::compute::call("multiply", arguments); return 0; } // case parser::FT_DIVIDES: int arrow_divides(std::vector& children, pb::Function* fn, const pb::PrimitiveType& return_type, arrow::compute::Expression& out) { for (int i = 0; i < children.size(); ++i) { - if (0 != build_arrow_expr_with_cast(children[i], nullptr, i)) { + if (0 != build_arrow_expr_with_cast(children[i], pb::DOUBLE)) { return -1; } } // 需要特殊处理除数是0, 返回NULL arrow::Datum null_datum = std::make_shared(); - out = arrow::compute::call("divide_checked", { - arrow::compute::call("cast", {children[0]->arrow_expr()}, arrow::compute::CastOptions::Unsafe(arrow::float64())), + out = arrow::compute::call("divide", { + children[0]->arrow_expr(), arrow::compute::call("if_else", { arrow::compute::call("equal", {children[1]->arrow_expr(), arrow::compute::literal(0)}), arrow::compute::literal(null_datum), - arrow::compute::call("cast", {children[1]->arrow_expr()}, arrow::compute::CastOptions::Unsafe(arrow::float64())) + children[1]->arrow_expr() }) }); return 0; @@ -483,7 +547,7 @@ int arrow_least(std::vector& children, pb::Function* fn, const pb::Pr } for (auto i = 0; i < arguments.size(); ++i) { if (is_string(children[i]->col_type())) { - arguments[i] = arrow::compute::call("cast", {arguments[i]}, arrow::compute::CastOptions::Unsafe(arrow::float64())); + arguments[i] = arrow_cast(arguments[i], children[i]->col_type(), pb::DOUBLE); } } arrow::compute::ElementWiseAggregateOptions options(/*skip_nulls*/false); @@ -498,7 +562,7 @@ int arrow_greatest(std::vector& children, pb::Function* fn, const pb: } for (auto i = 0; i < arguments.size(); ++i) { if (is_string(children[i]->col_type())) { - arguments[i] = arrow::compute::call("cast", {arguments[i]}, arrow::compute::CastOptions::Unsafe(arrow::float64())); + arguments[i] = arrow_cast(arguments[i], children[i]->col_type(), pb::DOUBLE); } } arrow::compute::ElementWiseAggregateOptions options(/*skip_nulls*/false); @@ -826,7 +890,7 @@ static arrow::Status ExecStringInput(arrow::compute::KernelContext* ctx, const a if (type != pb::STRING) { return arrow::Status::TypeError("not support"); } - arrow::compute::internal::VisitArrayValuesInline( + arrow::compute::internal::VisitArrayValuesInline( input, [&](std::string_view v) { *out_values++ = datetime_to_date(str_to_datetime(v.data(), v.length())); @@ -885,7 +949,7 @@ static arrow::Status ExecStringInput(arrow::compute::KernelContext* ctx, const a if (type != pb::STRING) { return arrow::Status::TypeError("not support"); } - arrow::compute::internal::VisitArrayValuesInline( + arrow::compute::internal::VisitArrayValuesInline( input, [&](std::string_view v) { *out_values++ = str_to_time(v.data(), v.length()); @@ -944,7 +1008,7 @@ static arrow::Status ExecStringInput(arrow::compute::KernelContext* ctx, const a if (type != pb::STRING) { return arrow::Status::TypeError("not support"); } - arrow::compute::internal::VisitArrayValuesInline( + arrow::compute::internal::VisitArrayValuesInline( input, [&](std::string_view v) { *out_values++ = str_to_datetime(v.data(), v.length()); @@ -1003,7 +1067,7 @@ static arrow::Status ExecStringInput(arrow::compute::KernelContext* ctx, const a if (type != pb::STRING) { return arrow::Status::TypeError("not support"); } - arrow::compute::internal::VisitArrayValuesInline( + arrow::compute::internal::VisitArrayValuesInline( input, [&](std::string_view v) { *out_values++ = datetime_to_timestamp(str_to_datetime(v.data(), v.length())); @@ -1098,6 +1162,85 @@ static arrow::Status Exec(arrow::compute::KernelContext* ctx, const arrow::compu } }; +template +inline OutputValueCType string_view_to_intergal(const std::string_view& v) { + // 不用stollu, 避免string_view转string + // + : uint64; - : int64; 兼容前面空格和正负号，判断溢出 + char* pos = const_cast(v.data()); + const char* end = v.data() + v.length(); + while (pos < end && *pos == ' ') { + pos++; + } + if (pos < end) { + if (*pos == '-') { + int64_t result = 0; + auto res = std::from_chars(pos, end, result); + if (res.ec == std::errc::result_out_of_range) { + // 溢出 + return std::numeric_limits::max(); + } + return result; + } else { + if (*pos == '+' && pos + 1 < end) { + pos++; + } + uint64_t result = 0; + auto res = std::from_chars(pos, end, result); + if (res.ec == std::errc::result_out_of_range) { + // 溢出 + return std::numeric_limits::max(); + } + return result; + } + } + return 0; +} + +template +struct ExecStringCastNumberic { + using OutputValueCType = typename arrow::TypeTraits::CType; + static arrow::Status Exec(arrow::compute::KernelContext* ctx, const arrow::compute::ExecSpan& batch, arrow::compute::ExecResult* out) { + const arrow::ArraySpan& input = batch[0].array; + arrow::ArraySpan* out_data = out->array_span_mutable(); + OutputValueCType* out_values = out_data->GetValues(1); + if (std::is_floating_point::value) { + arrow::compute::internal::VisitArrayValuesInline( + input, + [&](std::string_view v) { + double result = 0.0; + using ::arrow_vendored::fast_float::from_chars; + from_chars(v.data(), v.data() + v.length(), result); + *out_values++ = OutputValueCType(result); + }, + [&]() { + // null + *out_values++ = OutputValueCType{}; + }); + } else if (std::is_integral::value) { + arrow::compute::internal::VisitArrayValuesInline( + input, + [&](std::string_view v) { + *out_values++ = string_view_to_intergal(v); + }, + [&]() { + // null + *out_values++ = OutputValueCType{}; + }); + } else { + return arrow::Status::TypeError("StringCastNumbericFunctor type error"); + } + return arrow::Status::OK(); + } +}; + +struct ParseStringToBoolean { + template + static OutValue Call(arrow::compute::KernelContext*, Arg0Value val, arrow::Status* st) { + auto result = string_view_to_intergal(val); + return (bool)result; + } +}; + arrow::Status ExecMurmurHash(arrow::compute::KernelContext* ctx, const arrow::compute::ExecSpan& batch, arrow::compute::ExecResult* out) { const arrow::ArraySpan& input = batch[0].array; arrow::ArraySpan* out_data = out->array_span_mutable(); @@ -1232,9 +1375,18 @@ arrow::Status ArrowFunctionManager::RegisterAllDefinedFunction() { arrow::compute::internal::GenerateNumeric(*in_ty), InitExprValueCast)); } - arrow::compute::ScalarKernel k_string_to_date({arrow::large_binary()}, arrow::uint32(), + arrow::compute::ScalarKernel k_binary_to_date({arrow::binary()}, arrow::uint32(), + ExecExprValueToDate::ExecStringInput, InitExprValueCast); + arrow::compute::ScalarKernel k_large_binary_to_date({arrow::large_binary()}, arrow::uint32(), ExecExprValueToDate::ExecStringInput, InitExprValueCast); - ARROW_RETURN_NOT_OK(expr_value_to_date_func->AddKernel(k_string_to_date)); + arrow::compute::ScalarKernel k_utf8_to_date({arrow::utf8()}, arrow::uint32(), + ExecExprValueToDate::ExecStringInput, InitExprValueCast); + arrow::compute::ScalarKernel k_large_utf8_to_date({arrow::large_utf8()}, arrow::uint32(), + ExecExprValueToDate::ExecStringInput, InitExprValueCast); + ARROW_RETURN_NOT_OK(expr_value_to_date_func->AddKernel(k_binary_to_date)); + ARROW_RETURN_NOT_OK(expr_value_to_date_func->AddKernel(k_large_binary_to_date)); + ARROW_RETURN_NOT_OK(expr_value_to_date_func->AddKernel(k_utf8_to_date)); + ARROW_RETURN_NOT_OK(expr_value_to_date_func->AddKernel(k_large_utf8_to_date)); ARROW_RETURN_NOT_OK(registry->AddFunction(expr_value_to_date_func)); } @@ -1250,9 +1402,18 @@ arrow::Status ArrowFunctionManager::RegisterAllDefinedFunction() { arrow::compute::internal::GenerateNumeric(*in_ty), InitExprValueCast)); } - arrow::compute::ScalarKernel k_string_to_datetime({arrow::large_binary()}, arrow::uint64(), + arrow::compute::ScalarKernel k_binary_to_datetime({arrow::binary()}, arrow::uint64(), + ExecExprValueToDateTime::ExecStringInput, InitExprValueCast); + arrow::compute::ScalarKernel k_large_binary_to_datetime({arrow::large_binary()}, arrow::uint64(), ExecExprValueToDateTime::ExecStringInput, InitExprValueCast); - ARROW_RETURN_NOT_OK(expr_value_to_datetime_func->AddKernel(k_string_to_datetime)); + arrow::compute::ScalarKernel k_utf8_to_datetime({arrow::utf8()}, arrow::uint64(), + ExecExprValueToDateTime::ExecStringInput, InitExprValueCast); + arrow::compute::ScalarKernel k_large_utf8_to_datetime({arrow::large_utf8()}, arrow::uint64(), + ExecExprValueToDateTime::ExecStringInput, InitExprValueCast); + ARROW_RETURN_NOT_OK(expr_value_to_datetime_func->AddKernel(k_binary_to_datetime)); + ARROW_RETURN_NOT_OK(expr_value_to_datetime_func->AddKernel(k_large_binary_to_datetime)); + ARROW_RETURN_NOT_OK(expr_value_to_datetime_func->AddKernel(k_utf8_to_datetime)); + ARROW_RETURN_NOT_OK(expr_value_to_datetime_func->AddKernel(k_large_utf8_to_datetime)); ARROW_RETURN_NOT_OK(registry->AddFunction(expr_value_to_datetime_func)); } @@ -1268,9 +1429,18 @@ arrow::Status ArrowFunctionManager::RegisterAllDefinedFunction() { arrow::compute::internal::GenerateNumeric(*in_ty), InitExprValueCast)); } - arrow::compute::ScalarKernel k_string_to_time({arrow::large_binary()}, arrow::int32(), + arrow::compute::ScalarKernel k_binary_to_time({arrow::binary()}, arrow::int32(), + ExecExprValueToTime::ExecStringInput, InitExprValueCast); + arrow::compute::ScalarKernel k_large_binary_to_time({arrow::large_binary()}, arrow::int32(), ExecExprValueToTime::ExecStringInput, InitExprValueCast); - ARROW_RETURN_NOT_OK(expr_value_to_time_func->AddKernel(k_string_to_time)); + arrow::compute::ScalarKernel k_utf8_to_time({arrow::utf8()}, arrow::int32(), + ExecExprValueToTime::ExecStringInput, InitExprValueCast); + arrow::compute::ScalarKernel k_large_utf8_to_time({arrow::large_utf8()}, arrow::int32(), + ExecExprValueToTime::ExecStringInput, InitExprValueCast); + ARROW_RETURN_NOT_OK(expr_value_to_time_func->AddKernel(k_binary_to_time)); + ARROW_RETURN_NOT_OK(expr_value_to_time_func->AddKernel(k_large_binary_to_time)); + ARROW_RETURN_NOT_OK(expr_value_to_time_func->AddKernel(k_utf8_to_time)); + ARROW_RETURN_NOT_OK(expr_value_to_time_func->AddKernel(k_large_utf8_to_time)); ARROW_RETURN_NOT_OK(registry->AddFunction(expr_value_to_time_func)); } @@ -1280,17 +1450,58 @@ arrow::Status ArrowFunctionManager::RegisterAllDefinedFunction() { { auto expr_value_to_timestamp_func = std::make_shared("expr_value_to_timestamp", arrow::compute::Arity::Unary(), /*doc=*/arrow::compute::FunctionDoc::Empty()); - for (const std::shared_ptr& in_ty : arrow::IntTypes()) { + for (const std::shared_ptr& in_ty : arrow::IntTypes()) { ARROW_RETURN_NOT_OK(expr_value_to_timestamp_func->AddKernel({in_ty}, // 输入类型 arrow::uint32(), // 输出类型 arrow::compute::internal::GenerateNumeric(*in_ty), InitExprValueCast)); } - arrow::compute::ScalarKernel k_string_to_timestamp({arrow::large_binary()}, arrow::uint32(), + arrow::compute::ScalarKernel k_binary_to_timestamp({arrow::binary()}, arrow::uint32(), + ExecExprValueToTimeStamp::ExecStringInput, InitExprValueCast); + arrow::compute::ScalarKernel k_large_binary_to_timestamp({arrow::large_binary()}, arrow::uint32(), ExecExprValueToTimeStamp::ExecStringInput, InitExprValueCast); - ARROW_RETURN_NOT_OK(expr_value_to_timestamp_func->AddKernel(k_string_to_timestamp)); + arrow::compute::ScalarKernel k_utf8_to_timestamp({arrow::utf8()}, arrow::uint32(), + ExecExprValueToTimeStamp::ExecStringInput, InitExprValueCast); + arrow::compute::ScalarKernel k_large_utf8_to_timestamp({arrow::large_utf8()}, arrow::uint32(), + ExecExprValueToTimeStamp::ExecStringInput, InitExprValueCast); + ARROW_RETURN_NOT_OK(expr_value_to_timestamp_func->AddKernel(k_binary_to_timestamp)); + ARROW_RETURN_NOT_OK(expr_value_to_timestamp_func->AddKernel(k_large_binary_to_timestamp)); + ARROW_RETURN_NOT_OK(expr_value_to_timestamp_func->AddKernel(k_utf8_to_timestamp)); + ARROW_RETURN_NOT_OK(expr_value_to_timestamp_func->AddKernel(k_large_utf8_to_timestamp)); ARROW_RETURN_NOT_OK(registry->AddFunction(expr_value_to_timestamp_func)); } + + /* + * string to numberic(兼容如'null','123a'不失败) + */ + { + // string cast numeric + for (const std::shared_ptr& out_ty : arrow::NumericTypes()) { + std::string func_name = "string_cast_" + out_ty->ToString(); + auto string_cast_numberic_func = std::make_shared(func_name, arrow::compute::Arity::Unary(), + /*doc=*/arrow::compute::FunctionDoc::Empty()); + ARROW_RETURN_NOT_OK(string_cast_numberic_func->AddKernel({arrow::binary()}, out_ty, + arrow::compute::internal::GenerateNumeric(*out_ty))); + ARROW_RETURN_NOT_OK(string_cast_numberic_func->AddKernel({arrow::large_binary()}, out_ty, + arrow::compute::internal::GenerateNumeric(*out_ty))); + ARROW_RETURN_NOT_OK(string_cast_numberic_func->AddKernel({arrow::utf8()}, out_ty, + arrow::compute::internal::GenerateNumeric(*out_ty))); + ARROW_RETURN_NOT_OK(string_cast_numberic_func->AddKernel({arrow::large_utf8()}, out_ty, + arrow::compute::internal::GenerateNumeric(*out_ty))); + ARROW_RETURN_NOT_OK(registry->AddFunction(string_cast_numberic_func)); + } + // 特殊处理string cast bool + std::string func_name = "string_cast_" + arrow::boolean()->ToString(); + auto string_cast_bool_func = std::make_shared(func_name, arrow::compute::Arity::Unary(), + /*doc=*/arrow::compute::FunctionDoc::Empty()); + for (const auto& in_ty : arrow::BaseBinaryTypes()) { + arrow::compute::ArrayKernelExec exec = arrow::compute::internal::GenerateVarBinaryBase(*in_ty); + ARROW_RETURN_NOT_OK(string_cast_bool_func->AddKernel({in_ty}, arrow::boolean(), exec)); + } + ARROW_RETURN_NOT_OK(registry->AddFunction(string_cast_bool_func)); + } + /* * murmur_hash */ @@ -1357,12 +1568,11 @@ arrow::Status ArrowFunctionManager::RegisterAllInteralFunction() { register_object("right_shift", arrow_right_shift); // 选择 - if (FLAGS_enable_arrow_complex_func) { - register_object("case_when", arrow_case_when); - register_object("case_expr_when", arrow_case_expr_when); - register_object("if", arrow_if); - register_object("ifnull", arrow_if_null); - } + register_object("case_when", arrow_case_when); + register_object("case_expr_when", arrow_case_expr_when); + register_object("if", arrow_if); + register_object("ifnull", arrow_if_null); + // string register_object("concat", arrow_concat); register_object("concat_ws", arrow_concat_ws); @@ -1374,6 +1584,9 @@ arrow::Status ArrowFunctionManager::RegisterAllInteralFunction() { register_object("lower", arrow_lower); register_object("reverse", arrow_reverse); register_object("repeat", arrow_repeat); + register_object("substring_index", arrow_substring_index); + register_object("replace", arrow_replace); + // 类型转换 register_object("cast_to_string", arrow_cast_to_string); register_object("cast_to_date", arrow_cast_to_date); @@ -1406,6 +1619,7 @@ arrow::Status ArrowFunctionManager::RegisterAllInteralFunction() { register_object("week", arrow_week); register_object("yearweek", arrow_yearweek); register_object("timestampdiff", arrow_timestampdiff); + register_object("datediff", arrow_datediff); // 其他 register_object("murmur_hash", arrow_murmur_hash); diff --git a/src/expr/arrow_string_function.cpp b/src/expr/arrow_string_function.cpp index 9fff424d1..0c3c087dc 100644 --- a/src/expr/arrow_string_function.cpp +++ b/src/expr/arrow_string_function.cpp @@ -154,6 +154,44 @@ int arrow_repeat(std::vector& children, pb::Function* fn, return 0; } +int arrow_substring_index(std::vector& children, pb::Function* fn, + const pb::PrimitiveType& return_type, arrow::compute::Expression& out) { + RETURN_NULL_IF_COLUMN_SATISFY_COND(children.size() != 3); + RETURN_NULL_IF_COLUMN_SATISFY_COND(children[1]->col_type() == pb::NULL_TYPE || children[2]->col_type() == pb::NULL_TYPE); + + std::string index = children[1]->get_value(nullptr).get_string(); + int64_t pos = children[2]->get_value(nullptr).get_numberic(); + if (0 != build_arrow_expr_with_cast(children[0], pb::STRING)) { + return -1; + } + if (index.empty() || pos == 0) { + out = arrow::literal(""); + return 0; + } + CommonFunctionOptions option(index, pos); + out = arrow::compute::call("baikal_substring_index", {children[0]->arrow_expr()}, std::move(option)); + return 0; +} + +int arrow_replace(std::vector& children, pb::Function* fn, + const pb::PrimitiveType& return_type, arrow::compute::Expression& out) { + RETURN_NULL_IF_COLUMN_SATISFY_COND(children.size() != 3); + RETURN_NULL_IF_COLUMN_SATISFY_COND(children[0]->col_type() == pb::NULL_TYPE); + + if (0 != build_arrow_expr_with_cast(children[0], pb::STRING)) { + return -1; + } + std::string pattern = children[1]->get_value(nullptr).get_string(); + std::string replacement = children[2]->get_value(nullptr).get_string(); + if (pattern.empty()) { + out = children[0]->arrow_expr(); + return 0; + } + arrow::compute::ReplaceSubstringOptions option(pattern, replacement); + out = arrow::compute::call("replace_substring", {children[0]->arrow_expr()}, std::move(option)); + return 0; +} + template struct UpperTransfer { using BuilderType = typename arrow::TypeTraits::BuilderType; @@ -200,6 +238,57 @@ struct LowerTransfer { } }; +template +struct SubstringIndex { + using BuilderType = typename arrow::TypeTraits::BuilderType; + static arrow::Status Exec(arrow::compute::KernelContext* ctx, const arrow::compute::ExecSpan& batch, arrow::compute::ExecResult* out) { + BuilderType builder; + const arrow::ArraySpan& input = batch[0].array; + CommonState* state = static_cast(ctx->state()); + const std::string& sub = state->str_value; + int64_t pos = state->int_value; + + RETURN_NOT_OK(arrow::compute::internal::VisitArrayValuesInline( + input, + [&](std::string_view v) { + std::vector pos_vec; + size_t last_pos = 0; + while (true) { + size_t find_pos = v.find(sub, last_pos); + if (find_pos != std::string::npos) { + pos_vec.emplace_back(find_pos); + last_pos = find_pos + sub.size(); + } else { + break; + } + } + if (pos > 0) { + if (pos <= pos_vec.size()) { + return builder.Append(v.substr(0, pos_vec[pos - 1])); + } else { + return builder.Append(v); + } + } else { + pos = -pos; + if (pos <= pos_vec.size()) { + pos = pos_vec.size() - pos; + return builder.Append(v.substr(pos_vec[pos] + sub.size())); + } else { + return builder.Append(v); + } + } + return builder.Append(v); + }, + [&]() { + return builder.AppendNull(); + })); + std::shared_ptr output_array; + RETURN_NOT_OK(builder.Finish(&output_array)); + out->value = std::move(output_array->data()); + return arrow::Status::OK(); + } +}; + arrow::Status ArrowFunctionManager::RegisterAllStringFunction() { auto registry = arrow::compute::GetFunctionRegistry(); { @@ -222,6 +311,17 @@ arrow::Status ArrowFunctionManager::RegisterAllStringFunction() { } ARROW_RETURN_NOT_OK(registry->AddFunction(lower)); } + { + auto substring_index = std::make_shared("baikal_substring_index", arrow::compute::Arity::Unary(), + /*doc=*/arrow::compute::FunctionDoc::Empty()); + for (const std::shared_ptr& in_ty : arrow::BaseBinaryTypes()) { + ARROW_RETURN_NOT_OK( + substring_index->AddKernel({in_ty}, arrow::large_binary(), + arrow::compute::internal::GenerateVarBinaryToVarBinary(*in_ty), + InitCommonState)); + } + ARROW_RETURN_NOT_OK(registry->AddFunction(substring_index)); + } return arrow::Status::OK(); } } diff --git a/src/expr/arrow_time_function.cpp b/src/expr/arrow_time_function.cpp index 8cbc964fb..43f2619f7 100644 --- a/src/expr/arrow_time_function.cpp +++ b/src/expr/arrow_time_function.cpp @@ -27,50 +27,6 @@ namespace baikaldb { static const int32_t DATE_FORMAT_LENGTH = 128; -/* - * 通用config - */ -class CommonTimeOptionsType : public arrow::compute::FunctionOptionsType { -public: - static const arrow::compute::FunctionOptionsType* GetInstance() { - static std::unique_ptr instance(new CommonTimeOptionsType()); - return instance.get(); - } - const char* type_name() const override { return "CommonTimeOptionType"; } - std::string Stringify(const arrow::compute::FunctionOptions& options) const override { - return type_name(); - } - bool Compare(const arrow::compute::FunctionOptions& options, - const arrow::compute::FunctionOptions& other) const override { - const auto& lop = static_cast(options); - const auto& rop = static_cast(other); - return lop.str_value == rop.str_value && lop.int_value == rop.int_value;; - } - std::unique_ptr Copy(const arrow::compute::FunctionOptions& options) const override { - const auto& opts = static_cast(options); - return std::make_unique(opts.str_value, opts.int_value); - } -}; - -struct CommonTimeState : public arrow::compute::KernelState { - std::string str_value; - int64_t int_value = 0; - CommonTimeState(const std::string& conf, int64_t value) : str_value(conf), int_value(value) {} -}; - -arrow::Result> InitCommonTimeState(arrow::compute::KernelContext*, - const arrow::compute::KernelInitArgs& args) { - auto func_options = static_cast(args.options); - if (func_options == nullptr) { - return std::make_unique("", 0); - } - return std::make_unique(func_options->str_value, func_options->int_value); -} - -CommonTimeFunctionOptions::CommonTimeFunctionOptions(const std::string& value, int64_t int_value) - : arrow::compute::FunctionOptions(CommonTimeOptionsType::GetInstance()), str_value(value), int_value(int_value) {} - - /* * 以下三个实际上不会走到, 直接转成常量表达式了 */ @@ -124,7 +80,7 @@ int arrow_date_format(std::vector& children, pb::Function* fn, const if (0 != build_arrow_expr_with_cast(children[0], pb::TIMESTAMP)) { return -1; } - CommonTimeFunctionOptions option(children[1]->get_value(nullptr).get_string()); + CommonFunctionOptions option(children[1]->get_value(nullptr).get_string()); out = arrow::compute::call("baikal_date_format", {children[0]->arrow_expr()}, std::move(option)); return 0; } @@ -134,7 +90,7 @@ int arrow_time_format(std::vector& children, pb::Function* fn, const if (0 != build_arrow_expr_with_cast(children[0], pb::TIME)) { return -1; } - CommonTimeFunctionOptions option(children[1]->get_value(nullptr).get_string()); + CommonFunctionOptions option(children[1]->get_value(nullptr).get_string()); out = arrow::compute::call("baikal_time_format", {children[0]->arrow_expr()}, std::move(option)); return 0; } @@ -186,7 +142,7 @@ int arrow_date_sub(std::vector& children, pb::Function* fn, const pb: } int32_t interval = children[1]->get_value(nullptr).get_numberic(); std::string unit = children[2]->get_value(nullptr).get_string(); - CommonTimeFunctionOptions option(unit, interval); + CommonFunctionOptions option(unit, interval); out = arrow::compute::call("baikal_date_sub", {children[0]->arrow_expr()}, std::move(option)); return 0; } @@ -209,7 +165,7 @@ int arrow_date_add(std::vector& children, pb::Function* fn, const pb: } int32_t interval = children[1]->get_value(nullptr).get_numberic(); std::string unit = children[2]->get_value(nullptr).get_string(); - CommonTimeFunctionOptions option(unit, interval); + CommonFunctionOptions option(unit, interval); out = arrow::compute::call("baikal_date_add", {children[0]->arrow_expr()}, std::move(option)); return 0; } @@ -286,7 +242,7 @@ int arrow_week(std::vector& children, pb::Function* fn, const pb::Pri if (children.size() > 1) { mode = children[1]->get_value(nullptr).get_numberic() % 8; } - CommonTimeFunctionOptions option("", mode); + CommonFunctionOptions option("", mode); out = arrow::compute::call("baikal_week", {children[0]->arrow_expr()}, std::move(option)); return 0; } @@ -300,7 +256,7 @@ int arrow_yearweek(std::vector& children, pb::Function* fn, const pb: if (children.size() > 1) { mode = children[1]->get_value(nullptr).get_numberic() % 8; } - CommonTimeFunctionOptions option("", mode); + CommonFunctionOptions option("", mode); out = arrow::compute::call("baikal_yearweek", {children[0]->arrow_expr()}, std::move(option)); return 0; } @@ -328,6 +284,32 @@ int arrow_timestampdiff(std::vector& children, pb::Function* fn, cons } return 0; } + +int arrow_datediff(std::vector& children, pb::Function* fn, const pb::PrimitiveType& return_type, arrow::compute::Expression& out) { + RETURN_NULL_IF_COLUMN_SATISFY_COND(children.size() != 2); + BUILD_ARROW_EXPR_RET(children[0]); + BUILD_ARROW_EXPR_RET(children[1]); + auto left = children[0]->arrow_expr(); + auto right = children[1]->arrow_expr(); + if (children[0]->col_type() == pb::INT64) { + left = arrow_cast(left, children[0]->col_type(), pb::STRING); + left = arrow_cast(left, pb::STRING, pb::DATE); + } else { + left = arrow_cast(left, children[0]->col_type(), pb::DATE); + } + if (children[1]->col_type() == pb::INT64) { + right = arrow_cast(right, children[1]->col_type(), pb::STRING); + right = arrow_cast(right, pb::STRING, pb::DATE); + } else { + right = arrow_cast(right, children[1]->col_type(), pb::DATE); + } + left = arrow_cast(left, pb::DATE, pb::TIMESTAMP); + right = arrow_cast(right, pb::DATE, pb::TIMESTAMP); + out = arrow::compute::call("subtract_checked", {arrow_cast(left, pb::TIMESTAMP, pb::INT64), + arrow_cast(right, pb::TIMESTAMP, pb::INT64)}); + out = arrow::compute::call("divide_checked", {out, arrow::compute::literal(3600 * 24)}); + return 0; +} /* * date_format arrow实现 */ @@ -337,7 +319,7 @@ struct ExecDateFormat { using BuilderType = typename arrow::TypeTraits::BuilderType; static arrow::Status Exec(arrow::compute::KernelContext* ctx, const arrow::compute::ExecSpan& batch, arrow::compute::ExecResult* out) { - CommonTimeState* state = static_cast(ctx->state()); + CommonState* state = static_cast(ctx->state()); std::string format = state->str_value; BuilderType builder; const arrow::ArraySpan& input = batch[0].array; @@ -347,7 +329,8 @@ struct ExecDateFormat { // 核心 struct tm t_result; time_t t = (uint32_t)v; - localtime_r(&t, &t_result); + localtime_fixed_r(&t, &t_result); + char s[DATE_FORMAT_LENGTH]; date_format_internal(s, sizeof(s), format.data(), &t_result); return builder.Append(std::string(s)); @@ -370,7 +353,7 @@ struct ExecTimeFormat { using BuilderType = typename arrow::TypeTraits::BuilderType; static arrow::Status Exec(arrow::compute::KernelContext* ctx, const arrow::compute::ExecSpan& batch, arrow::compute::ExecResult* out) { - CommonTimeState* state = static_cast(ctx->state()); + CommonState* state = static_cast(ctx->state()); std::string format = state->str_value; BuilderType builder; const arrow::ArraySpan& input = batch[0].array; @@ -404,7 +387,7 @@ struct ExecHour { using OutputValueCType = typename arrow::TypeTraits::CType; static arrow::Status Exec(arrow::compute::KernelContext* ctx, const arrow::compute::ExecSpan& batch, arrow::compute::ExecResult* out) { - CommonTimeState* state = static_cast(ctx->state()); + CommonState* state = static_cast(ctx->state()); std::string format = state->str_value; const arrow::ArraySpan& input = batch[0].array; arrow::ArraySpan* out_data = out->array_span_mutable(); @@ -431,7 +414,7 @@ struct ExecDateSub { using OutputValueCType = typename arrow::TypeTraits::CType; static arrow::Status Exec(arrow::compute::KernelContext* ctx, const arrow::compute::ExecSpan& batch, arrow::compute::ExecResult* out) { - CommonTimeState* state = static_cast(ctx->state()); + CommonState* state = static_cast(ctx->state()); int32_t interval = state->int_value; const std::string& unit = state->str_value; const arrow::ArraySpan& input = batch[0].array; @@ -473,7 +456,7 @@ struct ExecDateAdd { using OutputValueCType = typename arrow::TypeTraits::CType; static arrow::Status Exec(arrow::compute::KernelContext* ctx, const arrow::compute::ExecSpan& batch, arrow::compute::ExecResult* out) { - CommonTimeState* state = static_cast(ctx->state()); + CommonState* state = static_cast(ctx->state()); int32_t interval = state->int_value; const std::string& unit = state->str_value; const arrow::ArraySpan& input = batch[0].array; @@ -550,7 +533,7 @@ struct ExecWeek { using OutputValueCType = typename arrow::TypeTraits::CType; static arrow::Status Exec(arrow::compute::KernelContext* ctx, const arrow::compute::ExecSpan& batch, arrow::compute::ExecResult* out) { - CommonTimeState* state = static_cast(ctx->state()); + CommonState* state = static_cast(ctx->state()); int64_t mode = state->int_value; int year = 0; int weeks = 0; @@ -589,7 +572,7 @@ struct ExecYearWeek { using OutputValueCType = typename arrow::TypeTraits::CType; static arrow::Status Exec(arrow::compute::KernelContext* ctx, const arrow::compute::ExecSpan& batch, arrow::compute::ExecResult* out) { - CommonTimeState* state = static_cast(ctx->state()); + CommonState* state = static_cast(ctx->state()); int64_t mode = state->int_value; int year = 0; int weeks = 0; @@ -631,7 +614,7 @@ arrow::Status ArrowFunctionManager::RegisterAllTimeFunction() { ARROW_RETURN_NOT_OK( arrow_date_format->AddKernel({in_ty}, arrow::large_binary(), arrow::compute::internal::GenerateNumeric(*in_ty), - InitCommonTimeState)); + InitCommonState)); } ARROW_RETURN_NOT_OK(registry->AddFunction(arrow_date_format)); } @@ -645,7 +628,7 @@ arrow::Status ArrowFunctionManager::RegisterAllTimeFunction() { ARROW_RETURN_NOT_OK( arrow_time_format->AddKernel({in_ty}, arrow::large_binary(), arrow::compute::internal::GenerateNumeric(*in_ty), - InitCommonTimeState)); + InitCommonState)); } ARROW_RETURN_NOT_OK(registry->AddFunction(arrow_time_format)); } @@ -659,7 +642,7 @@ arrow::Status ArrowFunctionManager::RegisterAllTimeFunction() { ARROW_RETURN_NOT_OK( arrow_hour->AddKernel({in_ty}, arrow::uint32(), arrow::compute::internal::GenerateNumeric(*in_ty), - InitCommonTimeState)); + InitCommonState)); } ARROW_RETURN_NOT_OK(registry->AddFunction(arrow_hour)); } @@ -673,7 +656,7 @@ arrow::Status ArrowFunctionManager::RegisterAllTimeFunction() { ARROW_RETURN_NOT_OK( arrow_date_sub->AddKernel({in_ty}, arrow::uint64(), arrow::compute::internal::GenerateNumeric(*in_ty), - InitCommonTimeState)); + InitCommonState)); } ARROW_RETURN_NOT_OK(registry->AddFunction(arrow_date_sub)); } @@ -687,7 +670,7 @@ arrow::Status ArrowFunctionManager::RegisterAllTimeFunction() { ARROW_RETURN_NOT_OK( arrow_date_add->AddKernel({in_ty}, arrow::uint64(), arrow::compute::internal::GenerateNumeric(*in_ty), - InitCommonTimeState)); + InitCommonState)); } ARROW_RETURN_NOT_OK(registry->AddFunction(arrow_date_add)); } @@ -701,7 +684,7 @@ arrow::Status ArrowFunctionManager::RegisterAllTimeFunction() { ARROW_RETURN_NOT_OK( time_to_sec->AddKernel({in_ty}, arrow::uint32(), arrow::compute::internal::GenerateNumeric(*in_ty), - InitCommonTimeState)); + InitCommonState)); } ARROW_RETURN_NOT_OK(registry->AddFunction(time_to_sec)); } @@ -714,7 +697,7 @@ arrow::Status ArrowFunctionManager::RegisterAllTimeFunction() { for (const std::shared_ptr& in_ty : arrow::NumericTypes()) { arrow::compute::ScalarKernel kernel({in_ty}, arrow::uint32(), arrow::compute::internal::GenerateNumeric(*in_ty), - InitCommonTimeState); + InitCommonState); /// Kernel expects a pre-allocated buffer to write the result bitmap /// into. The preallocated memory is not zeroed (except for the last byte), /// so the kernel should ensure to completely populate the bitmap. @@ -733,7 +716,7 @@ arrow::Status ArrowFunctionManager::RegisterAllTimeFunction() { for (const std::shared_ptr& in_ty : arrow::NumericTypes()) { arrow::compute::ScalarKernel kernel({in_ty}, arrow::uint32(), arrow::compute::internal::GenerateNumeric(*in_ty), - InitCommonTimeState); + InitCommonState); kernel.null_handling = arrow::compute::NullHandling::COMPUTED_PREALLOCATE; kernel.mem_allocation = arrow::compute::MemAllocation::PREALLOCATE; ARROW_RETURN_NOT_OK(yearweek_func->AddKernel(kernel)); diff --git a/src/expr/expr_node.cpp b/src/expr/expr_node.cpp index d374e90af..ed0cf8b67 100644 --- a/src/expr/expr_node.cpp +++ b/src/expr/expr_node.cpp @@ -822,7 +822,7 @@ int ExprNode::serialize_tree(uint64_t& serialize_slot_id) { if (slot_id > 63) { return -1; } - serialize_slot_id |= 1 << slot_id; + serialize_slot_id |= 1LL << slot_id; } } else if (is_function_eq()) { if (_children[0] == nullptr) { @@ -833,7 +833,7 @@ int ExprNode::serialize_tree(uint64_t& serialize_slot_id) { if (slot_id > 63) { return -1; } - serialize_slot_id |= 1 << slot_id; + serialize_slot_id |= 1LL << slot_id; } return 0; } diff --git a/src/expr/fn_manager.cpp b/src/expr/fn_manager.cpp index c3e3f5aa5..6d32dde51 100644 --- a/src/expr/fn_manager.cpp +++ b/src/expr/fn_manager.cpp @@ -304,7 +304,7 @@ void FunctionManager::register_operators() { register_object_ret("cast_to_datetime", cast_to_datetime, pb::DATETIME); register_object_ret("cast_to_string", cast_to_string, pb::STRING); register_object_ret("cast_to_signed", cast_to_signed, pb::INT64); - register_object_ret("cast_to_unsigned", cast_to_unsigned, pb::INT64); + register_object_ret("cast_to_unsigned", cast_to_unsigned, pb::UINT64); register_object_ret("cast_to_double", cast_to_double, pb::DOUBLE); } diff --git a/src/expr/internal_functions.cpp b/src/expr/internal_functions.cpp index 2996850a3..dfaa6e296 100644 --- a/src/expr/internal_functions.cpp +++ b/src/expr/internal_functions.cpp @@ -1534,7 +1534,7 @@ ExprValue date_format(const std::vector& input) { ExprValue tmp = input[0]; time_t t = tmp.cast_to(pb::TIMESTAMP)._u.uint32_val; struct tm t_result; - localtime_r(&t, &t_result); + localtime_fixed_r(&t, &t_result); char s[DATE_FORMAT_LENGTH]; date_format_internal(s, sizeof(s), input[1].str_val.c_str(), &t_result); ExprValue format_result(pb::STRING); @@ -2250,7 +2250,7 @@ ExprValue extract(const std::vector& input) { ExprValue tmp(pb::UINT32); time_t t = arg2.cast_to(pb::TIMESTAMP)._u.uint32_val; struct tm tm; - localtime_r(&t, &tm); + localtime_fixed_r(&t, &tm); if (input[0].str_val == "year") { tmp._u.uint32_val = tm.tm_year + 1900; } else if (input[0].str_val == "month") { diff --git a/src/expr/predicate.cpp b/src/expr/predicate.cpp index ea6845bb0..665987b90 100644 --- a/src/expr/predicate.cpp +++ b/src/expr/predicate.cpp @@ -594,8 +594,14 @@ int LikePredicate::transfer_to_arrow_expression() { } else { args.emplace_back(arrow::compute::call("cast", {_children[0]->arrow_expr()}, arrow::compute::CastOptions::Unsafe(arrow::large_binary()))); } - arrow::compute::MatchSubstringOptions opt(children(1)->get_value(nullptr).get_string(), /*ignore_case*/false); - _arrow_expr = arrow::compute::call("match_like", args, std::move(opt)); + if (_fn.fn_op() == parser::FT_EXACT_LIKE) { + covent_exact_pattern(children(1)->get_value(nullptr).get_string()); + arrow::compute::MatchSubstringOptions opt(_regex_pattern, /*ignore_case*/false); + _arrow_expr = arrow::compute::call("match_substring_regex", args, std::move(opt)); + } else { + arrow::compute::MatchSubstringOptions opt(children(1)->get_value(nullptr).get_string(), /*ignore_case*/false); + _arrow_expr = arrow::compute::call("match_like", args, std::move(opt)); + } return 0; } diff --git a/src/expr/scalar_fn_call.cpp b/src/expr/scalar_fn_call.cpp index 402a1ef25..972fc672b 100644 --- a/src/expr/scalar_fn_call.cpp +++ b/src/expr/scalar_fn_call.cpp @@ -33,7 +33,9 @@ int ScalarFnCall::init(const pb::ExprNode& node) { return -1; } _fn = node.fn(); - _origin_fn_name = _fn.name(); + if (_fn.origin_name().empty()) { + _fn.set_origin_name(_fn.name()); + } // rand不是const if (node_type() == pb::FUNCTION_CALL && _fn.name() == "rand") { _is_constant = false; @@ -256,7 +258,9 @@ const std::unordered_map ARROW_FUNC_SLOT_REF_COUNT = { {"repeat", 1}, {"substr", 1}, {"week", 1}, - {"yearweek", 1} + {"yearweek", 1}, + {"substring_index", 1}, + {"replace", 1} }; bool ScalarFnCall::can_use_arrow_vector() { @@ -329,7 +333,7 @@ std::string ScalarFnCall::to_sql(const std::unordered_map& return ""; } } - std::string fn_name = _origin_fn_name; + std::string fn_name = _fn.origin_name(); // 一元负数运算符特殊处理 if (_fn.fn_op() == parser::FT_UMINUS) { fn_name = "uminus"; diff --git a/src/expr/window_fn_call.cpp b/src/expr/window_fn_call.cpp index 7925c0c28..b90d20e3a 100644 --- a/src/expr/window_fn_call.cpp +++ b/src/expr/window_fn_call.cpp @@ -89,7 +89,7 @@ int WindowFnCall::open() { case DENSE_RANK: case PERCENT_RANK: case CUME_DIST: { - WindowRankIntermediate* window_intermediate = (WindowRankIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->mem_row_compare = _mem_row_compare; break; } @@ -103,11 +103,11 @@ int WindowFnCall::open() { return -1; } if (_children[0]->node_type() != pb::INT_LITERAL) { - DB_WARNING("_window_type: %d, invalid children node_type: %d", + DB_WARNING("_window_type: %d, invalid children node_type: %d", _window_type, _children[0]->node_type()); return -1; } - WindowNtileIntermediate* window_intermediate = (WindowNtileIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->n = _children[0]->get_value(nullptr).cast_to(pb::INT64).get_numberic(); if (window_intermediate->n <= 0) { DB_WARNING("Invalid n: %ld", window_intermediate->n); @@ -131,11 +131,11 @@ int WindowFnCall::open() { return -1; } if (_children[1]->node_type() != pb::INT_LITERAL) { - DB_WARNING("_window_type: %d, invalid children node_type: %d", + DB_WARNING("_window_type: %d, invalid children node_type: %d", _window_type, _children[1]->node_type()); return -1; } - WindowLeadLagIntermediate* window_intermediate = (WindowLeadLagIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->offset = _children[1]->get_value(nullptr).cast_to(pb::INT64).get_numberic(); if (window_intermediate->offset < 0) { DB_WARNING("Invalid offset: %ld", window_intermediate->offset); @@ -177,11 +177,11 @@ int WindowFnCall::open() { return -1; } if (_children[1]->node_type() != pb::INT_LITERAL) { - DB_WARNING("_window_type: %d, invalid children node_type: %d", + DB_WARNING("_window_type: %d, invalid children node_type: %d", _window_type, _children[1]->node_type()); return -1; } - WindowValueIntermediate* window_intermediate = (WindowValueIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->n = _children[1]->get_value(nullptr).cast_to(pb::INT64).get_numberic(); if (window_intermediate->n <= 0) { DB_WARNING("Invalid n: %ld", window_intermediate->n); @@ -200,7 +200,6 @@ void WindowFnCall::close() { ExprNode::close(); _tuple_id = -1; _slot_id = -1; - _window_intermediate.clear(); _mem_row_compare = nullptr; } @@ -234,7 +233,7 @@ int WindowFnCall::type_inferer() { _col_type = pb::DOUBLE; break; } - case MIN: + case MIN: case MAX: { if (_children.size() == 0) { DB_WARNING("has no child"); @@ -244,12 +243,12 @@ int WindowFnCall::type_inferer() { break; } case ROW_NUMBER: - case RANK: + case RANK: case DENSE_RANK: { _col_type = pb::INT64; break; } - case PERCENT_RANK: + case PERCENT_RANK: case CUME_DIST: { _col_type = pb::DOUBLE; break; @@ -312,55 +311,46 @@ int WindowFnCall::initialize() { switch (_window_type) { case COUNT_STAR: case COUNT: { - WindowCountIntermediate window_intermediate; - _window_intermediate.assign((char*)&window_intermediate, sizeof(WindowCountIntermediate)); + _window_intermediate.emplace(); break; } case SUM: { - WindowSumIntermediate window_intermediate; - _window_intermediate.assign((char*)&window_intermediate, sizeof(WindowSumIntermediate)); + _window_intermediate.emplace(); break; } case AVG: { - WindowAvgIntermediate window_intermediate; - _window_intermediate.assign((char*)&window_intermediate, sizeof(WindowAvgIntermediate)); + _window_intermediate.emplace(); break; } case MIN: case MAX: { - WindowMinMaxIntermediate window_intermediate; - _window_intermediate.assign((char*)&window_intermediate, sizeof(WindowMinMaxIntermediate)); + _window_intermediate.emplace(); break; } case ROW_NUMBER: { - WindowRowNumberIntermediate window_intermediate; - _window_intermediate.assign((char*)&window_intermediate, sizeof(WindowRowNumberIntermediate)); + _window_intermediate.emplace(); break; } case RANK: case DENSE_RANK: - case PERCENT_RANK: + case PERCENT_RANK: case CUME_DIST: { - WindowRankIntermediate window_intermediate; - _window_intermediate.assign((char*)&window_intermediate, sizeof(WindowRankIntermediate)); + _window_intermediate.emplace(); break; } case NTILE: { - WindowNtileIntermediate window_intermediate; - _window_intermediate.assign((char*)&window_intermediate, sizeof(WindowNtileIntermediate)); + _window_intermediate.emplace(); break; } case LEAD: case LAG: { - WindowLeadLagIntermediate window_intermediate; - _window_intermediate.assign((char*)&window_intermediate, sizeof(WindowLeadLagIntermediate)); + _window_intermediate.emplace(); break; } case FIRST_VALUE: case LAST_VALUE: case NTH_VALUE: { - WindowValueIntermediate window_intermediate; - _window_intermediate.assign((char*)&window_intermediate, sizeof(WindowValueIntermediate)); + _window_intermediate.emplace(); break; } default: { @@ -382,12 +372,12 @@ int WindowFnCall::call(RowBatch* batch, const int start, const int end) { } switch (_window_type) { case COUNT_STAR: { - WindowCountIntermediate* window_intermediate = (WindowCountIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->cnt += (end - start); break; } case COUNT: { - WindowCountIntermediate* window_intermediate = (WindowCountIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); for (int i = start; i < end; ++i) { bool can_count = true; for (auto child : _children) { @@ -403,7 +393,7 @@ int WindowFnCall::call(RowBatch* batch, const int start, const int end) { break; } case SUM: { - WindowSumIntermediate* window_intermediate = (WindowSumIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); for (int i = start; i < end; ++i) { ExprValue value = _children[0]->get_value(batch->get_row(i).get()).cast_to(_col_type); if (!value.is_null()) { @@ -417,7 +407,7 @@ int WindowFnCall::call(RowBatch* batch, const int start, const int end) { break; } case AVG: { - WindowAvgIntermediate* window_intermediate = (WindowAvgIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); for (int i = start; i < end; ++i) { ExprValue value = _children[0]->get_value(batch->get_row(i).get()).cast_to(_col_type); if (!value.is_null()) { @@ -428,7 +418,7 @@ int WindowFnCall::call(RowBatch* batch, const int start, const int end) { break; } case MIN: { - WindowMinMaxIntermediate* window_intermediate = (WindowMinMaxIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); for (int i = start; i < end; ++i) { ExprValue value = _children[0]->get_value(batch->get_row(i).get()).cast_to(_col_type); if (!value.is_null()) { @@ -441,7 +431,7 @@ int WindowFnCall::call(RowBatch* batch, const int start, const int end) { break; } case MAX: { - WindowMinMaxIntermediate* window_intermediate = (WindowMinMaxIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); for (int i = start; i < end; ++i) { ExprValue value = _children[0]->get_value(batch->get_row(i).get()).cast_to(_col_type); if (!value.is_null()) { @@ -458,14 +448,14 @@ int WindowFnCall::call(RowBatch* batch, const int start, const int end) { } case RANK: case DENSE_RANK: - case PERCENT_RANK: + case PERCENT_RANK: case CUME_DIST: { - WindowRankIntermediate* window_intermediate = (WindowRankIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->cur_batch = batch; break; } case NTILE: { - WindowNtileIntermediate* window_intermediate = (WindowNtileIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); if (window_intermediate->n <= 0) { DB_WARNING("invalid n: %ld", window_intermediate->n); return -1; @@ -477,22 +467,22 @@ int WindowFnCall::call(RowBatch* batch, const int start, const int end) { } case LEAD: case LAG: { - WindowLeadLagIntermediate* window_intermediate = (WindowLeadLagIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->cur_batch = batch; break; } case FIRST_VALUE: { - WindowValueIntermediate* window_intermediate = (WindowValueIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->value = _children[0]->get_value(batch->get_row(start).get()).cast_to(_col_type); break; } case LAST_VALUE: { - WindowValueIntermediate* window_intermediate = (WindowValueIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->value = _children[0]->get_value(batch->get_row(end - 1).get()).cast_to(_col_type); break; } case NTH_VALUE: { - WindowValueIntermediate* window_intermediate = (WindowValueIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); int idx = start + window_intermediate->n - 1; if (idx >= start && idx < end) { window_intermediate->value = _children[0]->get_value(batch->get_row(idx).get()).cast_to(_col_type); @@ -511,18 +501,18 @@ int WindowFnCall::get_result(ExprValue& expr_value) { switch (_window_type) { case COUNT_STAR: case COUNT: { - WindowCountIntermediate* window_intermediate = (WindowCountIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); expr_value.type = _col_type; expr_value._u.int64_val = window_intermediate->cnt; break; } case SUM: { - WindowSumIntermediate* window_intermediate = (WindowSumIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); expr_value = window_intermediate->sum; break; } case AVG: { - WindowAvgIntermediate* window_intermediate = (WindowAvgIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); if (window_intermediate->cnt != 0) { expr_value.type = _col_type; expr_value._u.double_val = window_intermediate->sum / window_intermediate->cnt; @@ -533,12 +523,12 @@ int WindowFnCall::get_result(ExprValue& expr_value) { } case MIN: case MAX: { - WindowMinMaxIntermediate* window_intermediate = (WindowMinMaxIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); expr_value = window_intermediate->min_max; break; } case ROW_NUMBER: { - WindowRowNumberIntermediate* window_intermediate = (WindowRowNumberIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); expr_value.type = _col_type; window_intermediate->row_number++; expr_value._u.int64_val = window_intermediate->row_number; @@ -551,14 +541,14 @@ int WindowFnCall::get_result(ExprValue& expr_value) { expr_value.type = _col_type; expr_value._u.int64_val = 1; } else { - WindowRankIntermediate* window_intermediate = (WindowRankIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->cur_idx++; if (window_intermediate->cur_idx == 1) { window_intermediate->last_rank = 1; } else { - if (window_intermediate->cur_idx < 2 || + if (window_intermediate->cur_idx < 2 || window_intermediate->cur_idx > window_intermediate->cur_batch->size()) { - DB_WARNING("Invalid window index : %ld, batch size: %lu", + DB_WARNING("Invalid window index : %ld, batch size: %lu", window_intermediate->cur_idx, window_intermediate->cur_batch->size()); return -1; } @@ -578,13 +568,13 @@ int WindowFnCall::get_result(ExprValue& expr_value) { break; } case PERCENT_RANK: { - // 计算「当前行排名 - 1」/ 「总行数 - 1」 + // 计算「当前行排名 - 1」/ 「总行数 - 1」 if (_mem_row_compare->need_not_compare()) { // 没有排序列，直接返回0 expr_value.type = _col_type; expr_value._u.double_val = 0; } else { - WindowRankIntermediate* window_intermediate = (WindowRankIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); int32_t num_rows = window_intermediate->cur_batch->size(); if (num_rows == 0) { DB_WARNING("empty batch"); @@ -594,9 +584,9 @@ int WindowFnCall::get_result(ExprValue& expr_value) { if (window_intermediate->cur_idx == 1) { window_intermediate->last_rank = 1; } else { - if (window_intermediate->cur_idx < 2 || + if (window_intermediate->cur_idx < 2 || window_intermediate->cur_idx > window_intermediate->cur_batch->size()) { - DB_WARNING("Invalid window index : %ld, batch size: %lu", + DB_WARNING("Invalid window index : %ld, batch size: %lu", window_intermediate->cur_idx, window_intermediate->cur_batch->size()); return -1; } @@ -622,21 +612,21 @@ int WindowFnCall::get_result(ExprValue& expr_value) { expr_value.type = _col_type; expr_value._u.double_val = 1; } else { - WindowRankIntermediate* window_intermediate = (WindowRankIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); int32_t num_rows = window_intermediate->cur_batch->size(); if (num_rows == 0) { DB_WARNING("empty batch"); return -1; } - if (window_intermediate->cur_idx < 0 || + if (window_intermediate->cur_idx < 0 || window_intermediate->cur_idx >= window_intermediate->cur_batch->size()) { - DB_WARNING("Invalid window index : %ld, batch size: %lu", + DB_WARNING("Invalid window index : %ld, batch size: %lu", window_intermediate->cur_idx, window_intermediate->cur_batch->size()); return -1; } while (true) { if (window_intermediate->last_rank < 0) { - DB_WARNING("Invalid window last rank: %ld, batch size: %lu", + DB_WARNING("Invalid window last rank: %ld, batch size: %lu", window_intermediate->last_rank, window_intermediate->cur_batch->size()); return -1; } @@ -664,7 +654,7 @@ int WindowFnCall::get_result(ExprValue& expr_value) { // - 2号桶: 3行 // - 3号桶: 2行 // - 4号桶: 2行 - WindowNtileIntermediate* window_intermediate = (WindowNtileIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); expr_value.type = _col_type; expr_value._u.int64_val = window_intermediate->cur_group_idx; ++window_intermediate->cur_idx; @@ -682,7 +672,7 @@ int WindowFnCall::get_result(ExprValue& expr_value) { case LAG: { // LEAD: 返回当前行下方指定偏移量的表达式值 // LAG: 返回当前行上方指定偏移量的表达式值 - WindowLeadLagIntermediate* window_intermediate = (WindowLeadLagIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); int64_t idx = window_intermediate->cur_idx; if (_window_type == LEAD) { idx += window_intermediate->offset; @@ -690,7 +680,7 @@ int WindowFnCall::get_result(ExprValue& expr_value) { idx -= window_intermediate->offset; } if (idx >= 0 && idx < window_intermediate->cur_batch->size()) { - expr_value = + expr_value = _children[0]->get_value(window_intermediate->cur_batch->get_row(idx).get()).cast_to(_col_type); } else { if (window_intermediate->default_expr != nullptr) { @@ -707,9 +697,9 @@ int WindowFnCall::get_result(ExprValue& expr_value) { break; } case FIRST_VALUE: - case LAST_VALUE: + case LAST_VALUE: case NTH_VALUE: { - WindowValueIntermediate* window_intermediate = (WindowValueIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); expr_value = window_intermediate->value; break; } @@ -725,31 +715,31 @@ int WindowFnCall::frame_reset() { switch (_window_type) { case COUNT_STAR: case COUNT: { - WindowCountIntermediate* window_intermediate = (WindowCountIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->cnt = 0; break; } case SUM: { - WindowSumIntermediate* window_intermediate = (WindowSumIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->sum = ExprValue::Null(); break; } case AVG: { - WindowAvgIntermediate* window_intermediate = (WindowAvgIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->sum = 0; window_intermediate->cnt = 0; break; } case MIN: case MAX: { - WindowMinMaxIntermediate* window_intermediate = (WindowMinMaxIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->min_max = ExprValue::Null(); break; } case ROW_NUMBER: case RANK: case DENSE_RANK: - case PERCENT_RANK: + case PERCENT_RANK: case CUME_DIST: case NTILE: case LEAD: @@ -758,9 +748,9 @@ int WindowFnCall::frame_reset() { break; } case FIRST_VALUE: - case LAST_VALUE: + case LAST_VALUE: case NTH_VALUE: { - WindowValueIntermediate* window_intermediate = (WindowValueIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->value = ExprValue::Null(); break; } @@ -776,29 +766,29 @@ int WindowFnCall::partition_reset() { switch (_window_type) { case COUNT_STAR: case COUNT: { - WindowCountIntermediate* window_intermediate = (WindowCountIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->cnt = 0; break; } case SUM: { - WindowSumIntermediate* window_intermediate = (WindowSumIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->sum = ExprValue::Null(); break; } case AVG: { - WindowAvgIntermediate* window_intermediate = (WindowAvgIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->sum = 0; window_intermediate->cnt = 0; break; } case MIN: case MAX: { - WindowMinMaxIntermediate* window_intermediate = (WindowMinMaxIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->min_max = ExprValue::Null(); break; } case ROW_NUMBER: { - WindowRowNumberIntermediate* window_intermediate = (WindowRowNumberIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->row_number = 0; break; } @@ -806,14 +796,14 @@ int WindowFnCall::partition_reset() { case DENSE_RANK: case PERCENT_RANK: case CUME_DIST: { - WindowRankIntermediate* window_intermediate = (WindowRankIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->cur_idx = 0; window_intermediate->last_rank = 0; window_intermediate->cur_batch = nullptr; break; } case NTILE: { - WindowNtileIntermediate* window_intermediate = (WindowNtileIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->cur_idx = 0; window_intermediate->cur_group_idx = 1; window_intermediate->quotient = 0; @@ -822,7 +812,7 @@ int WindowFnCall::partition_reset() { } case LEAD: case LAG: { - WindowLeadLagIntermediate* window_intermediate = (WindowLeadLagIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->cur_idx = 0; window_intermediate->cur_batch = nullptr; break; @@ -830,7 +820,7 @@ int WindowFnCall::partition_reset() { case FIRST_VALUE: case LAST_VALUE: case NTH_VALUE: { - WindowValueIntermediate* window_intermediate = (WindowValueIntermediate*)_window_intermediate.c_str(); + auto window_intermediate = &std::get(_window_intermediate); window_intermediate->value = ExprValue::Null(); break; } diff --git a/src/logical_plan/ddl_planner.cpp b/src/logical_plan/ddl_planner.cpp index e53a9a178..082d792e7 100644 --- a/src/logical_plan/ddl_planner.cpp +++ b/src/logical_plan/ddl_planner.cpp @@ -21,8 +21,8 @@ #include namespace baikaldb { -DEFINE_bool(unique_index_default_global, true, "unique_index_default_global"); -DEFINE_bool(normal_index_default_global, false, "normal_index_default_global"); +DEFINE_bool(unique_index_default_global, true, "Default unique index as global, default: true"); +DEFINE_bool(normal_index_default_global, false, "Default normal index as global, default: false"); int DDLPlanner::plan() { pb::MetaManagerRequest request; if (!_ctx->user_info->allow_ddl()) { @@ -878,11 +878,9 @@ int DDLPlanner::parse_create_table(pb::SchemaInfo& table) { } } else if (constraint->type == parser::CONSTRAINT_FULLTEXT) { index->set_index_type(pb::I_FULLTEXT); - can_support_ttl = false; index->set_storage_type(pb::ST_ARROW); } else if (constraint->type == parser::CONSTRAINT_VECTOR) { index->set_index_type(pb::I_VECTOR); - can_support_ttl = false; } else if (constraint->type == parser::CONSTRAINT_ROLLUP) { index->set_index_type(pb::I_ROLLUP); can_support_ttl = false; @@ -1109,6 +1107,34 @@ int DDLPlanner::parse_create_table(pb::SchemaInfo& table) { table.set_online_ttl_expire_time_us(0); DB_WARNING("ttl_duration: %ld", ttl_duration); } + // 新格式ttl配置优先 + json_iter = root.FindMember("ttl"); + if (json_iter != root.MemberEnd()) { + if (!can_support_ttl) { + DB_FATAL("fulltext/engine!=rocksdb can not create ttl table"); + return -1; + } + auto duration_iter = json_iter->value.FindMember("duration"); + if (duration_iter == json_iter->value.MemberEnd()) { + DB_FATAL("ttl need specify duration"); + return -1; + } + std::string ttl_field_name; + auto field_iter = json_iter->value.FindMember("field"); + if (field_iter != json_iter->value.MemberEnd()) { + ttl_field_name = field_iter->value.GetString(); + std::vector splits; + boost::split(splits, ttl_field_name, boost::is_any_of(".")); + ttl_field_name = splits.back(); + } + int64_t ttl_duration = duration_iter->value.GetInt64(); + table.set_ttl_duration(ttl_duration); + if (!ttl_field_name.empty()) { + table.mutable_ttl_field()->set_field_name(ttl_field_name); + } + table.set_online_ttl_expire_time_us(0); + DB_WARNING("ttl_duration: %ld, ttl field: %s", ttl_duration, ttl_field_name.c_str()); + } json_iter = root.FindMember("storage_compute_separate"); if (json_iter != root.MemberEnd()) { int64_t separate = json_iter->value.GetInt64(); @@ -1293,6 +1319,48 @@ int DDLPlanner::parse_create_table(pb::SchemaInfo& table) { dblink_info->mutable_mysql_info()->set_charset(mysql_info_iter->value.GetString()); } } + iter = value.FindMember("file_info"); + if (iter != value.MemberEnd()) { + const rapidjson::Value& file_info_value = iter->value; + auto file_info_iter = file_info_value.FindMember("cluster"); + if (file_info_iter != file_info_value.MemberEnd() && file_info_iter->value.IsString()) { + dblink_info->mutable_file_info()->set_cluster(file_info_iter->value.GetString()); + } + file_info_iter = file_info_value.FindMember("path"); + if (file_info_iter != file_info_value.MemberEnd() && file_info_iter->value.IsString()) { + dblink_info->mutable_file_info()->set_path(file_info_iter->value.GetString()); + } + file_info_iter = file_info_value.FindMember("username"); + if (file_info_iter != file_info_value.MemberEnd() && file_info_iter->value.IsString()) { + dblink_info->mutable_file_info()->set_username(file_info_iter->value.GetString()); + } + file_info_iter = file_info_value.FindMember("password"); + if (file_info_iter != file_info_value.MemberEnd() && file_info_iter->value.IsString()) { + dblink_info->mutable_file_info()->set_password(file_info_iter->value.GetString()); + } + file_info_iter = file_info_value.FindMember("partition_fields"); + if (file_info_iter != file_info_value.MemberEnd() && file_info_iter->value.IsArray()) { + for (size_t i = 0; i < file_info_iter->value.Size(); ++i) { + if (file_info_iter->value[i].IsString()) { + dblink_info->mutable_file_info()->add_partition_fields(file_info_iter->value[i].GetString()); + } + } + } + file_info_iter = file_info_value.FindMember("format"); + if (file_info_iter != file_info_value.MemberEnd() && file_info_iter->value.IsString()) { + const std::string& file_format_str = file_info_iter->value.GetString(); + pb::FileFormat file_format; + if (!pb::FileFormat_Parse(file_format_str, &file_format)) { + DB_WARNING("Invalid file_format_str: %s", file_format_str.c_str()); + return -1; + } + dblink_info->mutable_file_info()->set_format(file_format); + } + file_info_iter = file_info_value.FindMember("delimiter"); + if (file_info_iter != file_info_value.MemberEnd() && file_info_iter->value.IsString()) { + dblink_info->mutable_file_info()->set_delimiter(file_info_iter->value.GetString()); + } + } DB_WARNING("dblink_info: %s", dblink_info->ShortDebugString().c_str()); if (check_dblink_table_valid(table) != 0) { DB_WARNING("invalid dblink table, dblink_info: %s", dblink_info->ShortDebugString().c_str()); @@ -1758,10 +1826,12 @@ int DDLPlanner::parse_alter_table(pb::MetaManagerRequest& alter_request) { return -1; } if (tbl_ptr->engine == pb::DBLINK) { - DB_WARNING("dblink table not support alter"); - _ctx->stat_info.error_code = ER_BAD_TABLE_ERROR; - _ctx->stat_info.error_msg << "dblink table not support alter"; - return -1; + if (check_alter_dblink_table_valid(*tbl_ptr, *stmt) != 0) { + DB_WARNING("dblink table not support alter"); + _ctx->stat_info.error_code = ER_BAD_TABLE_ERROR; + _ctx->stat_info.error_msg << "dblink table not support alter"; + return -1; + } } if (spec->spec_type == parser::ALTER_SPEC_TABLE_OPTION) { if (spec->table_options.size() > 1) { @@ -1880,8 +1950,18 @@ int DDLPlanner::parse_alter_table(pb::MetaManagerRequest& alter_request) { _ctx->stat_info.error_code = ER_ALTER_OPERATION_NOT_SUPPORTED; _ctx->stat_info.error_msg << "field_name is empty"; return -1; - } - field->set_field_name(spec->column_name.value); + } + std::string drop_field_name = spec->column_name.value; + if (tbl_ptr->get_ttl_field() != nullptr) { + std::vector splits; + boost::split(splits, drop_field_name, boost::is_any_of(".")); + if (splits.back() == tbl_ptr->get_ttl_field()->short_name) { + _ctx->stat_info.error_code = ER_ALTER_OPERATION_NOT_SUPPORTED; + _ctx->stat_info.error_msg << "forbid drop ttl field"; + return -1; + } + } + field->set_field_name(drop_field_name); } else if (spec->spec_type == parser::ALTER_SPEC_MODIFY_COLUMN && spec->new_columns.size() > 0) { alter_request.set_op_type(pb::OP_MODIFY_FIELD); int column_len = spec->new_columns.size(); @@ -1915,7 +1995,7 @@ int DDLPlanner::parse_alter_table(pb::MetaManagerRequest& alter_request) { } } if (table->indexs_size() != 0) { - _ctx->stat_info.error_code = ER_ALTER_OPERATION_NOT_SUPPORTED;; + _ctx->stat_info.error_code = ER_ALTER_OPERATION_NOT_SUPPORTED; _ctx->stat_info.error_msg << "modify table column with index is not supported"; return -1; } @@ -2356,10 +2436,10 @@ int DDLPlanner::add_constraint_def(pb::SchemaInfo& table, parser::Constraint* co DB_WARNING("fulltext index only support one field."); return -1; } - if (has_ttl) { - DB_WARNING("fulltext index can't support ttl."); - return -1; - } + // if (has_ttl) { + // DB_WARNING("fulltext index can't support ttl."); + // return -1; + // } break; case parser::CONSTRAINT_VECTOR: index_type = pb::I_VECTOR; @@ -2393,6 +2473,7 @@ int DDLPlanner::add_constraint_def(pb::SchemaInfo& table, parser::Constraint* co } } + std::unordered_set index_fields_set; for (int32_t column_index = 0; column_index < constraint->columns.size(); ++column_index) { std::string column_name = constraint->columns[column_index]->name.value; if (_column_can_null[column_name] && index->index_type() != pb::I_FULLTEXT) { @@ -2401,6 +2482,13 @@ int DDLPlanner::add_constraint_def(pb::SchemaInfo& table, parser::Constraint* co _ctx->stat_info.error_msg << "index column : " << column_name << " should NOT NULL"; return -1; } + if (index_fields_set.count(column_name) != 0) { + DB_WARNING("index column : %s duplicate", column_name.c_str()); + _ctx->stat_info.error_code = ER_DUP_FIELDNAME; + _ctx->stat_info.error_msg << "Duplicate column name '" << column_name << "'"; + return -1; + } + index_fields_set.emplace(column_name); index->add_field_names(column_name); } if (constraint->index_option != nullptr) { @@ -2430,6 +2518,35 @@ int DDLPlanner::add_constraint_def(pb::SchemaInfo& table, parser::Constraint* co StorageType_Parse(storage_type, &pb_storage_type); } index->set_storage_type(pb_storage_type); + + // vector index + auto iter = root.FindMember("vector_description"); + if (iter != root.MemberEnd()) { + index->set_vector_description(iter->value.GetString()); + } + iter = root.FindMember("dimension"); + if (iter != root.MemberEnd()) { + index->set_dimension(iter->value.GetInt()); + } + iter = root.FindMember("nprobe"); + if (iter != root.MemberEnd()) { + index->set_nprobe(iter->value.GetInt()); + } + iter = root.FindMember("efsearch"); + if (iter != root.MemberEnd()) { + index->set_efsearch(iter->value.GetInt()); + } + iter = root.FindMember("efconstruction"); + if (iter != root.MemberEnd()) { + index->set_efconstruction(iter->value.GetInt()); + } + iter = root.FindMember("metric_type"); + pb::MetricType metric_type = pb::METRIC_L2; + if (iter != root.MemberEnd()) { + std::string metric_type_str = iter->value.GetString(); + MetricType_Parse(metric_type_str, &metric_type); + index->set_metric_type(metric_type); + } } } catch (...) { DB_WARNING("parse create table json comments error [%s]", value); @@ -3045,7 +3162,6 @@ int DDLPlanner::add_default_partition_info(pb::SchemaInfo& table) { return 0; } -// 请求对应MetaServer检查外部映射表是否存在 int DDLPlanner::check_dblink_table_valid(const pb::SchemaInfo& table) { if (!table.has_dblink_info()) { DB_WARNING("table has no dblink_info, table: %s", table.ShortDebugString().c_str()); @@ -3059,10 +3175,28 @@ int DDLPlanner::check_dblink_table_valid(const pb::SchemaInfo& table) { return -1; } return 0; + } else if (dblink_info.type() == pb::LT_FILE) { + std::unordered_set field_names; + for (const auto& field : table.fields()) { + field_names.insert(field.field_name()); + } + for (const auto& partition_field_name : dblink_info.file_info().partition_fields()) { + if (field_names.find(partition_field_name) == field_names.end()) { + DB_WARNING("DBLink File partition field should be in table fields, table: %s", + table.ShortDebugString().c_str()); + return -1; + } + } + if (dblink_info.file_info().delimiter().size() > 1) { + DB_WARNING("DBLink File only support single char delimiter"); + return -1; + } + return 0; } else if (dblink_info.type() != pb::LT_BAIKALDB) { DB_WARNING("Invalid dblink type: %d", dblink_info.type()); return -1; } + // 请求对应MetaServer检查外部映射表是否存在 if (_ctx == nullptr) { DB_WARNING("_ctx is nullptr"); return -1; @@ -3110,4 +3244,37 @@ int DDLPlanner::check_dblink_table_valid(const pb::SchemaInfo& table) { return 0; } +int DDLPlanner::check_alter_dblink_table_valid(const TableInfo& table, const parser::AlterTableStmt& stmt) { + const pb::DBLinkInfo& dblink_info = table.dblink_info; + if (dblink_info.type() == pb::LT_FILE) { + // 只允许添加或删除字段，且删除字段不能为分区字段 + if (stmt.alter_specs.size() < 1) { + DB_WARNING("No alter_specs"); + return -1; + } + parser::AlterTableSpec* spec = stmt.alter_specs[0]; + if (spec == nullptr) { + DB_WARNING("spec is nullptr"); + return -1; + } + if (spec->spec_type != parser::ALTER_SPEC_ADD_COLUMN && + spec->spec_type != parser::ALTER_SPEC_DROP_COLUMN) { + DB_WARNING("Invalid spec_type: %d", spec->spec_type); + return -1; + } + if (spec->spec_type == parser::ALTER_SPEC_DROP_COLUMN) { + const std::string& drop_column_name = spec->column_name.c_str(); + for (const auto& partition_field_name : dblink_info.file_info().partition_fields()) { + if (drop_column_name == partition_field_name) { + DB_WARNING("Can not drop partition field, %s", partition_field_name.c_str()); + return -1; + } + } + } + } else { + return -1; + } + return 0; +} + } // end of namespace baikaldb diff --git a/src/logical_plan/ddl_work_planner.cpp b/src/logical_plan/ddl_work_planner.cpp index 267ca0b56..c23deb489 100755 --- a/src/logical_plan/ddl_work_planner.cpp +++ b/src/logical_plan/ddl_work_planner.cpp @@ -166,7 +166,6 @@ int DDLWorkPlanner::create_txn_dml_node(std::unique_ptr& t secondary_node_ptr->init(plan_node); index_ddl_manager_node->add_child(secondary_node_ptr); } - if (create_single_txn(std::move(index_ddl_manager_node), txn_node) != 0) { DB_WARNING("create signele txn error."); return -1; diff --git a/src/logical_plan/kill_planner.cpp b/src/logical_plan/kill_planner.cpp index 85b578721..0c39a2bd7 100755 --- a/src/logical_plan/kill_planner.cpp +++ b/src/logical_plan/kill_planner.cpp @@ -19,7 +19,7 @@ #include "state_machine.h" namespace baikaldb { - +DEFINE_bool(kill_select_out_txn_without_raft, false, "kill select out txn without raft"); int KillPlanner::plan() { create_packet_node(pb::OP_KILL); auto client = _ctx->client_conn; @@ -46,6 +46,12 @@ int KillPlanner::plan() { DB_WARNING("conn_id equal %ld is_query:%d", k->conn_id, k->is_query); _ctx->kill_ctx = sock->query_ctx; _ctx->kill_ctx->kill_all_ctx(); + if (FLAGS_kill_select_out_txn_without_raft + && sock->txn_id == 0 + && sock->query_ctx != nullptr + && (sock->query_ctx->stmt_type == parser::NT_SELECT || sock->query_ctx->stmt_type == parser::NT_UNION)) { + _ctx->kill_without_raft = true; + } if (!k->is_query) { sock->state = STATE_ERROR; //client_free 会core diff --git a/src/logical_plan/logical_planner.cpp b/src/logical_plan/logical_planner.cpp index d5ceb4526..27703bf4d 100644 --- a/src/logical_plan/logical_planner.cpp +++ b/src/logical_plan/logical_planner.cpp @@ -42,6 +42,7 @@ DECLARE_int32(bthread_concurrency); //bthread.cpp namespace baikaldb { DEFINE_bool(enable_plan_cache, false, "enable plan cache"); DEFINE_bool(enable_convert_charset, false, "enable convert charset"); +DEFINE_bool(sign_with_resouce_tag, true, "sign without resource tag"); DECLARE_string(log_plat_name); DECLARE_bool(enable_dblink); @@ -524,10 +525,16 @@ int LogicalPlanner::generate_sql_sign(QueryContext* ctx, parser::StmtNode* stmt) std::string str; if (stat_info->sign == 0) { stmt->set_print_sample(true); - stat_info->sample_sql << "family_table_tag_optype_plat=[" << stat_info->family << "\t" - << stat_info->table << "\t" << stat_info->resource_tag << "\t" << op_type << "\t" - << FLAGS_log_plat_name << "] sql=[" << stmt << "]"; uint64_t out[2]; + if (FLAGS_sign_with_resouce_tag) { + stat_info->sample_sql << "family_table_tag_optype_plat=[" << stat_info->family << "\t" + << stat_info->table << "\t" << stat_info->resource_tag << "\t" << op_type << "\t" + << FLAGS_log_plat_name << "] sql=[" << stmt << "]"; + } else { + stat_info->sample_sql << "family_table_tag_optype_plat=[" << stat_info->family << "\t" + << stat_info->table << "\t" << "" << "\t" << op_type << "\t" + << FLAGS_log_plat_name << "] sql=[" << stmt << "]"; + } str = stat_info->sample_sql.str(); butil::MurmurHash3_x64_128(str.c_str(), str.size(), 0x1234, out); stat_info->sign = out[0]; @@ -612,6 +619,7 @@ int LogicalPlanner::gen_subquery_plan(parser::DmlNode* subquery, SmartPlanTableC _cur_sub_ctx->table_with_clause_mapping = _ctx->table_with_clause_mapping; _cur_sub_ctx->is_union_subquery = expr_params.is_union_subquery; _cur_sub_ctx->efsearch = _ctx->efsearch; + _cur_sub_ctx->nprobe = _ctx->nprobe; // from子查询完全ctx完全独立 if (expr_params.is_from_subquery || expr_params.is_union_subquery) { plan_state.reset(new (std::nothrow)PlanTableContext); @@ -831,6 +839,8 @@ int LogicalPlanner::add_table(const std::string& database, const std::string& ta _plan_table_ctx->dblink_table_mapping[tableid] = orig_tbl_ptr; } else if (tbl_ptr->dblink_info.type() == pb::LT_MYSQL) { _ctx->has_dblink_mysql = true; + } else if (tbl_ptr->dblink_info.type() == pb::LT_FILE) { + _ctx->has_dblink_file = true; } else { DB_WARNING("unknown dblink type: %d", tbl_ptr->dblink_info.type()); return -1; @@ -2955,9 +2965,12 @@ int LogicalPlanner::create_common_subquery_expr(const parser::SubqueryExpr* item pb::ExprNode* node = expr.add_nodes(); construct_literal_expr(row, node); } - } else { + } else if (rows.size() == 1) { pb::ExprNode* node = expr.add_nodes(); construct_literal_expr(rows[0], node); + } else { + DB_WARNING("not data row_filed_number"); + return -1; } } } else { @@ -4432,8 +4445,8 @@ int LogicalPlanner::can_use_dblink(SmartTable table) { _ctx->stat_info.error_msg << "dblink table only supprt INSERT/DELETE/UPDATE/SELECT/UNION"; return -1; } - } else if (table->dblink_info.type() == pb::LT_MYSQL){ - // dblink MySQL表只支持SELECT/UNION + } else if (table->dblink_info.type() == pb::LT_MYSQL || table->dblink_info.type() == pb::LT_FILE){ + // dblink MySQL/File表只支持SELECT/UNION if (_ctx->stmt_type != parser::NT_SELECT && _ctx->stmt_type != parser::NT_UNION) { DB_WARNING("dblink table not support stmt type: %d", _ctx->stmt_type); diff --git a/src/logical_plan/select_planner.cpp b/src/logical_plan/select_planner.cpp index b53d6dc0c..5869c21de 100644 --- a/src/logical_plan/select_planner.cpp +++ b/src/logical_plan/select_planner.cpp @@ -183,7 +183,7 @@ bool SelectPlanner::is_full_export() { return false; } // 包含DBLink Mysql的表不支持全量导出 - if (_ctx->has_dblink_mysql) { + if (_ctx->has_dblink_mysql || _ctx->has_dblink_file) { return false; } //代价信息统计时不走full export流程 diff --git a/src/logical_plan/update_planner.cpp b/src/logical_plan/update_planner.cpp index c27c7f2a3..622fcdc13 100644 --- a/src/logical_plan/update_planner.cpp +++ b/src/logical_plan/update_planner.cpp @@ -270,7 +270,7 @@ int UpdatePlanner::parse_kv_list() { if (update_field_ids.count(field.id) != 0) { continue; } - if (field.on_update_value == "(current_timestamp())") { + if (field.on_update_value == "(current_timestamp())" && !_ctx->disable_on_update) { pb::Expr value_expr; auto node = value_expr.add_nodes(); node->set_num_children(0); diff --git a/src/meta_server/common_state_machine.cpp b/src/meta_server/common_state_machine.cpp index 60fa309e8..6fd030b74 100644 --- a/src/meta_server/common_state_machine.cpp +++ b/src/meta_server/common_state_machine.cpp @@ -198,8 +198,9 @@ void CommonStateMachine::start_check_bns() { } } void CommonStateMachine::on_leader_start() { - start_check_bns(); _is_leader.store(true); + start_check_bns(); + } void CommonStateMachine::on_leader_start(int64_t term) { @@ -241,14 +242,15 @@ void CommonStateMachine::start_check_migrate() { while (_node.is_leader()) { int time = 0; while (time < sleep_time_count) { - if (!_node.is_leader()) { + if (!_node.is_leader() || !_is_leader.load()) { return; } bthread_usleep(1000 * 1000LL); ++time; } - SELF_TRACE("start check migrate, count: %ld", count); - ++count; + if (++count % 10 == 0) { + DB_WARNING("start check migrate, region_id: %ld, count: %ld", _dummy_region_id, count); + } check_migrate(); } } diff --git a/src/meta_server/ddl_manager.cpp b/src/meta_server/ddl_manager.cpp index 9d3dab665..e2cbf94d1 100755 --- a/src/meta_server/ddl_manager.cpp +++ b/src/meta_server/ddl_manager.cpp @@ -31,6 +31,7 @@ DEFINE_int32(max_ddl_retry_time, 30, "max ddl retry time"); DECLARE_int32(baikal_heartbeat_interval_us); DEFINE_bool(all_rollup_region_need_execute, false, "all rollup region need execute"); DEFINE_int32(single_store_max_ddlwork_num, 3, "store max ddlwork num"); +DEFINE_int32(single_store_max_delete_ddl_num, 1, "store max delete ddlwork num"); DEFINE_bool(cold_data_rollup_done, true, "cold data rollup done"); std::string construct_ddl_work_key(const std::string& identify, const std::initializer_list& ids) { @@ -214,7 +215,7 @@ void DBManager::process_baikal_heartbeat(const pb::BaikalHeartBeatRequest* reque DB_DEBUG("dll_response : %s address %s", response->ShortDebugString().c_str(), address.c_str()); } -bool DBManager::round_robin_select(std::string* selected_address, bool is_column_ddl) { +bool DBManager::round_robin_select(std::string* selected_address, bool is_column_update_ddl) { BAIDU_SCOPED_LOCK(_address_instance_mutex); auto iter = _address_instance_map.find(_last_rolling_instance); if (iter == _address_instance_map.end() || (++iter) == _address_instance_map.end()) { @@ -234,7 +235,7 @@ bool DBManager::round_robin_select(std::string* selected_address, bool is_column auto find_task_map = _common_task_map.init_if_not_exist_else_update(iter->first, false, [¤t_task_number](CommonTaskMap& db_task_map){ current_task_number = db_task_map.doing_task_map.size() + db_task_map.to_do_task_map.size(); }); - int32_t max_concurrent = is_column_ddl ? FLAGS_baikaldb_max_concurrent * 5 : FLAGS_baikaldb_max_concurrent; + int32_t max_concurrent = is_column_update_ddl ? FLAGS_baikaldb_max_concurrent * 5 : FLAGS_baikaldb_max_concurrent; if (!find_task_map || current_task_number < max_concurrent) { _last_rolling_instance = iter->first; *selected_address = iter->first; @@ -246,8 +247,8 @@ bool DBManager::round_robin_select(std::string* selected_address, bool is_column return false; } -bool DBManager::select_instance(std::string* selected_address, bool is_column_ddl) { - return round_robin_select(selected_address, is_column_ddl); +bool DBManager::select_instance(std::string* selected_address, bool is_column_update_ddl) { + return round_robin_select(selected_address, is_column_update_ddl); } int DBManager::execute_task(MemRegionDdlWork& work) { @@ -269,7 +270,12 @@ int DBManager::execute_task(MemRegionDdlWork& work) { ++iter2; } }); - int32_t max_concurrent = work.region_info.op_type() == pb::OP_MODIFY_FIELD ? + + bool is_column_update = work.region_info.op_type() == pb::OP_MODIFY_FIELD + && !work.region_info.column_ddl_info().update_slots().empty(); + bool is_column_delete = work.region_info.op_type() == pb::OP_MODIFY_FIELD + && work.region_info.column_ddl_info().update_slots().empty(); + int32_t max_concurrent = is_column_update ? FLAGS_single_table_ddl_max_concurrent * 10 : FLAGS_single_table_ddl_max_concurrent; if (all_task_count_by_table_id > max_concurrent) { DB_NOTICE("table %s ddl task count %d reach max concurrency %d", table_id_prefix.c_str(), @@ -281,13 +287,15 @@ int DBManager::execute_task(MemRegionDdlWork& work) { auto& region_ddl_info = work.region_info; work.update_timestamp = butil::gettimeofday_us(); std::string address; - if (select_instance(&address, work.region_info.op_type() == pb::OP_MODIFY_FIELD)) { + if (select_instance(&address, is_column_update)) { auto task_id = std::to_string(region_ddl_info.table_id()) + "_" + std::to_string(region_ddl_info.region_id()); // 一个store同一时间只能执行single_store_max_ddlwork_num个ddl任务 SmartRegionInfo region_info = RegionManager::get_instance()->get_region_info(region_ddl_info.region_id()); - if (region_info != nullptr && work.region_info.op_type() != pb::OP_MODIFY_FIELD) { + if (region_info != nullptr && !is_column_update) { BAIDU_SCOPED_LOCK(_task_store_mutex); - if (store_ddlwork_cnt_map[region_info->leader()] >= FLAGS_single_store_max_ddlwork_num) { + int single_store_max_ddlwork_num = + is_column_delete ? FLAGS_single_store_max_delete_ddl_num : FLAGS_single_store_max_ddlwork_num; + if (store_ddlwork_cnt_map[region_info->leader()] >= single_store_max_ddlwork_num) { DB_NOTICE("store_check address_%s is doing %d ddl_work", region_info->leader().c_str(), store_ddlwork_cnt_map[region_info->leader()]); return -2; diff --git a/src/meta_server/meta_server.cpp b/src/meta_server/meta_server.cpp index 06fae939d..a7392d1cc 100644 --- a/src/meta_server/meta_server.cpp +++ b/src/meta_server/meta_server.cpp @@ -188,6 +188,7 @@ void MetaServer::distribute_leader_thread() { int META_TYPE = 0; int AUTO_INCR_TYPE = 1; int TSO_TYPE = 2; + TimeCost time_cost; while (!_shutdown) { int sleep_s = 60; if (FLAGS_distribute_leader_thread_s > 0) { @@ -205,6 +206,9 @@ void MetaServer::distribute_leader_thread() { if (FLAGS_distribute_leader_thread_s <= 0) { continue; } + if (time_cost.get_time() < 2 * 3600ULL * 1000 * 1000) { + continue; + } if (!_meta_state_machine->is_leader()) { continue; } diff --git a/src/meta_server/meta_state_machine.cpp b/src/meta_server/meta_state_machine.cpp index 5c6b20869..22f57701d 100644 --- a/src/meta_server/meta_state_machine.cpp +++ b/src/meta_server/meta_state_machine.cpp @@ -399,7 +399,7 @@ void MetaStateMachine::on_apply(braft::Iterator& iter) { break; } case pb::OP_UPDATE_TTL_DURATION: { - TableManager::get_instance()->update_ttl_duration(request, iter.index(), done); + TableManager::get_instance()->update_ttl_info(request, iter.index(), done); break; } case pb::OP_UPDATE_BYTE_SIZE: { @@ -459,7 +459,7 @@ void MetaStateMachine::on_apply(braft::Iterator& iter) { break; } case pb::OP_SPLIT_REGION: { - RegionManager::get_instance()->split_region(request, done); + RegionManager::get_instance()->split_region(request, iter.index(), done); break; } case pb::OP_MODIFY_RESOURCE_TAG: { diff --git a/src/meta_server/region_manager.cpp b/src/meta_server/region_manager.cpp index 4a9096976..846568454 100644 --- a/src/meta_server/region_manager.cpp +++ b/src/meta_server/region_manager.cpp @@ -86,10 +86,11 @@ void RegionManager::update_region(const pb::MetaManagerRequest& request, } if (region_info.version() == master_region_ptr->version() && region_info.conf_version() <= master_region_ptr->conf_version()) { - if (peer_is_equal(region_info, *master_region_ptr)) { + if (peer_is_equal(region_info, *master_region_ptr, true)) { DB_WARNING("region_id: %ld, request conf_version %ld <= master conf_version %ld", region_id, region_info.conf_version(), master_region_ptr->conf_version()); IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "equal conf version"); + return; } } region_info.set_conf_version(master_region_ptr->conf_version() + 1); @@ -297,7 +298,7 @@ void RegionManager::drop_region(const pb::MetaManagerRequest& request, DB_NOTICE("drop region success, request:%s", request.ShortDebugString().c_str()); } -void RegionManager::split_region(const pb::MetaManagerRequest& request, braft::Closure* done) { +void RegionManager::split_region(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done) { auto& region_split_info = request.region_split(); int64_t region_id = region_split_info.region_id(); int new_region_num = region_split_info.new_region_num(); @@ -320,7 +321,7 @@ void RegionManager::split_region(const pb::MetaManagerRequest& request, braft::C IF_DONE_SET_RESPONSE(done, pb::INTERNAL_ERROR, "write db fail"); return; } - DB_WARNING("generate %d region_id: [%ld, %ld]", new_region_num, new_region_start_id, new_region_end_id); + DB_WARNING("apply_index: %ld, generate %d region_id: [%ld, %ld]", apply_index, new_region_num, new_region_start_id, new_region_end_id); //更新内存 set_max_region_id(new_region_end_id); if (done && ((MetaServerClosure*)done)->response) { @@ -1947,7 +1948,7 @@ void RegionManager::leader_heartbeat_for_region(const pb::StoreHeartBeatRequest* } bool peer_changed = false; if (!leader_region.simple()) { - peer_changed = !peer_is_equal(leader_region_info, *master_region_info); + peer_changed = !peer_is_equal(leader_region_info, *master_region_info, false); } check_whether_update_region(region_id, peer_changed, instance, leader_region, master_region_info, response); if (!peer_changed) { diff --git a/src/meta_server/schema_manager.cpp b/src/meta_server/schema_manager.cpp index b388be736..485994544 100644 --- a/src/meta_server/schema_manager.cpp +++ b/src/meta_server/schema_manager.cpp @@ -206,10 +206,15 @@ void SchemaManager::process_schema_info(google::protobuf::RpcController* control auto mutable_request = const_cast(request); mutable_request->mutable_ddlwork_info()->set_begin_timestamp(butil::gettimeofday_s()); } + if (request->op_type() == pb::OP_ADD_INDEX) { + auto ret = TableManager::get_instance()->pre_process_for_add_index(request, response, log_id, done); + if (ret < 0) { + return; + } + } } if (request->op_type() == pb::OP_UPDATE_TTL_DURATION && !request->table_info().has_ttl_duration()) { - // 只能修改有ttl的表 ERROR_SET_RESPONSE(response, pb::INPUT_PARAM_ERROR, "ttl_duration must > 0", request->op_type(), log_id); return; @@ -221,6 +226,12 @@ void SchemaManager::process_schema_info(google::protobuf::RpcController* control "no partition_info", request->op_type(), log_id); return; } + if (request->op_type() == pb::OP_ADD_PARTITION) { + auto ret = TableManager::get_instance()->pre_process_for_add_partition(request, response, log_id, done); + if (ret < 0) { + return; + } + } if (request->op_type() == pb::OP_UPDATE_CHARSET && !request->table_info().has_charset()) { ERROR_SET_RESPONSE(response, pb::INPUT_PARAM_ERROR, @@ -912,6 +923,92 @@ int SchemaManager::pre_process_for_create_table(const pb::MetaManagerRequest* re return -1; } + // 对主表和全局索引在状态机外选实例 + std::unordered_set need_pick_instance_index; + for (const auto& index : table_info.indexs()) { + if (index.index_type() == pb::I_PRIMARY || index.is_global()) { + need_pick_instance_index.insert(index.index_name()); + } + } + if (!table_info.has_partition_num()) { + table_info.set_partition_num(1); + } + // 先处理split key + std::unordered_set processed_index_name; + for (auto idx = 0; + idx < table_info.partition_num() && TableManager::get_instance()->is_create_table_support_engine(table_info.engine()); + ++idx) { + int64_t partition_id = idx; + std::string partition_resource_tag = resource_tag; + std::string partition_main_logical_room = main_logical_room; + if (table_info.has_partition_info() && table_info.partition_info().type() == pb::PT_RANGE) { + if (table_info.partition_info().range_partition_infos(idx).has_resource_tag()) { + partition_resource_tag = table_info.partition_info().range_partition_infos(idx).resource_tag(); + partition_main_logical_room = ""; + } + if (table_info.partition_info().range_partition_infos(idx).has_partition_id()) { + partition_id = table_info.partition_info().range_partition_infos(idx).partition_id(); + } + } + auto partition_index_instances = mutable_request->add_partition_index_instance(); + partition_index_instances->set_partition_id(partition_id); + for (auto i = 0; i < table_info.split_keys_size(); ++i) { + auto split_key = table_info.mutable_split_keys(i); + std::string index_name = split_key->index_name(); + auto index_instances = partition_index_instances->add_index_instances(); + index_instances->set_index_name(index_name); + for (auto j = 0; j <= split_key->split_keys_size(); ++j) { + std::string instance; + int ret = ClusterManager::get_instance()->select_instance_rolling( + {partition_resource_tag, partition_main_logical_room, ""}, {}, instance); + if (ret < 0) { + DB_WARNING("select instance fail"); + ERROR_SET_RESPONSE(response, pb::INTERNAL_ERROR, "select instance fail", request->op_type(), log_id); + return -1; + } + index_instances->add_instances(instance); + } + processed_index_name.insert(index_name); + } + } + for (const auto& index_name : processed_index_name) { + need_pick_instance_index.erase(index_name); + } + //没有指定split_key的索引 + for (auto i = 0; + i < table_info.partition_num() && TableManager::get_instance()->is_create_table_support_engine(table_info.engine()); + ++i) { + int64_t partition_id = i; + std::string partition_resource_tag = resource_tag; + std::string partition_main_logical_room = main_logical_room; + if (table_info.has_partition_info() && table_info.partition_info().type() == pb::PT_RANGE) { + if (table_info.partition_info().range_partition_infos(i).has_resource_tag()) { + partition_resource_tag = table_info.partition_info().range_partition_infos(i).resource_tag(); + partition_main_logical_room = ""; + } + if (table_info.partition_info().range_partition_infos(i).has_partition_id()) { + partition_id = table_info.partition_info().range_partition_infos(i).partition_id(); + } + } + auto partition_index_instances = mutable_request->mutable_partition_index_instance(i); + for (auto& index : need_pick_instance_index) { + std::string instance; + if (namespace_name == "INTERNAL" && table_info.database() == "baikaldb" && table_info.table_name() == "__baikaldb_instance") { + // for baikaldb instance id, do nothing + } else { + int ret = ClusterManager::get_instance()->select_instance_rolling( + {partition_resource_tag, partition_main_logical_room, ""}, {}, instance); + if (ret < 0) { + DB_WARNING("select instance fail"); + ERROR_SET_RESPONSE(response, pb::INTERNAL_ERROR, "select instance fail", request->op_type(), log_id); + return -1; + } + } + auto index_instances = partition_index_instances->add_index_instances(); + index_instances->set_index_name(index); + index_instances->add_instances(instance); + } + } return 0; } diff --git a/src/meta_server/table_manager.cpp b/src/meta_server/table_manager.cpp index 80c2c6572..ea23b0ac2 100644 --- a/src/meta_server/table_manager.cpp +++ b/src/meta_server/table_manager.cpp @@ -29,7 +29,6 @@ namespace baikaldb { DECLARE_int32(concurrency_num); DEFINE_int32(region_replica_num, 3, "region replica num, default:3"); -DEFINE_int32(learner_region_replica_num, 1, "learner region replica num, default:1"); DEFINE_int32(region_region_size, 100 * 1024 * 1024, "region size, default:100M"); DEFINE_int64(table_tombstone_gc_time_s, 3600 * 24 * 5, "time interval to clear table_tombstone. default(5d)"); DEFINE_uint64(statistics_heart_beat_bytesize, 256 * 1024 * 1024, "default(256M)"); @@ -39,6 +38,20 @@ DEFINE_bool(use_partition_split_key, false, "add dynamic partition use get_parti DEFINE_int32(dynamic_partition_change_time_s, 3000, "dynamic partition change time, default: 3000s"); DEFINE_int32(dynamic_partition_change_cnt, 10, "dynamic partition change count, default: 10"); +const std::set TableManager::_need_check_wordrank_types = { +#ifdef BAIDU_INTERNAL + pb::S_DEFAULT, +#endif + pb::S_WORDRANK, + pb::S_WORDSEG_BASIC, + pb::S_WORDRANK_Q2B_ICASE, + pb::S_WORDRANK_Q2B_ICASE_UNLIMIT +}; + +const std::set TableManager::_need_check_wordweight_types = { + pb::S_WORDWEIGHT, pb::S_WORDWEIGHT_NO_FILTER, pb::S_WORDWEIGHT_NO_FILTER_SAME_WEIGHT +}; + void TableTimer::run() { DB_NOTICE("Table Timer run."); std::vector delete_schemas; @@ -526,6 +539,7 @@ void TableManager::create_table(const pb::MetaManagerRequest& request, const int std::string namespace_name = table_info.namespace_name(); std::string database_name = namespace_name + "\001" + table_info.database(); std::string table_name = database_name + "\001" + table_info.table_name(); + std::string resource_tag = table_info.resource_tag(); TableMem table_mem; //校验合法性, 准备数据 @@ -579,10 +593,21 @@ void TableManager::create_table(const pb::MetaManagerRequest& request, const int IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "field not illegal"); return; } - ret = alloc_index_id(table_info, table_mem, max_table_id_tmp); + ret = alloc_index_id(table_info, table_mem, max_table_id_tmp, resource_tag); if (ret < 0) { DB_WARNING("table:%s 's index info not illegal", table_name.c_str()); - IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "index not illegal"); + std::string err_msg = "index not legal"; + if (ret == -2) { + err_msg += ", cluster doesn't support workrank/workweight"; + } + IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, err_msg); + return; + } + std::string err_msg; + ret = check_ttl_info(table_info, err_msg); + if (ret != 0) { + DB_WARNING("table:%s 's ttl info not illegal", table_name.c_str()); + IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, err_msg); return; } if (table_info.engine() == pb::BINLOG) { @@ -669,7 +694,7 @@ void TableManager::create_table(const pb::MetaManagerRequest& request, const int } DB_WARNING("paritition %s.", table_info.partition_info().ShortDebugString().c_str()); table_mem.schema_pb = table_info; - ret = write_schema_for_not_level(table_mem, done, max_table_id_tmp, has_auto_increment); + ret = write_schema_for_not_level(request, table_mem, done, max_table_id_tmp, has_auto_increment, apply_index); if (ret != 0) { DB_WARNING("write rocksdb fail when create table, table:%s", table_name.c_str()); IF_DONE_SET_RESPONSE(done, pb::INTERNAL_ERROR, "write db fail"); @@ -799,7 +824,7 @@ void TableManager::drop_table(const pb::MetaManagerRequest& request, const int64 } if (check_table_is_linked(drop_table_id)) { DB_WARNING("table is linked, request:%s", request.ShortDebugString().c_str()); - IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "table is linked binlog table"); + IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "table is linked binlog table, watt data stream may need to be restarted after operation"); return; } std::vector delete_rocksdb_keys; @@ -1057,11 +1082,16 @@ void TableManager::rename_table(const pb::MetaManagerRequest& request, return; } - if (check_table_has_ddlwork(table_id) || check_table_is_linked(table_id)) { + if (check_table_has_ddlwork(table_id)) { DB_WARNING("table is doing ddl, request:%s", request.ShortDebugString().c_str()); IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "table is doing ddl"); return; } + if (check_table_is_linked(table_id)) { + DB_WARNING("table is linked binlog table, request:%s", request.ShortDebugString().c_str()); + IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "table is linked binlog table, watt data stream may need to be restarted after operation"); + return; + } std::string namespace_name = request.table_info().namespace_name(); std::string database_name = namespace_name + "\001" + request.table_info().database(); std::string old_table_name = database_name + "\001" + request.table_info().table_name(); @@ -1116,12 +1146,17 @@ void TableManager::swap_table(const pb::MetaManagerRequest& request, return; } - if (check_table_has_ddlwork(table_id) || check_table_is_linked(table_id) || - check_table_has_ddlwork(new_table_id) || check_table_is_linked(new_table_id)) { + if (check_table_has_ddlwork(table_id) || check_table_has_ddlwork(new_table_id)) { DB_WARNING("table is doing ddl, request:%s", request.ShortDebugString().c_str()); IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "table is doing ddl"); return; } + + if (check_table_is_linked(table_id) || check_table_is_linked(new_table_id)) { + DB_WARNING("table is linked binlog table, request:%s", request.ShortDebugString().c_str()); + IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "table is linked binlog table, watt data stream may need to be restarted after operation"); + return; + } pb::SchemaInfo mem_schema_pb; get_table_info(table_id, mem_schema_pb); pb::SchemaInfo new_mem_schema_pb; @@ -1239,6 +1274,120 @@ void TableManager::update_split_lines(const pb::MetaManagerRequest& request, }); } +// 状态机外选store实例 +int TableManager::pre_process_for_add_partition(const pb::MetaManagerRequest* request, pb::MetaManagerResponse* response, uint64_t log_id, google::protobuf::Closure* done) { + if (!request->has_table_info() || + !request->table_info().has_partition_info()) { + DB_WARNING("request has no valid partition info, request: %s", request->ShortDebugString().c_str()); + ERROR_SET_RESPONSE(response, pb::INPUT_PARAM_ERROR, + "request has no valid partition info", request->op_type(), log_id); + return -1; + } + + int64_t table_id = -1; + if (check_table_exist(request->table_info(), table_id) != 0) { + DB_WARNING("check table exist fail, request:%s", request->ShortDebugString().c_str()); + ERROR_SET_RESPONSE(response, pb::INPUT_PARAM_ERROR, + "table not exist", request->op_type(), log_id); + return -1; + } + pb::SchemaInfo mem_schema_pb; + if (get_table_info(table_id, mem_schema_pb) != 0) { + DB_WARNING("Fail to get_table_info, request:%s", request->ShortDebugString().c_str()); + ERROR_SET_RESPONSE(response, pb::INPUT_PARAM_ERROR, + "Fail to get_table_info", request->op_type(), log_id); + return -1; + } + if (!mem_schema_pb.has_partition_info()) { + DB_WARNING("mem_schema_pb has no partition_info, request: %s", request->ShortDebugString().c_str()); + ERROR_SET_RESPONSE(response, pb::INPUT_PARAM_ERROR, + "table has no partition_info", request->op_type(), log_id); + return -1; + } + if (mem_schema_pb.partition_info().type() != pb::PT_RANGE) { + DB_WARNING("not range partition table, partition type: %d", mem_schema_pb.partition_info().type()); + ERROR_SET_RESPONSE(response, pb::INPUT_PARAM_ERROR, + "not range partition table", request->op_type(), log_id); + return -1; + } + if (mem_schema_pb.mutable_partition_info() == nullptr) { + DB_WARNING("partition_info is nullptr, request: %s", request->ShortDebugString().c_str()); + ERROR_SET_RESPONSE(response, pb::INTERNAL_ERROR, + "table partition_info is nullptr", request->op_type(), log_id); + return -1; + } + std::unordered_map global_indexes_template; + for (const auto& index : mem_schema_pb.indexs()) { + if (index.index_type() == pb::I_PRIMARY || index.is_global()) { + global_indexes_template[index.index_name()] = index.index_id(); + } + } + auto mutable_request = const_cast(request); + for (auto& range_partition_info : request->table_info().partition_info().range_partition_infos()) { + // 获取分区预分裂值 + ::google::protobuf::RepeatedPtrField split_keys; + split_keys.CopyFrom(range_partition_info.split_keys()); + if (range_partition_info.is_pre_split() && split_keys.empty()) { + split_keys.CopyFrom(mem_schema_pb.split_keys()); + } + + // 如果partition设置resource_tag，则使用partition的resource_tag + std::string partition_resource_tag = mem_schema_pb.resource_tag(); + std::string partition_main_logical_room = mem_schema_pb.main_logical_room(); + if (range_partition_info.has_resource_tag()) { + partition_resource_tag = range_partition_info.resource_tag(); + partition_main_logical_room = ""; + } + boost::trim(partition_resource_tag); + + std::unordered_map global_indexes = global_indexes_template; + std::unordered_set processed_index_name; + auto partition_index_instance = mutable_request->add_partition_index_instance(); + // 有split_key的索引先处理 + for (const auto& split_key : split_keys) { + const std::string& index_name = split_key.index_name(); + if (global_indexes.find(index_name) == global_indexes.end()) { + continue; + } + auto index_instances = partition_index_instance->add_index_instances(); + index_instances->set_index_name(index_name); + for (size_t i = 0; i <= split_key.split_keys_size(); ++i) { + std::string instance; + int ret = ClusterManager::get_instance()->select_instance_rolling( + {partition_resource_tag, partition_main_logical_room, ""}, {}, instance); + if (ret < 0) { + DB_WARNING("select instance fail, request: %s", request->ShortDebugString().c_str()); + ERROR_SET_RESPONSE(response, pb::INTERNAL_ERROR, + "select instance failr", request->op_type(), log_id); + return -1; + } + index_instances->add_instances(instance); + } + processed_index_name.insert(index_name); + } + for (const auto& index_name : processed_index_name) { + global_indexes.erase(index_name); + } + for (const auto& index : global_indexes) { + std::string index_name = index.first; + std::string instance; + int ret = ClusterManager::get_instance()->select_instance_rolling( + {partition_resource_tag, partition_main_logical_room, ""}, {}, instance); + if (ret < 0) { + DB_WARNING("select instance failed, request: %s", request->ShortDebugString().c_str()); + ERROR_SET_RESPONSE(response, pb::INTERNAL_ERROR, + "select instance failr", request->op_type(), log_id); + return -1; + } + auto index_instances = partition_index_instance->add_index_instances(); + index_instances->set_index_name(index_name); + index_instances->add_instances(instance); + } + } + return 0; +} + +// 状态机内 void TableManager::add_partition(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done) { @@ -1282,6 +1431,7 @@ void TableManager::add_partition(const pb::MetaManagerRequest& request, const pb::PrimitiveType partition_col_type = mem_schema_pb.partition_info().field_info().mysql_type(); int64_t partition_id = mem_schema_pb.partition_info().max_range_partition_id(); int64_t tmp_max_region_id = RegionManager::get_instance()->get_max_region_id(); + int64_t start_region_id = tmp_max_region_id + 1; // 创建region std::shared_ptr> init_regions(new std::vector{}); @@ -1308,8 +1458,19 @@ void TableManager::add_partition(const pb::MetaManagerRequest& request, pb::PartitionInfo* p_read_add_partition_info = real_add_schema_pb.mutable_partition_info(); p_read_add_partition_info->clear_range_partition_infos(); + if (request.table_info().partition_info().range_partition_infos_size() + != request.partition_index_instance_size()) { + DB_FATAL("request range_partition_infos_size: %d withpartition_index_instance_size: %d not match", + request.table_info().partition_info().range_partition_infos_size(), request.partition_index_instance_size()); + IF_DONE_SET_RESPONSE(done, pb::INTERNAL_ERROR, "partition_index_instance_size not match"); + return; + } + + int idx = 0; std::unordered_set partition_ids; for (auto& range_partition_info : request.table_info().partition_info().range_partition_infos()) { + auto& partition_index_instance = request.partition_index_instance(idx); + idx++; const std::string& partition_name = range_partition_info.partition_name(); if (partition_names.find(partition_name) != partition_names.end()) { if (request.has_is_dynamic_change() && request.is_dynamic_change()) { @@ -1403,15 +1564,25 @@ void TableManager::add_partition(const pb::MetaManagerRequest& request, if (global_indexes.find(index_name) == global_indexes.end()) { continue; } - for (size_t i = 0; i <= split_key.split_keys_size(); ++i) { - std::string instance; - int ret = ClusterManager::get_instance()->select_instance_rolling( - {partition_resource_tag, partition_main_logical_room, ""}, {}, instance); - if (ret < 0) { - DB_WARNING("select instance fail, request: %s", request.ShortDebugString().c_str()); - IF_DONE_SET_RESPONSE(done, pb::INTERNAL_ERROR, "select instance fail"); - return; + pb::IndexInstance* match_index_instance = nullptr; + for (auto& index_instances : partition_index_instance.index_instances()) { + if (index_name == index_instances.index_name()) { + match_index_instance = const_cast(&index_instances); } + } + if (match_index_instance == nullptr) { + DB_FATAL("index_instance is null, index_name: %s", index_name.c_str()); + IF_DONE_SET_RESPONSE(done, pb::INTERNAL_ERROR, "match_index_instance is null"); + return; + } + if (match_index_instance->instances().size() != split_key.split_keys_size() + 1) { + DB_FATAL("picked_instance size(%d) not equal to split_keys_size(%d)", + match_index_instance->instances().size(), split_key.split_keys_size()); + IF_DONE_SET_RESPONSE(done, pb::INTERNAL_ERROR, "match_index_instance size not match"); + return; + } + for (size_t i = 0; i <= split_key.split_keys_size(); ++i) { + std::string instance = match_index_instance->instances().at(i); pb::InitRegion init_region_request; pb::RegionInfo* region_info = init_region_request.mutable_region_info(); if (region_info == nullptr) { @@ -1444,14 +1615,25 @@ void TableManager::add_partition(const pb::MetaManagerRequest& request, global_indexes.erase(index_name); } for (const auto& index : global_indexes) { - std::string instance; - int ret = ClusterManager::get_instance()->select_instance_rolling( - {partition_resource_tag, partition_main_logical_room, ""}, {}, instance); - if (ret < 0) { - DB_WARNING("select instance failed, request: %s", request.ShortDebugString().c_str()); - IF_DONE_SET_RESPONSE(done, pb::INTERNAL_ERROR, "select instance failed"); + std::string index_name = index.first; + pb::IndexInstance* match_index_instance = nullptr; + for (auto& index_instances : partition_index_instance.index_instances()) { + if (index_name == index_instances.index_name()) { + match_index_instance = const_cast(&index_instances); + } + } + if (match_index_instance == nullptr) { + DB_FATAL("index_instance is null, index_name: %s", index_name.c_str()); + IF_DONE_SET_RESPONSE(done, pb::INTERNAL_ERROR, "match_index_instance is null"); + return; + } + if (match_index_instance->instances().size() != 1) { + DB_FATAL("picked_instance size(%d) not equal to split_keys_size(1)", + match_index_instance->instances().size()); + IF_DONE_SET_RESPONSE(done, pb::INTERNAL_ERROR, "match_index_instance size not match"); return; } + std::string instance = match_index_instance->instances().at(0); pb::InitRegion init_region_request; pb::RegionInfo* region_info = init_region_request.mutable_region_info(); if (region_info == nullptr) { @@ -1518,7 +1700,8 @@ void TableManager::add_partition(const pb::MetaManagerRequest& request, IF_DONE_SET_RESPONSE(done, pb::INTERNAL_ERROR, "write db fail"); return; } - + DB_WARNING("apply_index: %ld, generate %ld region_id: [%ld, %ld]", + apply_index, tmp_max_region_id - start_region_id + 1, start_region_id, tmp_max_region_id); // 更新内存 RegionManager::get_instance()->set_max_region_id(tmp_max_region_id); set_table_pb(mem_schema_pb); @@ -2061,6 +2244,34 @@ void TableManager::update_statistics(const pb::MetaManagerRequest& request, DB_NOTICE("update table statistics success, request:%s", stat_pb.ShortDebugString().c_str()); } +// network error: -1; eq: 1, neq: 0 +int TableManager::check_gflag(const IdcInfo& idc, const std::string& gflag_name, const std::string& expected_value) { + int max_retry_times = 3; + for (int retry_times = 0; retry_times < max_retry_times; retry_times++) { + std::string new_instance; + auto ret = ClusterManager::get_instance()->select_instance_rolling(idc, {},new_instance); + std::string response; + ret = brpc_with_http(new_instance, new_instance + "/flags/" + gflag_name, response); + if (ret != 0) { + DB_WARNING("query gflag failed, instance_addr: %s, gflag_name: %s, retry_times: %d", + new_instance.c_str(), gflag_name.c_str(), retry_times); + continue; + } + auto start_pos = response.find(gflag_name + " | "); + if (start_pos == std::string::npos) { + // should never happen, treated as neq + DB_WARNING("query gflag failed, response unexpected format, instance_addr: %s, gflag_name: %s, response: %s", + new_instance.c_str(), gflag_name.c_str(), response.c_str()); + return 0; + } + // 多取一位，最后一位应当为空格，防止是前缀 + int substr_len = expected_value.size() + 1; + std::string actual_value = response.substr(start_pos + gflag_name.size() + 3, substr_len); + return actual_value == expected_value + " " ? 1 : 0; + } + return -1; +} + void TableManager::update_resource_tag(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done) { @@ -2078,6 +2289,45 @@ void TableManager::update_resource_tag(const pb::MetaManagerRequest& request, IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "resource_tag not exist"); return ; } + + bool need_check_wordrank = std::any_of( + mem_schema_pb.indexs().begin(), mem_schema_pb.indexs().end(), + [&](const auto& idx) { + return idx.index_type() == pb::I_FULLTEXT && + _need_check_wordrank_types.count(idx.segment_type()) != 0; + }); + + bool need_check_wordweight = std::any_of( + mem_schema_pb.indexs().begin(), mem_schema_pb.indexs().end(), + [&](const auto& idx) { + return idx.index_type() == pb::I_FULLTEXT && + _need_check_wordweight_types.count(idx.segment_type()) != 0; + }); + + if (need_check_wordrank) { + IdcInfo idc; + idc.resource_tag = resource_tag; + int cmp_value = check_gflag(idc, "use_fulltext_wordseg_wordrank_segment", "true"); + if (cmp_value < 0) { + DB_WARNING("check gflag network error, pass check"); + } else if (cmp_value == 0) { + IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "new resource tag does not support wordrank"); + return ; + } + } + + if (need_check_wordweight) { + IdcInfo idc; + idc.resource_tag = resource_tag; + int cmp_value = check_gflag(idc, "use_fulltext_wordweight_segment", "true"); + if (cmp_value < 0) { + DB_WARNING("check gflag network error, pass check"); + } else if (cmp_value == 0) { + IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "new resource tag does not support wordweight"); + return ; + } + } + mem_schema_pb.set_resource_tag(request.table_info().resource_tag()); mem_schema_pb.set_version(mem_schema_pb.version() + 1); auto ret = update_schema_for_rocksdb(table_id, mem_schema_pb, done); @@ -2121,35 +2371,134 @@ void TableManager::update_dists(const pb::MetaManagerRequest& request, }); } -void TableManager::update_ttl_duration(const pb::MetaManagerRequest& request, - const int64_t apply_index, +int TableManager::check_ttl_info(const pb::SchemaInfo& table_schema, std::string& err_msg) { + if (!table_schema.has_ttl_duration() || table_schema.ttl_duration() <= 0) { + return 0; + } + if (table_schema.engine() != pb::ROCKSDB && table_schema.engine() != pb::ROCKSDB_CSTORE) { + DB_WARNING("table engine doesn't support ttl"); + err_msg = "table engine doesn't support ttl"; + return -1; + } + + std::string ttl_field_name = table_schema.has_ttl_field() ? table_schema.ttl_field().field_name() : ""; + + if (table_schema.has_partition_info() && !ttl_field_name.empty()) { + DB_WARNING("field ttl doesn't support partition table, table name:%s", table_schema.table_name().c_str()); + err_msg = "field ttl doesn't support partition table"; + return -1; + } + + if (!ttl_field_name.empty()) { + bool field_exist = false; + for (const auto& field : table_schema.fields()) { + if (field.field_name() == ttl_field_name) { + field_exist = true; + if (field.can_null()) { + DB_WARNING("ttl_field can null, field_name[%s]", ttl_field_name.c_str()); + err_msg = "ttl_field should be NOT NULL"; + return -1; + } + if (field.mysql_type() != pb::PrimitiveType::DATETIME + && field.mysql_type() != pb::PrimitiveType::TIMESTAMP + && field.mysql_type() != pb::PrimitiveType::DATE) { + DB_WARNING("ttl_field is not supported type, type[%s]", pb::PrimitiveType_Name(field.mysql_type()).c_str()); + err_msg = "ttl_field type not supported"; + return -1; + } + break; + } + } + if (!field_exist) { + DB_WARNING("ttl_field not found, ttl field name: %s", ttl_field_name.c_str()); + err_msg = "ttl_field not found"; + return -1; + } + } + + bool ttl_field_in_pk = if_pk_contains_ttl_field(table_schema); + for (const auto& index: table_schema.indexs()) { + if (!is_ttl_support_index_type(index)) { + DB_WARNING("index with type[%s], is_global[%d] doesn't support field ttl", + pb::IndexType_Name(index.index_type()).c_str(), index.is_global()); + err_msg = "index does support field ttl"; + return -1; + } + if (!index.is_global() || ttl_field_name.empty()) { + continue; + } + + bool find_field = ttl_field_in_pk || index.index_type() == pb::I_PRIMARY; + for (const auto& field: index.field_names()) { + if (field == ttl_field_name) { + find_field = true; + break; + } + } + if (!find_field) { + DB_WARNING("global index doesn't contain ttl field"); + err_msg = "global index doesn't contain ttl field"; + return -1; + } + } + return 0; +} + +void TableManager::update_ttl_info(const pb::MetaManagerRequest& request, + const int64_t apply_index, braft::Closure* done) { - update_table_internal(request, apply_index, done, + if (!request.table_info().has_ttl_duration()) { + IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "ttl_duration must > 0"); + return; + } + int64_t table_id = 0; + if (check_table_exist(request.table_info(), table_id) != 0) { + DB_WARNING("check table exist fail, request:%s", request.ShortDebugString().c_str()); + IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "table not exist"); + return; + } + pb::SchemaInfo old_schema; + get_table_info(table_id, old_schema); + // 更新ttl_duration + if (old_schema.has_ttl_duration() && old_schema.ttl_duration() > 0) { + if (request.table_info().has_ttl_field() && !request.table_info().ttl_field().field_name().empty()) { + // 更新时不指定ttl field就默认不变 + std::string old_field_name = old_schema.has_ttl_field() ? old_schema.ttl_field().field_name() : ""; + std::string new_field_name = request.table_info().ttl_field().field_name(); + if (old_field_name != new_field_name) { + DB_WARNING("new_ttl_field_name[%s] doesn't match old_ttl_field_name[%s]", + old_field_name.c_str(), new_field_name.c_str()); + IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "can not change ttl_field"); + return; + } + } + DB_WARNING("update table ttl info, table_id:%ld", table_id); + } else { + // 添加ttl + old_schema.set_ttl_duration(request.table_info().ttl_duration()); + if (request.table_info().has_ttl_field()) { + old_schema.mutable_ttl_field()->CopyFrom(request.table_info().ttl_field()); + } + std::string err_msg; + if (check_ttl_info(old_schema, err_msg) != 0) { + DB_WARNING("add ttl to table[%ld] failed, err_msg: %s", table_id, err_msg.c_str()); + IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, err_msg); + return; + } + } + update_table_internal(request, apply_index, done, [](const pb::MetaManagerRequest& request, pb::SchemaInfo& mem_schema_pb, braft::Closure* done) { if (mem_schema_pb.ttl_duration() > 0 && request.table_info().ttl_duration() > 0) { // 只修改ttl mem_schema_pb.set_ttl_duration(request.table_info().ttl_duration()); } else if (mem_schema_pb.ttl_duration() <= 0 && request.table_info().ttl_duration() > 0) { // online ttl - bool can_support_ttl = true; - for (const auto& index : mem_schema_pb.indexs()) { - if (index.index_type() == pb::I_FULLTEXT) { - can_support_ttl = false; - break; - } - } - - if (mem_schema_pb.engine() == pb::REDIS || mem_schema_pb.engine() == pb::BINLOG) { - can_support_ttl = false; - } - - if (!can_support_ttl) { - DB_WARNING("can't support ttl, req: %s", request.ShortDebugString().c_str()); - return; - } int64_t online_ttl_expire_time_us = butil::gettimeofday_us() + request.table_info().ttl_duration() * 1000000LL; mem_schema_pb.set_ttl_duration(request.table_info().ttl_duration()); mem_schema_pb.set_online_ttl_expire_time_us(online_ttl_expire_time_us); + if (request.table_info().has_ttl_field()) { + mem_schema_pb.mutable_ttl_field()->CopyFrom(request.table_info().ttl_field()); + } } else { DB_WARNING("update fail, resuest.ttl_duration:%ld mem_schema_pb.ttl_duration:%ld", request.table_info().ttl_duration(), mem_schema_pb.ttl_duration()); @@ -2942,10 +3291,12 @@ int TableManager::load_statistics_snapshot(const std::string& value) { return 0; } -int TableManager::write_schema_for_not_level(TableMem& table_mem, +int TableManager::write_schema_for_not_level(const pb::MetaManagerRequest& request, + TableMem& table_mem, braft::Closure* done, int64_t max_table_id_tmp, - bool has_auto_increment) { + bool has_auto_increment, + const int64_t apply_index) { //如果创建成功，则不需要做任何操作 //如果失败，则需要报错，手工调用删除table的接口 std::vector rocksdb_keys; @@ -2978,10 +3329,6 @@ int TableManager::write_schema_for_not_level(TableMem& table_mem, } // Range Partition - std::string main_logical_room; - get_main_logical_room(table_mem.schema_pb, main_logical_room); - std::string resource_tag = table_mem.schema_pb.resource_tag(); - boost::trim(resource_tag); std::string namespace_name = table_mem.schema_pb.namespace_name(); std::string database = table_mem.schema_pb.database(); std::string table_name = table_mem.schema_pb.table_name(); @@ -2994,34 +3341,40 @@ int TableManager::write_schema_for_not_level(TableMem& table_mem, return -1; } } - + if (is_create_table_support_engine(table_mem.schema_pb.engine()) + && request.partition_index_instance_size() != table_mem.schema_pb.partition_num()) { + DB_FATAL("partition_index_instance_size(%d) not equal to partition_num(%ld)", + request.partition_index_instance_size(), table_mem.schema_pb.partition_num()); + return -1; + } //有split_key的索引先处理 std::vector processed_index_name; for (auto i = 0; i < table_mem.schema_pb.partition_num() && is_create_table_support_engine(table_mem.schema_pb.engine()); ++i) { - int64_t partition_id = i; - std::string partition_resource_tag = resource_tag; - std::string partition_main_logical_room = main_logical_room; - if (table_mem.schema_pb.has_partition_info() && table_mem.schema_pb.partition_info().type() == pb::PT_RANGE) { - if (table_mem.schema_pb.partition_info().range_partition_infos(i).has_resource_tag()) { - partition_resource_tag = table_mem.schema_pb.partition_info().range_partition_infos(i).resource_tag(); - partition_main_logical_room = ""; - } - if (table_mem.schema_pb.partition_info().range_partition_infos(i).has_partition_id()) { - partition_id = table_mem.schema_pb.partition_info().range_partition_infos(i).partition_id(); - } - } + auto& partition_index_instances = request.partition_index_instance(i); + int64_t partition_id = partition_index_instances.partition_id(); for (auto& split_key : table_mem.schema_pb.split_keys()) { std::string index_name = split_key.index_name(); - for (auto j = 0; j <= split_key.split_keys_size(); ++j) { - std::string instance; - int ret = ClusterManager::get_instance()->select_instance_rolling( - {partition_resource_tag, partition_main_logical_room, ""}, {}, instance); - if (ret < 0) { - DB_WARNING("select instance fail"); - return -1; + pb::IndexInstance* match_index_instance = nullptr; + for (auto& index_instances : partition_index_instances.index_instances()) { + if (index_name == index_instances.index_name()) { + match_index_instance = const_cast(&index_instances); } + } + if (match_index_instance == nullptr) { + DB_FATAL("index_instance is null, index_name: %s", index_name.c_str()); + IF_DONE_SET_RESPONSE(done, pb::PARSE_TO_PB_FAIL, "pick instance fail"); + return -1; + } + if (match_index_instance->instances().size() != split_key.split_keys_size() + 1) { + DB_FATAL("picked_instance size(%d) not equal to split_keys_size(%d)", + match_index_instance->instances().size(), split_key.split_keys_size()); + IF_DONE_SET_RESPONSE(done, pb::PARSE_TO_PB_FAIL, "pick instance fail"); + return -1; + } + for (auto j = 0; j <= split_key.split_keys_size(); ++j) { + std::string instance = match_index_instance->instances().at(j); pb::InitRegion init_region_request; pb::RegionInfo* region_info = init_region_request.mutable_region_info(); region_info->set_region_id(++tmp_max_region_id); @@ -3056,30 +3409,28 @@ int TableManager::write_schema_for_not_level(TableMem& table_mem, for (auto i = 0; i < table_mem.schema_pb.partition_num() && is_create_table_support_engine(table_mem.schema_pb.engine()); ++i) { - int64_t partition_id = i; - std::string partition_resource_tag = resource_tag; - std::string partition_main_logical_room = main_logical_room; - if (table_mem.schema_pb.has_partition_info() && table_mem.schema_pb.partition_info().type() == pb::PT_RANGE) { - if (table_mem.schema_pb.partition_info().range_partition_infos(i).has_resource_tag()) { - partition_resource_tag = table_mem.schema_pb.partition_info().range_partition_infos(i).resource_tag(); - partition_main_logical_room = ""; - } - if (table_mem.schema_pb.partition_info().range_partition_infos(i).has_partition_id()) { - partition_id = table_mem.schema_pb.partition_info().range_partition_infos(i).partition_id(); - } - } + auto partition_index_instances = request.partition_index_instance(i); + int64_t partition_id = partition_index_instances.partition_id(); for (auto& index : global_index) { - std::string instance; - if (namespace_name == "INTERNAL" && database == "baikaldb" && table_name == "__baikaldb_instance") { - // for baikaldb instance id, do nothing - } else { - int ret = ClusterManager::get_instance()->select_instance_rolling( - {partition_resource_tag, partition_main_logical_room, ""}, {}, instance); - if (ret < 0) { - DB_WARNING("select instance fail"); - return -1; + auto index_name = index.first; + pb::IndexInstance* match_index_instance = nullptr; + for (auto& index_instances : partition_index_instances.index_instances()) { + if (index_name == index_instances.index_name()) { + match_index_instance = const_cast(&index_instances); } } + if (match_index_instance == nullptr) { + DB_FATAL("picked_instance is null, index_name: %s", index_name.c_str()); + IF_DONE_SET_RESPONSE(done, pb::PARSE_TO_PB_FAIL, "picked_instance fail"); + return -1; + } + if (match_index_instance->instances().size() != 1) { + DB_FATAL("picked_instance size(%d) not equal to split_keys_size(1)", + match_index_instance->instances().size()); + IF_DONE_SET_RESPONSE(done, pb::PARSE_TO_PB_FAIL, "pick instance fail"); + return -1; + } + std::string instance = match_index_instance->instances().at(0); pb::InitRegion init_region_request; pb::RegionInfo* region_info = init_region_request.mutable_region_info(); region_info->set_region_id(++tmp_max_region_id); @@ -3125,6 +3476,8 @@ int TableManager::write_schema_for_not_level(TableMem& table_mem, IF_DONE_SET_RESPONSE(done, pb::INTERNAL_ERROR, "write db fail"); return -1; } + DB_WARNING("apply_index: %ld, generate %ld region_id: [%ld, %ld]", + apply_index, tmp_max_region_id - start_region_id + 1, start_region_id, tmp_max_region_id); RegionManager::get_instance()->set_max_region_id(tmp_max_region_id); if (done) { ((MetaServerClosure*)done)->op_type = pb::OP_CREATE_TABLE; @@ -3516,11 +3869,32 @@ int TableManager::alloc_field_id(pb::SchemaInfo& table_info, bool& has_auto_incr return 0; } -int TableManager::alloc_index_id(pb::SchemaInfo& table_info, TableMem& table_mem, int64_t& max_table_id_tmp) { +bool TableManager::if_pk_contains_ttl_field(const pb::SchemaInfo& table_info){ + if (table_info.has_ttl_field() && !table_info.ttl_field().field_name().empty()) { + std::string ttl_field_name = table_info.ttl_field().field_name(); + for (const auto& index: table_info.indexs()) { + if (index.index_type() != pb::I_PRIMARY) { + continue; + } + for (const auto& field_name: index.field_names()) { + if (field_name == ttl_field_name) { + return true; + } + } + } + } + return false; +} + +// -1 normal error, -2 don't support workrank or workweight +int TableManager::alloc_index_id(pb::SchemaInfo& table_info, TableMem& table_mem, + int64_t& max_table_id_tmp, std::string resource_tag) { bool has_primary_key = false; std::string table_name = table_info.table_name(); std::unordered_set primary_key_field_id_set; std::vector vector_idxes; + std::vector fulltext_idxes; + fulltext_idxes.reserve(3); //分配index_id， 序列与table_id共享, 必须有primary_key for (auto i = 0; i < table_info.indexs_size(); ++i) { std::string index_name = table_info.indexs(i).index_name(); @@ -3548,6 +3922,9 @@ int TableManager::alloc_index_id(pb::SchemaInfo& table_info, TableMem& table_mem if (table_info.indexs(i).index_type() == pb::I_VECTOR) { vector_idxes.emplace_back(i); } + if (table_info.indexs(i).index_type() == pb::I_FULLTEXT) { + fulltext_idxes.emplace_back(i); + } if (table_info.indexs(i).index_type() != pb::I_PRIMARY) { table_info.mutable_indexs(i)->set_index_id(++max_table_id_tmp); continue; @@ -3639,10 +4016,47 @@ int TableManager::alloc_index_id(pb::SchemaInfo& table_info, TableMem& table_mem return -1; } } + + if (fulltext_idxes.size() != 0) { + bool support_wordrank = false, support_wordweight = false; + IdcInfo idc; + idc.resource_tag = resource_tag; + int cmp_value = check_gflag(idc, "use_fulltext_wordseg_wordrank_segment", "true"); + if (cmp_value < 0) { + DB_WARNING("check gflag network error, pass check"); + support_wordrank = true; + } else if (cmp_value == 0) { + DB_WARNING("store cluster does not support wordrank, resource_tag: %s", resource_tag.c_str()); + support_wordrank = false; + } else { + support_wordrank = true; + } + + cmp_value = check_gflag(idc, "use_fulltext_wordweight_segment", "true"); + if (cmp_value < 0) { + DB_WARNING("check gflag network error, pass check"); + support_wordweight = true; + } else if (cmp_value == 0) { + DB_WARNING("store cluster does not support wordweight, resource_tag: %s", resource_tag.c_str()); + support_wordweight = false; + } else { + support_wordweight = true; + } + for (int idx : fulltext_idxes) { + if (0 != _need_check_wordrank_types.count(table_info.indexs(idx).segment_type()) && !support_wordrank) { + DB_WARNING("resource_tag don't support workrank"); + return -2; + } + if (0 != _need_check_wordweight_types.count(table_info.indexs(idx).segment_type()) && !support_wordweight) { + DB_WARNING("resource_tag don't support workweight"); + return -2; + } + } + } return 0; } -int64_t TableManager::get_pre_regionid(int64_t table_id, +int64_t TableManager::get_pre_regionid(int64_t table_id, const std::string& start_key, int64_t partition) { if (!exist_table_id(table_id)) { DB_WARNING("table_id: %ld not exist", table_id); @@ -4489,7 +4903,7 @@ void TableManager::drop_index(const pb::MetaManagerRequest& request, const int64 // 忽略大小写 return boost::algorithm::iequals(info.index_name(), index_req.index_name()) && (info.index_type() == pb::I_UNIQ || info.index_type() == pb::I_KEY || - info.index_type() == pb::I_FULLTEXT || info.index_type() == pb::I_ROLLUP); + info.index_type() == pb::I_FULLTEXT || info.index_type() == pb::I_ROLLUP || info.index_type() == pb::I_VECTOR); }); if (index_to_del != std::end(schema_info.indexs())) { if (index_req.hint_status() == pb::IHS_VIRTUAL || index_to_del->hint_status() == pb::IHS_VIRTUAL) { @@ -4525,6 +4939,113 @@ void TableManager::drop_index(const pb::MetaManagerRequest& request, const int64 } } +// 状态机外, 只有leader处理 +int TableManager::pre_process_for_add_index(const pb::MetaManagerRequest* request, + pb::MetaManagerResponse* response, + uint64_t log_id, + google::protobuf::Closure* done) { + int64_t table_id = 0; + if (check_table_exist(request->table_info(), table_id) != 0 && + request->table_info().table_id() == table_id) { + DB_WARNING("DDL_LOG[add_index] check table exist fail, request:%s", request->ShortDebugString().c_str()); + response->set_errcode(pb::INPUT_PARAM_ERROR); + response->set_errmsg("table not exist"); + return -1; + } + pb::SchemaInfo schema_info; + int ret = get_table_info(table_id, schema_info); + if (ret != 0) { + DB_WARNING("DDL_LOG[add_index] table not in table_info_map, request:%s", request->DebugString().c_str()); + response->set_errcode(pb::INPUT_PARAM_ERROR); + response->set_errmsg("table not in table_info_map"); + return -1; + } + const pb::IndexInfo& index_info = request->table_info().indexs(0); + std::string index_name = index_info.index_name(); + // 状态机外选global index的机器 + std::string main_logical_room; + if (schema_info.has_main_logical_room()) { + main_logical_room = schema_info.main_logical_room(); + } + std::string resource_tag = schema_info.resource_tag(); + boost::trim(resource_tag); + + if (index_info.has_index_type() && index_info.index_type() == pb::I_FULLTEXT && index_info.has_segment_type()) { + if (0 != _need_check_wordrank_types.count(index_info.segment_type())) { + IdcInfo idc; + idc.resource_tag = resource_tag; + int cmp_value = check_gflag(idc, "use_fulltext_wordseg_wordrank_segment", "true"); + if (cmp_value < 0) { + DB_WARNING("check gflag network error, pass check"); + } else if (cmp_value == 0) { + DB_WARNING("DDL_LOG[add_index] store cluster does not support wordrank, resource_tag: %s", resource_tag.c_str()); + response->set_errcode(pb::INPUT_PARAM_ERROR); + response->set_errmsg("store cluster does not support wordrank, resource_tag: " + resource_tag); + return -1; + } + } else if (0 != _need_check_wordweight_types.count(index_info.segment_type())) { + IdcInfo idc; + idc.resource_tag = resource_tag; + int cmp_value = check_gflag(idc, "use_fulltext_wordweight_segment", "true"); + if (cmp_value < 0) { + DB_WARNING("check gflag network error, pass check"); + } else if (cmp_value == 0) { + DB_WARNING("DDL_LOG[add_index] store cluster does not support wordweight, resource_tag: %s", resource_tag.c_str()); + response->set_errcode(pb::INPUT_PARAM_ERROR); + response->set_errmsg("store cluster does not support wordweight, resource_tag: " + resource_tag); + return -1; + } + } + } + + if (!index_info.is_global()) { + return 0; + } + + if (schema_info.has_partition_info() && schema_info.partition_info().type() == pb::PT_RANGE) { + if (schema_info.partition_num() != schema_info.partition_info().range_partition_infos_size()) { + DB_WARNING("Invalid partition_num %ld or range_partition_infos_size: %d", + schema_info.partition_num(), + (int)schema_info.partition_info().range_partition_infos_size()); + response->set_errcode(pb::INTERNAL_ERROR); + response->set_errmsg("Invalid partition_num or range_partition_infos_size"); + return -1; + } + } + auto mutable_request = const_cast(request); + for (auto i = 0; + i < schema_info.partition_num() && is_create_table_support_engine(schema_info.engine()); + ++i) { + int64_t partition_id = i; + std::string partition_resource_tag = resource_tag; + std::string partition_main_logical_room = main_logical_room; + if (schema_info.has_partition_info() && schema_info.partition_info().type() == pb::PT_RANGE) { + if (schema_info.partition_info().range_partition_infos(i).has_resource_tag()) { + partition_resource_tag = schema_info.partition_info().range_partition_infos(i).resource_tag(); + partition_main_logical_room = ""; + } + if (schema_info.partition_info().range_partition_infos(i).has_partition_id()) { + partition_id = schema_info.partition_info().range_partition_infos(i).partition_id(); + } + } + auto partition_index_instances = mutable_request->add_partition_index_instance(); + partition_index_instances->set_partition_id(partition_id); + auto index_instances = partition_index_instances->add_index_instances(); + std::string instance; + int ret = ClusterManager::get_instance()->select_instance_rolling( + {partition_resource_tag, partition_main_logical_room, ""}, {}, instance); + if (ret < 0) { + response->set_errcode(pb::INTERNAL_ERROR); + response->set_errmsg("no valid store address for global index"); + return -1; + } + index_instances->set_index_name(index_name); + index_instances->add_instances(instance); + } + return 0; +} + +// 状态机内 void TableManager::add_index(const pb::MetaManagerRequest& request, const int64_t apply_index, braft::Closure* done) { @@ -4561,7 +5082,7 @@ void TableManager::add_index(const pb::MetaManagerRequest& request, IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "fields info fail"); return; } - DB_DEBUG("DDL_LOG[add_index] check field success."); + pb::SchemaInfo mem_schema_pb; ret = get_table_info(table_id, mem_schema_pb); if (ret != 0) { @@ -4569,6 +5090,31 @@ void TableManager::add_index(const pb::MetaManagerRequest& request, IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "table not in table_info_map"); return; } + + if (request.table_info().indexs(0).is_global()) { + std::string ttl_field_name; + ret = get_ttl_field_name(table_id, ttl_field_name); + if (ret != 0) { + DB_WARNING("DDL_LOG[add_index] check fields info fail, request:%s", request.ShortDebugString().c_str()); + IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "get ttl fields info failed"); + return; + } + if (!ttl_field_name.empty()) { + auto any_fields_equal = std::any_of( + std::begin(first_index_fields), + std::end(first_index_fields), + [&ttl_field_name](const std::string& field_name) -> bool { + return field_name == ttl_field_name; + } + ); + if (!any_fields_equal && !if_pk_contains_ttl_field(mem_schema_pb)) { + DB_WARNING("DDL_LOG[add_index] check fields info fail, request:%s", request.ShortDebugString().c_str()); + IF_DONE_SET_RESPONSE(done, pb::INPUT_PARAM_ERROR, "global index doesn't contain ttl field"); + return; + } + } + } + DB_DEBUG("DDL_LOG[add_index] check field success."); if (request.table_info().indexs(0).index_type() == pb::I_ROLLUP) { if (!check_rollup_field_all_in_pk(table_id, request.table_info().indexs(0))) { DB_WARNING("rollup key not in pk, request:%s", request.ShortDebugString().c_str()); @@ -4641,7 +5187,7 @@ void TableManager::add_index(const pb::MetaManagerRequest& request, index_info.set_state(pb::IS_PUBLIC); _just_add_virtual_index_info.insert(index_info.index_id());//保存虚拟索引id，后续drop_index的流程中删除相应的id } else { - ret = do_add_index(mem_schema_pb, apply_index, done, table_id, index_info); + ret = do_add_index(request, mem_schema_pb, apply_index, done, table_id, index_info); } if (ret != 0) { DB_WARNING("add global|local index error."); @@ -4683,18 +5229,11 @@ void TableManager::add_index(const pb::MetaManagerRequest& request, IF_DONE_SET_RESPONSE(done, pb::SUCCESS, "success"); } -int TableManager::init_global_index_region(const pb::SchemaInfo& schema_info, int64_t table_id, braft::Closure* done, pb::IndexInfo& index_info) { +int TableManager::init_global_index_region(const pb::MetaManagerRequest& request, const pb::SchemaInfo& schema_info, + int64_t table_id, braft::Closure* done, pb::IndexInfo& index_info, const int64_t apply_index) { std::vector rocksdb_keys; std::vector rocksdb_values; - // Range Partition - std::string main_logical_room; - if (schema_info.has_main_logical_room()) { - main_logical_room = schema_info.main_logical_room(); - } - std::string resource_tag = schema_info.resource_tag(); - boost::trim(resource_tag); - if (schema_info.has_partition_info() && schema_info.partition_info().type() == pb::PT_RANGE) { if (schema_info.partition_num() != schema_info.partition_info().range_partition_infos_size()) { DB_WARNING("Invalid partition_num %ld or range_partition_infos_size: %d", @@ -4708,48 +5247,54 @@ int TableManager::init_global_index_region(const pb::SchemaInfo& schema_info, in //与store交互 //准备partition_num个数的regionInfo int64_t tmp_max_region_id = RegionManager::get_instance()->get_max_region_id(); + int64_t start_region_id = tmp_max_region_id + 1; std::shared_ptr> init_regions(new std::vector{}); //没有指定split_key的索引 - for (auto i = 0; - i < schema_info.partition_num() && is_create_table_support_engine(schema_info.engine()); - ++i) { - int64_t partition_id = i; - std::string partition_resource_tag = resource_tag; - std::string partition_main_logical_room = main_logical_room; - if (schema_info.has_partition_info() && schema_info.partition_info().type() == pb::PT_RANGE) { - if (schema_info.partition_info().range_partition_infos(i).has_resource_tag()) { - partition_resource_tag = schema_info.partition_info().range_partition_infos(i).resource_tag(); - partition_main_logical_room = ""; + if (is_create_table_support_engine(schema_info.engine())) { + if (schema_info.partition_num() != request.partition_index_instance_size()) { + DB_FATAL("has partition num: %ld, but only %d partition_index_instance", schema_info.partition_num(), request.partition_index_instance_size()); + return -1; + } + for (auto i = 0; i < schema_info.partition_num(); ++i) { + int64_t partition_id = i; + if (schema_info.has_partition_info() && schema_info.partition_info().type() == pb::PT_RANGE) { + if (schema_info.partition_info().range_partition_infos(i).has_partition_id()) { + partition_id = schema_info.partition_info().range_partition_infos(i).partition_id(); + } } - if (schema_info.partition_info().range_partition_infos(i).has_partition_id()) { - partition_id = schema_info.partition_info().range_partition_infos(i).partition_id(); + auto& partition_index_instances = request.partition_index_instance(i); + if (partition_index_instances.partition_id() != partition_id) { + DB_FATAL("partition_id not match, partition_id: %ld, partition_index_instances.partition_id(): %ld", partition_id, partition_index_instances.partition_id()); + return -1; + } + if (partition_index_instances.index_instances_size() == 0 + || partition_index_instances.index_instances(0).instances_size() == 0) { + DB_WARNING("no instance address"); + return -1; } - } - std::string instance; - int ret = ClusterManager::get_instance()->select_instance_rolling( - {partition_resource_tag, partition_main_logical_room, ""}, {}, instance); - if (ret < 0) { - DB_WARNING("select instance fail"); - return -1; + std::string instance = partition_index_instances.index_instances(0).instances(0); + pb::InitRegion init_region_request; + pb::RegionInfo* region_info = init_region_request.mutable_region_info(); + region_info->set_region_id(++tmp_max_region_id); + region_info->set_table_id(index_info.index_id()); + region_info->set_main_table_id(table_id); + region_info->set_table_name(schema_info.table_name()); + construct_common_region(region_info, schema_info.replica_num()); + region_info->set_partition_id(partition_id); + region_info->add_peers(instance); + region_info->set_leader(instance); + region_info->set_can_add_peer(false);// 简化理解，让raft addpeer必须发送snapshot + // region_info->set_partition_num(schema_info.partition_num()); // 当前系统未使用region partition_num，range分区不易维护 + region_info->set_is_binlog_region(false); // binlog表不能有全局索引 + pb::SchemaInfo* request_schema_info = init_region_request.mutable_schema_info(); + *request_schema_info = schema_info; + request_schema_info->set_version(schema_info.version() + 1); + request_schema_info->add_indexs()->CopyFrom(index_info); + init_region_request.set_snapshot_times(2); + init_regions->emplace_back(init_region_request); + DB_WARNING("init_region_request: %s", init_region_request.DebugString().c_str()); } - pb::InitRegion init_region_request; - pb::RegionInfo* region_info = init_region_request.mutable_region_info(); - region_info->set_region_id(++tmp_max_region_id); - region_info->set_table_id(index_info.index_id()); - region_info->set_main_table_id(table_id); - region_info->set_table_name(schema_info.table_name()); - construct_common_region(region_info, schema_info.replica_num()); - region_info->set_partition_id(partition_id); - region_info->add_peers(instance); - region_info->set_leader(instance); - region_info->set_can_add_peer(false);// 简化理解，让raft addpeer必须发送snapshot - // region_info->set_partition_num(schema_info.partition_num()); // 当前系统未使用region partition_num，range分区不易维护 - region_info->set_is_binlog_region(false); // binlog表不能有全局索引 - *(init_region_request.mutable_schema_info()) = schema_info; - init_region_request.set_snapshot_times(2); - init_regions->emplace_back(init_region_request); - DB_WARNING("init_region_request: %s", init_region_request.DebugString().c_str()); } //持久化region_id std::string max_region_id_key = RegionManager::get_instance()->construct_max_region_id_key(); @@ -4774,6 +5319,8 @@ int TableManager::init_global_index_region(const pb::SchemaInfo& schema_info, in IF_DONE_SET_RESPONSE(done, pb::INTERNAL_ERROR, "write db fail"); return -1; } + DB_WARNING("apply_index: %ld, generate %ld region_id: [%ld, %ld]", + apply_index, tmp_max_region_id - start_region_id + 1, start_region_id, tmp_max_region_id); RegionManager::get_instance()->set_max_region_id(tmp_max_region_id); //leader发送请求 if (done && is_create_table_support_engine(schema_info.engine())) { @@ -4790,11 +5337,12 @@ int TableManager::init_global_index_region(const pb::SchemaInfo& schema_info, in return 0; } -int TableManager::do_add_index(const pb::SchemaInfo& mem_schema_pb, +int TableManager::do_add_index(const pb::MetaManagerRequest& request, + const pb::SchemaInfo& mem_schema_pb, const int64_t apply_index, braft::Closure* done, const int64_t table_id, pb::IndexInfo& index_info) { int64_t start_region_id = RegionManager::get_instance()->get_max_region_id(); - if (index_info.is_global() && init_global_index_region(mem_schema_pb, table_id, done, index_info) != 0) { + if (index_info.is_global() && init_global_index_region(request, mem_schema_pb, table_id, done, index_info, apply_index) != 0) { DB_WARNING("table_id[%ld] add global index init global region failed.", table_id); if (done) { IF_DONE_SET_RESPONSE(done, pb::INTERNAL_ERROR, "init global region failed"); @@ -4821,6 +5369,26 @@ int TableManager::do_add_index(const pb::SchemaInfo& mem_schema_pb, return 0; } +int TableManager::get_ttl_field_name(int64_t table_id, std::string& ttl_field_name) { + ttl_field_name.clear(); + DoubleBufferedTableMemMapping::ScopedPtr info; + if (_table_mem_infos.Read(&info) != 0) { + DB_WARNING("read double_buffer_table error."); + return -1; + } + auto table_mem_iter = info->table_info_map.find(table_id); + if (table_mem_iter == info->table_info_map.end()) { + DB_WARNING("table_id:[%ld] not exist.", table_id); + return -1; + } + if (table_mem_iter->second.schema_pb.ttl_duration() <= 0 || !table_mem_iter->second.schema_pb.has_ttl_field()) { + DB_WARNING("table_id:[%ld] does not exist ttl field", table_id); + return 0; + } + ttl_field_name = table_mem_iter->second.schema_pb.ttl_field().field_name(); + return 0; +} + bool TableManager::check_field_exist(const std::string &field_name, int64_t table_id) { DoubleBufferedTableMemMapping::ScopedPtr info; diff --git a/src/meta_server/tso_state_machine.cpp b/src/meta_server/tso_state_machine.cpp index 00bf0f276..eb6753400 100755 --- a/src/meta_server/tso_state_machine.cpp +++ b/src/meta_server/tso_state_machine.cpp @@ -373,7 +373,6 @@ void TSOStateMachine::update_timestamp() { } void TSOStateMachine::on_leader_start() { - start_check_bns(); DB_WARNING("tso leader start"); int64_t now = tso::clock_realtime_ms(); pb::TsoTimestamp current; @@ -395,6 +394,7 @@ void TSOStateMachine::on_leader_start() { DB_WARNING("sync timestamp ok"); _is_leader.store(true); _tso_update_timer.start(); + start_check_bns(); }; Bthread bth; bth.run(func); diff --git a/src/physical_plan/exec_type_analyzer.cpp b/src/physical_plan/exec_type_analyzer.cpp index cdb36b21f..64e72db73 100644 --- a/src/physical_plan/exec_type_analyzer.cpp +++ b/src/physical_plan/exec_type_analyzer.cpp @@ -20,18 +20,25 @@ namespace baikaldb { DECLARE_bool(enable_plan_cache); DECLARE_bool(use_arrow_vector); DEFINE_int32(mpp_hash_partition_num, 0, "mpp hash partition num"); +DEFINE_int32(file_scan_concurrency, 4, "file scan concurrency"); DEFINE_bool(enable_decide_mpp_by_db_statistics, false, "enable decide mpp by db statistics"); DECLARE_int64(mpp_min_statistics_rows); DECLARE_int64(mpp_min_statistics_bytes); int ExecTypeAnalyzer::analyze(QueryContext* ctx) { if (ctx->is_from_subquery - || ctx->is_union_subquery - || ctx->is_insert_select_subquery) { + || ctx->is_union_subquery + || ctx->is_insert_select_subquery) { return 0; } bool can_vectorize = can_use_arrow_vector(ctx); if (!can_vectorize) { + if (ctx->must_vectorize) { + ctx->stat_info.error_code = ER_NOT_SUPPORTED_YET; + ctx->stat_info.error_msg << "This SQL only support in vectorize mode"; + DB_WARNING("This SQL only support in vectorize mode"); + return -1; + } ctx->use_mpp = false; return 0; } @@ -65,7 +72,7 @@ bool ExecTypeAnalyzer::can_use_arrow_vector(QueryContext* ctx) { state.reset_vectorize_info(); // full join必须走向量化, 即使集群没开向量化开关 - bool must_vectorize = false; + bool& must_vectorize = ctx->must_vectorize; // sql注释或者sign配置指定走向量化 bool defined_vectorize = false; std::vector joins; @@ -76,6 +83,16 @@ bool ExecTypeAnalyzer::can_use_arrow_vector(QueryContext* ctx) { break; } } + if (!must_vectorize) { + std::vector scans; + ctx->root->get_node_pass_subquery(pb::SCAN_NODE, scans); + for (auto& scan : scans) { + if (static_cast(scan)->is_file_scan_node()) { + must_vectorize = true; + break; + } + } + } if (ctx->client_conn != nullptr && ctx->client_conn->txn_id != 0) { if (joins.size() > 0 || ctx->root->get_node_pass_subquery(pb::UNION_NODE) != nullptr) { @@ -128,9 +145,7 @@ bool ExecTypeAnalyzer::can_use_arrow_vector(QueryContext* ctx) { } } else { if (must_vectorize) { - ctx->stat_info.error_code = ER_NOT_SUPPORTED_YET; - ctx->stat_info.error_msg << "FULL JOIN only support in vectorize mode."; - return -1; + return false; } state.execute_type = pb::EXEC_ROW; } @@ -146,10 +161,19 @@ bool ExecTypeAnalyzer::can_use_mpp(QueryContext* ctx) { // 事务中不能使用mpp return false; } + std::vector scan_nodes; + ctx->root->get_node_pass_subquery(pb::SCAN_NODE, scan_nodes); + bool has_file_scan_node = false; + for (auto* scan_node : scan_nodes) { + if (scan_node != nullptr && static_cast(scan_node)->is_file_scan_node()) { + has_file_scan_node = true; + break; + } + } ExecNode* join_node = ctx->root->get_node_pass_subquery(pb::JOIN_NODE); ExecNode* agg_node = ctx->root->get_node_pass_subquery(pb::AGG_NODE); - if (join_node == nullptr && agg_node == nullptr) { - // 没有join/agg的没必要走mpp + if (join_node == nullptr && agg_node == nullptr && !has_file_scan_node) { + // 没有join/agg/file_scan的没必要走mpp return false; } if (ctx->is_explain && ctx->explain_type != SHOW_PLAN) { diff --git a/src/physical_plan/index_selector.cpp b/src/physical_plan/index_selector.cpp index b2b0e61c9..d3c4ed705 100644 --- a/src/physical_plan/index_selector.cpp +++ b/src/physical_plan/index_selector.cpp @@ -11,6 +11,8 @@ // 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 "file_system.h" +#include "file_scan_node.h" #include "index_selector.h" #include "slot_ref.h" #include "scalar_fn_call.h" @@ -20,7 +22,7 @@ #include "parser.h" namespace baikaldb { -DEFINE_bool(use_column_storage, false, "whether use column storage"); +DECLARE_int32(file_scan_concurrency); using namespace range; DEFINE_bool(use_index_merge, false, "if use index merge in index select"); @@ -67,7 +69,8 @@ int IndexSelector::analyze(QueryContext* ctx) { sort_node->set_limit(limit_node->other_limit()); } for (auto& scan_node_ptr : scan_nodes) { - if (!static_cast(scan_node_ptr)->is_rocksdb_scan_node()) { + ScanNode* scan_node = static_cast(scan_node_ptr); + if (!scan_node->is_rocksdb_scan_node() && !scan_node->is_file_scan_node()) { continue; } ExecNode* parent_node_ptr = scan_node_ptr->get_parent(); @@ -83,11 +86,12 @@ int IndexSelector::analyze(QueryContext* ctx) { int ret = 0; std::map field_range_type; bool index_has_null = false; + if (join_node != NULL || agg_node != NULL) { IndexSelectorOptions options; options.execute_type = ctx->runtime_state == nullptr ? pb::EXEC_ROW : ctx->runtime_state->execute_type; ret = index_selector(ctx->tuple_descs(), - static_cast(scan_node_ptr), + scan_node, filter_node, NULL, join_node, @@ -102,7 +106,7 @@ int IndexSelector::analyze(QueryContext* ctx) { IndexSelectorOptions options; options.execute_type = ctx->runtime_state == nullptr ? pb::EXEC_ROW : ctx->runtime_state->execute_type; ret = index_selector(ctx->tuple_descs(), - static_cast(scan_node_ptr), + scan_node, filter_node, sort_node, join_node, @@ -157,7 +161,7 @@ void IndexSelector::analyze_join_index(QueryContext* ctx, ScanNode* scan_node, E } _ctx = ctx; // 和joiner真正下推in一样的逻辑 - if (!scan_node->is_rocksdb_scan_node()) { + if (!scan_node->is_rocksdb_scan_node() && !scan_node->is_file_scan_node()) { return; } ExecNode* parent_node_ptr = scan_node->get_parent(); @@ -183,6 +187,7 @@ void IndexSelector::analyze_join_index(QueryContext* ctx, ScanNode* scan_node, E IndexSelectorOptions options; options.join_on_conditions = in_condition; options.execute_type = ctx->runtime_state == nullptr ? pb::EXEC_ROW : ctx->runtime_state->execute_type; + index_selector(ctx->tuple_descs(), scan_node, filter_node, @@ -406,7 +411,7 @@ void IndexSelector::hit_field_or_like_range(ExprNode* expr, std::mapget_table_info_ptr(table_id); + auto table_info_ptr = SchemaFactory::get_instance()->get_table_info_ptr(table_id); if (table_info_ptr == nullptr) { return; } @@ -549,7 +554,7 @@ void IndexSelector::hit_field_range(ExprNode* expr, } RangeType tmp_type; - auto try_add_into_fulltext = [this, field_id](FulltextInfoNode* index_node, range::FieldRange&& range, int64_t index_id) { + auto try_add_into_fulltext = [field_id](FulltextInfoNode* index_node, range::FieldRange&& range, int64_t index_id) { if (index_node != nullptr && index_id > 0) { auto& inner_node = boost::get(index_node->info); inner_node.children.emplace_back(new FulltextInfoNode); @@ -732,6 +737,10 @@ int64_t IndexSelector::index_selector(const std::vector& tu std::map& field_range_type, const std::string& sample_sql, const IndexSelectorOptions& options) { + if (scan_node == nullptr) { + DB_WARNING("scan_node is nullptr "); + return -1; + } int64_t table_id = scan_node->table_id(); int32_t tuple_id = scan_node->tuple_id(); auto table_info = _factory->get_table_info_ptr(table_id); @@ -739,11 +748,19 @@ int64_t IndexSelector::index_selector(const std::vector& tu DB_WARNING("table info not found:%ld", table_id); return -1; } + if (scan_node->is_file_scan_node()) { + TimeCost tm; + if (select_partition_files(table_info, scan_node, filter_node, options.join_on_conditions) != 0) { + DB_WARNING("Fail to select_partition_files, table_id: %ld", table_id); + return -1; + } + DB_WARNING("select_partition_files tm: %ld", tm.get_time()); + return 0; + } bool can_use_column_storage = false; - if (FLAGS_use_column_storage && options.execute_type == pb::EXEC_ARROW_ACERO && table_info->schema_conf.use_column_storage()) { + if (options.execute_type == pb::EXEC_ARROW_ACERO && table_info->schema_conf.use_column_storage()) { // db侧判断是否使用列存: - // 1. FLAGS_use_column_storage是否为true // 2. 是否走列式执行 // 3. 表是否支持列存 // store侧判断是否使用列存: @@ -917,6 +934,9 @@ int64_t IndexSelector::index_selector(const std::vector& tu if (_ctx != nullptr && _ctx->efsearch != -1) { sort_property.efsearch = _ctx->efsearch; } + if (_ctx != nullptr && _ctx->nprobe != -1) { + sort_property.nprobe = _ctx->nprobe; + } access_path->calc_index_match(sort_property); if (index_info.type == pb::I_FULLTEXT && index_info.index_hint_status == pb::IHS_NORMAL && access_path->is_possible) { @@ -1103,6 +1123,257 @@ int IndexSelector::select_partition(SmartTable& table_info, ScanNode* scan_node, return 0; } +int IndexSelector::select_partition_files( + SmartTable& table_info, ScanNode* scan_node, FilterNode* filter_node, ExprNode* join_on_conditions) { + if (table_info == nullptr) { + DB_WARNING("table_info is nullptr"); + return -1; + } + FileScanNode* file_scan_node = static_cast(scan_node); + if (file_scan_node == nullptr) { + DB_WARNING("file_scan_node is nullptr"); + return -1; + } + const pb::FileInfo& pb_file_info = table_info->dblink_info.file_info(); + std::shared_ptr fs = create_filesystem(pb_file_info.cluster(), + pb_file_info.username(), + pb_file_info.password(), + AFS_CLIENT_CONF_PATH); + ScopeGuard guard([&fs] () { + destroy_filesystem(fs); + }); + if (fs == nullptr) { + DB_WARNING("Fail to create_filesystem, %s", pb_file_info.ShortDebugString().c_str()); + return -1; + } + std::vector partition_files; + partition_files.reserve(100); + const std::string& table_file_path = pb_file_info.path(); + FileInfo file_info; + int ret = fs->get_file_info(table_file_path, file_info, nullptr); + if (ret < 0) { + DB_WARNING("Fail to get_file_info, path: %s", table_file_path.c_str()); + return -1; + } + // 非分区表 + // - 如果file_info path是文件，则选择该文件； + // - 如果file_info path是目录，则选择该目录下的文件集合（不递归子目录）； + if (pb_file_info.partition_fields().empty()) { + if (file_info.mode == FileMode::I_FILE) { + pb::PartitionFile partition_file; + partition_file.set_file_path(table_file_path); + partition_files.emplace_back(partition_file); + } else if (file_info.mode == FileMode::I_DIR) { + std::vector file_paths; + ret = ReadDirImpl::get_all_files(fs.get(), table_file_path, file_paths); + if (ret < 0) { + DB_WARNING("Fail to get_all_files"); + return -1; + } + for (const auto& file_path : file_paths) { + pb::PartitionFile partition_file; + partition_file.set_file_path(table_file_path + "/" + file_path); + partition_files.emplace_back(partition_file); + } + } else { + DB_WARNING("Invalid mode: %d, %s", (int)file_info.mode, table_file_path.c_str()); + return -1; + } + if (partition_files.empty()) { + file_scan_node->partition_property()->has_no_input_data = true; + } else { + file_scan_node->set_files(partition_files); + } + return 0; + } + // 分区表 + if (file_info.mode != FileMode::I_DIR) { + DB_WARNING("Invalid file mode: %d", (int)file_info.mode); + return -1; + } + int32_t tuple_id = file_scan_node->tuple_id(); + pb::TupleDescriptor* tuple_desc = nullptr; + if (_ctx != nullptr) { + tuple_desc = _ctx->get_tuple_desc(tuple_id); + } + if (tuple_desc == nullptr) { + DB_WARNING("tuple_desc is nullptr"); + return -1; + } + std::vector field_id2slot(table_info->fields.back().id + 1, 0); + for (const auto& slot : tuple_desc->slots()) { + if (slot.field_id() >= field_id2slot.size()) { + DB_WARNING("vector out of range, field_id: %d", slot.field_id()); + continue; + } + field_id2slot[slot.field_id()] = slot.slot_id(); + } + size_t partition_fields_size = pb_file_info.partition_fields().size(); + std::vector partition_slot_ids(partition_fields_size, 0); + std::vector partition_field_types(partition_fields_size); + for (int i = 0; i < partition_fields_size; ++i) { + const auto& partition_field = pb_file_info.partition_fields(i); + FieldInfo* partition_field_info = table_info->get_field_ptr(partition_field); + if (partition_field_info == nullptr) { + DB_WARNING("Fail to get_field_ptr, %s", partition_field.c_str()); + return -1; + } + int32_t partition_field_id = partition_field_info->id; + if (partition_field_id < 0 || partition_field_id >= field_id2slot.size()) { + DB_WARNING("Invalid partition_field_id: %d, %s", partition_field_id, partition_field.c_str()); + return -1; + } + int32_t partition_slot_id = field_id2slot[partition_field_id]; + if (partition_slot_id <= 0) { + continue; + } + partition_slot_ids[i] = partition_slot_id; + partition_field_types[i] = partition_field_info->type; + } + // 获取每个分区字段涉及的条件 + std::unordered_map> slot_conjuncts_map; + ScopeGuard conjunct_guard([&slot_conjuncts_map] () { + for (auto& [_, conjuncts] : slot_conjuncts_map) { + for (auto* conjunct : conjuncts) { + conjunct->close(); + } + } + }); + if (filter_node != nullptr && filter_node->mutable_conjuncts() != nullptr) { + for (auto* conjunct : *filter_node->mutable_conjuncts()) { + if (conjunct != nullptr) { + // TODO - 一个条件涉及多个分区字段不会进行过滤，后续有需要可以升级 + std::set> tuple_slot_ids; + conjunct->get_all_tuple_slot_ids(tuple_slot_ids); + if (tuple_slot_ids.size() == 1) { + const auto& tuple_slot_id = *tuple_slot_ids.begin(); + if (tuple_slot_id.first == tuple_id && + std::find(partition_slot_ids.begin(), partition_slot_ids.end(), tuple_slot_id.second) != partition_slot_ids.end()) { + slot_conjuncts_map[tuple_slot_id.second].emplace_back(conjunct); + } + } + } + } + } + if (join_on_conditions != nullptr) { + std::set> tuple_slot_ids; + join_on_conditions->get_all_tuple_slot_ids(tuple_slot_ids); + if (tuple_slot_ids.size() == 1) { + const auto& tuple_slot_id = *tuple_slot_ids.begin(); + if (tuple_slot_id.first == tuple_id && + std::find(partition_slot_ids.begin(), partition_slot_ids.end(), tuple_slot_id.second) != partition_slot_ids.end()) { + slot_conjuncts_map[tuple_slot_id.second].emplace_back(join_on_conditions); + } + } + } + for (auto& [_, conjuncts] : slot_conjuncts_map) { + for (auto& conjunct : conjuncts) { + ret = conjunct->open(); + if (ret < 0) { + DB_WARNING("expr open fail, ret: %d", ret); + return ret; + } + } + } + // 获取符合条件的分区文件集合 + std::shared_ptr mem_row_desc = std::make_shared(); + std::vector tuple_descs; + tuple_descs.emplace_back(*tuple_desc); + ret = mem_row_desc->init(tuple_descs); + if (ret < 0) { + DB_WARNING("_mem_row_desc init fail"); + return -1; + } + std::vector partition_vals; + partition_vals.reserve(2); + // 存储各个分区的分区值，[["20250610","12"], ["20250611","13"]] + std::vector> partition_vals_vec; + partition_vals_vec.reserve(100); + // 存储各个分区的分区目录，eventday=20250610/hour=12、eventday=20250611/hour=13 + // 长度与partition_vals_vec相等 + std::vector partition_dirs; + partition_dirs.reserve(100); + // @breif 递归处理每一层目录，获取符合条件的分区目录 + // @param dir_path 当前目录 + // @param idx 当前目录层级 + std::function + recurse_calc_partition_dirs = [&] (const std::string& dir_path, const int idx) { + if (idx >= partition_slot_ids.size()) { + partition_vals_vec.emplace_back(partition_vals); + partition_dirs.emplace_back(dir_path); + return 0; + } + std::vector child_dirs; + TimeCost tm; + int ret = ReadDirImpl::get_all_dirs(fs.get(), dir_path, child_dirs); + if (ret < 0) { + DB_WARNING("Fail to get_all_files"); + return -1; + } + int32_t partition_slot_id = partition_slot_ids[idx]; + pb::PrimitiveType field_type = partition_field_types[idx]; + for (const auto& child_dir : child_dirs) { + const std::string& pattern = pb_file_info.partition_fields(idx) + "="; + // 是否以pattern开头 + if (child_dir.find(pattern) == 0) { + const std::string& value = child_dir.substr(pattern.size()); + bool is_valid = true; + if (partition_slot_id > 0 && + slot_conjuncts_map.find(partition_slot_id) != slot_conjuncts_map.end()) { + ExprValue expr_value(field_type, value); + std::unique_ptr row = mem_row_desc->fetch_mem_row(); + row->set_value(tuple_id, partition_slot_id, expr_value); + for (auto* conjunct : slot_conjuncts_map[partition_slot_id]) { + ExprValue conjunct_value = conjunct->get_value(row.get()); + if (conjunct_value.is_null() || conjunct_value.get_numberic() == false) { + is_valid = false; + break; + } + } + } + if (is_valid) { + partition_vals.emplace_back(value); + ret = recurse_calc_partition_dirs(dir_path + "/" + child_dir, idx + 1); + partition_vals.pop_back(); + if (ret < 0) { + DB_WARNING("Fail to calc_partition_files"); + return ret; + } + } + } + } + return 0; + }; + ret = recurse_calc_partition_dirs(table_file_path, 0); + if (ret < 0) { + DB_WARNING("Fail to calc_partition_files"); + return -1; + } + for (int i = 0; i < partition_dirs.size(); ++i) { + std::vector file_paths; + TimeCost tm; + ret = ReadDirImpl::get_all_files(fs.get(), partition_dirs[i], file_paths); + if (ret < 0) { + DB_WARNING("Fail to get_all_files"); + return -1; + } + for (const auto& file_path : file_paths) { + pb::PartitionFile partition_file; + for (const auto& partition_val : partition_vals_vec[i]) { + partition_file.add_partition_vals(partition_val); + } + partition_file.set_file_path(partition_dirs[i] + "/" + file_path); + partition_files.emplace_back(partition_file); + } + } + if (partition_files.empty()) { + file_scan_node->partition_property()->has_no_input_data = true; + } else { + file_scan_node->set_files(partition_files); + } + return 0; +} + // 检查ROLLUP索引是否要选择 bool IndexSelector::check_rollup_index_valid(SmartTable& table_info, const IndexInfo& index_info, diff --git a/src/physical_plan/limit_calc.cpp b/src/physical_plan/limit_calc.cpp index f45fa4178..5a21e424d 100644 --- a/src/physical_plan/limit_calc.cpp +++ b/src/physical_plan/limit_calc.cpp @@ -46,6 +46,10 @@ void LimitCalc::_analyze_limit(QueryContext* ctx, ExecNode* node, int64_t limit) // 没有agg_fn时, 在agg_node的open阶段可以使用limit。 case pb::MERGE_AGG_NODE: if (static_cast(node)->mutable_agg_fn_calls()->empty()) { + if (node->children(0)->node_type() == pb::SELECT_MANAGER_NODE) { + //跳过SELECT_MANAGER_NODE，否则select distinct f from test limit 10;有bug，会少返回数据 + node = node->children(0); + } break; } else { return; diff --git a/src/physical_plan/mpp_analyzer.cpp b/src/physical_plan/mpp_analyzer.cpp index 57ba969e0..9b9bccbe5 100644 --- a/src/physical_plan/mpp_analyzer.cpp +++ b/src/physical_plan/mpp_analyzer.cpp @@ -19,9 +19,12 @@ #include "vectorize_helpper.h" #include "filter_node.h" #include "agg_node.h" +#include "file_scan_node.h" +#include "db_service.h" namespace baikaldb { DECLARE_int32(mpp_hash_partition_num); +DECLARE_int32(file_scan_concurrency); int MppAnalyzer::analyze(QueryContext* ctx) { if (ctx->get_runtime_state()->execute_type == pb::EXEC_ARROW_ACERO) { @@ -30,10 +33,6 @@ int MppAnalyzer::analyze(QueryContext* ctx) { if (!ctx->use_mpp) { return 0; } - // 包含DBLink Mysql表的查询不支持MPP执行 - if (ctx->has_dblink_mysql) { - return 0; - } // 设置每个节点的partition_property和schema if (0 != ctx->root->set_partition_property_and_schema(ctx)) { return -1; @@ -168,6 +167,8 @@ int MppAnalyzer::create_exchange_node_pair(QueryContext* ctx, pb_es_node->set_log_id(ctx->stat_info.log_id); pb_es_node->mutable_partition_property()->set_type(parent->type); if (use_broadcast_shuffle) { + er_node->partition_property()->type = pb::BroadcastPartitionType; + pb_er_node->mutable_partition_property()->set_type(pb::BroadcastPartitionType); pb_es_node->mutable_partition_property()->set_type(pb::BroadcastPartitionType); } ExchangeSenderNode* es_node = new (std::nothrow) ExchangeSenderNode; @@ -227,13 +228,16 @@ int MppAnalyzer::add_exchange_and_separate_for_index_join_inner_node(QueryContex NodePartitionProperty* parent_property, std::set& need_runtime_filter_select_manager_nodes, bool need_use_broadcast_exchange) { - DB_DEBUG("node: %s, parent_property(%p): %s, node_property(%p): %s, need_runtime_filter_select_manager_nodes size: %ld", + DB_DEBUG("node: %s, parent_property(%p): %s, node_property(%p): %s, " + "need_runtime_filter_select_manager_nodes size: %ld, need_use_broadcast_exchange: %d", pb::PlanNodeType_Name(node->node_type()).c_str(), fragment->partition_property, parent_property->print().c_str(), node->partition_property(), node->partition_property()->print().c_str(), - need_runtime_filter_select_manager_nodes.size()); + need_runtime_filter_select_manager_nodes.size(), + need_use_broadcast_exchange); + if (node->partition_property()->has_no_input_data) { return 0; } @@ -460,6 +464,7 @@ int MppAnalyzer::seperate_store_fragment(QueryContext* ctx, ExecNode* select_manager, NodePartitionProperty* parent_property, bool need_use_broadcast_exchange) { + ScanNode* scan_node = static_cast(select_manager->get_node(pb::SCAN_NODE)); // parent -> SelectManager -> store (xxx) ExchangeReceiverNode* receiver = nullptr; ExchangeSenderNode* sender = nullptr; @@ -475,6 +480,8 @@ int MppAnalyzer::seperate_store_fragment(QueryContext* ctx, store_fragment->root = sender; store_fragment->parent = fragment.get(); store_fragment->partition_property = sender->partition_property(); + store_fragment->runtime_state = fragment->last_runtime_state; + store_fragment->last_runtime_state = fragment->last_runtime_state; fragment->receivers.emplace_back(receiver); fragment->children.emplace_back(store_fragment); @@ -492,11 +499,7 @@ int MppAnalyzer::seperate_store_fragment(QueryContext* ctx, receiver->set_sender_fragment_id(store_fragment->fragment_id); sender->set_receiver_fragment_id(fragment->fragment_id); - // sortnode - SortNode* sort_node = static_cast(sender->get_node(pb::SORT_NODE)); - if (sort_node != nullptr) { - receiver->init_sort_info(sort_node); - } + receiver->init_condition_and_sort_info(select_manager); // 如果使用全局索引, 在这里串行访问全局索引, 生成主表的region信息, 供exchange使用, 否则exchange获取到的是全局索引的region信息 // TODO 是否有更好的方法 @@ -577,14 +580,26 @@ int MppAnalyzer::check_need_add_exchange(NodePartitionProperty* parent_propety, *need_add_exchange = true; return 0; } + if (parent_propety->hash_partition_propertys.empty()) { + DB_FATAL("parent hash_partition_propertys is empty, but is HashPartitionType"); + return -1; + } for (auto& hash_columns : node_property->hash_partition_propertys) { - if (parent_propety->hash_partition_propertys[0]->hash_partition_is_same(hash_columns.get())) { + if (hash_columns.get() == parent_propety->hash_partition_propertys[0].get() + || (hash_columns->type == pb::HashPartitionType + && parent_propety->hash_partition_propertys[0]->hash_partition_is_same(hash_columns.get()))) { *need_add_exchange = false; return 0; } } *need_add_exchange = true; - break; + return 0; + case pb::BroadcastPartitionType: + *need_add_exchange = (node_property->type != pb::BroadcastPartitionType); + return 0; + case pb::RandomPartitionType: + *need_add_exchange = (node_property->type != pb::RandomPartitionType); + return 0; default: DB_FATAL("not support partition type: %d", node_property->type); return -1; @@ -681,8 +696,8 @@ int MppAnalyzer::build_fragment(QueryContext* ctx, SmartFragment& fragment) { int fragment_id = fragment->fragment_id; ctx->fragments[fragment_id] = fragment; if (fragment->receivers.size() == 0) { - // store fragment - return 0; + // scan fragment + return build_scan_fragment(ctx, fragment); } if (fragment->root->node_type() == pb::PACKET_NODE) { // 需要特殊处理fragment0, 因为root是pack_node, 不是exchange sender @@ -702,7 +717,8 @@ int MppAnalyzer::build_fragment(QueryContext* ctx, SmartFragment& fragment) { ExchangeSenderNode* sender = static_cast(fragment->root); // 先添加主db sender->add_address(SchemaFactory::get_instance()->get_address()); - if (fragment->partition_property->type == pb::HashPartitionType) { + if (fragment->partition_property->type == pb::HashPartitionType + || fragment->partition_property->type == pb::RandomPartitionType) { // rolling N-1 db int mpp_hash_partition_num = FLAGS_mpp_hash_partition_num; if (ctx->mpp_hash_num > 0) { @@ -726,7 +742,9 @@ int MppAnalyzer::build_fragment(QueryContext* ctx, SmartFragment& fragment) { SmartFragment& next_fragment = fragment->children[i]; ExchangeSenderNode* next_fragment_sender = static_cast(next_fragment->root); if (fragment->partition_property->type == pb::SinglePartitionType - || fragment->partition_property->type == pb::HashPartitionType) { + || fragment->partition_property->type == pb::HashPartitionType + || fragment->partition_property->type == pb::BroadcastPartitionType + || fragment->partition_property->type == pb::RandomPartitionType) { // 本fragment sender对应的下游sender, 确定了本fragment分布地址, 就可以指定下游sender的destination next_fragment_sender->set_destination(sender->get_fragment_address()); } else { @@ -739,6 +757,91 @@ int MppAnalyzer::build_fragment(QueryContext* ctx, SmartFragment& fragment) { } return 0; } + +// 生成db侧执行scan的fragment +int MppAnalyzer::build_scan_fragment(QueryContext* ctx, SmartFragment& fragment) { + if (ctx == nullptr) { + DB_WARNING("ctx is nullptr"); + return -1; + } + if (fragment == nullptr) { + DB_WARNING("fragment is nullptr"); + return -1; + } + ExchangeSenderNode* sender = static_cast(fragment->root); + if (sender == nullptr) { + DB_WARNING("sender is nullptr"); + return -1; + } + ScanNode* scan_node = static_cast(sender->get_node(pb::SCAN_NODE)); + if (scan_node == nullptr) { + return 0; + } + if (scan_node->partition_property()->has_no_input_data) { + return 0; + } + ExchangeReceiverNode* receiver = static_cast(sender->get_exchange_receiver_node()); + if (receiver == nullptr) { + DB_WARNING("receiver is nullptr"); + return -1; + } + SelectManagerNode* select_manager = + static_cast(receiver->get_relate_select_manager_node()); + if (select_manager == nullptr) { + DB_WARNING("select_manager is nullptr"); + return -1; + } + if (scan_node->is_file_scan_node()) { + // 将离线文件平均分配到多个实例执行 + FileScanNode* file_scan_node = static_cast(scan_node); + const auto& partition_files = file_scan_node->get_files(); + if (partition_files.empty()) { + return 0; + } + std::vector db_instances; + int32_t file_scan_concurrency = + std::min(FLAGS_file_scan_concurrency, static_cast(partition_files.size())); + if (file_scan_concurrency > 1 && + SchemaFactory::get_instance()->rolling_pick_db(file_scan_concurrency - 1, db_instances) != 0) { + DB_WARNING("pick db instances failed"); + return -1; + } + db_instances.emplace_back(SchemaFactory::get_instance()->get_address()); + std::map> db_partition_files; + for (int i = 0; i < partition_files.size(); ++i) { + const std::string& db_instance = db_instances[i % db_instances.size()]; + db_partition_files[db_instance].emplace_back(partition_files[i]); + } + std::map db_request_map; + for (const auto& [ith_db, files] : db_partition_files) { + sender->add_address(ith_db); + file_scan_node->set_files(files); + pb::DAGFragmentRequest request; + int ret = DBInteract::get_instance()->construct_mpp_dag_request(ctx, {fragment->fragment_id}, request); + if (ret != 0) { + DB_WARNING("construct mpp dag request failed"); + return -1; + } + db_request_map[ith_db] = request; + } + receiver->set_destination(sender->get_fragment_address()); + select_manager->set_db_request_map(db_request_map); + } else if (scan_node->is_mysql_scan_node()) { + sender->add_address(SchemaFactory::get_instance()->get_address()); + receiver->set_destination(sender->get_fragment_address()); + std::map db_request_map; + pb::DAGFragmentRequest request; + int ret = DBInteract::get_instance()->construct_mpp_dag_request(ctx, {fragment->fragment_id}, request); + if (ret != 0) { + DB_WARNING("construct mpp dag request failed"); + return -1; + } + db_request_map[SchemaFactory::get_instance()->get_address()] = request; + select_manager->set_db_request_map(db_request_map); + } + return 0; +} + } /* vim: set ts=4 sw=4 sts=4 tw=100 */ diff --git a/src/physical_plan/physical_planner.cpp b/src/physical_plan/physical_planner.cpp index cd78b258b..e64558750 100644 --- a/src/physical_plan/physical_planner.cpp +++ b/src/physical_plan/physical_planner.cpp @@ -118,6 +118,10 @@ int PhysicalPlanner::analyze(QueryContext* ctx) { if (ctx->return_empty) { ctx->root->set_return_empty(); } + ret = ConditionOptimizer().analyze(ctx); + if (ret < 0) { + return ret; + } // mpp添加Exchange等 ret = MppAnalyzer().analyze(ctx); if (ret < 0) { @@ -279,33 +283,19 @@ int PhysicalPlanner::send_fragment_to_other_db(QueryContext* ctx) { int64_t send_time = 0; for (auto iter : ctx->db_to_fragments) { const std::string& ith_db = iter.first; - if (ith_db == SchemaFactory::get_instance()->get_address()) { continue; } time_cost.reset(); pb::DAGFragmentRequest request; - request.set_op(pb::OP_FRAGMENT_START); - request.set_log_id(cur_state.log_id()); - request.set_sql_sign(cur_state.sign); - // add fragment infos - for (auto& fragment_id : iter.second) { - pb::Plan fragment_plan; - std::shared_ptr fragment_info = ctx->fragments[fragment_id]; - pb::FragmentInfo* fragment_info_pb = request.add_fragments(); - fragment_info_pb->set_fragment_id(fragment_id); - pb::RuntimeState* pb_rs = fragment_info_pb->mutable_runtime_state(); - fragment_info->runtime_state->to_proto(pb_rs); - ExecNode::create_pb_plan(0, &fragment_plan, fragment_info->root); - fragment_info_pb->mutable_plan()->CopyFrom(fragment_plan); - } - // add userinfo, 副db planrouter需要 - if (ctx->user_info != nullptr) { - request.set_username(ctx->user_info->username); + int ret = DBInteract::get_instance()->construct_mpp_dag_request(ctx, iter.second, request); + if (ret != 0) { + DB_WARNING("Fail to construct_mpp_dag_fragment"); + return -1; } pack_time += time_cost.get_time(); time_cost.reset(); - int ret = DBInteract::get_instance()->handle_mpp_dag_fragment(request, ith_db); + ret = DBInteract::get_instance()->handle_mpp_dag_fragment(request, ith_db); if (0 != ret) { DB_FATAL("logid: %lu, send request to other db fail, ret: %d", cur_state.log_id(), ret); return -1; @@ -322,14 +312,20 @@ int PhysicalPlanner::send_fragment_to_other_db(QueryContext* ctx) { // 只有mpp主db执行失败/超时调用，停止所有db的fragment执行 int PhysicalPlanner::stop_mpp(QueryContext* ctx) { RuntimeState& cur_state = *ctx->get_runtime_state(); - for (auto iter : ctx->db_to_fragments) { + std::vector db_list; + db_list.reserve(ctx->db_to_fragments.size() + 1); + db_list.emplace_back(SchemaFactory::get_instance()->get_address()); + for (auto& iter : ctx->db_to_fragments) { + db_list.emplace_back(iter.first); + } + for (auto& db : db_list) { pb::DAGFragmentRequest request; request.set_op(pb::OP_FRAGMENT_STOP); request.set_log_id(cur_state.log_id()); if (ctx->user_info != nullptr) { request.set_username(ctx->user_info->username); } - int ret = DBInteract::get_instance()->handle_mpp_dag_fragment(request, iter.first); + int ret = DBInteract::get_instance()->handle_mpp_dag_fragment(request, db); if (0 != ret) { DB_FATAL("send request to other db for stop fail, ret: %d", ret); return -1; @@ -368,6 +364,18 @@ int PhysicalPlanner::execute(QueryContext* ctx, DataBuffer* send_buf) { state.statistics_types->insert(pb::StatisticType::ST_HISTOGRAM); state.statistics_types->insert(pb::StatisticType::ST_CMSKETCH); state.statistics_types->insert(pb::StatisticType::ST_HYPERLOGLOG); + } else if (boost::istarts_with(type, "max_sample_regions")) { + auto iter = type.find(":"); + std::string max_region_string; + if (iter != std::string::npos) { + max_region_string = type.substr(iter + 1); + } + int64_t max_sample_regions = 0; + if (boost::conversion::try_lexical_convert(max_region_string, max_sample_regions)) { + state.max_sample_regions = max_sample_regions; + } else { + DB_WARNING("convert max_sample_region to int failed"); + } } } // 兜底，如果state.statistics_types 为空就默认选hist和cms @@ -399,12 +407,20 @@ int PhysicalPlanner::execute(QueryContext* ctx, DataBuffer* send_buf) { } } if (explain_is_trace(ctx->explain_type)) { - ctx->trace_node.set_node_type(ctx->root->node_type()); - ctx->root->set_trace(&ctx->trace_node); - ctx->root->create_trace(); - for (auto& subquery : ctx->sub_query_plans) { - subquery->root->set_trace(&ctx->trace_node); - subquery->root->create_trace(); + if (ctx->use_mpp) { + for (auto& [id, fragment] : ctx->fragments) { + fragment->trace_node.set_node_type(fragment->root->node_type()); + fragment->root->set_trace(&fragment->trace_node); + fragment->root->create_trace(); + } + } else { + ctx->trace_node.set_node_type(ctx->root->node_type()); + ctx->root->set_trace(&ctx->trace_node); + ctx->root->create_trace(); + for (auto& subquery : ctx->sub_query_plans) { + subquery->root->set_trace(&ctx->trace_node); + subquery->root->create_trace(); + } } } if (ctx->use_mpp) { diff --git a/src/physical_plan/plan_router.cpp b/src/physical_plan/plan_router.cpp index 43cf2636d..20fb7191f 100644 --- a/src/physical_plan/plan_router.cpp +++ b/src/physical_plan/plan_router.cpp @@ -148,8 +148,15 @@ int PlanRouter::insert_node_analyze(T* node, QueryContext* ctx) { return 0; } -int PlanRouter::scan_node_analyze(RocksdbScanNode* scan_node, QueryContext* ctx, bool has_join, +int PlanRouter::scan_node_analyze(ScanNode* scan_node, QueryContext* ctx, bool has_join, const std::set& escape_get_region_infos) { + if (scan_node == nullptr) { + DB_WARNING("scan_node is null"); + return -1; + } + if (!scan_node->is_rocksdb_scan_node()) { + return 0; + } SchemaFactory* schema_factory = SchemaFactory::get_instance(); if (ctx->debug_region_id != -1) { pb::RegionInfo info; @@ -190,11 +197,18 @@ int PlanRouter::scan_node_analyze(RocksdbScanNode* scan_node, QueryContext* ctx, return scan_plan_router(scan_node, get_slot_id, get_tuple_desc, has_join, escape_get_region_infos); } -int PlanRouter::scan_plan_router(RocksdbScanNode* scan_node, +int PlanRouter::scan_plan_router(ScanNode* scan_node, const std::function& get_slot_id, const std::function& get_tuple_desc, bool has_join, const std::set& escape_get_region_infos) { + if (scan_node == nullptr) { + DB_WARNING("scan_node is null"); + return -1; + } + if (!scan_node->is_rocksdb_scan_node()) { + return 0; + } //pb::ScanNode* pb_scan_node = scan_node->mutable_pb_node()->mutable_derive_node()->mutable_scan_node(); int64_t main_table_id = scan_node->table_id(); SchemaFactory* schema_factory = SchemaFactory::get_instance(); diff --git a/src/physical_plan/separate.cpp b/src/physical_plan/separate.cpp index b085efa96..4892dc43d 100644 --- a/src/physical_plan/separate.cpp +++ b/src/physical_plan/separate.cpp @@ -320,7 +320,8 @@ int Separate::separate_join(QueryContext* ctx, const std::vector& joi for (auto& scan_node_ptr : scan_nodes) { // INFORMATION_SCHEMA if (!static_cast(scan_node_ptr)->is_rocksdb_scan_node() && - !static_cast(scan_node_ptr)->is_mysql_scan_node()) { + !static_cast(scan_node_ptr)->is_mysql_scan_node() && + !static_cast(scan_node_ptr)->is_file_scan_node()) { continue; } bool is_rocksdb_scan_node = static_cast(scan_node_ptr)->is_rocksdb_scan_node(); diff --git a/src/protocol/db_service.cpp b/src/protocol/db_service.cpp index cf9f5555f..816f8e26c 100644 --- a/src/protocol/db_service.cpp +++ b/src/protocol/db_service.cpp @@ -5,6 +5,8 @@ #include #include #include +#include +#include #include "dual_scan_node.h" #include "arrow_io_excutor.h" #include "network_socket.h" @@ -302,27 +304,40 @@ int DbService::handle_fragment_scan_nodes(RuntimeState* state, std::vector(node); - // 不clear select_index_common 会core - scan_node->mutable_pb_node()->mutable_derive_node()->mutable_scan_node()->clear_use_indexes(); - - // 关联scan node和对应的select manager node - auto sm_node = static_cast(scan_node->get_parent_node(pb::SELECT_MANAGER_NODE)); - if (sm_node == nullptr) { - DB_FATAL("mpp execute fail: fragment get select manager node fail, log_id: %lu, tuple_id: %d", state->log_id(), scan_node->tuple_id()); + // 3. dblink scan_node场景 + ScanNode* scan_node = static_cast(node); + if (scan_node == nullptr) { + DB_WARNING("scan_node is nullptr"); return -1; } - static_cast(scan_node)->set_related_manager_node(sm_node); + auto sm_node = static_cast(scan_node->get_parent_node(pb::SELECT_MANAGER_NODE)); + if (sm_node != nullptr) { + // 不clear select_index_common 会core + scan_node->mutable_pb_node()->mutable_derive_node()->mutable_scan_node()->clear_use_indexes(); - // index selector and plan router - DB_DEBUG("scannode tuple_id: %d, need planrouter", scan_node->tuple_id()); - bool index_has_null = false; - if (0 != Joiner::do_plan_router(state, {scan_node}, index_has_null, false)) { - DB_FATAL("mpp execute fail: fragment plan router fail, log_id: %lu, tuple_id: %d", state->log_id(), scan_node->tuple_id()); - return -1; - } - if (index_has_null) { - sm_node->set_return_empty(); + // 关联scan node和对应的select manager node + if (sm_node == nullptr) { + DB_FATAL("mpp fragment get select manager node fail, log_id: %lu, tuple_id: %d", state->log_id(), scan_node->tuple_id()); + return -1; + } + static_cast(scan_node)->set_related_manager_node(sm_node); + + // index selector and plan router + DB_DEBUG("scannode tuple_id: %d, need planrouter", scan_node->tuple_id()); + bool index_has_null = false; + if (0 != Joiner::do_plan_router(state, {scan_node}, index_has_null, false, false)) { + DB_FATAL("mpp execute fail: fragment plan router fail, log_id: %lu, tuple_id: %d", state->log_id(), scan_node->tuple_id()); + return -1; + } + if (index_has_null) { + sm_node->set_return_empty(); + } + } else { + // 发送到副db的file_scan_node/mysql_scan_node可能没有select manager node，用于执行扫描。 + if (!scan_node->is_file_scan_node() && !scan_node->is_mysql_scan_node()) { + DB_WARNING("scan node has no select manager node, log_id: %lu, tuple_id: %d", state->log_id(), scan_node->tuple_id()); + return -1; + } } } return 0; @@ -414,7 +429,6 @@ int DbService::fragment_internal_open(google::protobuf::RpcController* controlle return -1; } } - ret = root->open(&state); if (ret < 0) { err_msg = "open plan fail"; diff --git a/src/protocol/handle_helper.cpp b/src/protocol/handle_helper.cpp index 96606c5ba..5626991e9 100644 --- a/src/protocol/handle_helper.cpp +++ b/src/protocol/handle_helper.cpp @@ -719,7 +719,7 @@ bool HandleHelper::_handle_split_lines(const SmartSocket& client, const std::vec } bool HandleHelper::_handle_ttl_duration(const SmartSocket& client, const std::vector& split_vec) { - if(!client || !client->user_info || split_vec.size() != 4) { + if(!client || !client->user_info || (split_vec.size() != 4 && split_vec.size() != 6)) { DB_FATAL("param invalid"); return false; } @@ -736,6 +736,10 @@ bool HandleHelper::_handle_ttl_duration(const SmartSocket& client, const std::ve return false; } info->set_ttl_duration(ttl); + if (split_vec.size() == 6) { + pb::FieldInfo* ttl_field = info->mutable_ttl_field(); + ttl_field->set_field_name(split_vec[5]); + } MetaServerInteract::get_instance()->send_request("meta_manager", request, response); DB_WARNING("req:%s res:%s", request.ShortDebugString().c_str(), response.ShortDebugString().c_str()); if(!_make_response_packet(client, response.ShortDebugString())) { @@ -1837,7 +1841,24 @@ bool HandleHelper::_handle_schema_conf(const SmartSocket& client, const std::vec schema_conf->set_use_column_storage(is_open); } else if (key == "force_column_storage") { schema_conf->set_force_column_storage(is_open); + } else if (key == "column_only_read_base") { + schema_conf->set_column_only_read_base(is_open); } else if (key == "enable_column_engine") { + if (is_open) { + auto table = factory->get_table_info_ptr(table_id); + if (table == nullptr) { + DB_FATAL("table null table name: %s, table_id: %ld", full_name.c_str(), table_id); + client->state = STATE_ERROR; + return false; + } + if (!(table->schema_conf.storage_compute_separate())) { + if (!table->has_version) { + DB_FATAL("table name: %s, table_id: %ld need use column storage", full_name.c_str(), table_id); + client->state = STATE_ERROR; + return false; + } + } + } schema_conf->set_enable_column_engine(is_open); } else if (key == "olap_pre_split_cnt") { int32_t olap_pre_split_cnt = strtol(split_vec[4].c_str(), NULL, 10); diff --git a/src/protocol/main.cpp b/src/protocol/main.cpp index 496e46356..bda56743c 100644 --- a/src/protocol/main.cpp +++ b/src/protocol/main.cpp @@ -53,7 +53,7 @@ void crash_handler(int sig) { int main(int argc, char **argv) { // Initail signal handlers. signal(SIGPIPE, SIG_IGN); - signal(SIGSEGV, (sighandler_t)baikaldb::crash_handler); + //signal(SIGSEGV, (sighandler_t)baikaldb::crash_handler); signal(SIGINT, (sighandler_t)baikaldb::handle_exit_signal); signal(SIGTERM, (sighandler_t)baikaldb::handle_exit_signal); #ifdef BAIKALDB_REVISION @@ -175,6 +175,7 @@ int main(int argc, char **argv) { server->stop(); baikaldb::MemoryGCHandler::get_instance()->close(); baikaldb::MemTrackerPool::get_instance()->close(); + baikaldb::GlobalArrowExecutor::shutdown(); dummy_server.Stop(0); dummy_server.Join(); diff --git a/src/protocol/network_server.cpp b/src/protocol/network_server.cpp index 8527b8bbb..16b050213 100644 --- a/src/protocol/network_server.cpp +++ b/src/protocol/network_server.cpp @@ -54,19 +54,19 @@ DEFINE_bool(fetch_instance_id, false, "fetch baikaldb instace id, used for gener DEFINE_string(hostname, "HOSTNAME", "matrix instance name"); DEFINE_bool(insert_agg_sql, false, "whether insert agg_sql"); DEFINE_int32(batch_insert_agg_sql_size, 50, "batch size for insert"); -DEFINE_int32(batch_insert_sign_sql_interval_us, 10 * 60 * 1000 * 1000, "batch_insert_sign_sql_interval_us default 10min"); DEFINE_bool(enable_tcp_keep_alive, false, "enable tcp keepalive flag"); DECLARE_int32(baikal_heartbeat_interval_us); DEFINE_bool(open_to_collect_slow_query_infos, false, "open to collect slow_query_infos, default: false"); DEFINE_uint64(limit_slow_sql_size, 50, "each sign to slow query sql counts, default: 50"); DEFINE_int32(slow_query_batch_size, 100, "slow query sql batch size, default: 100"); DECLARE_bool(auto_update_meta_list); -DEFINE_string(afs_gc_hostname, "", "afs_gc_hostname"); +DEFINE_string(afs_gc_hostname, "", "AFS garbage collection hostname, default: empty"); DECLARE_string(baikal_resource_tag); DEFINE_bool(dump_slow_sqls, false, "whether dump slow sqls"); DEFINE_bool(auto_kill_timeout_query, false, "auto kill timeout query"); DEFINE_uint64(load_mpp_sign_interval_min, 0, "load mpp sign interval min, default: 0, disable load mpp sign"); DECLARE_string(baikal_resource_tag); +DECLARE_string(log_plat_name); static const std::string instance_table_name = "INTERNAL.baikaldb.__baikaldb_instance"; @@ -555,10 +555,12 @@ void NetworkServer::print_agg_sql() { static std::map> parent_sign_to_subquery_signs; static std::set sign_to_counts; static std::set parent_sign_to_counts; + static std::unordered_map special_signs; TimeCost cost; TimeCost reset_counter_cost; TimeCost degrade_cost; TimeCost load_mpp_signs_time; + TimeCost handle_special_sings_time; while (!_shutdown) { bool need_reset_counter = false; if (reset_counter_cost.get_time() > 24 * 3600 * 1000 * 1000LL) { @@ -599,7 +601,7 @@ void NetworkServer::print_agg_sql() { time_t timep; struct tm tm; time(&timep); - localtime_r(&timep, &tm); + localtime_fixed_r(&timep, &tm); struct CountErr { int64_t count = 0; @@ -622,6 +624,42 @@ void NetworkServer::print_agg_sql() { sql_values.reserve(4096); } + if (handle_special_sings_time.get_time() > 600 * 1000 * 1000ULL) { + handle_special_sings_time.reset(); + auto func = [](const SmartTable& table) -> bool { + for (auto sign : table->sign_blacklist) { + special_signs.insert({sign, ""}); + } + for (auto sign : table->sign_forcelearner) { + special_signs.insert({sign, ""}); + } + for (auto sign : table->sign_rolling) { + special_signs.insert({sign, ""}); + } + for (auto& sign_index : table->sign_forceindex) { + std::vector vec; + boost::split(vec, sign_index, boost::is_any_of(":")); + if (vec.size() != 2) { + continue; + } + uint64_t sign_num = strtoull(vec[0].c_str(), nullptr, 10); + special_signs.insert({sign_num, ""}); + } + for (auto& sign_index : table->sign_exec_type) { + std::vector vec; + boost::split(vec, sign_index, boost::is_any_of(":")); + if (vec.size() != 2) { + continue; + } + uint64_t sign_num = strtoull(vec[0].c_str(), nullptr, 10); + special_signs.insert({sign_num, ""}); + } + return false; + }; + std::vector database_table; + factory->get_table_by_filter(database_table, func); + } + for (auto& pair : sample.internal_map) { if (pair.first.empty()) { continue; @@ -725,6 +763,15 @@ void NetworkServer::print_agg_sql() { std::string plat = sql_agg.substr(pos, sql_agg.find_first_of(']', pos) - pos); pos = sql_agg.find_first_of('[') + 1; std::string sql_text = sql_agg.substr(pos, sql_agg.find_first_of(']', pos) - pos); + resource_tag = pair2.second.resource_tag; + + if (special_signs.count(out_sign) > 0) { + std::string sample_sql_without_sign = "family_table_tag_optype_plat=[" + family + "\t" + + tbl + "\t" + "" + "\t" + op_type + "\t" + + FLAGS_log_plat_name + "] sql=[" + sql_text + "]"; + special_signs[out_sign] = sample_sql_without_sign; + } + sql_text = boost::replace_all_copy(sql_text, "'", "\\'"); // 避免REPLACE INTO的执行发生递归 if (family == "BaikalStat" && tbl == "baikaldb_trace_info" && @@ -791,6 +838,7 @@ void NetworkServer::print_agg_sql() { DB_WARNING("insert agg_sql_by_sign: %s failed", sql_values.c_str()); } } + SchemaFactory::get_instance()->set_special_signs(special_signs); for (auto& pair : table_count_err) { // 10s pv>50 出错率>0.5则读路由去备库 if (pair.second.count > FLAGS_backup_pv_threshold && @@ -1337,8 +1385,8 @@ void NetworkServer::process_slow_query_map() { struct tm local_end_time; time_t start_time = slow_query_info.start_time; time_t end_time = slow_query_info.end_time; - localtime_r(&start_time, &local_start_time); - localtime_r(&end_time, &local_end_time); + localtime_fixed_r(&start_time, &local_start_time); + localtime_fixed_r(&end_time, &local_end_time); //localtime => format("%Y-%m-%d %H:%M:%S") char start_time_str[40], end_time_str[40]; snprintf(start_time_str, 40, "%04d-%02d-%02d %02d:%02d:%02d:%03d", @@ -1744,7 +1792,7 @@ int NetworkServer::make_worker_process() { } _conn_bvars_update_bth.run([this](){client_conn_bvars_update();}); -#ifdef BAIDU_INTERNAL +#if defined(BAIDU_INTERNAL) _ext_fs_gc_bth.run([this]() { while (!_shutdown) { if (FLAGS_afs_gc_hostname == FLAGS_hostname) { @@ -1984,6 +2032,10 @@ void NetworkServer::client_conn_bvars_update() { int32_t client_cnt = 0; int32_t max_running_time = 0; EpollInfo* epoll_info = NetworkServer::get_instance()->get_epoll_info(); + if (epoll_info == NULL) { + DB_WARNING("epoll_info not initialized yet."); + continue; + } for (int32_t idx = 0; idx < CONFIG_MPL_EPOLL_MAX_SIZE; ++idx) { const SmartSocket& sock = epoll_info->get_fd_mapping(idx); if (sock == NULL || sock->is_free || sock->fd == -1 || sock->ip == "") { diff --git a/src/protocol/show_helper.cpp b/src/protocol/show_helper.cpp index 5d925d107..2050716cd 100644 --- a/src/protocol/show_helper.cpp +++ b/src/protocol/show_helper.cpp @@ -30,6 +30,8 @@ void ShowHelper::init() { this, std::placeholders::_1, std::placeholders::_2); _calls[SQL_SHOW_DATABASES] = std::bind(&ShowHelper::_show_databases, this, std::placeholders::_1, std::placeholders::_2); + _calls[SQL_SHOW_SCHEMAS] = std::bind(&ShowHelper::_show_databases, + this, std::placeholders::_1, std::placeholders::_2); _calls[SQL_SHOW_TABLES] = std::bind(&ShowHelper::_show_tables, this, std::placeholders::_1, std::placeholders::_2); _calls[SQL_SHOW_CREATE_TABLE] = std::bind(&ShowHelper::_show_create_table, @@ -972,8 +974,28 @@ bool ShowHelper::_show_create_table(const SmartSocket& client, const std::vector } oss << (field.auto_inc ? "AUTO_INCREMENT " : ""); oss << (field.is_unique_indicator ? "UNIQUE " : ""); + auto comment_escape_string = [] (const std::string& origin_comment) { + // 简单转义规则，按照mysql，'为''其他为backslash转义 + static std::unordered_map escape_map { + {'\'', "''"}, + {'\\', "\\\\"}, + {'\n', "\\n"}, + {'\0', "\\0"} + }; + std::string escape_comment; + escape_comment.reserve(origin_comment.size() + 10); + for (auto c: origin_comment) { + if (escape_map.count(c) != 0) { + escape_comment.append(escape_map[c]); + } else { + escape_comment.push_back(c); + } + } + return escape_comment; + }; + if (!field.comment.empty()) { - oss << "COMMENT '" << field.comment << "'"; + oss << "COMMENT '" << comment_escape_string(field.comment) << "'"; } oss << ",\n"; } @@ -1085,7 +1107,7 @@ bool ShowHelper::_show_create_table(const SmartSocket& client, const std::vector oss << " DEFAULT CHARSET=" << charset_map[info.charset]; oss <<" AVG_ROW_LENGTH=" << info.byte_size_per_record; } - oss << " COMMENT='{\"resource_tag\":\"" << info.resource_tag << "\""; + oss << R"( COMMENT='{"resource_tag":")" << info.resource_tag << "\""; if (!info.comment.empty()) { oss << ", \"comment\":\"" << info.comment << "\""; } @@ -1094,7 +1116,11 @@ bool ShowHelper::_show_create_table(const SmartSocket& client, const std::vector oss << ", \"region_split_lines\":" << info.region_split_lines; } if (info.ttl_info.ttl_duration_s > 0) { - oss << ", \"ttl_duration\":" << info.ttl_info.ttl_duration_s; + std::string ttl_field_info = ""; + if (info.ttl_info.ttl_field != nullptr) { + ttl_field_info = "\"field\": \"" + info.ttl_info.ttl_field->short_name + "\", "; + } + oss << ", \"ttl\":{" << ttl_field_info << "\"duration\": " << info.ttl_info.ttl_duration_s << "}"; } if (info.learner_resource_tags.size() > 0) { oss << ", \"learner_resource_tag\": ["; @@ -1190,6 +1216,43 @@ bool ShowHelper::_show_create_table(const SmartSocket& client, const std::vector oss << ",\"table_name\":\"" << info.dblink_info.mysql_info().table_name() << "\""; oss << ",\"charset\":\"" << info.dblink_info.mysql_info().charset() << "\""; oss << "}"; + } else if (info.dblink_info.has_file_info()) { + oss << ", \"file_info\": {"; + oss << "\"cluster\":\"" << info.dblink_info.file_info().cluster() << "\""; + oss << ",\"path\":\"" << info.dblink_info.file_info().path() << "\""; + oss << ",\"username\":\"" << info.dblink_info.file_info().username() << "\""; + // oss << ",\"password\":\"" << info.dblink_info.file_info().password() << "\""; // 隐去password + if (info.dblink_info.file_info().partition_fields().size() > 0) { + oss << ",\"partition_fields\":\"["; + std::string partition_fields_str; + for (const auto& partition_field : info.dblink_info.file_info().partition_fields()) { + partition_fields_str += "\"" + partition_field + "\","; + } + if (!partition_fields_str.empty()) { + partition_fields_str.pop_back(); + } + oss << partition_fields_str << "]"; + } + oss << ",\"format\":\"" << pb::FileFormat_Name(info.dblink_info.file_info().format()) << "\""; + if (!info.dblink_info.file_info().delimiter().empty()) { + static auto process_slash_func = [] (const std::string& str) -> std::string { + static std::unordered_map slash_map = { + {'\t', "\\\\t"}, + {'\u0001', "\\\\u0001"} + }; + std::string out; + for (char c : str) { + if (slash_map.find(c) != slash_map.end()) { + out += slash_map[c]; + } else { + out += c; + } + } + return out; + }; + oss << ",\"delimiter\":\"" << process_slash_func(info.dblink_info.file_info().delimiter()) << "\""; + } + oss << "}"; } else { oss << ",\"meta_name\":\"" << info.dblink_info.meta_name() << "\""; oss << ",\"namespace_name\":\"" << info.dblink_info.namespace_name() << "\""; @@ -2519,7 +2582,9 @@ bool ShowHelper::_show_schema_conf(const SmartSocket& client, const std::vector< "enable_column_engine", "olap_pre_split_cnt", "cold_use_column_only", - "force_column_storage"}; + "cold_use_column_only", + "force_column_storage", + "column_only_read_base"}; // 前三个conf按照bool解析, pk_prefix_balance按照int32来解析 if (split_vec.size() != 3 || allowed_conf.find(split_vec[2]) == allowed_conf.end()) { client->state = STATE_ERROR; @@ -2553,7 +2618,8 @@ bool ShowHelper::_show_schema_conf(const SmartSocket& client, const std::vector< || split_vec[2] == "enable_column_engine" || split_vec[2] == "olap_pre_split_cnt" || split_vec[2] == "cold_use_column_only" - || split_vec[2] == "force_column_storage") { + || split_vec[2] == "force_column_storage" + || split_vec[2] == "column_only_read_base") { names.emplace_back("value"); } @@ -2855,6 +2921,25 @@ bool ShowHelper::_show_all_tables(const SmartSocket& client, const std::vectorengine == pb::DBLINK; }; + type_func_map["statistics"] = [] (const SmartTable& table) { + return table != nullptr && table->have_statistics; + }; + type_func_map["cost"] = [] (const SmartTable& table) { + auto& pb_conf = table->schema_conf; + const google::protobuf::Reflection* reflection = pb_conf.GetReflection(); + const google::protobuf::Descriptor* descriptor = pb_conf.GetDescriptor(); + const google::protobuf::FieldDescriptor* field = nullptr; + field = descriptor->FindFieldByName(TABLE_SWITCH_COST); + if (field == nullptr) { + return false; + } + bool has_field = reflection->HasField(pb_conf, field); + if (!has_field) { + return false; + } + return reflection->GetBool(pb_conf, field); + }; + type_func_map["type_timestamp"] = [](const SmartTable& table) { if (table != nullptr) { for (auto& f : table->fields) { @@ -3299,7 +3384,12 @@ bool ShowHelper::_show_user(const SmartSocket& client, const std::vectoruser_info; + if (user == nullptr || !user->is_super) { + _wrapper->make_err_packet(client, ER_NO_SUCH_USER, "Only Support Super User To Query"); + client->state = STATE_READ_QUERY_RESULT; + return false; + } auto info = factory->get_user_info(split_vec[2]); if (info == nullptr) { DB_WARNING("user name not exist [%s]", split_vec[2].c_str()); diff --git a/src/protocol/state_machine.cpp b/src/protocol/state_machine.cpp index 61259c668..40155f757 100644 --- a/src/protocol/state_machine.cpp +++ b/src/protocol/state_machine.cpp @@ -383,7 +383,8 @@ void StateMachine::_print_query_time(SmartSocket client) { sql_agg_cost << BvarMap(stat_info->sample_sql.str(), index_id, stat_info->table_id, stat_info->total_time, err_count * stat_info->total_time, rows, stat_info->num_scan_rows, stat_info->read_disk_size, stat_info->num_filter_rows, stat_info->region_count, - field_range_type, err_count, stat_info->sign, subquery_signs); + field_range_type, err_count, stat_info->sign, subquery_signs, + stat_info->resource_tag); } if (op_type == pb::OP_SELECT || op_type == pb::OP_UNION) { @@ -1299,6 +1300,11 @@ int StateMachine::_get_json_attributes(std::shared_ptr ctx) { ctx->no_binlog = json_iter->value.GetBool(); DB_WARNING("no_binlog: %d", ctx->no_binlog); } + json_iter = root.FindMember("disable_on_update"); + if (json_iter != root.MemberEnd() && root["disable_on_update"].IsBool()) { + ctx->disable_on_update = json_iter->value.GetBool(); + DB_WARNING("disable_on_update: %d", ctx->disable_on_update); + } json_iter = root.FindMember("mpp"); if (json_iter != root.MemberEnd() && root["mpp"].IsBool()) { ctx->use_mpp = json_iter->value.GetBool(); @@ -1313,6 +1319,11 @@ int StateMachine::_get_json_attributes(std::shared_ptr ctx) { ctx->efsearch = json_iter->value.GetInt(); DB_WARNING("efsearch: %d", ctx->efsearch); } + json_iter = root.FindMember("nprobe"); + if (json_iter != root.MemberEnd()) { + ctx->nprobe = json_iter->value.GetInt(); + DB_WARNING("nprobe: %d", ctx->nprobe); + } } catch (...) { DB_WARNING("parse extra file error [%s]", json_str.c_str()); continue; diff --git a/src/raft/log_entry_reader.cpp b/src/raft/log_entry_reader.cpp index 4c38d42e3..8193c6325 100644 --- a/src/raft/log_entry_reader.cpp +++ b/src/raft/log_entry_reader.cpp @@ -30,6 +30,7 @@ int LogEntryReader::read_log_entry(int64_t region_id, int64_t log_index, std::st DB_FATAL("read log entry fail, region_id: %ld, log_index: %ld", region_id, log_index); return -1; } + RocksdbVars::get_instance()->raft_log_scan_times_count << 1; rocksdb::Slice slice(log_value); LogHead head(slice); if (head.type != braft::ENTRY_TYPE_DATA) { @@ -66,6 +67,7 @@ int LogEntryReader::read_log_entry(int64_t region_id, int64_t start_log_index, i options.fill_cache = false; std::unique_ptr iter(_rocksdb->new_iterator(options, _log_cf)); iter->Seek(log_data_key.data()); + RocksdbVars::get_instance()->raft_log_scan_times_count << 1; for (; iter->Valid(); iter->Next()) { if (!iter->key().starts_with(prefix.data())) { DB_WARNING("read end info, region_id: %ld, key:%s", region_id, iter->key().ToString(true).c_str()); @@ -129,6 +131,7 @@ int LogEntryReader::read_txn_last_log_entry(int64_t region_id, int64_t start_log std::map log_index_txn_map; std::unique_ptr iter(_rocksdb->new_iterator(options, _log_cf)); iter->Seek(log_data_key.data()); + RocksdbVars::get_instance()->raft_log_scan_times_count << 1; for (; iter->Valid(); iter->Next()) { if (!iter->key().starts_with(prefix.data())) { DB_WARNING("read end info, region_id: %ld, key:%s", region_id, iter->key().ToString(true).c_str()); diff --git a/src/raft/my_raft_log_storage.cpp b/src/raft/my_raft_log_storage.cpp index bb3d60686..c4bddf378 100644 --- a/src/raft/my_raft_log_storage.cpp +++ b/src/raft/my_raft_log_storage.cpp @@ -281,6 +281,7 @@ braft::LogEntry* MyRaftLogStorage::get_entry(const int64_t index) { index, _region_id); return NULL; } + RocksdbVars::get_instance()->raft_log_scan_times_count << 1; if (value.size() < LOG_HEAD_SIZE) { DB_FATAL("value of log index:%ld of region id:%ld is corrupted", index, _region_id); diff --git a/src/reverse/reverse_common.cpp b/src/reverse/reverse_common.cpp index 022fca7b4..38fa4fc9c 100644 --- a/src/reverse/reverse_common.cpp +++ b/src/reverse/reverse_common.cpp @@ -19,11 +19,11 @@ #include #include "proto/reverse.pb.h" namespace baikaldb { -DEFINE_string(q2b_utf8_path, "./conf/q2b_utf8.dic", "q2b_utf8_path"); -DEFINE_string(q2b_gbk_path, "./conf/q2b_gbk.dic", "q2b_gbk_path"); -DEFINE_string(punctuation_path, "./conf/punctuation.dic", "punctuation_path"); -DEFINE_bool(reverse_print_log, false, "reverse_print_log"); -DEFINE_bool(enable_print_convert_log, false, "enable_print_convert_log"); +DEFINE_string(q2b_utf8_path, "./conf/q2b_utf8.dic", "Path to q2b UTF-8 dictionary, default: ./conf/q2b_utf8.dic"); +DEFINE_string(q2b_gbk_path, "./conf/q2b_gbk.dic", "Path to q2b GBK dictionary, default: ./conf/q2b_gbk.dic"); +DEFINE_string(punctuation_path, "./conf/punctuation.dic", "Path to punctuation dictionary, default: ./conf/punctuation.dic"); +DEFINE_bool(reverse_print_log, false, "Enable reverse print logging, default: false"); +DEFINE_bool(enable_print_convert_log, false, "Enable print convert logging, default: false"); std::atomic_long g_statistic_insert_key_num = {0}; std::atomic_long g_statistic_delete_key_num = {0}; @@ -70,7 +70,7 @@ int Tokenizer::init() { return 0; } -#ifdef BAIDU_INTERNAL +#if defined(BAIDU_INTERNAL) && !defined(__aarch64__) drpc::NLPCClient* wordrank_client; drpc::NLPCClient* wordseg_client; drpc::NLPCClient* wordweight_client; diff --git a/src/reverse/reverse_interface.cpp b/src/reverse/reverse_interface.cpp index 40c09821a..8c6d36c9a 100644 --- a/src/reverse/reverse_interface.cpp +++ b/src/reverse/reverse_interface.cpp @@ -47,7 +47,7 @@ int segment( case pb::S_ES_STANDARD: ret = Tokenizer::get_instance()->es_standard(word, term_map, charset); break; -#ifdef BAIDU_INTERNAL +#if defined(BAIDU_INTERNAL) && !defined(__aarch64__) case pb::S_WORDRANK: ret = Tokenizer::get_instance()->wordrank(word, term_map, charset); break; diff --git a/src/runtime/arrow_io_excutor.cpp b/src/runtime/arrow_io_excutor.cpp index dc5fb40cf..d68f30e18 100644 --- a/src/runtime/arrow_io_excutor.cpp +++ b/src/runtime/arrow_io_excutor.cpp @@ -23,6 +23,9 @@ #include #include "exchange_sender_node.h" +namespace arrow { +DECLARE_int64(refuse_if_exceed_max_malloc_gb); +} namespace baikaldb { static bool validate_arrow_multi_threads(const char*, int32_t val) { return val >= 2; @@ -262,11 +265,12 @@ int GlobalArrowExecutor::init() { void GlobalArrowExecutor::execute(RuntimeState* state, arrow::Result>* result) { if (state->vectorlized_parallel_execution == false) { // 不开启pipeline并行 - *result = arrow::acero::DeclarationToTable(arrow::acero::Declaration::Sequence(std::move(state->acero_declarations)), false); - state->vectorlized_parallel_execution = false; + *result = arrow::acero::DeclarationToTable(arrow::acero::Declaration::Sequence(std::move(state->acero_declarations)), + /*use_threads=*/false, + GetMemoryPoolForRead()); } else { // 开启pipeline并行, 异步模式 - arrow::compute::ExecContext exec_context(arrow::default_memory_pool(), arrow::internal::GetCpuThreadPool()); + arrow::compute::ExecContext exec_context(GetMemoryPoolForRead(), arrow::internal::GetCpuThreadPool()); exec_context.set_use_threads(true); bthread::Mutex mu; bthread::ConditionVariable cond; @@ -284,7 +288,20 @@ void GlobalArrowExecutor::execute(RuntimeState* state, arrow::Resultvectorlized_parallel_execution = true; } } + +void GlobalArrowExecutor::shutdown() { + DB_WARNING("GlobalArrowExecutor begin shutdown"); + // wait = false: 等待当前正在执行的task结束, 清空当前pending的task + arrow::internal::GetCpuThreadPool()->Shutdown(/*wait=*/false); + DB_WARNING("GlobalArrowExecutor shutdown finish"); +} + +arrow::MemoryPool* GetMemoryPoolForRead() { + if (arrow::FLAGS_refuse_if_exceed_max_malloc_gb > 0) { + return arrow::system_memory_pool_with_limit(); + } + return arrow::default_memory_pool(); +} } // namespace baikaldb diff --git a/src/runtime/chunk.cpp b/src/runtime/chunk.cpp index d0739c369..c463a93f5 100644 --- a/src/runtime/chunk.cpp +++ b/src/runtime/chunk.cpp @@ -22,7 +22,7 @@ #include "mem_row_descriptor.h" #include #include - +#include "arrow_io_excutor.h" namespace baikaldb { DEFINE_int32(chunk_max_size_mb, 100, "chunk max size mb"); @@ -42,45 +42,45 @@ int Chunk::init_tuple_info(const pb::TupleDescriptor* tuple) { { case FieldDescriptorProto::TYPE_BOOL: // pb::NULL_TYPE, pb::BOOL _fields.emplace_back(std::make_shared(name, arrow::boolean())); - _builders.emplace_back(std::make_shared()); + _builders.emplace_back(std::make_shared(GetMemoryPoolForRead())); _size_per_row += sizeof(bool); break; case FieldDescriptorProto::TYPE_SINT32: // pb::INT8, pb::INT16, pb::INT32 case FieldDescriptorProto::TYPE_SFIXED32: // pb::TIME _fields.emplace_back(std::make_shared(name, arrow::int32())); - _builders.emplace_back(std::make_shared()); + _builders.emplace_back(std::make_shared(GetMemoryPoolForRead())); _size_per_row += sizeof(int32_t); break; case FieldDescriptorProto::TYPE_SINT64: // pb::INT64 _fields.emplace_back(std::make_shared(name, arrow::int64())); - _builders.emplace_back(std::make_shared()); + _builders.emplace_back(std::make_shared(GetMemoryPoolForRead())); _size_per_row += sizeof(int64_t); break; case FieldDescriptorProto::TYPE_UINT32: // pb::UINT8, pb::UINT16, pb::UINT32 case FieldDescriptorProto::TYPE_FIXED32: // pb::TIMESTAMP, pb::DATE _fields.emplace_back(std::make_shared(name, arrow::uint32())); - _builders.emplace_back(std::make_shared()); + _builders.emplace_back(std::make_shared(GetMemoryPoolForRead())); _size_per_row += sizeof(int32_t); break; case FieldDescriptorProto::TYPE_UINT64: // pb::UINT64 case FieldDescriptorProto::TYPE_FIXED64: // pb::DATETIME _fields.emplace_back(std::make_shared(name, arrow::uint64())); - _builders.emplace_back(std::make_shared()); + _builders.emplace_back(std::make_shared(GetMemoryPoolForRead())); _size_per_row += sizeof(int64_t); break; case FieldDescriptorProto::TYPE_FLOAT: // pb::FLOAT _fields.emplace_back(std::make_shared(name, arrow::float32())); - _builders.emplace_back(std::make_shared()); + _builders.emplace_back(std::make_shared(GetMemoryPoolForRead())); _size_per_row += sizeof(float); break; case FieldDescriptorProto::TYPE_DOUBLE: // pb::DOUBLE _fields.emplace_back(std::make_shared(name, arrow::float64())); - _builders.emplace_back(std::make_shared()); + _builders.emplace_back(std::make_shared(GetMemoryPoolForRead())); _size_per_row += sizeof(double); break; case FieldDescriptorProto::TYPE_BYTES: // pb::STRING, pb::HLL, pb::BITMAP, pb::TDIGEST _fields.emplace_back(std::make_shared(name, arrow::large_binary())); - _builders.emplace_back(std::make_shared()); + _builders.emplace_back(std::make_shared(GetMemoryPoolForRead())); break; default: DB_FATAL("unkown mysql type: %d", pb_type); diff --git a/src/session/binlog_context.cpp b/src/session/binlog_context.cpp index 97e361b5a..06c062ccd 100644 --- a/src/session/binlog_context.cpp +++ b/src/session/binlog_context.cpp @@ -25,7 +25,7 @@ #endif namespace baikaldb { -DEFINE_bool(meta_tso_autoinc_degrade, false, "meta_tso_autoinc_degrade"); +DEFINE_bool(meta_tso_autoinc_degrade, false, "Enable meta TSO autoinc degrade, default: false"); BRPC_VALIDATE_GFLAG(meta_tso_autoinc_degrade, brpc::PassValidate); DECLARE_int64(print_time_us); DECLARE_int64(retry_interval_us); diff --git a/src/store/backup.cpp b/src/store/backup.cpp index 014a285f7..546980e43 100755 --- a/src/store/backup.cpp +++ b/src/store/backup.cpp @@ -99,7 +99,7 @@ int Backup::backup_datainfo_to_file(const std::string& path, int64_t& file_size) rocksdb::Options options = db->get_options(db->get_data_handle()); options.bottommost_compression = rocksdb::kZSTD; options.bottommost_compression_opts = rocksdb::CompressionOptions(); - std::unique_ptr writer(new SstFileWriter(options, false)); + std::unique_ptr writer(new SstFileWriter(options)); rocksdb::ExternalSstFileInfo sst_file_info; auto ret = writer->open(path); if (!ret.ok()) { diff --git a/src/store/main.cpp b/src/store/main.cpp index 799ec98d8..2de219987 100644 --- a/src/store/main.cpp +++ b/src/store/main.cpp @@ -41,7 +41,7 @@ DECLARE_bool(use_fulltext_wordweight_segment); DECLARE_bool(use_fulltext_wordseg_wordrank_segment); DEFINE_string(wordrank_conf, "./config/drpc_client.xml", "wordrank conf path"); } // namespace baikaldb -DEFINE_bool(stop_server_before_core, true, "stop_server_before_core"); +DEFINE_bool(stop_server_before_core, true, "Stop server before core dump, default: true"); DEFINE_int32(compaction_sst_cache_capacity, 200000, "compaction_sst_cache_capacity"); brpc::Server server; @@ -105,7 +105,7 @@ int main(int argc, char **argv) { baikaldb::register_myraft_extension(); int ret = 0; baikaldb::Tokenizer::get_instance()->init(); -#ifdef BAIDU_INTERNAL +#if defined(BAIDU_INTERNAL) && !defined(__aarch64__) //init wordrank_client ret = ::drpc::init_env(baikaldb::FLAGS_wordrank_conf); if (ret < 0) { @@ -136,7 +136,7 @@ int main(int argc, char **argv) { DB_WARNING("init nlpc success"); #endif - /* + /* auto call = []() { std::ifstream extra_fs("test_file"); std::string word((std::istreambuf_iterator(extra_fs)), @@ -238,6 +238,7 @@ int main(int argc, char **argv) { DB_WARNING("store qos close success"); baikaldb::MemoryGCHandler::get_instance()->close(); baikaldb::MemTrackerPool::get_instance()->close(); + baikaldb::GlobalArrowExecutor::shutdown(); // exit if server.join is blocked baikaldb::Bthread bth; bth.run([]() { diff --git a/src/store/meta_writer.cpp b/src/store/meta_writer.cpp index 48d57d7fa..34f046e0e 100644 --- a/src/store/meta_writer.cpp +++ b/src/store/meta_writer.cpp @@ -70,6 +70,8 @@ const std::string MetaWriter::COLUMN_HOT_FILE_INDENTIFY(1, 0x13); const std::string MetaWriter::COLUMN_COLD_FILE_INDENTIFY(1, 0x14); // 用于加速 show binlog_info detail binlog_table 获取data_cf_oldest_datetime const std::string MetaWriter::BINLOG_DATA_CF_OLDEST_IDENTIFY(1, 0x15); +// 列存需要记录事务的begin_index +const std::string MetaWriter::COLUMN_TXN_LOG_INDEX_IDENTIFY(1, 0x16); int MetaWriter::init_meta_info(const pb::RegionInfo& region_info) { std::vector keys; @@ -257,15 +259,43 @@ int MetaWriter::clear_error_pre_commit(int64_t region_id, uint64_t txn_id) { return write_batch(&batch, region_id); } -int MetaWriter::write_meta_begin_index(int64_t region_id, int64_t log_index, int64_t data_index, uint64_t txn_id) { +int MetaWriter::delete_column_txn_log_index(int64_t region_id, std::vector txn_ids) { + rocksdb::WriteBatch batch; + if (txn_ids.empty()) { + return 0; + } + for (auto txn_id : txn_ids) { + batch.Delete(_meta_cf, column_txn_log_index_key(region_id, txn_id)); + } + return write_batch(&batch, region_id); +} + +int MetaWriter::write_meta_begin_index(int64_t region_id, int64_t log_index, int64_t data_index, uint64_t txn_id, bool has_column_engine) { if (log_index == 0) { return 0; } rocksdb::WriteBatch batch; batch.Put(_meta_cf, applied_index_key(region_id), encode_applied_index(log_index, data_index)); batch.Put(_meta_cf, transcation_log_index_key(region_id, txn_id), encode_transcation_log_index_value(log_index)); + if (has_column_engine) { + batch.Put(_meta_cf, column_txn_log_index_key(region_id, txn_id), encode_rollup_region_init_index(log_index)); + } return write_batch(&batch, region_id); } + +int64_t MetaWriter::read_meta_begin_index(int64_t region_id, uint64_t txn_id) { + std::string value; + rocksdb::ReadOptions options; + auto status = _rocksdb->get(options, _meta_cf, rocksdb::Slice(column_txn_log_index_key(region_id, txn_id)), &value); + if (!status.ok()) { + DB_WARNING("Error while read applied index, Error %s, region_id: %ld, txn_id: %lu", + status.ToString().c_str(), region_id, txn_id); + return -1; + } + TableKey tk(value); + return tk.extract_i64(0); +} + int MetaWriter::write_meta_index_and_num_table_lines(int64_t region_id, int64_t log_index, int64_t data_index, int64_t num_table_lines, SmartTransaction txn) { if (log_index == 0) { @@ -389,6 +419,15 @@ int MetaWriter::clear_txn_log_index(int64_t region_id) { status.code(), status.ToString().c_str(), region_id); return -1; } + start_key = column_txn_log_index_key(region_id, 0); + end_key = column_txn_log_index_key(region_id, UINT64_MAX); + status = _rocksdb->remove_range(MetaWriter::write_options, _meta_cf, + start_key, end_key, false); + if (!status.ok()) { + DB_WARNING("remove_range error: code=%d, msg=%s, region_id: %ld", + status.code(), status.ToString().c_str(), region_id); + return -1; + } return 0; } int MetaWriter::clear_txn_infos(int64_t region_id) { @@ -705,6 +744,14 @@ std::string MetaWriter::log_index_key_prefix(int64_t region_id) const { key.append_char(MetaWriter::PREPARED_TXN_LOG_INDEX_IDENTIFY.c_str(), 1); return key.data(); } +std::string MetaWriter::column_txn_log_index_key(int64_t region_id, uint64_t txn_id) const { + MutTableKey key; + key.append_char(MetaWriter::META_IDENTIFY.c_str(), 1); + key.append_i64(region_id); + key.append_char(MetaWriter::COLUMN_TXN_LOG_INDEX_IDENTIFY.c_str(), 1); + key.append_u64(txn_id); + return key.data(); +} std::string MetaWriter::transcation_pb_key(int64_t region_id, uint64_t txn_id, int64_t log_index) const { MutTableKey key; key.append_char(MetaWriter::META_IDENTIFY.c_str(), 1); diff --git a/src/store/region.cpp b/src/store/region.cpp index 285ad3981..1c6aad616 100644 --- a/src/store/region.cpp +++ b/src/store/region.cpp @@ -34,6 +34,7 @@ #include "qos.h" #include "arrow_io_excutor.h" #include "exchange_sender_node.h" +#include "ttl_delete_node.h" #ifdef BAIDU_INTERNAL #include #else @@ -55,6 +56,8 @@ DECLARE_bool(raft_enable_leader_lease); namespace baikaldb { DEFINE_bool(use_fulltext_wordweight_segment, true, "load wordweight dict"); DEFINE_bool(use_fulltext_wordseg_wordrank_segment, true, "load wordseg wordrank dict"); +DEFINE_int32(binlog_cache_size_for_single_region, 10, "binlog_cache_size_for_single_region"); +DEFINE_bool(use_binlog_cache, false, "read binlog use_binlog_cache"); DEFINE_int32(election_timeout_ms, 1000, "raft election timeout(ms)"); DEFINE_int32(skew, 5, "split skew, default : 45% - 55%"); DEFINE_int32(reverse_level2_len, 5000, "reverse index level2 length, default : 5000"); @@ -68,7 +71,6 @@ DEFINE_int64(disable_write_wait_timeout_us, 1000 * 1000, DEFINE_int32(snapshot_interval_s, 600, "raft snapshot interval(s)"); DEFINE_int32(fetch_log_timeout_s, 60, "raft learner fetch log time out(s)"); DEFINE_int32(fetch_log_interval_ms, 10, "raft learner fetch log interval(ms)"); -DEFINE_int32(snapshot_timed_wait, 120 * 1000 * 1000LL, "snapshot timed wait default 120S"); DEFINE_int64(snapshot_diff_lines, 10000, "save_snapshot when num_table_lines diff"); DEFINE_int64(snapshot_diff_logs, 2000, "save_snapshot when log entries diff"); DEFINE_int64(snapshot_log_exec_time_s, 60, "save_snapshot when log entries apply time"); @@ -91,12 +93,12 @@ DEFINE_int32(no_op_timer_timeout_ms, 100, "no op timer timeout(ms)"); DEFINE_int32(follow_read_timeout_s, 10, "follow read timeout(s)"); DEFINE_bool(apply_partial_rollback, true, "apply partial rollback"); DEFINE_bool(demotion_read_index_without_leader, true, "demotion read index without leader"); -DEFINE_bool(report_all_to_meta, true, "report_all_to_meta"); +DEFINE_bool(report_all_to_meta, true, "Report all to meta server, default: true"); // 并发控制 DEFINE_int64(sign_concurrency_timeout_rate, 5, "sign_concurrency_timeout_rate, default: 5. (0 means without timeout)"); DEFINE_int64(min_sign_concurrency_timeout_ms,1000, "min_sign_concurrency_timeout_ms, default: 1s"); DEFINE_int64(max_sign_concurrency_wait_cnt, 2000, "max_sign_concurrency_wait_cnt, default: 2k"); -DEFINE_bool(open_sign_concurrency, true, "open_sign_concurrency"); +DEFINE_bool(open_sign_concurrency, true, "Enable sign concurrency, default: true"); // 向量化执行 DEFINE_bool(only_use_arrow, false, "only use arrow(no row pb), for automated case tests, default(false)"); DEFINE_bool(vectorlized_parallel_execution, false, "vectorlized parallel execution"); @@ -147,6 +149,9 @@ int Region::init(bool new_region, int32_t snapshot_times) { ON_SCOPE_EXIT([this]() { _can_heartbeat = true; }); + if (FLAGS_use_binlog_cache) { + _binlog_cache.init(FLAGS_binlog_cache_size_for_single_region); + } MutTableKey start; MutTableKey end; start.append_i64(_region_id); @@ -198,7 +203,7 @@ int Region::init(bool new_region, int32_t snapshot_times) { _region_info.table_id(), _region_id); return -1; } - + const auto charset = table_info.charset; for (int64_t index_id : table_info.indices) { IndexInfo info = _factory->get_index_info(index_id); @@ -222,7 +227,7 @@ int Region::init(bool new_region, int32_t snapshot_times) { segment_type = pb::S_UNIGRAMS; #endif } -#ifdef BAIDU_INTERNAL +#if defined(BAIDU_INTERNAL) && !defined(__aarch64__) if (segment_type == pb::S_WORDRANK || segment_type == pb::S_WORDSEG_BASIC || segment_type == pb::S_WORDRANK_Q2B_ICASE || segment_type == pb::S_WORDRANK_Q2B_ICASE_UNLIMIT) { if (!FLAGS_use_fulltext_wordseg_wordrank_segment) { @@ -230,7 +235,8 @@ int Region::init(bool new_region, int32_t snapshot_times) { open flag use_fulltext_workseg_wordrank_segment", _region_id); return -1; } - } else if ((segment_type == pb::S_WORDWEIGHT || segment_type == pb::S_WORDWEIGHT_NO_FILTER) && + } else if ((segment_type == pb::S_WORDWEIGHT || segment_type == pb::S_WORDWEIGHT_NO_FILTER + || segment_type == pb::S_WORDWEIGHT_NO_FILTER_SAME_WEIGHT) && !FLAGS_use_fulltext_wordweight_segment) { DB_FATAL("region %ld store not support wordweight segment, \ open flag use_fulltext_wordweight_segment", _region_id); @@ -277,11 +283,15 @@ int Region::init(bool new_region, int32_t snapshot_times) { TTLInfo ttl_info = _factory->get_ttl_duration(get_table_id()); if (ttl_info.ttl_duration_s > 0) { - _use_ttl = true; - if (ttl_info.online_ttl_expire_time_us > 0) { - // online TTL - _online_ttl_base_expire_time_us = ttl_info.online_ttl_expire_time_us; - } + SmartIndex global_index = SchemaFactory::get_instance()->get_index_info_ptr(get_global_index_id()); + if (global_index == nullptr) { + DB_FATAL("init region failed, missing region global index info. region_id:%ld, global_index_id:%ld", _region_id, _global_index_id); + return -1; + } + if (update_ttl_info() != 0) { + DB_FATAL("init region failed, update ttl info failed. region_id:%ld, global_index_id:%ld", _region_id, _global_index_id); + return -1; + } } _storage_compute_separate = _factory->get_separate_switch(get_table_id()); bool is_cold = false; @@ -325,7 +335,7 @@ int Region::init(bool new_region, int32_t snapshot_times) { boost::lexical_cast(_region_id); options.snapshot_file_system_adaptor = &_snapshot_adaptor; - _txn_pool.init(_region_id, _use_ttl, _online_ttl_base_expire_time_us); + _txn_pool.init(_region_id, use_normal_ttl(), _online_ttl_base_expire_time_us); bool is_restart = _restart; if (_is_learner) { DB_DEBUG("init learner."); @@ -385,7 +395,7 @@ int Region::init(bool new_region, int32_t snapshot_times) { } else { SplitCompactionFilter::get_instance()->set_filter_region_info( _region_id, _resource->region_info.end_key(), - _use_ttl, _online_ttl_base_expire_time_us); + use_normal_ttl(), _online_ttl_base_expire_time_us); } // follower read bthread::ExecutionQueueOptions opt; @@ -1691,7 +1701,7 @@ void Region::query(google::protobuf::RpcController* controller, } const auto& remote_side_tmp = butil::endpoint2str(cntl->remote_side()); const char* remote_side = remote_side_tmp.c_str(); - if (!is_leader()) { + if (!is_leader() && !request->extra_req().no_raft_log()) { if (!is_learner()) { _not_leader_alarm.not_leader_alarm(_node.leader_id()); // 非leader才返回 @@ -1715,6 +1725,11 @@ void Region::query(google::protobuf::RpcController* controller, _region_id, log_id, remote_side); return; } + if (_raft_status_error) { + response->set_errcode(is_learner() ? pb::LEARNER_NOT_READY : pb::NOT_LEADER); + response->set_errmsg("raft status error"); + return; + } if (request->extra_req().use_read_idx()) { if (!_ready_for_follower_read) { response->set_errcode(is_learner() ? pb::LEARNER_NOT_READY : pb::NOT_LEADER); @@ -1798,7 +1813,23 @@ void Region::query(google::protobuf::RpcController* controller, // TimeCost cost; switch (request->op_type()) { case pb::OP_KILL: - exec_out_txn_query(controller, request, response, done_guard.release()); + if (request->extra_req().no_raft_log()) { + DMLClosure* c = new DMLClosure; + c->op_type = request->op_type(); + c->log_id = log_id; + c->response = response; + c->region = this; + c->remote_side = remote_side; + dml_1pc(*request, request->op_type(), request->plan(), request->tuples(), + *response, 0, 0, static_cast(c)); + if (response->errcode() != pb::SUCCESS) { + DB_FATAL("dml exec failed, region_id: %ld log_id:%lu", _region_id, log_id); + delete c; + return; + } + } else { + exec_out_txn_query(controller, request, response, done_guard.release()); + } break; case pb::OP_TXN_QUERY_PRIMARY_REGION: exec_txn_query_primary_region(controller, request, response, done_guard.release()); @@ -2856,7 +2887,8 @@ int Region::select_vectorized(RuntimeState& state, ExecNode* root, const pb::Sto if (limit > 0) { table = table->Slice(0, std::min(limit, table->num_rows())); } - arrow::Result> record_batch_result = table->CombineChunksToBatch(); + arrow::MemoryPool* pool = GetMemoryPoolForRead(); + arrow::Result> record_batch_result = table->CombineChunksToBatch(pool); if (!record_batch_result.ok()) { DB_FATAL("arrow execute fail: CombineChunksToBatch fail, region_id: %ld, logid: %lu, status: %s", _region_id, request.log_id(), record_batch_result.status().ToString().c_str()); @@ -2870,7 +2902,7 @@ int Region::select_vectorized(RuntimeState& state, ExecNode* root, const pb::Sto } int64_t combine_time = t.get_time(); t.reset(); - arrow::Result> schema_ret = arrow::ipc::SerializeSchema(*(record_batch->schema()), arrow::default_memory_pool()); + arrow::Result> schema_ret = arrow::ipc::SerializeSchema(*(record_batch->schema()), pool); if (!schema_ret.ok()) { DB_FATAL("arrow execute fail: serialize schema fail, region_id: %ld, logid: %lu, status: %s", _region_id, request.log_id(), schema_ret.status().ToString().c_str()); @@ -3561,6 +3593,7 @@ void Region::do_apply(int64_t term, int64_t index, const pb::StoreReq& request, _meta_writer->write_olap_info(_region_id, olap_info); _olap_state.store(pb::OLAP_ACTIVE); DB_WARNING("olap region_id: %ld state change to OLAP_ACTIVE", _region_id); + _column_mgr.remove_column_data(pb::CS_INVALID, 0); } uint64_t txn_id = request.txn_infos_size() > 0 ? request.txn_infos(0).txn_id():0; //事务流程中DML处理 @@ -3959,7 +3992,7 @@ void Region::apply_kv_in_txn(const pb::StoreReq& request, braft::Closure* done, write_begin_index = true; } if (write_begin_index) { - auto ret = _meta_writer->write_meta_begin_index(_region_id, index, _data_index, txn_id); + auto ret = _meta_writer->write_meta_begin_index(_region_id, index, _data_index, txn_id, txn->has_column_engine()); //DB_WARNING("write meta info when prepare, region_id: %ld, applied_index: %ld, txn_id: %ld", // _region_id, index, txn_id); if (ret < 0) { @@ -4683,6 +4716,8 @@ void Region::adjustkey_and_add_version(const pb::StoreReq& request, str_to_hex(request.end_key()).c_str(), applied_index, term); set_region_with_update_range(region_info_mem); + // 调整startkey endkey之后重置列存 + _column_mgr.remove_column_data(pb::CS_INVALID, 0); _last_split_time_cost.reset(); } @@ -4750,6 +4785,8 @@ void Region::validate_and_add_version(const pb::StoreReq& request, txn_infos.push_back(txn_info); } _txn_pool.update_txn_num_rows_after_split(txn_infos); + // 分裂后重置列存 + _column_mgr.remove_column_data(pb::CS_INVALID, 0); // 分裂后主动执行compact DB_WARNING("region_id: %ld, new_region_id: %ld, split do compact in queue", _region_id, _split_param.new_region_id); @@ -4863,6 +4900,7 @@ void Region::add_version_for_split_region(const pb::StoreReq& request, braft::Cl _region_id, _applied_index, term); _region_control.reset_region_status(); set_region_with_update_range(region_info_mem); + _column_mgr.remove_column_data(pb::CS_INVALID, 0); if (!compare_and_set_legal()) { DB_FATAL("split timeout, region was set split fail, region_id: %ld", _region_id); if (done != nullptr) { @@ -5043,6 +5081,10 @@ void Region::on_snapshot_save(braft::SnapshotWriter* writer, braft::Closure* don DB_WARNING("vector index is nullptr, region_id: %ld, index_id: %ld", _region_id, index_id); return; } + auto info = SchemaFactory::get_instance()->get_index_info_ptr(index_id); + if (info != nullptr && info->state == pb::IS_DELETE_LOCAL) { + continue; + } DB_WARNING("begin dump index:%s, index_id: %ld", _snapshot_path.c_str(), index_id); if (idx->is_separate()) { std::vector faiss_name_vec; @@ -5228,12 +5270,11 @@ void Region::on_snapshot_load_for_restart(braft::SnapshotReader* reader, _applied_index = applied_index; } } else { - //系统在执行commit之前重启 + //系统在执行commit之前重启 if (log_index < start_log_index) { start_log_index = log_index; } txn_ids.insert(txn_id); - } } int64_t max_applied_index = std::max(snapshot_index, _applied_index); @@ -5636,7 +5677,7 @@ void Region::compact_data_in_queue() { RegionControl::compact_data_in_queue(_region_id); } -void Region::vector_schema_change() { +void Region::vector_schema_change(std::unordered_map& vector_indexs) { if (_shutdown) { return; } @@ -5644,15 +5685,27 @@ void Region::vector_schema_change() { ON_SCOPE_EXIT([this]() { _multi_thread_cond.decrease_signal(); }); - std::map vector_index_map; + std::map cur_vector_index_map; { BAIDU_SCOPED_LOCK(_reverse_index_map_lock); - if (_vector_index_map.empty()) { - return; + cur_vector_index_map = _vector_index_map; + } + + // index change + for (auto& [index_id, index_state] : vector_indexs) { + if (cur_vector_index_map.count(index_id) == 0 && index_state != pb::IS_PUBLIC) { + auto new_vec_index = new VectorIndex(); + if (new_vec_index->init(_region_info, _region_id, index_id, _table_id) != 0) { + DB_FATAL("vec index init fail"); + continue; + } + BAIDU_SCOPED_LOCK(_reverse_index_map_lock); + _vector_index_map[index_id] = new_vec_index; } - vector_index_map = _vector_index_map; } - for (auto& pair : vector_index_map) { + + // field change + for (auto& pair : cur_vector_index_map) { auto vec_idx = pair.second; int64_t index_id = pair.first; if (vec_idx == nullptr) { @@ -5660,8 +5713,7 @@ void Region::vector_schema_change() { return; } if (vec_idx->schema_change() != 0) { - DB_WARNING("Fail to compact, region_id: %ld, index_id: %ld", _region_id, index_id); - return; + DB_FATAL("Fail to compact, region_id: %ld, index_id: %ld", _region_id, index_id); } } } @@ -6647,7 +6699,7 @@ void Region::write_local_rocksdb_for_split() { } } else if (index_info.type == pb::I_UNIQ || index_info.type == pb::I_KEY) { rocksdb::Slice tmp_value = iter->value(); - if (_use_ttl) { + if (use_normal_ttl()) { ttl_decode(tmp_value, &index_info, _online_ttl_base_expire_time_us); } if (!Transaction::fits_region_range(key_slice, tmp_value, @@ -6845,9 +6897,17 @@ void Region::write_local_rocksdb_for_split() { return; } bool has_vector_index = false; + std::set invalid_vec_indexes; { BAIDU_SCOPED_LOCK(_reverse_index_map_lock); - has_vector_index = _vector_index_map.size() > 0; + for (auto& [index_idx, idx] : _vector_index_map) { + auto info = SchemaFactory::get_instance()->get_index_info_ptr(index_idx); + if (info != nullptr && info->state == pb::IS_DELETE_LOCAL) { + invalid_vec_indexes.insert(std::to_string(index_idx)); + } else { + has_vector_index = true; + } + } } if (has_vector_index) { // 复制faiss文件并同步增量 @@ -6858,12 +6918,29 @@ void Region::write_local_rocksdb_for_split() { dir_iter iter(_snapshot_path); dir_iter end; for (; iter != end; ++iter) { + bool is_valid_file = true; std::string child_path = iter->path().c_str(); std::vector split_vec; boost::split(split_vec, child_path, boost::is_any_of("/")); - std::string out_path = "/" + split_vec.back(); + std::string filename = split_vec.back(); + std::string out_path = "/" + filename; if (boost::istarts_with(out_path, "/faissindex_")) { - files.emplace_back(out_path); + // faissindex_273_0_0 faissindex_delbitmap_273 faissindex_notcachefields_273 faissindex_scalardata_273 + // faissindex_277_cacheinfo faissindex_277_111 faissindex_delbitmap_277_111 faissindex_notcachefields_277_111 faissindex_scalardata_277_111 + std::vector name_split; + boost::split(name_split, filename, boost::is_any_of("_")); + for (auto s_idx = 1; s_idx < name_split.size(); ++s_idx) { + if (is_digits(name_split[s_idx])) { + if (invalid_vec_indexes.count(name_split[s_idx]) > 0) { + DB_WARNING("ignore file: %s", filename.c_str()); + is_valid_file = false; + } + break; + } + } + if (is_valid_file) { + files.emplace_back(out_path); + } } } } catch (boost::filesystem::filesystem_error& e) { @@ -7751,6 +7828,7 @@ void Region::print_log_entry(const int64_t start_index, const int64_t end_index) read_options.fill_cache = false; std::unique_ptr iter(_rocksdb->new_iterator(read_options, RocksWrapper::RAFT_LOG_CF)); iter->Seek(log_data_key.data()); + RocksdbVars::get_instance()->raft_log_scan_times_count << 1; for (int i = 0; iter->Valid() && i < 100; iter->Next(), i++) { TableKey key(iter->key()); int64_t log_index = key.extract_i64(sizeof(int64_t) + 1); @@ -7795,6 +7873,7 @@ int Region::exec_rollup_region_finish_request(const pb::StoreReq& request, read_options.fill_cache = false; std::unique_ptr iter(_rocksdb->new_iterator(read_options, RocksWrapper::RAFT_LOG_CF)); iter->Seek(begin_log_data_key.data()); + RocksdbVars::get_instance()->raft_log_scan_times_count << 1; for (int i = 0; iter->Valid() && (unappiled_begin_index + i) < unappiled_end_index; iter->Next(), i++) { rocksdb::Slice value_slice(iter->value()); LogHead head(iter->value()); @@ -8202,7 +8281,7 @@ int Region::add_reverse_index(int64_t table_id, const std::set& index_i segment_type = pb::S_NO_SEGMENT; } if (segment_type == pb::S_DEFAULT) { - #ifdef BAIDU_INTERNAL + #if defined(BAIDU_INTERNAL) && !defined(__aarch64__) segment_type = pb::S_WORDRANK; #else segment_type = pb::S_UNIGRAMS; @@ -8252,6 +8331,13 @@ void Region::remove_local_index_data() { if (table_info == nullptr) { return; } + + std::map vector_index_map; + { + BAIDU_SCOPED_LOCK(_reverse_index_map_lock); + vector_index_map = _vector_index_map; + } + bool need_snapshot = false; for (auto index_id: table_info->indices) { auto index_info = _factory->get_index_info_ptr(index_id); if (index_info == nullptr) { @@ -8259,8 +8345,23 @@ void Region::remove_local_index_data() { } if (index_info->state == pb::IS_DELETE_LOCAL) { delete_local_rocksdb_for_ddl(main_table_id, index_id); + + if (index_info->type == pb::I_VECTOR + && vector_index_map.count(index_id) > 0 + && !vector_index_map[index_id]->is_removed()) { + // clear mem data + need_snapshot = true; + vector_index_map[index_id]->reset(); + vector_index_map[index_id]->set_removed(true); + DB_WARNING("table_id: %ld, region_id: %ld, index_id: %ld, deleted faiss index in memory", + main_table_id, _region_id, index_id); + } } } + if (need_snapshot) { + // 写no op, 下次snapshot删除faiss文件 + do_snapshot(); + } } void Region::delete_local_rocksdb_for_ddl(int64_t table_id, int64_t index_id) { @@ -8397,14 +8498,187 @@ bool Region::can_use_approximate_split() { } } +int Region::update_ttl_info() { + BAIDU_SCOPED_LOCK(_ttl_mutex); + if (_shutdown) { + return 0; + } + bool ttl_info_updated = false; + // 拒绝修改ttl字段的逻辑放在meta + TTLInfo ttl_info = _factory->get_ttl_duration(get_table_id()); + if (ttl_info.ttl_duration_s <= 0) { + return 0; + } + + // 初始化ttl信息，只会更新一次 + if (!use_ttl()) { + _ttl_field_type = TTLFieldType::T_NORMAL_TTL; + std::shared_ptr ttl_field = ttl_info.ttl_field; + if (ttl_field != nullptr) { + // 只允许修改一次ttl字段 + ttl_info_updated = true; + _ttl_field_type = TTLFieldType::T_NON_TTL; + _ttl_field = std::make_shared(*ttl_field); + + SmartIndex global_index = SchemaFactory::get_instance()->get_index_info_ptr(get_global_index_id()); + SmartTable table_info_ptr = SchemaFactory::get_instance()->get_table_info_ptr(_table_id); + if (global_index == nullptr || table_info_ptr == nullptr) { + DB_FATAL("get global_index_info or table_info failed"); + return -1; + } + // 主表的话需要从table_info内获取，全局索引从index_info内获取 + const std::vector& value_fields = + global_index->id == global_index->pk ? table_info_ptr->fields : global_index->pk_fields; + for (const auto& field: value_fields) { + if (field.id == ttl_field->id) { + _ttl_field_type = global_index->id == global_index->pk + ? TTLFieldType::T_FIELD_PK_VALUE + : global_index->type == pb::I_UNIQ + ? TTLFieldType::T_FIELD_GUI_VALUE + : TTLFieldType::T_FIELD_GI_KEY; + if (global_index->id == global_index->pk) { + _decode_ttl_field_map[_ttl_field->id] = ttl_field.get(); + } + break; + } + } + // 由于table_info_ptr->fields内是所有字段，包含主键和普通字段 + // _ttl_field_in_rocksdb_key = true 优先级要高于_ttl_field_in_rocksdb_key = false + for (const auto& field: global_index->fields) { + if (field.id == ttl_field->id) { + _ttl_field_type = TTLFieldType::T_FIELD_KEY; + break; + } + } + + if (_ttl_field_type == TTLFieldType::T_NON_TTL) { + // should never reach here + DB_FATAL("No ttl field find in main_table or global index, region_id: %ld, index_id: %ld", _region_id, _global_index_id); + return -1; + } + } + } + + // ttl_duration可能修改，每次都需要更新 + if (ttl_info.ttl_duration_s > 0) { + if (_ttl_duration_s != ttl_info.ttl_duration_s) { + ttl_info_updated = true; + _ttl_duration_s = ttl_info.ttl_duration_s; + } + // online TTL + if (ttl_info.online_ttl_expire_time_us > 0 && ttl_info.online_ttl_expire_time_us != _online_ttl_base_expire_time_us) { + ttl_info_updated = true; + _online_ttl_base_expire_time_us = ttl_info.online_ttl_expire_time_us; + } + } + if (ttl_info_updated && _init_success) { + _txn_pool.update_ttl_info(use_normal_ttl(), _online_ttl_base_expire_time_us); + DB_WARNING("table_id: %ld, region_id: %ld, ttl_duration_s: %ld, online_ttl_expire_time_us: %ld, %s", + get_table_id(), _region_id, ttl_info.ttl_duration_s, + ttl_info.online_ttl_expire_time_us, timestamp_to_str(ttl_info.online_ttl_expire_time_us / 1000000).c_str()); + } + return 0; +} + +int Region::get_ttl_timestamp(IndexInfo& index_info, + const rocksdb::Slice& rocksdb_key, + const rocksdb::Slice& rocksdb_value, + const SmartRecord& record_template, + uint64_t& ttl_timestamp) const { + if (need_decode_ttl_field()) { + SmartRecord record = record_template->clone(true); + int ret = 0; + if (_ttl_field_type == TTLFieldType::T_FIELD_KEY) { + // main table或者全局索引的key内 + int pos = KEY_PREFIX_LENGTH; + ret = record->decode_key(index_info, rocksdb_key, pos); + if (ret != 0) { + DB_WARNING("ttl decode field value from key failed, index_id:%ld, field_id:%d", index_info.id, _ttl_field->id); + return ret; + } + } else if (_ttl_field_type == TTLFieldType::T_FIELD_PK_VALUE) { + // main table的value内 + // 为空会有默认值 + TupleRecord tuple_record(rocksdb_value); + ret = tuple_record.decode_fields(_decode_ttl_field_map, record); + if (ret != 0) { + DB_WARNING("ttl decode field value from pk value failed, index_id:%ld, field_id:%d", index_info.id, _ttl_field->id); + return ret; + } + } else if (_ttl_field_type == TTLFieldType::T_FIELD_GUI_VALUE) { + // 全局唯一索引，且ttl字段在主键内不在索引内 + int pos = 0; + ret = record->decode_primary_key(index_info, rocksdb_value, pos); + if (ret != 0) { + DB_WARNING("ttl decode field value from guidx value failed, index_id:%ld, field_id:%d", index_info.id, _ttl_field->id); + return ret; + } + } else if (_ttl_field_type == TTLFieldType::T_FIELD_GI_KEY) { + // 全局索引，且ttl字段在主键内不在索引内 + int pos = KEY_PREFIX_LENGTH; + ret = record->decode_key(index_info, rocksdb_key, pos); + if (ret != 0) { + DB_WARNING("decode field value from key failed, index_id:%ld", index_info.id); + return ret; + } + ret = record->decode_primary_key(index_info, rocksdb_key, pos); + if (ret != 0) { + DB_WARNING("ttl decode field value from guidx value failed, index_id:%ld, field_id:%d", index_info.id, _ttl_field->id); + return ret; + } + } + const FieldDescriptor* field_descriptor = record->get_field_by_idx(_ttl_field->pb_idx); + if (field_descriptor == nullptr) { + DB_WARNING("no ttl field found!"); + return -1; + } + ExprValue ttl_field_value = record->get_value(field_descriptor); + ttl_field_value.cast_to(_ttl_field->type); + ttl_field_value.cast_to(pb::TIMESTAMP); + ttl_timestamp = (ttl_field_value.get_numberic() + _ttl_duration_s) * 1000000LL; + } else { + rocksdb::Slice rocksdb_value_cpy = rocksdb_value; + ttl_timestamp = ttl_decode(rocksdb_value_cpy, &index_info, is_field_ttl() ? 0 : _online_ttl_base_expire_time_us); + } + return 0; +} + +SmartState Region::create_ttl_delete_runtime_status() { + if (_is_global_index) { + return nullptr; + } + pb::StoreReq req; + req.set_region_id(_region_id); + req.set_region_version(_version); + + pb::Plan plan; + plan.add_nodes()->set_limit(0); + + SmartState state_ptr = std::make_shared(); + RuntimeState& state = *state_ptr; + state.set_resource(get_resource()); + state.set_remote_side("region_" + std::to_string(_region_id) + "_ttl_thread"); + StateOption option; + option.store_compute_separate = false; + + int ret = state.init(req, plan, RepeatedPtrField(), &_txn_pool, option); + if (ret != 0) { + DB_WARNING("RuntimeStatue init failed, region_id: %ld", _region_id); + return nullptr; + } + state.set_reverse_index_map(_reverse_index_map); + state.set_vector_index_map(_vector_index_map); + return state_ptr; +} + // 后续要用compaction filter 来维护，现阶段主要有num_table_lines维护问题 void Region::ttl_remove_expired_data() { - if (!_use_ttl && !is_binlog_region()) { + if (!use_ttl() && !is_binlog_region()) { return; } if (_shutdown) { return; - } + } _multi_thread_cond.increase(); ON_SCOPE_EXIT([this]() { _multi_thread_cond.decrease_signal(); @@ -8424,7 +8698,12 @@ void Region::ttl_remove_expired_data() { int64_t global_index_id = get_global_index_id(); int64_t main_table_id = get_table_id(); int64_t read_timestamp_us = butil::gettimeofday_us(); + + bool is_cstore = _factory->get_table_engine(main_table_id) == pb::ROCKSDB_CSTORE; + auto resource = get_resource(); + std::vector indices; + bool delete_by_pk = is_field_ttl() && !_is_global_index; // 包含field ttl或者fulltext、vector index时，需要按照主键删除 if (_is_global_index) { indices.push_back(global_index_id); } else { @@ -8434,14 +8713,55 @@ void Region::ttl_remove_expired_data() { continue; } indices.push_back(index_id); + SmartIndex index_info_ptr = SchemaFactory::get_instance()->get_index_info_ptr(index_id); + if ((index_info_ptr->type == pb::I_FULLTEXT || index_info_ptr->type == pb::I_VECTOR) + && !is_cstore) { + delete_by_pk = true; + break; + } + } + if (delete_by_pk) { + indices.clear(); + indices.emplace_back(main_table_id); } } + + SmartRecord record_template = _factory->new_record(_table_id); + if (record_template == nullptr) { + DB_WARNING("new_record from table[%ld] failed, skip ttl.", _table_id); + return; + } + SmartIndex pk_info = _factory->get_index_info_ptr(_table_id); + SmartIndex global_index_info = _factory->get_index_info_ptr(_global_index_id); + + std::map local_secondary_index_fields_map; + std::set pk_ids; + for (const auto& iter: pk_info->fields) { + pk_ids.emplace(iter.id); + } + + SmartTable table_info_ptr = _factory->get_table_info_ptr(_table_id); + for (auto& field: table_info_ptr->fields) { + if (pk_ids.count(field.id) == 0) { + local_secondary_index_fields_map[field.id] = &field; + } + } + std::atomic write_sst_lines(0); - IndexInfo pk_info = _factory->get_index_info(main_table_id); - auto resource = get_resource(); - bool is_cstore = _factory->get_table_engine(main_table_id) == pb::ROCKSDB_CSTORE; + SmartState ttl_runtime_state = nullptr; + std::unique_ptr ttl_delete_node = nullptr; + if (delete_by_pk) { + ttl_runtime_state = create_ttl_delete_runtime_status(); + ttl_delete_node = std::make_unique(); + pb::PlanNode node; + node.mutable_derive_node()->mutable_delete_node()->set_table_id(main_table_id); + ttl_delete_node->init(node); + } + int64_t del_pk_num = 0; + + // start process index for (int64_t index_id : indices) { MutTableKey table_prefix; table_prefix.append_i64(_region_id).append_i64(index_id); @@ -8466,11 +8786,16 @@ void Region::ttl_remove_expired_data() { region_oldest_ts = _meta_writer->read_binlog_oldest_ts(_region_id); } // 内部txn，不提交出作用域自动析构 - SmartTransaction txn(new Transaction(0, nullptr)); + SmartTransaction txn(new Transaction(0, &_txn_pool)); + if (ttl_runtime_state != nullptr) { + txn->set_resource(resource); + ttl_runtime_state->set_txn(txn); + } txn->begin(txn_opt); rocksdb::Status s; + int commit_batch = 100; for (iter->Seek(table_prefix.data()); iter->Valid(); iter->Next()) { - if (FLAGS_stop_ttl_data || _shutdown) { + if (FLAGS_stop_ttl_data || _shutdown) { break; } ++count; @@ -8483,12 +8808,12 @@ void Region::ttl_remove_expired_data() { oldest_tso = timestamp_to_ts(tso::get_timestamp_internal(oldest_tso) - 7200); // data cf 比binlog cf延迟两小时删除 if (commit_tso > oldest_tso) { // 未过期，后边的ts都不会过期，直接跳出 - DB_WARNING("commit_tso: %ld, %s, oldest_tso: %ld, %s", - commit_tso, ts_to_datetime_str(commit_tso).c_str(), + DB_WARNING("commit_tso: %ld, %s, oldest_tso: %ld, %s", + commit_tso, ts_to_datetime_str(commit_tso).c_str(), oldest_tso, ts_to_datetime_str(oldest_tso).c_str()); break; } - } + } } else { if (index_info.type == pb::I_PRIMARY || _is_global_index) { // check end_key @@ -8496,8 +8821,12 @@ void Region::ttl_remove_expired_data() { break; } } - rocksdb::Slice value_slice1(iter->value()); - if (ttl_decode(value_slice1, &index_info, _online_ttl_base_expire_time_us) > read_timestamp_us) { + uint64_t ttl_timestamp = std::numeric_limits::max(); + int ret = get_ttl_timestamp(index_info, iter->key(), iter->value(), record_template, ttl_timestamp); + if (ret < 0) { + continue; + } + if (ttl_timestamp > read_timestamp_us) { //未过期 continue; } @@ -8506,24 +8835,51 @@ void Region::ttl_remove_expired_data() { if (!is_binlog_region()) { s = txn->get_txn()->GetForUpdate(read_opt, _data_cf, iter->key(), &value); if (!s.ok()) { - DB_WARNING("index %ld, region_id: %ld GetForUpdate failed, status: %s", + DB_WARNING("index %ld, region_id: %ld GetForUpdate failed, status: %s", index_id, _region_id, s.ToString().c_str()); continue; } - - rocksdb::Slice value_slice2(value); - if (ttl_decode(value_slice2, &index_info, _online_ttl_base_expire_time_us) > read_timestamp_us) { + uint64_t ttl_timestamp = std::numeric_limits::max(); + int ret = get_ttl_timestamp(index_info, iter->key(), rocksdb::Slice(value), record_template, ttl_timestamp); + if (ret < 0) { + continue; + } + if (ttl_timestamp > read_timestamp_us) { //加锁校验未过期 continue; } } else { // do nothing, binlog region不需要加锁再次校验 } - s = txn->get_txn()->Delete(_data_cf, iter->key()); - if (!s.ok()) { - DB_FATAL("index %ld, region_id: %ld Delete failed, status: %s", - index_id, _region_id, s.ToString().c_str()); - continue; + if (delete_by_pk) { + // 只有主键能进到这里 + SmartRecord record = record_template->clone(true); + int pos = KEY_PREFIX_LENGTH; + int ret = record->decode_key(*pk_info, iter->key(), pos); + if (ret != 0) { + DB_WARNING("decode main table keys failed, table id: %ld", _table_id); + continue; + } + rocksdb::Slice value_slice(value); + if (_ttl_field == nullptr) { + // 普通ttl需要移除value头部的时间戳 + ttl_decode(value_slice, &index_info, _online_ttl_base_expire_time_us); + } + TupleRecord tuple_record(value_slice); + // 删除索引需要全部的索引字段 + ret = tuple_record.decode_fields(local_secondary_index_fields_map, record); + if (ret != 0) { + DB_WARNING("decode main table value failed, table_id: %ld", _table_id); + continue; + } + ttl_delete_node->add_delete_records(record); + } else { + s = txn->get_txn()->Delete(_data_cf, iter->key()); + if (!s.ok()) { + DB_FATAL("index %ld, region_id: %ld Delete failed, status: %s", + index_id, _region_id, s.ToString().c_str()); + continue; + } } // for cstore only, remove_columns if (is_cstore && index_info.type == pb::I_PRIMARY && !_is_global_index) { @@ -8539,15 +8895,28 @@ void Region::ttl_remove_expired_data() { ++_num_delete_lines; ++del_pk_num; } - if (++num_remove_lines % 100 == 0) { - // 批量提交，减少内部锁冲突 - s = txn->commit(); - if (!s.ok()) { - DB_FATAL("index %ld, region_id: %ld commit failed, status: %s", - index_id, _region_id, s.ToString().c_str()); - continue; + if (++num_remove_lines % commit_batch == 0) { + if (delete_by_pk) { + // ttl_delete_node内部会commit或rollback + if (ttl_delete_node->open(ttl_runtime_state.get()) < 0) { + DB_FATAL("index %ld, region_id: %ld ttl_delete_node open failed", + index_id, _region_id); + del_pk_num -= commit_batch; + } + } else { + // 批量提交，减少内部锁冲突 + s = txn->commit(); + if (!s.ok()) { + DB_FATAL("index %ld, region_id: %ld commit failed, status: %s", + index_id, _region_id, s.ToString().c_str()); + continue; + } + } + txn.reset(new Transaction(0, &_txn_pool)); + if (ttl_runtime_state != nullptr) { + txn->set_resource(resource); + ttl_runtime_state->set_txn(txn); } - txn.reset(new Transaction(0, nullptr)); txn->begin(txn_opt); } } @@ -8568,14 +8937,24 @@ void Region::ttl_remove_expired_data() { _binlog_param.data_cf_oldest_ts = ts; } } - if (num_remove_lines % 100 != 0) { - s = txn->commit(); - if (!s.ok()) { - DB_FATAL("index %ld, region_id: %ld commit failed, status: %s", - index_id, _region_id, s.ToString().c_str()); + if (num_remove_lines % commit_batch != 0) { + if (delete_by_pk) { + // ttl_delete_node内部会commit或rollback + if (ttl_delete_node->open(ttl_runtime_state.get()) < 0) { + DB_FATAL("index %ld, region_id: %ld ttl_delete_node open failed", + index_id, _region_id); + del_pk_num -= (num_remove_lines % commit_batch); + } + } else { + // 批量提交，减少内部锁冲突 + s = txn->commit(); + if (!s.ok()) { + DB_FATAL("index %ld, region_id: %ld commit failed, status: %s", + index_id, _region_id, s.ToString().c_str()); + } } } - DB_WARNING("scan index:%ld, cost: %ld, scan count: %ld, remove lines: %ld, region_id: %ld", + DB_WARNING("scan index:%ld, cost: %ld, scan count: %ld, remove lines: %ld, region_id: %ld", index_id, cost.get_time(), count, num_remove_lines, _region_id); if (index_id == global_index_id) { // num_table_lines维护不准，用来空region merge @@ -8736,6 +9115,11 @@ void Region::check_peer_latency() { } else { _node.get_status(&status); } + if (status.state == raft::STATE_ERROR) { + _raft_status_error = true; + } else { + _raft_status_error = false; + } int64_t dml_latency = get_dml_latency(); int64_t latency = dml_latency * (status.committed_index - _applied_index); diff --git a/src/store/region_binlog.cpp b/src/store/region_binlog.cpp index bc68d9d2d..0be36fe32 100755 --- a/src/store/region_binlog.cpp +++ b/src/store/region_binlog.cpp @@ -43,6 +43,7 @@ DEFINE_int64(binlog_seek_batch, 10000, "10000"); DEFINE_int64(binlog_use_seek_interval_min, 60, "1h"); DEFINE_int64(offline_binlog_size_peer_sst, 1073741824LL, "defualt 1GB, -1 means no limit"); DEFINE_bool(binlog_force_get, false, "false"); +DECLARE_bool(use_binlog_cache); DECLARE_int64(print_time_us); DECLARE_string(meta_server_bns); DECLARE_string(db_path); @@ -403,9 +404,12 @@ std::string Region::binlog_get_str_val(const std::string& name, const std::mapsecond.get_string(); } -void Region::binlog_get_scan_fields(std::map& field_ids, std::vector& field_slot, - SmartTable& binlog_table, SmartIndex& binlog_pri) { +void Region::binlog_get_scan_fields(std::map& field_ids, BinlogReadFields& binlog_fields, std::vector& field_slot, + SmartTable& binlog_table, SmartIndex& binlog_pri, bool read_all) { field_slot.resize(binlog_table->fields.back().id + 1); + static const std::unordered_set read_field = { + "binlog_type", "start_ts", "binlog_row_cnt" + }; std::set pri_field_ids; for (auto& field_info : binlog_pri->fields) { @@ -413,16 +417,27 @@ void Region::binlog_get_scan_fields(std::map& field_ids, st } for (auto& field : binlog_table->fields) { - field_slot[field.id] = field.id; - if (pri_field_ids.count(field.id) == 0) { - field_ids[field.id] = &field; + if (field.short_name == "ts") { + binlog_fields.ts_field = &field; + } else if (field.short_name == "binlog_type") { + binlog_fields.binlog_type_field = &field; + } else if (field.short_name == "start_ts") { + binlog_fields.start_ts_field = &field; + } else if (field.short_name == "binlog_row_cnt") { + binlog_fields.binlog_row_cnt_field = &field; + } + if (read_all || read_field.count(field.short_name) == 1) { + field_slot[field.id] = field.id; + if (pri_field_ids.count(field.id) == 0) { + field_ids[field.id] = &field; + } } } } void Region::binlog_get_field_values(std::map& field_value_map, SmartRecord& record, SmartTable& binlog_table) { for (auto& field : binlog_table->fields) { - auto f = record->get_field_by_tag(field.id); + auto f = record->get_field_by_idx(field.pb_idx); field_value_map[field.short_name] = record->get_value(f); } } @@ -515,7 +530,7 @@ int Region::binlog_scan_when_restart() { ExprValue value; value.type = pb::INT64; value._u.int64_val = begin_ts; - left_record->set_value(left_record->get_field_by_tag(1), value); + left_record->set_value(left_record->get_field_by_idx(0), value); right_record->decode(""); MutTableKey left_key, right_key; @@ -528,9 +543,10 @@ int Region::binlog_scan_when_restart() { &_region_info, 1, 0, false, false, false); std::map field_ids; + BinlogReadFields binlog_fields; std::vector field_slot; - binlog_get_scan_fields(field_ids, field_slot, binlog_table, binlog_pri); + binlog_get_scan_fields(field_ids, binlog_fields, field_slot, binlog_table, binlog_pri); TableIterator* table_iter = Iterator::scan_primary(nullptr, range, field_ids, field_slot, false, true); if (table_iter == nullptr) { @@ -869,12 +885,12 @@ int Region::write_binlog_value(const std::map& field_val default_value = ExprValue::Now(); default_value.cast_to(field.type); } - if (0 != record->set_value(record->get_field_by_tag(field.id), default_value)) { + if (0 != record->set_value(record->get_field_by_idx(field.pb_idx), default_value)) { DB_WARNING("fill insert value failed"); return -1; } } else { - if (0 != record->set_value(record->get_field_by_tag(field.id), iter->second)) { + if (0 != record->set_value(record->get_field_by_idx(field.pb_idx), iter->second)) { DB_WARNING("fill insert value failed"); return -1; } @@ -1341,65 +1357,6 @@ int BinlogReadMgr::get_binlog_finish(pb::StoreRes* response) { return 0; } -int64_t Region::read_data_cf_oldest_ts() { - int64_t begin_ts = 0; - SmartTable binlog_table = _factory->get_table_info_ptr(get_table_id()); - SmartIndex binlog_pri = _factory->get_index_info_ptr(get_table_id()); - SmartRecord left_record = _factory->new_record(*binlog_table); - SmartRecord right_record = _factory->new_record(*binlog_table); - if (left_record == nullptr || right_record == nullptr) { - return -1; - } - - ExprValue value; - value.type = pb::INT64; - value._u.int64_val = begin_ts; - left_record->set_value(left_record->get_field_by_tag(1), value); - right_record->decode(""); - - MutTableKey left_key, right_key; - if (left_record->encode_key(*binlog_pri.get(), left_key, binlog_pri.get()->fields.size(), false, false) != 0) { - DB_FATAL("Fail to encode_key left, table:%ld", binlog_table.get()->id); - return -1; - } - left_key.set_full(true); - IndexRange range(left_key, right_key, binlog_pri.get(), binlog_pri.get(), - &_region_info, 1, 0, false, false, false); - - std::map field_ids; - std::vector field_slot; - binlog_get_scan_fields(field_ids, field_slot, binlog_table, binlog_pri); - - TableIterator* table_iter = Iterator::scan_primary(nullptr, range, field_ids, field_slot, false, true); - if (table_iter == nullptr) { - DB_WARNING("open TableIterator fail, table_id:%ld", get_table_id()); - return -1; - } - - ON_SCOPE_EXIT(([this, table_iter]() { - delete table_iter; - })); - - SmartRecord record = _factory->new_record(*binlog_table); - int ret = 0; - record->clear(); - if (!table_iter->valid()) { - DB_WARNING("region_id: %ld table_iter is invalid", _region_id); - return -1; - } - - ret = table_iter->get_next(record); - if (ret < 0) { - DB_WARNING("region_id: %ld get_next failed", _region_id); - return -1; - } - - std::map field_value_map; - binlog_get_field_values(field_value_map, record, binlog_table); - int64_t ts = binlog_get_int64_val("ts", field_value_map); - return ts; -} - bool Region::flash_back_need_read(const pb::StoreReq* request, const std::map& field_value_map, const std::set& req_db_tables, @@ -1539,6 +1496,17 @@ void Region::read_binlog(const pb::StoreReq* request, return; } + if (FLAGS_use_binlog_cache && !is_read_offline_binlog && !request->binlog_desc().flash_back_read()) { + int ret = _binlog_cache.find(begin_ts, response); + if (ret == 0) { + int64_t select_cost = cost.get_time(); + Store::get_instance()->select_time_cost << select_cost; + static bvar::LatencyRecorder cache_binlog_time("cache_binlog_time", 60); + cache_binlog_time << cost.get_time(); + return; + } + } + SmartRecord left_record = _factory->new_record(*binlog_table); SmartRecord right_record = _factory->new_record(*binlog_table); if (left_record == nullptr || right_record == nullptr) { @@ -1550,7 +1518,7 @@ void Region::read_binlog(const pb::StoreReq* request, ExprValue value; value.type = pb::INT64; value._u.int64_val = begin_ts; - left_record->set_value(left_record->get_field_by_tag(1), value); + left_record->set_value(left_record->get_field_by_idx(0), value); right_record->decode(""); MutTableKey left_key, right_key; @@ -1564,9 +1532,10 @@ void Region::read_binlog(const pb::StoreReq* request, &_region_info, 1, 0, false, false, false); std::map field_ids; + BinlogReadFields binlog_fields; std::vector field_slot; - binlog_get_scan_fields(field_ids, field_slot, binlog_table, binlog_pri); + binlog_get_scan_fields(field_ids, binlog_fields, field_slot, binlog_table, binlog_pri, request->binlog_desc().flash_back_read()); TableIterator* table_iter = Iterator::scan_binlog_primary(range, field_ids, field_slot, is_read_offline_binlog); if (table_iter == nullptr) { @@ -1579,7 +1548,6 @@ void Region::read_binlog(const pb::StoreReq* request, })); int64_t max_fake_binlog = 0; - std::map field_value_map; SmartRecord record = _factory->new_record(*binlog_table); BinlogReadMgr binlog_reader(_region_id, begin_ts, remote_side, log_id, binlog_cnt, is_read_offline_binlog); int ret = 0; @@ -1587,7 +1555,7 @@ void Region::read_binlog(const pb::StoreReq* request, // SQL闪回预先插入 std::set req_db_tables; std::set req_signs; - if ((request->binlog_desc().flash_back_read() || request->binlog_desc().read_offline_binlog())) { + if (request->binlog_desc().flash_back_read()) { if (request->binlog_desc().db_tables_size() > 0) { for (const std::string& db_table : request->binlog_desc().db_tables()) { req_db_tables.insert(db_table); @@ -1613,11 +1581,13 @@ void Region::read_binlog(const pb::StoreReq* request, break; } std::map field_value_map; - binlog_get_field_values(field_value_map, record, binlog_table); - int64_t ts = binlog_get_int64_val("ts", field_value_map); // type 为 COMMIT 时，ts 为 commit_ts - BinlogType binlog_type = static_cast(binlog_get_int64_val("binlog_type", field_value_map)); - int64_t start_ts = binlog_get_int64_val("start_ts", field_value_map); - int64_t binlog_row_cnt = binlog_get_int64_val("binlog_row_cnt", field_value_map); + if (request->binlog_desc().flash_back_read()) { + binlog_get_field_values(field_value_map, record, binlog_table); + } + int64_t ts = binlog_fields.get_ts(record); // type 为 COMMIT 时，ts 为 commit_ts + BinlogType binlog_type = binlog_fields.get_binlog_type(record); + int64_t start_ts = binlog_fields.get_start_ts(record);; + int64_t binlog_row_cnt = binlog_fields.get_binlog_row_cnt(record); DB_DEBUG("ts:%ld,start_ts:%ld, binlog_type:%d", ts, start_ts, binlog_type); if (binlog_row_cnt <= 0) { binlog_row_cnt = 1; @@ -1645,7 +1615,7 @@ void Region::read_binlog(const pb::StoreReq* request, } // SQL闪回读取时过滤 - if ((request->binlog_desc().flash_back_read() || request->binlog_desc().read_offline_binlog()) + if (request->binlog_desc().flash_back_read() && !flash_back_need_read(request, field_value_map, req_db_tables, req_signs)) { continue; } @@ -1683,7 +1653,16 @@ void Region::read_binlog(const pb::StoreReq* request, response->set_errcode(pb::SUCCESS); response->set_errmsg("read binlog success"); + // 下发binlog才缓存，否则ts推进不了 + if (FLAGS_use_binlog_cache && response->binlogs_size() > 0 + && !is_read_offline_binlog && !request->binlog_desc().flash_back_read()) { + _binlog_cache.add(begin_ts, *response); + } int64_t select_cost = cost.get_time(); + if (response->binlogs_size() == 0) { + static bvar::LatencyRecorder no_binlog_time("no_binlog_time", 60); + no_binlog_time << cost.get_time(); + } Store::get_instance()->select_time_cost << select_cost; } @@ -1721,7 +1700,6 @@ void Region::query_binlog_ts(const pb::StoreReq* request, data_cf_oldest_ts = _binlog_param.data_cf_oldest_ts; } - int64_t begin_ts = 0; int64_t binlog_cf_oldest_ts = RocksWrapper::get_instance()->get_oldest_ts_in_binlog_cf(); auto binlog_info = response->mutable_binlog_info(); binlog_info->set_region_id(request->region_id()); @@ -1934,7 +1912,7 @@ class OfflineBinlogSstWriter { int OfflineBinlogSstWriter::open_writer() { rocksdb::Options option = RocksWrapper::get_instance()->get_cold_options(); option.env = rocksdb::Env::Default(); - _writer.reset(new SstFileWriter(option, false)); + _writer.reset(new SstFileWriter(option)); if (_writer == nullptr) { DB_FATAL("region_id: %ld backup task fail, SstFileWriter is nullptr", _region_id); return -1; @@ -2543,8 +2521,9 @@ int Region::write_offline_binlog_data() { */ std::map field_ids; + BinlogReadFields binlog_fields; std::vector field_slot; - binlog_get_scan_fields(field_ids, field_slot, binlog_table, binlog_pri); + binlog_get_scan_fields(field_ids, binlog_fields, field_slot, binlog_table, binlog_pri); MutTableKey prefix; prefix.append_i64(_region_id).append_i64(_table_id); @@ -2589,7 +2568,6 @@ int Region::write_offline_binlog_data() { int ret = 0; bool batch_finish = false; - std::map field_value_map; std::map start_binlog_map; SmartRecord record = _factory->new_record(*binlog_table); BinlogReadMgr binlog_reader(_region_id, _offline_binlog_task.backup_task_start_ts, "backup_task", 0, 0, false); diff --git a/src/store/region_control.cpp b/src/store/region_control.cpp index 2a0c79053..86d50c394 100644 --- a/src/store/region_control.cpp +++ b/src/store/region_control.cpp @@ -29,8 +29,8 @@ DECLARE_string(stable_uri); DECLARE_int64(transfer_leader_catchup_time_threshold); DECLARE_int64(store_heart_beat_interval_us); DEFINE_int32(compact_interval, 1, "compact_interval xx (s)"); -DEFINE_bool(allow_compact_range, true, "allow_compact_range"); -DEFINE_bool(allow_blocking_flush, true, "allow_blocking_flush"); +DEFINE_bool(allow_compact_range, true, "Allow compact range, default: true"); +DEFINE_bool(allow_blocking_flush, true, "Allow blocking flush, default: true"); DEFINE_bool(snapshot_consistency, true, "avoid bottommost compaction"); int RegionControl::remove_data(int64_t drop_region_id) { rocksdb::WriteOptions options; @@ -290,6 +290,7 @@ int RegionControl::remove_log_entry(int64_t drop_region_id) { opt.fill_cache = false; std::unique_ptr iter(rocksdb->new_iterator(opt, rocksdb->get_raft_log_handle())); iter->Seek(log_data_key.data()); + RocksdbVars::get_instance()->raft_log_scan_times_count << 1; if (iter->Valid()) { int64_t log_index = TableKey(iter->key()).extract_i64(sizeof(int64_t) + 1); rocksdb::Slice value(iter->value()); diff --git a/src/store/region_olap.cpp b/src/store/region_olap.cpp index 267202c66..2dc5b64e9 100644 --- a/src/store/region_olap.cpp +++ b/src/store/region_olap.cpp @@ -1216,7 +1216,7 @@ int ColdFileSstWriter::write_kv(const rocksdb::Slice& key, const rocksdb::Slice& if (_writer == nullptr) { rocksdb::Options option = RocksWrapper::get_instance()->get_cold_options(); option.env = rocksdb::Env::Default(); - _writer = new SstFileWriter(option, false); + _writer = new SstFileWriter(option); _write_count = 0; _cost.reset(); _tmp_file_name = FLAGS_db_path + "_tmp/" + std::to_string(_region_id) + "_" +std::to_string(butil::gettimeofday_us()); diff --git a/src/store/store.cpp b/src/store/store.cpp index eb1cff77e..d6bdba80e 100644 --- a/src/store/store.cpp +++ b/src/store/store.cpp @@ -54,8 +54,6 @@ DEFINE_int32(update_used_size_interval_us, 10 * 1000 * 1000, "update used size i DEFINE_int32(init_region_concurrency, 10, "init region concurrency when start"); DEFINE_int32(split_threshold , 150, "split_threshold, default: 150% * region_size / 100"); DEFINE_int64(min_split_lines, 200000, "min_split_lines, protected when wrong param put in table"); -DEFINE_int64(flush_region_interval_us, 10 * 60 * 1000 * 1000LL, - "flush region interval, default(10 min)"); DEFINE_int64(transaction_clear_interval_ms, 5000LL, "transaction clear interval, default(5s)"); DEFINE_int64(binlog_timeout_check_ms, 10 * 1000LL, @@ -68,6 +66,8 @@ DECLARE_int64(flush_memtable_interval_us); DEFINE_int32(max_split_concurrency, 2, "max split region concurrency, default:2"); DEFINE_int64(none_region_merge_interval_us, 5 * 60 * 1000 * 1000LL, "none region merge interval, default(5 min)"); +DEFINE_int64(region_merge_for_no_write_timeout_s, 0LL, + "region_merge_for_no_write_timeout_s, default not merge"); DEFINE_int64(region_delay_remove_timeout_s, 3600 * 24LL, "region_delay_remove_time_s, default(1d)"); DEFINE_bool(use_approximate_size, true, @@ -82,8 +82,8 @@ DEFINE_string(network_segment, "", "network segment of store set by user"); DEFINE_string(container_id, "", "container_id for zoombie instance"); DEFINE_int32(rocksdb_perf_level, rocksdb::kDisable, "rocksdb_perf_level"); DEFINE_bool(stop_ttl_data, false, "stop ttl data"); -DEFINE_bool(stop_cold_region_flush, false, "stop_cold_region_flush"); -DEFINE_bool(olap_region_split_enable, false, "olap_region_split_enable"); +DEFINE_bool(stop_cold_region_flush, false, "Stop cold region flush, default: false"); +DEFINE_bool(olap_region_split_enable, false, "Enable OLAP region split, default: false"); DEFINE_int64(check_peer_delay_min, 1, "check peer delay min"); DEFINE_bool(process_delete_regions_when_init, false, "process delete regions when init"); DECLARE_bool(store_rocks_hang_check); @@ -138,7 +138,7 @@ int Store::init_before_listen(std::vector& init_region_ids) { boost::trim(FLAGS_resource_tag); _resource_tag = FLAGS_resource_tag; -#ifdef BAIDU_INTERNAL +#if defined(BAIDU_INTERNAL) // 初始化外部文件系统，用于olap std::vector ugi_infos; ret = get_afs_infos(ugi_infos); @@ -1442,7 +1442,7 @@ void Store::query_file_system(google::protobuf::RpcController* controller, return; } if (ret != 0) { - DB_FATAL("remote_compaction_id: %s fail to do file operation, op: %s", + DB_WARNING("remote_compaction_id: %s fail to do file operation, op: %s", request->remote_compaction_id().c_str(), pb::CompactionOpType_Name(request->op_type()).c_str()); response->set_errcode(pb::COMPACTION_FILE_SYSTEM_ERROR); @@ -1792,9 +1792,13 @@ void Store::ttl_remove_thread() { if (_shutdown) { return; } - - traverse_copy_region_map([](const SmartRegion& region) { - region->update_ttl_info(); + + std::set update_failed_set; + traverse_copy_region_map([&update_failed_set](const SmartRegion& region) { + if (region->update_ttl_info() != 0) { + DB_FATAL("update ttl info failed, region_id: %ld", region->region_info().region_id()); + update_failed_set.emplace(region->get_region_id()); + } }); // 控制ttl在ttl_remove_interval_period指定时间范围,典型是晚上流量低峰 std::vector periods = string_split(FLAGS_ttl_remove_interval_period, '-'); @@ -1814,8 +1818,8 @@ void Store::ttl_remove_thread() { } if (time.get_time() > FLAGS_ttl_remove_interval_s * 1000 * 1000LL) { - traverse_copy_region_map([](const SmartRegion& region) { - if (!FLAGS_stop_ttl_data) { + traverse_copy_region_map([&update_failed_set](const SmartRegion& region) { + if (!FLAGS_stop_ttl_data && update_failed_set.count(region->get_region_id()) == 0) { region->ttl_remove_expired_data(); } }); @@ -2758,8 +2762,10 @@ void Store::whether_split_thread() { } } } - - if (!_factory->get_merge_switch(ptr_region->get_table_id())) { + // 开启merge开关，或者长时间不写（配置FLAGS_region_merge_for_no_write_timeout_s>0），进行merge操作，防止线上有太多空region + if (!_factory->get_merge_switch(ptr_region->get_table_id()) + && (FLAGS_region_merge_for_no_write_timeout_s <= 0 + || ptr_region->get_timecost() < FLAGS_region_merge_for_no_write_timeout_s * 1000 * 1000LL)) { continue; } //简化特殊处理，首尾region不merge @@ -3002,7 +3008,7 @@ int Store::get_used_size_per_region(const std::vector& region_ids, uint void Store::update_schema_info(const pb::SchemaInfo& table, std::map>* reverse_index_map, - std::unordered_set* vector_table_set) { + std::unordered_map>* vector_index_map) { //锁住的是update_table和table_info_mapping, table_info锁的位置不能改 _factory->update_table(table); if (table.has_deleted() && table.deleted()) { @@ -3016,8 +3022,8 @@ void Store::update_schema_info(const pb::SchemaInfo& table, } } if (index_info.index_type() == pb::I_VECTOR) { - if (vector_table_set != nullptr) { - vector_table_set->emplace(table.table_id()); + if (vector_index_map != nullptr) { + (*vector_index_map)[table.table_id()].emplace(std::make_pair(index_info.index_id(), index_info.state())); } } } @@ -3125,7 +3131,18 @@ void Store::construct_heart_beat_request( int64_t applied_index = 0; int64_t data_index = 0; _meta_writer->read_applied_index(region_id, &applied_index, &data_index); - Region::add_peer_info(request, region_info, applied_index, is_learner, true); + bool need_report = false; + // 刚add peer后，还没收到on_conf_change就重启，重启上报后会误删自己 + // 所以判断自己在不在peer中，保证收到on_conf_change才上报 + for (const auto& peer : region_info.peers()) { + if (peer == _address) { + need_report = true; + break; + } + } + if (need_report) { + Region::add_peer_info(request, region_info, applied_index, is_learner, true); + } } } @@ -3170,10 +3187,10 @@ void Store::process_heart_beat_response(const pb::StoreHeartBeatResponse& respon RocksWrapper::get_instance()->adjust_option(_param_map); } _meta_need_report = response.need_report(); - std::unordered_set vector_table_set; + std::unordered_map> vector_index_map; std::map> reverse_index_map; for (auto& schema_info : response.schema_change_info()) { - update_schema_info(schema_info, &reverse_index_map, &vector_table_set); + update_schema_info(schema_info, &reverse_index_map, &vector_index_map); } if (!reverse_index_map.empty()) { traverse_copy_region_map([this, &reverse_index_map](const SmartRegion& region) { @@ -3185,12 +3202,12 @@ void Store::process_heart_beat_response(const pb::StoreHeartBeatResponse& respon } }); } - if (!vector_table_set.empty()) { - traverse_copy_region_map([this, &vector_table_set](const SmartRegion& region) { + if (!vector_index_map.empty()) { + traverse_copy_region_map([this, &vector_index_map](const SmartRegion& region) { if (!region->removed()) { - auto iter = vector_table_set.find(region->get_table_id()); - if (iter != vector_table_set.end()) { - region->vector_schema_change(); + auto iter = vector_index_map.find(region->get_table_id()); + if (iter != vector_index_map.end()) { + region->vector_schema_change(iter->second); } } }); diff --git a/src/vector_index/vector_index.cpp b/src/vector_index/vector_index.cpp index 33527f0da..12433a7db 100755 --- a/src/vector_index/vector_index.cpp +++ b/src/vector_index/vector_index.cpp @@ -21,16 +21,21 @@ #include #include #include +#include +#include #include #include #include "arrow/vendored/fast_float/fast_float.h" +#include namespace baikaldb { DEFINE_int32(efsearch_retry_pow, 4, "efsearch_retry_pow"); -DEFINE_int32(max_efsearch, 16384, "max_efsearch"); DEFINE_int64(compaction_interval_s, 24 * 60 * 60, "vector compaction interval(s)"); DEFINE_bool(need_brute_force_when_insufficient, true, "need brute force when sufficient"); DEFINE_double(table_lines_compaction_threshold, 0.7, "compaction threshold"); +DEFINE_int32(first_train_vector_num, 10000, "first train vector num"); +DEFINE_int32(retraining_set_ratio, 10, "vector index retraining set ratio"); +DECLARE_int64(print_time_us); class IDSelectorBitmap : public faiss::IDSelectorBitmapUserDefine { public: @@ -74,6 +79,20 @@ class IDSelectorBitmap : public faiss::IDSelectorBitmapUserDefine { std::shared_ptr _del_bitmap; }; +int mres_to_int(const std::ssub_match& mr, int deflt = -1, int begin = 0) { + if (mr.length() == 0) { + return deflt; + } + int res = -1; + try { + res = std::stoi(mr.str().substr(begin)); + } catch (std::exception& e) { + DB_WARNING("Fail to mres_to_int: %s", e.what()); + return -1; + } + return res; +} + int VectorIndex::init(const pb::RegionInfo& region_info, int64_t region_id, int64_t index_id, int64_t table_id) { _rocksdb = RocksWrapper::get_instance(); _region_id = region_id; @@ -95,28 +114,89 @@ int VectorIndex::init(const pb::RegionInfo& region_info, int64_t region_id, int6 DB_WARNING("dimension fail:%d", _dimension); return -1; } - // TODO: 后续可以用正则匹配 - if (info.vector_description.empty()) { - _vector_description = "IDMap2,HNSW16"; - _is_hnsw = true; - } else if (info.vector_description == "Flat") { - _vector_description = "IDMap2,Flat"; - _is_flat = true; - } else if (info.vector_description == "L2norm,Flat") { - _vector_description = "IDMap2,L2norm,Flat"; - _is_flat = true; - _is_l2norm = true; - } else if (boost::istarts_with(info.vector_description, "HNSW")) { - _vector_description = "IDMap2," + info.vector_description; - _is_hnsw = true; - } else if (boost::istarts_with(info.vector_description, "L2norm,HNSW")) { - _vector_description = "IDMap2," + info.vector_description; + std::string desc = info.vector_description; + _vector_description = "IDMap2"; + // metric_type + if (boost::istarts_with(desc, "L2norm")) { _is_l2norm = true; + _vector_description += ",L2norm"; + desc = desc.substr(7); + } + // index type + if (desc.empty()) { + _vector_description += ",HNSW16"; _is_hnsw = true; } else { - DB_WARNING("not support description:%s", info.vector_description.c_str()); - return -1; + // SUPPORTED: + // decode by Flat: Flat \ HNSW \ IVF + // decode by PQ: PQmxn \ HNSW,PQmx8 \ IVF,PQmxn + // decode by PQFS: PQmx4fs \ IVF,PQmx4fs + std::smatch sm; + std::string code_string = desc; + if (std::regex_match(desc, sm, std::regex("HNSW([0-9]*)([,_].*)?"))) { + _is_hnsw = true; + int hnsw_M = mres_to_int(sm[1], 16); + if (hnsw_M <= 0) { + DB_WARNING("hnsw_M must be greater than 0: %d", hnsw_M); + return -1; + } + code_string = sm[2].length() > 0 ? sm[2].str().substr(1) : "Flat"; + _vector_description += ",HNSW" + std::to_string(hnsw_M); + } else if (std::regex_match(desc, sm, std::regex("IVF([0-9]*)([,_].*)?"))) { + _is_ivf = true; + _ivf_nlist = mres_to_int(sm[1], 16); + if (_ivf_nlist <= 0) { + DB_WARNING("ivf_nlist must be greater than 0: %d", _ivf_nlist); + return -1; + } + code_string = sm[2].length() > 0 ? sm[2].str().substr(1) : "Flat"; + _train_vec_size = std::max(_ivf_nlist, FLAGS_first_train_vector_num); + _vector_description += ",IVF" + std::to_string(_ivf_nlist); + } + // code, include PQ, PQFS, FLAT + if (std::regex_match(code_string, sm, std::regex("PQ([0-9]+)x4fs?"))) { + int M = mres_to_int(sm[1]); + if (M <= 0) { + DB_WARNING("PQ_M must be greater than 0: %d", M); + return -1; + } + _is_pqfs = true; + _train_vec_size = FLAGS_first_train_vector_num; + _vector_description += ",PQ" + std::to_string(M) + "x4fs"; + if (_is_hnsw) { + // faiss高版本支持 + DB_WARNING("not support HNSW+PQFS, desc: %s", _vector_description.c_str()); + return -1; + } + } else if (std::regex_match(code_string, sm, std::regex("PQ([0-9]+)(x[0-9]+)?"))) { + int M = mres_to_int(sm[1]); + int nbit = mres_to_int(sm[2], 8, 1); + if (M <= 0 || nbit <= 0) { + DB_WARNING("PQ_M and nbit must be greater than 0: %d, %d", M, nbit); + return -1; + } + _is_pq = true; + _train_vec_size = FLAGS_first_train_vector_num; + + if (_is_hnsw) { + if (nbit != 8) { + // faiss高版本支持 + DB_WARNING("faiss 1.7.4 only support PQx8 for hnsw, but defined nbit: %d, desc: %s", nbit, _vector_description.c_str()); + return -1; + } + _vector_description += ",PQ" + std::to_string(M); + } else { + _vector_description += ",PQ" + std::to_string(M) + "x" + std::to_string(nbit); + } + } else { + _is_flat = !_is_hnsw && !_is_ivf; + if (!_is_hnsw) { + // "HNSW"后不能加Flat, "IVF16"后必须带Flat + _vector_description += ",Flat"; + } + } } + DB_WARNING("info.vector_description: %s -> _vector_description: %s", info.vector_description.c_str(), _vector_description.c_str()); _metrix_type = static_cast(info.metric_type); if (_metrix_type == faiss::METRIC_INNER_PRODUCT) { _metric_compare = std::greater(); @@ -380,6 +460,12 @@ int VectorIndex::insert_vector( SmartRecord record) { int64_t cache_idx = 0; SmartFaissIndex faiss_index; + // if (_begin_train && _faiss_index->index != nullptr && !_faiss_index->index->is_trained) { + // // slow down for first vtrain -> [50ms, 1s] + // int64_t slow_down_us = std::max(_train_begin_time.get_time() / 10, (int64_t)50000); + // slow_down_us = std::min(slow_down_us, (int64_t)1000000); + // bthread_usleep(slow_down_us); + // } ScopeGuard auto_decrease([this, &faiss_index] () { if (_is_separate && faiss_index != nullptr) { faiss_index->dec_ref(); @@ -410,14 +496,15 @@ int VectorIndex::insert_vector( if (ret != 0) { return ret; } - return add_to_faiss(faiss_index, word, cache_idx, record); + return add_to_faiss(faiss_index, word, cache_idx, record, false); } int VectorIndex::delete_vector( SmartTransaction& txn, const std::string& word, const std::string& pk, - SmartRecord record) { + SmartRecord record, + pb::IndexState index_status) { SmartFaissIndex faiss_index; if (_is_separate) { uint64_t separate_value = 0; @@ -439,7 +526,7 @@ int VectorIndex::delete_vector( ScopeGuard write_auto_decrease([&faiss_index]() { faiss_index->real_writing_cond.decrease_signal(); }); - return del_to_rocksdb(txn->get_txn(), faiss_index, pk, V_DELETE); + return del_to_rocksdb(txn->get_txn(), faiss_index, pk, V_DELETE, index_status); } int VectorIndex::search_vector( @@ -449,7 +536,7 @@ int VectorIndex::search_vector( SmartTable& table_info, const std::string& search_data, int64_t topk, - int32_t efsearch, + VectorSearchParam& param, std::vector& records, std::vector& vector_filter_exprs, std::vector& scan_filter_exprs, @@ -464,7 +551,7 @@ int VectorIndex::search_vector( } else { faiss_index = _faiss_index; } - return search(txn, faiss_index, pk_info, table_info, search_data, topk, efsearch, + return search(txn, faiss_index, pk_info, table_info, search_data, topk, param, records, vector_filter_exprs, scan_filter_exprs, pre_filter_exprs); } @@ -525,7 +612,8 @@ int VectorIndex::del_to_rocksdb( myrocksdb::Transaction* txn, SmartFaissIndex faiss_index, const std::string& pk, - VFlag flag) { + VFlag flag, + pb::IndexState index_status) { if (faiss_index == nullptr) { DB_WARNING("faiss_index is nullptr"); return -1; @@ -543,6 +631,10 @@ int VectorIndex::del_to_rocksdb( std::string value; res = txn->Get(read_opt, data_cf, key, &value); if (!res.ok()) { + // TODO: 是否IsNotFound都能直接返回成功? 无论什么索引状态 + if (index_status != pb::IS_PUBLIC && res.IsNotFound()) { + return 0; + } DB_WARNING("rocksdb get error: code=%d, msg=%s", res.code(), res.ToString().c_str()); return -1; @@ -631,7 +723,8 @@ void from_chars_to_float_vec(const std::string& str, std::vector& vec) { int VectorIndex::add_to_faiss(SmartFaissIndex faiss_index, const std::string& word, int64_t cache_idx, - SmartRecord record) { + SmartRecord record, + bool enable_train) { if (faiss_index == nullptr) { DB_WARNING("faiss_index is nullptr"); return -1; @@ -660,12 +753,20 @@ int VectorIndex::add_to_faiss(SmartFaissIndex faiss_index, // 从标量缓存中删除不需要缓存的列，对这些列采用后过滤方式 std::shared_ptr scalar_data; std::shared_ptr flat_scalar_data; - if (VectorizeHelpper::change_arrow_record_batch_schema(arrow_record_batch->schema(), faiss_index->scalar_data, &scalar_data, true) != 0){ + if (VectorizeHelpper::change_arrow_record_batch_schema( + arrow_record_batch->schema(), + faiss_index->scalar_data, + &scalar_data, + /*need_check_momery_limit=*/false) != 0){ DB_WARNING("Fail to change arrow record batch schema"); return -1; } if (faiss_index->flat_scalar_data != nullptr && - VectorizeHelpper::change_arrow_record_batch_schema(arrow_record_batch->schema(), faiss_index->flat_scalar_data, &flat_scalar_data, true) != 0) { + VectorizeHelpper::change_arrow_record_batch_schema( + arrow_record_batch->schema(), + faiss_index->flat_scalar_data, + &flat_scalar_data, + /*need_check_momery_limit=*/false) != 0) { DB_WARNING("Fail to change arrow record batch schema"); return -1; } @@ -707,7 +808,27 @@ int VectorIndex::add_to_faiss(SmartFaissIndex faiss_index, DB_WARNING("concatenate_record_batches failed"); return -1; } - if (faiss_index->flat_index->ntotal >= 1000) { + + bool index_need_train = need_train(); + bool trigger_train = false; + if (enable_train + && index_need_train + && !faiss_index->index->is_trained + && _trigger_retrain_vector_num > 0 + && faiss_index->flat_index->ntotal > _trigger_retrain_vector_num) { + _begin_train = true; + trigger_train = true; + _train_begin_time.reset(); + } + + /* 三种情况: + * 1. 不需要train的向量类型, 如HNSW,Flat: + * a. 正常dml写入过程, 每满1000个向量，就将一级Flat添加到二级索引里 + * 2. 需要train的向量类型, 如IVF1000,Flat: + * a. 已经trained: 正常dml写入过程, 每满1000个向量，就将一级Flat添加到二级索引里 + * b. 还没trained: 索引重建过程, 到达_trigger_retrain_vector_num个向量时，触发train，然后将一级Flat添加到二级索引里 + */ + if (faiss_index->flat_index->ntotal >= 1000 && (!index_need_train || faiss_index->index->is_trained || trigger_train)) { // reconstruct from faiss_index, 减少缓存，降低内存占用 faiss::IndexIDMap2* idmap2_index = static_cast(faiss_index->flat_index); std::vector origin_cache_idxs; @@ -730,6 +851,19 @@ int VectorIndex::add_to_faiss(SmartFaissIndex faiss_index, } cache_idxs.insert(cache_idxs.end(), origin_cache_idxs.begin(), origin_cache_idxs.end()); cache_vectors.insert(cache_vectors.end(), origin_cache_vectors.begin(), origin_cache_vectors.end()); + if (trigger_train) { + try { + TimeCost cost; + faiss_index->index->train(cache_idxs.size(), &cache_vectors[0]); + _train_vec_size = cache_idxs.size(); + DB_WARNING("region_id: %ld, index: %s, dim: %d, index_size: %ld, train cost: %ld using vec size: %ld", + _region_id, _vector_description.c_str(), _dimension, cache_idxs.size(), cost.get_time(), _train_vec_size); + _begin_train = false; + } catch (std::exception& e) { + DB_WARNING("Fail to train faiss index: %s", e.what()); + return -1; + } + } // 合并flat层标量数据 record_batches = {faiss_index->scalar_data, flat_concat_record_batch}; std::shared_ptr concat_record_batch; @@ -777,7 +911,7 @@ int VectorIndex::search( SmartTable& table_info, const std::string& search_data, int64_t topk, - int32_t efsearch, + VectorSearchParam& param, std::vector& records, std::vector& vector_filter_exprs, std::vector& scan_filter_exprs, @@ -800,6 +934,7 @@ int VectorIndex::search( std::vector dis; int64_t least = 10; int64_t search_count = k; + TimeCost total_cost; { bthread::RWLockRdGuard lock(faiss_index->rwlock); @@ -876,92 +1011,123 @@ int VectorIndex::search( DB_WARNING("Fail to search, maybe not support search_params"); return -1; } - DB_WARNING("flat search:%ld, prefilter_cost:%ld", mink, prefilter_cost); - } else { - if (faiss_index->index == nullptr || faiss_index->flat_index == nullptr) { - DB_WARNING("faiss_index->index/faiss_index->flat_index is nullptr"); - return -1; + if (total_cost.get_time() > FLAGS_print_time_us) { + DB_WARNING("flat search:%ld, prefilter_cost:%ld", mink, prefilter_cost); } - int64_t mink = std::max(std::min(k, faiss_index->flat_index->ntotal), least); - flat_idxs.resize(mink, -1); - flat_dis.resize(mink, -1); - faiss::SearchParameters flat_search_params; - std::shared_ptr flat_bitmap_sel; - int64_t flat_prefilter_cost = 0; - if (has_filter) { - TimeCost tm; - auto arrow_res = arrow::compute::ExecuteScalarExpression( - arrow_filter_expr, *(faiss_index->flat_scalar_data->schema()), faiss_index->flat_scalar_data); - if (!arrow_res.status().ok()) { - DB_WARNING("Fail to ExecuteScalarExpression, %s", arrow_res.status().message().c_str()); + } else { + do { + if (faiss_index->index == nullptr || faiss_index->flat_index == nullptr) { + DB_WARNING("faiss_index->index/faiss_index->flat_index is nullptr"); return -1; } - flat_prefilter_cost = tm.get_time(); - auto flat_datum_array = (*arrow_res).array_as(); - flat_bitmap_sel = std::make_shared(flat_datum_array, faiss_index->flat_del_bitmap); - } else { - flat_bitmap_sel = std::make_shared(nullptr, faiss_index->flat_del_bitmap); - } - flat_search_params.sel = flat_bitmap_sel.get(); - try { - faiss_index->flat_index->search(1, &search_vector[0], mink, &flat_dis[0], &flat_idxs[0], &flat_search_params); - } catch (...) { - DB_WARNING("Fail to search, maybe not support search_params"); - return -1; - } - int64_t flat_search_cnt = 0; - for (int64_t i : flat_idxs) { - if (i == -1) { - break; + // search flat vector index first + int64_t mink = std::max(std::min(k, faiss_index->flat_index->ntotal), least); + flat_idxs.resize(mink, -1); + flat_dis.resize(mink, -1); + faiss::SearchParameters flat_search_params; + std::shared_ptr flat_bitmap_sel; + int64_t flat_prefilter_cost = 0; + if (has_filter) { + TimeCost tm; + auto arrow_res = arrow::compute::ExecuteScalarExpression( + arrow_filter_expr, *(faiss_index->flat_scalar_data->schema()), faiss_index->flat_scalar_data); + if (!arrow_res.status().ok()) { + DB_WARNING("Fail to ExecuteScalarExpression, %s", arrow_res.status().message().c_str()); + return -1; + } + flat_prefilter_cost = tm.get_time(); + auto flat_datum_array = (*arrow_res).array_as(); + flat_bitmap_sel = std::make_shared(flat_datum_array, faiss_index->flat_del_bitmap); } else { - ++flat_search_cnt; + flat_bitmap_sel = std::make_shared(nullptr, faiss_index->flat_del_bitmap); } - } - mink = std::max(std::min(k, faiss_index->index->ntotal), least); - idxs.resize(mink, -1); - dis.resize(mink, -1); - bool do_brute_force = false; - faiss::SearchParametersHNSW search_params; - search_params.efSearch = std::max(_efsearch, efsearch); - search_params.need_brute_force = true; // 如果过滤数量大于faiss索引数量的93%，则退化成暴搜； - search_params.need_brute_force_when_insufficient = FLAGS_need_brute_force_when_insufficient; // hnsw搜索结果不足topK时退化成暴搜； - search_params.do_brute_force = &do_brute_force; // hnsw搜索过程中是否退化暴搜 - std::shared_ptr bitmap_sel; - int64_t prefilter_cost = 0; - if (has_filter) { - TimeCost tm; - auto arrow_res = arrow::compute::ExecuteScalarExpression( - arrow_filter_expr, *(faiss_index->scalar_data->schema()), faiss_index->scalar_data); - if (!arrow_res.status().ok()) { - DB_WARNING("Fail to ExecuteScalarExpression, %s", arrow_res.status().message().c_str()); + flat_search_params.sel = flat_bitmap_sel.get(); + try { + faiss_index->flat_index->search(1, &search_vector[0], mink, &flat_dis[0], &flat_idxs[0], &flat_search_params); + } catch (...) { + DB_WARNING("Fail to search, maybe not support search_params"); return -1; } - prefilter_cost = tm.get_time(); - auto datum_array = (*arrow_res).array_as(); - bitmap_sel = std::make_shared(datum_array, faiss_index->del_bitmap); - } else { - bitmap_sel = std::make_shared(nullptr, faiss_index->del_bitmap); - } - search_params.sel = bitmap_sel.get(); - try { - faiss_index->index->search(1, &search_vector[0], mink, &dis[0], &idxs[0], &search_params); - } catch (...) { - DB_WARNING("Fail to search, maybe not support search_params"); - return -1; - } - int64_t search_cnt = 0; - for (int64_t i : idxs) { - if (i == -1) { + int64_t flat_search_cnt = 0; + for (int64_t i : flat_idxs) { + if (i == -1) { + break; + } else { + ++flat_search_cnt; + } + } + // search real vector index + if (faiss_index->index->ntotal == 0) { break; + } + mink = std::max(std::min(k, faiss_index->index->ntotal), least); + idxs.resize(mink, -1); + dis.resize(mink, -1); + std::shared_ptr bitmap_sel; + int64_t prefilter_cost = 0; + if (has_filter) { + TimeCost tm; + auto arrow_res = arrow::compute::ExecuteScalarExpression( + arrow_filter_expr, *(faiss_index->scalar_data->schema()), faiss_index->scalar_data); + if (!arrow_res.status().ok()) { + DB_WARNING("Fail to ExecuteScalarExpression, %s", arrow_res.status().message().c_str()); + return -1; + } + prefilter_cost = tm.get_time(); + auto datum_array = (*arrow_res).array_as(); + bitmap_sel = std::make_shared(datum_array, faiss_index->del_bitmap); } else { - ++search_cnt; + bitmap_sel = std::make_shared(nullptr, faiss_index->del_bitmap); } - } - DB_WARNING("region_id: %ld, index_id: %ld, separate_value: %lu, search:%ld, faiss.ntotal:%lu,%lu, " - "search_cnt: %ld,%ld, prefilter_time_cost: %ld, %ld, efsearch: %d, do_brute_force: %d", - _region_id, _index_id, faiss_index->separate_value, mink, faiss_index->flat_index->ntotal, faiss_index->index->ntotal, - flat_search_cnt, search_cnt, flat_prefilter_cost, prefilter_cost, search_params.efSearch, do_brute_force); - search_count = flat_search_cnt + search_cnt; + bool do_brute_force = false; + try { + if (_is_hnsw) { + // HNSW + Flat/PQ + faiss::SearchParametersHNSW search_params; + search_params.efSearch = std::max(_efsearch, param.efsearch); + search_params.need_brute_force = true; // 如果过滤数量大于faiss索引数量的93%，则退化成暴搜； + search_params.need_brute_force_when_insufficient = FLAGS_need_brute_force_when_insufficient; // hnsw搜索结果不足topK时退化成暴搜； + search_params.do_brute_force = &do_brute_force; // hnsw搜索过程中是否退化暴搜 + search_params.sel = bitmap_sel.get(); + faiss_index->index->search(1, &search_vector[0], mink, &dis[0], &idxs[0], &search_params); + } else if (_is_ivf) { + // IVF + Flat/PQ/PQFS + faiss::SearchParametersIVF search_params; + search_params.nprobe = std::max(_nprobe, param.nprobe); + search_params.sel = bitmap_sel.get(); + faiss_index->index->search(1, &search_vector[0], mink, &dis[0], &idxs[0], &search_params); + } else if (_is_pq) { + // PQ + faiss::SearchParametersPQ search_params; + search_params.search_type = faiss::IndexPQ::Search_type_t::ST_PQ; + search_params.sel = bitmap_sel.get(); + faiss_index->index->search(1, &search_vector[0], mink, &dis[0], &idxs[0], &search_params); + } else { + // PQFS + faiss::SearchParameters search_params; + search_params.sel = bitmap_sel.get(); + faiss_index->index->search(1, &search_vector[0], mink, &dis[0], &idxs[0], &search_params); + } + } catch (...) { + DB_WARNING("Fail to search, maybe not support search_params"); + return -1; + } + int64_t search_cnt = 0; + for (int64_t i : idxs) { + if (i == -1) { + break; + } else { + ++search_cnt; + } + } + if (total_cost.get_time() > FLAGS_print_time_us) { + DB_WARNING("region_id: %ld, index_id: %ld, separate_value: %lu, search:%ld, faiss.ntotal:%lu,%lu, " + "search_cnt: %ld,%ld, prefilter_time_cost: %ld, %ld, do_brute_force: %d", + _region_id, _index_id, faiss_index->separate_value, mink, faiss_index->flat_index->ntotal, faiss_index->index->ntotal, + flat_search_cnt, search_cnt, flat_prefilter_cost, prefilter_cost, do_brute_force); + } + search_count = flat_search_cnt + search_cnt; + } while (0); } } std::vector result_idxs; @@ -1012,8 +1178,11 @@ int VectorIndex::construct_records(myrocksdb::Transaction* txn, keys.reserve(idxs.size()); std::vector rocksdb_keys; rocksdb_keys.reserve(idxs.size()); + std::vector str_rocksdb_keys; + str_rocksdb_keys.reserve(idxs.size()); for (int64_t id : idxs) { - std::string key = _key_prefix; + std::string& key = str_rocksdb_keys.emplace_back(); + key = _key_prefix; uint8_t level = 0; key.append((char*)&level, sizeof(uint8_t)); if (_is_separate) { @@ -1282,7 +1451,13 @@ int VectorIndex::init_faiss_index(SmartFaissIndex faiss_index) { DB_WARNING("faiss_index is nullptr"); return -1; } - faiss_index->index = faiss::index_factory(_dimension, _vector_description.c_str(), _metrix_type); + try { + faiss_index->index = faiss::index_factory(_dimension, _vector_description.c_str(), _metrix_type); + } catch (std::exception& e) { + DB_FATAL("faiss_index->index init fail, _dimension: %d, _vector_description: %s, _metrix_type: %d, err: %s", + _dimension, _vector_description.c_str(), _metrix_type, e.what()); + return -1; + } if (faiss_index->index == nullptr) { DB_WARNING("faiss_index->index init fail"); return -1; @@ -1291,6 +1466,9 @@ int VectorIndex::init_faiss_index(SmartFaissIndex faiss_index) { faiss::ParameterSpace().set_index_parameter(faiss_index->index, "efSearch", _efsearch); faiss::ParameterSpace().set_index_parameter(faiss_index->index, "efConstruction", _efconstruction); } + if (_is_ivf) { + faiss::ParameterSpace().set_index_parameter(faiss_index->index, "nprobe", _nprobe); + } if (!_is_flat) { if (_is_l2norm) { faiss_index->flat_index = faiss::index_factory(_dimension, "IDMap2,L2norm,Flat", _metrix_type); @@ -1490,7 +1668,7 @@ int VectorIndex::restore_faiss_index(const pb::RegionInfo& region_info, SmartFai return ret; } bthread::RWLockWrGuard lock(faiss_index->rwlock); - ret = add_to_faiss(faiss_index, word, cache_idx, record); + ret = add_to_faiss(faiss_index, word, cache_idx, record, true); if (ret < 0) { DB_WARNING("add_to_faiss fail, index_id: %ld", _index_id); return ret; @@ -1562,9 +1740,9 @@ int VectorIndex::restore_faiss_index(const pb::RegionInfo& region_info, SmartFai } } DB_WARNING("restore_index success, region_id: %ld, faiss dump_idx: %ld, " - "faiss cache_idx:%ld, del_cnt:%ld, separate_value: %lu", + "faiss cache_idx:%ld, del_cnt:%ld, separate_value: %lu, train_vec_size: %ld", region_info.region_id(), faiss_index->dump_idx, faiss_index->cache_idx.load(), - del_cnt, faiss_index->separate_value); + del_cnt, faiss_index->separate_value, _train_vec_size); return 0; } @@ -1590,12 +1768,20 @@ int VectorIndex::schema_change_faiss_index(SmartFaissIndex faiss_index) { std::shared_ptr scalar_data; std::shared_ptr flat_scalar_data; if (faiss_index->scalar_data != nullptr && - VectorizeHelpper::change_arrow_record_batch_schema(arrow_schema, faiss_index->scalar_data, &scalar_data, true) != 0){ + VectorizeHelpper::change_arrow_record_batch_schema( + arrow_schema, + faiss_index->scalar_data, + &scalar_data, + /*need_check_momery_limit=*/false) != 0){ DB_WARNING("Fail to change arrow record batch schema"); return -1; } if (faiss_index->flat_scalar_data != nullptr && - VectorizeHelpper::change_arrow_record_batch_schema(arrow_schema, faiss_index->flat_scalar_data, &flat_scalar_data, true) != 0) { + VectorizeHelpper::change_arrow_record_batch_schema( + arrow_schema, + faiss_index->flat_scalar_data, + &flat_scalar_data, + /*need_check_momery_limit=*/false) != 0) { DB_WARNING("Fail to change arrow record batch schema"); return -1; } @@ -1605,6 +1791,19 @@ int VectorIndex::schema_change_faiss_index(SmartFaissIndex faiss_index) { return 0; } +void VectorIndex::cal_new_train_vec_size(int64_t valid_vec_size) { + if (!need_train()) { + _trigger_retrain_vector_num = -1; + return; + } + int64_t cal_train_vec_size = valid_vec_size * FLAGS_retraining_set_ratio / 100; + int64_t min_train_vec_size = FLAGS_first_train_vector_num; + if (_is_ivf) { + min_train_vec_size = std::max(min_train_vec_size, (int64_t)_ivf_nlist); + } + _trigger_retrain_vector_num = std::max(min_train_vec_size, cal_train_vec_size); +} + int VectorIndex::compact_faiss_index(const pb::RegionInfo& region_info, SmartFaissIndex faiss_index, bool is_force) { if (faiss_index == nullptr) { DB_WARNING("faiss_index is nullptr"); @@ -1686,23 +1885,46 @@ bool VectorIndex::need_compact_faiss_index(SmartFaissIndex faiss_index) { if (faiss_index->flat_index != nullptr) { ntotal += faiss_index->flat_index->ntotal; } + + int64_t vec_ntotal = ntotal - faiss_index->del_count; + if (faiss_index->index != nullptr && need_train()) { + int64_t trigger_train_vec_ntotal = _train_vec_size; // first train + if (faiss_index->index->is_trained) { + trigger_train_vec_ntotal = _train_vec_size * FLAGS_retraining_set_ratio * 2; + } + + if (vec_ntotal >= trigger_train_vec_ntotal) { + cal_new_train_vec_size(vec_ntotal); + DB_WARNING("triger train: vector_description: %s, vec_ntotal: %ld, trained: %d, _train_vec_size: %ld, trigger_train_vec_ntotal: %ld", + _vector_description.c_str(), + vec_ntotal, + faiss_index->index->is_trained, + _train_vec_size, + trigger_train_vec_ntotal); + return true; + } + } if (ntotal > 10 && faiss_index->del_count * 100 / ntotal > 30) { + cal_new_train_vec_size(vec_ntotal); DB_WARNING( - "region_id: %ld, _index_id: %ld, ntotal:%ld, _del_count:%ld", + "region_id: %ld, _index_id: %ld, ntotal:%ld, _del_count:%ld, trigger_retrain_vector_num: %ld", _region_id, _index_id, ntotal, - faiss_index->del_count.load()); + faiss_index->del_count.load(), + _trigger_retrain_vector_num); return true; } if (faiss_index->last_compaction_ts.get_time() > FLAGS_compaction_interval_s * 1000 * 1000LL && ntotal > 0 && faiss_index->del_count * 100 / ntotal > 5) { + cal_new_train_vec_size(vec_ntotal); DB_WARNING( - "region_id: %ld, _index_id: %ld, ntotal:%ld, _del_count:%ld", + "region_id: %ld, _index_id: %ld, ntotal:%ld, _del_count:%ld, trigger_retrain_vector_num: %ld", _region_id, _index_id, ntotal, - faiss_index->del_count.load()); + faiss_index->del_count.load(), + _trigger_retrain_vector_num); return true; } // 用于向量隔离索引删除空FaissIndex @@ -1964,7 +2186,20 @@ int VectorIndex::read_not_cache_fields_from_file(SmartFaissIndex faiss_index, co ScopeGuard auto_decrease([&extra_fs] () { extra_fs.close(); }); - std::string not_cache_fileds_str((std::istreambuf_iterator(extra_fs)), std::istreambuf_iterator()); + std::string file_str((std::istreambuf_iterator(extra_fs)), std::istreambuf_iterator()); + std::string not_cache_fileds_str = file_str; + std::vector vecs; + vecs.reserve(2); + boost::split(vecs, file_str, boost::is_any_of(";")); + if (vecs.size() == 2) { + try { + _train_vec_size = std::stoi(vecs[1]); + not_cache_fileds_str = vecs[0]; + } catch (...) { + DB_WARNING("Invalid _train_vec_size: %s", vecs[1].c_str()); + return -1; + } + } std::vector vec; vec.reserve(8); boost::split(vec, not_cache_fileds_str, boost::is_any_of(",")); @@ -1976,8 +2211,8 @@ int VectorIndex::read_not_cache_fields_from_file(SmartFaissIndex faiss_index, co return -1; } } - DB_WARNING("succ read_not_cache_fields_from_file, need_not_cache_fields: size: %lu, str: %s", - faiss_index->need_not_cache_fields.size(), not_cache_fileds_str.c_str()); + DB_WARNING("succ read_not_cache_fields_from_file, _train_vec_size: %ld, need_not_cache_fields: size: %lu, str: %s", + _train_vec_size, faiss_index->need_not_cache_fields.size(), file_str.c_str()); return 0; } @@ -2055,6 +2290,8 @@ int VectorIndex::write_not_cache_fields_to_file(SmartFaissIndex faiss_index, con if (!not_cache_fields_str.empty()) { not_cache_fields_str.pop_back(); } + // 减少文件数, 上次train向量数复用 + not_cache_fields_str += ";" + std::to_string(_train_vec_size); extra_fs << not_cache_fields_str; DB_WARNING("success write_not_cache_fields_to_file, file_name: %s, not_cache_fields_str: %s", full_file_name.c_str(), not_cache_fields_str.c_str()); diff --git a/test/test_arrow_compute.cpp b/test/test_arrow_compute.cpp index 2a769a734..dd266449b 100644 --- a/test/test_arrow_compute.cpp +++ b/test/test_arrow_compute.cpp @@ -19,6 +19,8 @@ #include #include #include "common.h" +#include "expr_value.h" +#include "arrow_function.h" #include #include #include @@ -36,7 +38,7 @@ #include "arrow/compute/registry.h" #include #include - +#include int main(int argc, char* argv[]) @@ -362,9 +364,135 @@ void run_test_case_when() { } } +template +struct test_string_cast_same_with_row { + using CType = typename arrow::TypeTraits::CType; + using BuilderType = typename arrow::TypeTraits::BuilderType; + + static void test(std::string func_name, std::shared_ptr& input_array, std::vector& row_values) { + std::cout << "func_name: " << func_name << std::endl; + + std::vector input_vector = {input_array}; + auto res = arrow::compute::CallFunction(func_name, input_vector); + EXPECT_TRUE(res.ok()); + + auto output_data = *res; + const CType* data = output_data.mutable_array()->GetValues(1); + for (int64_t i = 0; i < output_data.length(); ++i) { + CType expect = row_values[i].get_numberic(); + CType actual = *(data + i); + std::cout << "row: " << i << ", str: " << row_values[i].str_val << ", expect: " << expect << ", actual: " << actual << std::endl; + ASSERT_EQ(actual, expect); + } + } +}; + +struct test_string_cast_same_with_row_bool_type { + static void test(std::string func_name, std::shared_ptr& input_array, std::vector& row_values) { + std::cout << "func_name: " << func_name << std::endl; + + std::vector input_vector = {input_array}; + auto res = arrow::compute::CallFunction(func_name, input_vector); + EXPECT_TRUE(res.ok()); + + auto output_data = *res; + const uint8_t* bitmap_buffer = output_data.mutable_array()->buffers[1]->data(); + for (int64_t i = 0; i < output_data.length(); ++i) { + bool expect = row_values[i].get_numberic(); + bool actual = arrow::bit_util::GetBit(bitmap_buffer, i); + std::cout << "row: " << i << ", str: " << row_values[i].str_val << ", expect: " << expect << ", actual: " << actual << std::endl; + ASSERT_EQ(actual, expect); + } + } +}; + +void run_test_string_cast_numberic() { + auto input = std::make_shared(); + std::vector row_values; + std::vector test_strs = { + "", + "+", + "-", + "100", + "00101", + "+102", + "-103", + "104abc", + "-105abc", + "+106abc", + " 101003", + " +101004", + " -101004", + "+a1005", + "-a1006", + " +1007.999", + "-1008.111", + "255", // uint8_max + "256", // uint8_max + 1 + "65535", // uint16_max + "65536", // uint16_max + 1 + "4294967295", // uint32_max + "4294967296", // uint32_max + 1 + "18446744073709551615", // uint64_max + "18446744073709551616", // uint64_max + 1 + "127", // int8_max + "+127", // int8_max + "128", // int8_max + 1 + "+128", // int8_max + 1 + "-128", // int8_min + "-129", // int8_min - 1 + "32767", // int16_max + "+32767", // int16_max + "32768", // int16_max + 1 + "+32768", // int16_max + 1 + "-32768", // int16_min + "-32769", // int16_min - 1 + "2147483647", // int32_max + "+2147483647", // int32_max + "2147483648", // int32_max + 1 + "+2147483648", // int32_max + 1 + "-2147483648", // int32_min + "-2147483649", // int32_min - 1 + "9223372036854775807", // int64_max + "+9223372036854775807", // int64_max + "9223372036854775808", // int64_max + 1 + "+9223372036854775808", // int64_max + 1 + "-9223372036854775808", // int64_min + "-9223372036854775809", // int64_min - 1 + "null", + " +null" + }; + for (auto& str : test_strs) { + input->Append(str); + + ExprValue expr_val; + expr_val.type = pb::STRING; + expr_val.str_val = str; + row_values.emplace_back(expr_val); + } + std::shared_ptr input_array; + auto s = input->Finish(&input_array); + + test_string_cast_same_with_row::test("string_cast_uint8", input_array, row_values); + test_string_cast_same_with_row::test("string_cast_uint16", input_array, row_values); + test_string_cast_same_with_row::test("string_cast_uint32", input_array, row_values); + test_string_cast_same_with_row::test("string_cast_uint64", input_array, row_values); + test_string_cast_same_with_row::test("string_cast_int8", input_array, row_values); + test_string_cast_same_with_row::test("string_cast_int16", input_array, row_values); + test_string_cast_same_with_row::test("string_cast_int32", input_array, row_values); + test_string_cast_same_with_row::test("string_cast_int64", input_array, row_values); + test_string_cast_same_with_row::test("string_cast_float", input_array, row_values); + test_string_cast_same_with_row::test("string_cast_double", input_array, row_values); + test_string_cast_same_with_row_bool_type::test("string_cast_bool", input_array, row_values); +} TEST(test_arrow_vector_execute, case_all) { run_test_run_acero_sync(); run_test_case_when(); } + +TEST(test_string_cast_numberic, case_all) { + EXPECT_TRUE(baikaldb::ArrowFunctionManager::instance()->RegisterAllArrowFunction() == 0); + run_test_string_cast_numberic(); +} } // namespace baikal diff --git a/test/test_common.cpp b/test/test_common.cpp index 603fbb30e..374292db9 100644 --- a/test/test_common.cpp +++ b/test/test_common.cpp @@ -438,9 +438,9 @@ TEST(BvarMap, bvarmap) { uint64_t parent_sign = 2385825078143366794; std::set subquery_signs = {8394144613061275097, 8919421716185942419}; //bm << BvarMap(std::make_pair("abc", 1)); - bm << BvarMap("abc", 1, 101, 101, 10, 1, 5, 3, 100, 1, field_range_type, 1, parent_sign, subquery_signs); - bm << BvarMap("abc", 4, 102, 102, 20, 2, 6, 2, 100, 1, field_range_type, 1, parent_sign, subquery_signs); - bm << BvarMap("bcd", 5, 103, 103, 30, 3, 7, 1, 100, 1, field_range_type, 1, parent_sign, subquery_signs); + bm << BvarMap("abc", 1, 101, 101, 10, 1, 5, 3, 100, 1, field_range_type, 1, parent_sign, subquery_signs, ""); + bm << BvarMap("abc", 4, 102, 102, 20, 2, 6, 2, 100, 1, field_range_type, 1, parent_sign, subquery_signs, ""); + bm << BvarMap("bcd", 5, 103, 103, 30, 3, 7, 1, 100, 1, field_range_type, 1, parent_sign, subquery_signs, ""); std::cout << bm.get_value(); } @@ -483,10 +483,10 @@ TEST(test_gbk_regex, match) { std::cout << "wregex match result : " << boost::regex_match(wval1, wzhao_regex) << '\n'; } -DEFINE_int64(gflags_test_int64, 20000, ""); -DEFINE_double(gflags_test_double, 0.234, ""); -DEFINE_string(gflags_test_string, "abc", ""); -DEFINE_bool(gflags_test_bool, true, ""); +DEFINE_int64(gflags_test_int64, 20000, "Test int64 gflag with default value 20000"); +DEFINE_double(gflags_test_double, 0.234, "Test double gflag with default value 0.234"); +DEFINE_string(gflags_test_string, "abc", "Test string gflag with default value 'abc'"); +DEFINE_bool(gflags_test_bool, true, "Test boolean gflag with default value true"); TEST(gflags_test, case_all) { if (!google::SetCommandLineOption("gflags_test_int64", "1000").empty()) { @@ -529,7 +529,6 @@ TEST(gflags_test, case_all) { TEST(bns_to_meta_bns_test, case_all) { static std::map mapping = { - {"31.opera-adp-baikalStore-000-nj.FENGCHAO.njjs", "group.opera-ps-baikalMeta-000-bj.FENGCHAO.all"}, {"28.opera-atomkv-baikalStore-000-bj.FENGCHAO.bjhw", "group.opera-atomkv-baikalMeta-000-bj.FENGCHAO.all"}, {"2.opera-bigtree-baikalStore-000-bj.FENGCHAO.bjyz", "group.opera-atomkv-baikalMeta-000-bj.FENGCHAO.all"}, {"5.opera-coffline-baikalStore-000-mix.FENGCHAO.dbl", "group.opera-coffline-baikalMeta-000-bj.FENGCHAO.all"}, @@ -543,7 +542,7 @@ TEST(bns_to_meta_bns_test, case_all) { {"1.opera-aladdin-baikaldb-000-bj.FENGCHAO.dbl", "group.opera-aladdin-baikalMeta-000-bj.FENGCHAO.all"}, {"group.opera-aladdin-baikaldb-000-nj.FENGCHAO.all", "group.opera-aladdin-baikalMeta-000-bj.FENGCHAO.all"}, {"7.opera-aladdin-baikalStore-000-mix.FENGCHAO.gzhxy", "group.opera-aladdin-baikalMeta-000-bj.FENGCHAO.all"}, - {"0.opera-hmkv-baikalStore-000-bd.FENGCHAO.bddwd", "group.opera-holmes-baikalMeta-000-yq.FENGCHAO.all"}, + {"0.opera-hmkv-baikalStore-000-bdnew.FENGCHAO.bddwd", "group.opera-holmes-baikalMeta-000-yq.FENGCHAO.all"}, {"55.opera-hm-baikalStore-000-bd.FENGCHAO.bddwd", "group.opera-holmes-baikalMeta-000-yq.FENGCHAO.all"}, {"group.opera-hm-baikalStore-000-bd.FENGCHAO.all.serv", "group.opera-holmes-baikalMeta-000-yq.FENGCHAO.all"}, {"1.opera-adp-baikalBinlog-000-bj.FENGCHAO.bjhw", "group.opera-ps-baikalMeta-000-bj.FENGCHAO.all"}, @@ -574,9 +573,9 @@ TEST(brpc_http_get_info_test_db, case_all) { } } -DEFINE_int32(bvar_test_total_time_s, 0, ""); -DEFINE_int32(bvar_test_loop_time_ms, 100, ""); -DEFINE_int32(bvar_test_interaval_time_s, 60, ""); +DEFINE_int32(bvar_test_total_time_s, 0, "Total time in seconds for bvar testing with default value 0"); +DEFINE_int32(bvar_test_loop_time_ms, 100, "Loop time in milliseconds for bvar testing with default value 100"); +DEFINE_int32(bvar_test_interaval_time_s, 60, "Interval time in seconds for bvar testing with default value 60"); TEST(bvar_window_test, bvar) { bvar::Adder count; bvar::Window> window_count(&count, FLAGS_bvar_test_interaval_time_s); @@ -593,7 +592,7 @@ TEST(bvar_window_test, bvar) { bthread_usleep(FLAGS_bvar_test_loop_time_ms * 1000); } } -DEFINE_int64(ttl_time_us, 60, ""); +DEFINE_int64(ttl_time_us, 60, "Time to live in microseconds with default value 60"); TEST(ttl_test, ttl) { int64_t now_time = butil::gettimeofday_us(); uint64_t ttl_storage = ttl_encode(now_time); diff --git a/test/test_date_time.cpp b/test/test_date_time.cpp index 3c269db4c..3ff79bf21 100644 --- a/test/test_date_time.cpp +++ b/test/test_date_time.cpp @@ -37,6 +37,40 @@ namespace baikaldb { //extern uint64_t str_to_datetime(const char* str_time); //extern time_t datetime_to_timestamp(uint64_t datetime); //extern uint64_t timestamp_to_datetime(time_t timestamp); +TEST(test_localtime_mktime_fixed_r, case_all) { + time_t t; + std::cout << -1 % 86400 << " : " << -1 / 86400 << "\n"; + std::cout << -2000000000 % 86400 << " : " << -2000000000 / 86400 << "\n"; + for (t = -2000000000; t < 2000000000LL; t += 34567*345) { + struct tm tm_std, tm_fixed; + if (t < 0) { // 负数时间，localtime_r会差几分钟，不知道为啥，用gmtime_r来代替测试 + gmtime_r(&t, &tm_std); + localtime_fixed_r(&t, &tm_fixed, 0); + } else { + localtime_r(&t, &tm_std); + localtime_fixed_r(&t, &tm_fixed); + } + + EXPECT_EQ(tm_std.tm_year, tm_fixed.tm_year) << t; + EXPECT_EQ(tm_std.tm_mon, tm_fixed.tm_mon) << t; + EXPECT_EQ(tm_std.tm_mday, tm_fixed.tm_mday) << t; + // 不支持夏令时 + if (tm_std.tm_year < 1986 - 1900 || tm_std.tm_year > 1991 - 1900) { + EXPECT_EQ(tm_std.tm_hour, tm_fixed.tm_hour) << t; + } + EXPECT_EQ(tm_std.tm_min, tm_fixed.tm_min) << t; + EXPECT_EQ(tm_std.tm_sec, tm_fixed.tm_sec) << t; + EXPECT_EQ(tm_std.tm_yday, tm_fixed.tm_yday) << t; + EXPECT_EQ(tm_std.tm_wday, tm_fixed.tm_wday) << t; + if (t < 0) { // 负数时间，mktime会差几分钟，不知道为啥，用timegm来代替测试 + EXPECT_EQ(timegm(&tm_fixed), mktime_fixed_r(&tm_fixed, 0)); + } else { + EXPECT_EQ(mktime(&tm_fixed), mktime_fixed_r(&tm_fixed)); + } + + } +} + TEST(test_stamp_to_str, case_all) { uint64_t year = 0; uint64_t mon = 0; diff --git a/test/test_dms.cpp b/test/test_dms.cpp index 7f63bd3b9..4c339aaa5 100644 --- a/test/test_dms.cpp +++ b/test/test_dms.cpp @@ -13,7 +13,6 @@ // limitations under the License. #include - #include "dms.h" int main(int argc, char* argv[]) { diff --git a/test/test_partition_meta.cpp b/test/test_partition_meta.cpp index 762dc4d01..5da5731f2 100644 --- a/test/test_partition_meta.cpp +++ b/test/test_partition_meta.cpp @@ -74,6 +74,7 @@ void construct_schema(pb::SchemaInfo& schema_info) { schema_info.set_database("test_db"); schema_info.set_namespace_name("test_namespace"); schema_info.set_resource_tag("e0-nj"); + schema_info.set_partition_num(1); // Field schema_info.add_fields(); @@ -299,6 +300,7 @@ TEST_F(TestPartition, test_alter_partition) { ASSERT_NE(p_schema_info, nullptr); construct_schema(*p_schema_info); EXPECT_EQ(_schema_manager->pre_process_for_partition(&request_create_table, &response_create_table, 0), 0); + EXPECT_EQ(_schema_manager->pre_process_for_create_table(&request_create_table, &response_create_table, 0), 0); _table_manager->create_table(request_create_table, 1, NULL); { baikaldb::DoubleBufferedTableMemMapping::ScopedPtr info; diff --git a/test/test_privilege_manager.cpp b/test/test_privilege_manager.cpp index 55ddbddd1..af37d992a 100644 --- a/test/test_privilege_manager.cpp +++ b/test/test_privilege_manager.cpp @@ -235,6 +235,9 @@ TEST_F(PrivilegeManagerTest, test_create_drop_modify) { request_create_table_fc.mutable_table_info()->set_table_name("userinfo"); request_create_table_fc.mutable_table_info()->set_database("FC_Word"); request_create_table_fc.mutable_table_info()->set_namespace_name("FengChao"); + request_create_table_fc.mutable_table_info()->set_engine(baikaldb::pb::ROCKSDB); + request_create_table_fc.mutable_table_info()->set_partition_num(1); + // request_create_table_fc.mutable_table_info()->add_init_store("127.0.0.1:8010"); request_create_table_fc.mutable_table_info()->set_resource_tag("e0"); baikaldb::pb::FieldInfo* field = request_create_table_fc.mutable_table_info()->add_fields(); @@ -258,6 +261,8 @@ TEST_F(PrivilegeManagerTest, test_create_drop_modify) { index->set_index_type(baikaldb::pb::I_KEY); index->add_field_names("username"); index->add_field_names("type"); + baikaldb::pb::MetaManagerResponse response_create_table; + EXPECT_EQ(_schema_manager->pre_process_for_create_table(&request_create_table_fc, &response_create_table, 0), 0); _table_manager->create_table(request_create_table_fc, 1, NULL); ASSERT_EQ(2, _namespace_manager->_max_namespace_id); @@ -366,6 +371,8 @@ TEST_F(PrivilegeManagerTest, test_create_drop_modify) { request_create_table_fc_level.mutable_table_info()->set_namespace_name("FengChao"); // request_create_table_fc_level.mutable_table_info()->add_init_store("127.0.0.1:8010"); request_create_table_fc_level.mutable_table_info()->set_resource_tag("e0"); + request_create_table_fc_level.mutable_table_info()->set_engine(baikaldb::pb::ROCKSDB); + request_create_table_fc_level.mutable_table_info()->set_partition_num(1); //request_create_table_fc_level.mutable_table_info()->set_upper_table_name("userinfo"); field = request_create_table_fc_level.mutable_table_info()->add_fields(); field->set_field_name("userid"); @@ -389,6 +396,7 @@ TEST_F(PrivilegeManagerTest, test_create_drop_modify) { index->set_index_type(baikaldb::pb::I_KEY); index->add_field_names("planname"); index->add_field_names("type"); + EXPECT_EQ(_schema_manager->pre_process_for_create_table(&request_create_table_fc_level, &response_create_table, 0), 0); _table_manager->create_table(request_create_table_fc_level, 2, NULL); ASSERT_EQ(2, _namespace_manager->_max_namespace_id); ASSERT_EQ(3, _database_manager->_max_database_id); diff --git a/test/test_region_manager.cpp b/test/test_region_manager.cpp index b167c3d86..3a7ec89fe 100644 --- a/test/test_region_manager.cpp +++ b/test/test_region_manager.cpp @@ -230,6 +230,8 @@ TEST_F(TestManagerTest, test_create_drop_modify) { request_create_table_fc.mutable_table_info()->set_table_name("userinfo"); request_create_table_fc.mutable_table_info()->set_database("FC_Word"); request_create_table_fc.mutable_table_info()->set_namespace_name("FengChao"); + request_create_table_fc.mutable_table_info()->set_engine(baikaldb::pb::ROCKSDB); + request_create_table_fc.mutable_table_info()->set_partition_num(1); // request_create_table_fc.mutable_table_info()->add_init_store("127.0.0.1:8010"); request_create_table_fc.mutable_table_info()->set_resource_tag("e0"); baikaldb::pb::FieldInfo* field = request_create_table_fc.mutable_table_info()->add_fields(); @@ -253,6 +255,8 @@ TEST_F(TestManagerTest, test_create_drop_modify) { index->set_index_type(baikaldb::pb::I_KEY); index->add_field_names("username"); index->add_field_names("type"); + baikaldb::pb::MetaManagerResponse response_create_table; + EXPECT_EQ(_schema_manager->pre_process_for_create_table(&request_create_table_fc, &response_create_table, 0), 0); _table_manager->create_table(request_create_table_fc, 1, NULL); ASSERT_EQ(2, _namespace_manager->_max_namespace_id); @@ -354,6 +358,8 @@ TEST_F(TestManagerTest, test_create_drop_modify) { request_create_table_fc_level.mutable_table_info()->set_table_name("planinfo"); request_create_table_fc_level.mutable_table_info()->set_database("FC_Word"); request_create_table_fc_level.mutable_table_info()->set_namespace_name("FengChao"); + request_create_table_fc_level.mutable_table_info()->set_engine(baikaldb::pb::ROCKSDB); + request_create_table_fc_level.mutable_table_info()->set_partition_num(1); // request_create_table_fc_level.mutable_table_info()->add_init_store("127.0.0.1:8010"); request_create_table_fc_level.mutable_table_info()->set_resource_tag("e0"); field = request_create_table_fc_level.mutable_table_info()->add_fields(); @@ -378,6 +384,7 @@ TEST_F(TestManagerTest, test_create_drop_modify) { index->set_index_type(baikaldb::pb::I_KEY); index->add_field_names("planname"); index->add_field_names("type"); + EXPECT_EQ(_schema_manager->pre_process_for_create_table(&request_create_table_fc_level, &response_create_table, 0), 0); _table_manager->create_table(request_create_table_fc_level, 2, NULL); { @@ -572,7 +579,7 @@ TEST_F(TestManagerTest, test_create_drop_modify) { baikaldb::pb::MetaManagerRequest split_region_request; split_region_request.set_op_type(baikaldb::pb::OP_SPLIT_REGION); split_region_request.mutable_region_split()->set_region_id(1); - _region_manager->split_region(split_region_request, NULL); + _region_manager->split_region(split_region_request, 0, NULL); ASSERT_EQ(3, _region_manager->get_max_region_id()); _schema_manager->load_snapshot(); ASSERT_EQ(3, _region_manager->get_max_region_id()); diff --git a/test/test_rocksdb.cpp b/test/test_rocksdb.cpp index 90e5bcb80..8ab9efe03 100644 --- a/test/test_rocksdb.cpp +++ b/test/test_rocksdb.cpp @@ -16,28 +16,491 @@ #include "meta_rocksdb.h" #include #include "sst_file_writer.h" +DEFINE_bool(only_ingest_external, false, "enable log"); +DEFINE_string(ingest_file_name, "", "ingest file name"); using namespace baikaldb; class RocksdbTEST : public testing::Test { public: ~RocksdbTEST() {} protected: virtual void SetUp() { - _rocksdb = RocksWrapper::get_instance(); - if (!_rocksdb) { - DB_FATAL("create rocksdb handler failed"); - return; - } - int ret = _rocksdb->init("./rocks_db"); - if (ret != 0) { - DB_FATAL("rocksdb init failed: code:%d", ret); - return; - } + } virtual void TearDown() {} - baikaldb::RocksWrapper* _rocksdb; }; -TEST_F(RocksdbTEST, ingest_test) { +void print_metadata_info(baikaldb::RocksWrapper* rocksdb) { + std::vector metadatas; + rocksdb->get_db()->GetLiveFilesMetaData(&metadatas); + for (const auto& metadata : metadatas) { + DB_WARNING("LiveFileMetaData[cf_name: %s level: %d]; " + "smallestkey[%s]; largestkey[%s]; " + "SstFileMetaData[smallest_seqno: %lu largest_seqno: %lu num_reads_sampled: %lu being_compacted: %d " + "num_entries: %lu num_deletions: %lu file_creation_time: %lu name: %s db_path: %s]; " + "FileStorageInfo[relative_filename: %s directory: %s file_number: %lu size: %lu epoch_number: %lu file_checksum: %s file_checksum_func_name: %s].", + metadata.column_family_name.c_str(), metadata.level, + metadata.smallestkey.c_str(), metadata.largestkey.c_str(), + metadata.smallest_seqno, metadata.largest_seqno, metadata.num_reads_sampled, (int)metadata.being_compacted, + metadata.num_entries, metadata.num_deletions, metadata.file_creation_time, metadata.name.c_str(), metadata.db_path.c_str(), + metadata.relative_filename.c_str(), metadata.directory.c_str(), metadata.file_number, metadata.size, metadata.epoch_number, + metadata.file_checksum.c_str(), metadata.file_checksum_func_name.c_str()); + } +} + +void print_properties_info(baikaldb::RocksWrapper* rocksdb) { + rocksdb::TablePropertiesCollection props; + rocksdb->get_db()->GetPropertiesOfAllTables(rocksdb->get_data_handle(), &props); + for (const auto& item : props) { + DB_WARNING("TablePropertiesCollection[cf_name: %s] [props: %s]; ", item.first.c_str(), item.second->ToString("; ", ": ").c_str()); + } +} + +TEST_F(RocksdbTEST, ingest_test7) { + baikaldb::RocksWrapper* _rocksdb = RocksWrapper::get_instance(); + if (!_rocksdb) { + DB_FATAL("create rocksdb handler failed"); + return; + } + int ret = _rocksdb->init("./rocks_db7"); + if (ret != 0) { + DB_FATAL("rocksdb init failed: code:%d", ret); + return; + } + ON_SCOPE_EXIT([_rocksdb]() { + _rocksdb->close(); + }); + // auto dbimpl = static_cast(_rocksdb->get_db()->GetRootDB()); + for (int i = 0; i < 15; i++) { + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i+1))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i+1).c_str()); + } + rocksdb::Slice start("0"); + rocksdb::Slice end("999999"); + rocksdb::CompactRangeOptions compact_options; + compact_options.exclusive_manual_compaction = false; + DB_WARNING("begin compact"); + auto s = _rocksdb->compact_range(compact_options, _rocksdb->get_data_handle(), &start, &end); + if (!s.ok()) { + DB_WARNING("compact_range error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + print_metadata_info(_rocksdb); + + + int i = 1; + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i+2))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i+2).c_str()); + i = 3; + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i+2))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i+2).c_str()); + _rocksdb->remove(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice("4")); + DB_WARNING("remove key: %s", "4"); + // _rocksdb->remove_range(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice("6"), rocksdb::Slice("8"), false); + DB_WARNING("remove_range key: [6-8]"); + // _rocksdb->remove_range(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice("7"), rocksdb::Slice("9"), false); + DB_WARNING("remove_range key: [7-9]"); + // _rocksdb->remove_range(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice("4"), rocksdb::Slice("5"), false); + DB_WARNING("remove_range key: [4-5]"); + + rocksdb::FlushOptions flush_options; + auto status = _rocksdb->flush(flush_options, _rocksdb->get_data_handle()); + if (!status.ok()) { + DB_WARNING("flush data to rocksdb fail, err_msg:%s", status.ToString().c_str()); + } + + _rocksdb->remove_range(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice("1"), rocksdb::Slice("3"), false); + + print_metadata_info(_rocksdb); + std::string delete_range_str; + s = _rocksdb->get_db()->RangeTombstoneSummary(_rocksdb->get_data_handle(), 10, &delete_range_str); + if (!s.ok()) { + DB_WARNING("RangeTombstoneSummary error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + DB_WARNING("delete_range_str: %s", delete_range_str.c_str()); + rocksdb::ReadOptions read_options; + read_options.ignore_range_deletions = true; + s = _rocksdb->get_db()->InternalTravel(read_options, _rocksdb->get_data_handle(), rocksdb::kMaxSequenceNumber); + if (!s.ok()) { + DB_WARNING("InternalTravel error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + DB_WARNING("ignore_range_deletions: true"); + read_options.ignore_range_deletions = false; + s = _rocksdb->get_db()->InternalTravel(read_options, _rocksdb->get_data_handle(), rocksdb::kMaxSequenceNumber); + if (!s.ok()) { + DB_WARNING("InternalTravel error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + bthread_usleep(60*1000*1000); +} + +TEST_F(RocksdbTEST, ingest_test6) { + baikaldb::RocksWrapper* _rocksdb = RocksWrapper::get_instance(); + if (!_rocksdb) { + DB_FATAL("create rocksdb handler failed"); + return; + } + int ret = _rocksdb->init("./rocks_db6"); + if (ret != 0) { + DB_FATAL("rocksdb init failed: code:%d", ret); + return; + } + ON_SCOPE_EXIT([_rocksdb]() { + _rocksdb->close(); + }); + // auto dbimpl = static_cast(_rocksdb->get_db()->GetRootDB()); + for (int i = 0; i < 15; i++) { + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i+1))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i+1).c_str()); + } + rocksdb::Slice start("0"); + rocksdb::Slice end("999999"); + rocksdb::CompactRangeOptions compact_options; + compact_options.exclusive_manual_compaction = false; + DB_WARNING("begin compact"); + auto s = _rocksdb->compact_range(compact_options, _rocksdb->get_data_handle(), &start, &end); + if (!s.ok()) { + DB_WARNING("compact_range error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + print_metadata_info(_rocksdb); + FLAGS_enable_remote_compaction = true; + s = _rocksdb->compact_files(rocksdb::CompactionOptions(), _rocksdb->get_data_handle(), {"./rocks_db6/000020.sst"}, 6); + if (!s.ok()) { + DB_WARNING("compact_files error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + + FLAGS_enable_remote_compaction = false; + rocksdb::ReadOptions read_options; + std::unique_ptr iter2(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + for (iter2->SeekToFirst(); iter2->Valid(); iter2->Next()) { + rocksdb::Slice key_slice(iter2->key()); + rocksdb::Slice value_slice(iter2->value()); + DB_WARNING("scan1 key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } + + for (int i = 0; i < 10; i++) { + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i+2))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i+2).c_str()); + } + const rocksdb::Snapshot* snapshot = _rocksdb->get_snapshot(); + rocksdb::FlushOptions flush_options; + auto status = _rocksdb->flush(flush_options, _rocksdb->get_data_handle()); + if (!status.ok()) { + DB_WARNING("flush data to rocksdb fail, err_msg:%s", status.ToString().c_str()); + } + print_metadata_info(_rocksdb); + + iter2.reset(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + for (iter2->SeekToFirst(); iter2->Valid(); iter2->Next()) { + rocksdb::Slice key_slice(iter2->key()); + rocksdb::Slice value_slice(iter2->value()); + DB_WARNING("scan2 key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } + + s = _rocksdb->get_db()->InternalTravel(read_options, _rocksdb->get_data_handle(), rocksdb::kMaxSequenceNumber); + if (!s.ok()) { + DB_WARNING("InternalTravel error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + +} + +TEST_F(RocksdbTEST, ingest_test5) { + baikaldb::RocksWrapper* _rocksdb = RocksWrapper::get_instance(); + if (!_rocksdb) { + DB_FATAL("create rocksdb handler failed"); + return; + } + int ret = _rocksdb->init("./rocks_db5"); + if (ret != 0) { + DB_FATAL("rocksdb init failed: code:%d", ret); + return; + } + ON_SCOPE_EXIT([_rocksdb]() { + _rocksdb->close(); + }); + // auto dbimpl = static_cast(_rocksdb->get_db()->GetRootDB()); + for (int i = 0; i < 15; i++) { + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i).c_str()); + } + rocksdb::Slice start("0"); + rocksdb::Slice end("999999"); + rocksdb::CompactRangeOptions compact_options; + compact_options.exclusive_manual_compaction = false; + DB_WARNING("begin compact"); + auto s = _rocksdb->compact_range(compact_options, _rocksdb->get_data_handle(), &start, &end); + if (!s.ok()) { + DB_WARNING("compact_range error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + print_metadata_info(_rocksdb); + for (int i = 0; i < 10; i++) { + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i+1))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i+1).c_str()); + } + const rocksdb::Snapshot* snapshot = _rocksdb->get_snapshot(); + rocksdb::FlushOptions flush_options; + auto status = _rocksdb->flush(flush_options, _rocksdb->get_data_handle()); + if (!status.ok()) { + DB_WARNING("flush data to rocksdb fail, err_msg:%s", status.ToString().c_str()); + } + print_metadata_info(_rocksdb); + DB_WARNING("222222222222222"); + for (int i = 3; i < 7; i++) { + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i+2))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i+2).c_str()); + } + + status = _rocksdb->flush(flush_options, _rocksdb->get_data_handle()); + if (!status.ok()) { + DB_WARNING("flush data to rocksdb fail, err_msg:%s", status.ToString().c_str()); + } + + rocksdb::ReadOptions read_options; + std::unique_ptr iter2(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + for (iter2->SeekToFirst(); iter2->Valid(); iter2->Next()) { + rocksdb::Slice key_slice(iter2->key()); + rocksdb::Slice value_slice(iter2->value()); + DB_WARNING("scan key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } + print_metadata_info(_rocksdb); + read_options.table_filter = [&](const rocksdb::TableProperties& props) { + if (props.orig_file_number == 20) { + // DB_WARNING("RETRUN FALSE file_number: %lu TableProperties: %s", props.orig_file_number, props.ToString("; ", ": ").c_str()); + return false; + } else { + // DB_WARNING("RETRUN TRUE file_number: %lu TableProperties: %s", props.orig_file_number, props.ToString("; ", ": ").c_str()); + return true; + } + }; + s = _rocksdb->compact_files(rocksdb::CompactionOptions(), _rocksdb->get_data_handle(), {"./rocks_db5/000022.sst","./rocks_db5/000024.sst"}, 1); + if (!s.ok()) { + DB_WARNING("compact_files error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + print_metadata_info(_rocksdb); + s = _rocksdb->get_db()->InternalTravel(read_options, _rocksdb->get_data_handle(), rocksdb::kMaxSequenceNumber); + if (!s.ok()) { + DB_WARNING("InternalTravel error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + s = _rocksdb->compact_files(rocksdb::CompactionOptions(), _rocksdb->get_data_handle(), {"./rocks_db5/000025.sst"}, 5); + if (!s.ok()) { + DB_WARNING("compact_files error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + print_metadata_info(_rocksdb); + s = _rocksdb->get_db()->InternalTravel(read_options, _rocksdb->get_data_handle(), rocksdb::kMaxSequenceNumber); + if (!s.ok()) { + DB_WARNING("InternalTravel error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + _rocksdb->release_snapshot(snapshot); + bthread_usleep(10*1000*1000); + s = _rocksdb->compact_files(rocksdb::CompactionOptions(), _rocksdb->get_data_handle(), {"./rocks_db5/000026.sst"}, 5); + if (!s.ok()) { + DB_WARNING("compact_files error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + print_metadata_info(_rocksdb); + s = _rocksdb->get_db()->InternalTravel(read_options, _rocksdb->get_data_handle(), rocksdb::kMaxSequenceNumber); + if (!s.ok()) { + DB_WARNING("InternalTravel error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + bthread_usleep(10*1000*1000); + s = _rocksdb->compact_files(rocksdb::CompactionOptions(), _rocksdb->get_data_handle(), {"./rocks_db5/000027.sst"}, 5); + if (!s.ok()) { + DB_WARNING("compact_files error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + print_metadata_info(_rocksdb); + s = _rocksdb->get_db()->InternalTravel(read_options, _rocksdb->get_data_handle(), rocksdb::kMaxSequenceNumber); + if (!s.ok()) { + DB_WARNING("InternalTravel error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + int i = 5; + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i+3))); + i = 6; + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i+3))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i+3).c_str()); + // snapshot = _rocksdb->get_snapshot(); + // DB_WARNING("seq: %lu", snapshot->GetSequenceNumber()); + status = _rocksdb->flush(flush_options, _rocksdb->get_data_handle()); + if (!status.ok()) { + DB_WARNING("flush data to rocksdb fail, err_msg:%s", status.ToString().c_str()); + } + print_metadata_info(_rocksdb); + bthread_usleep(10*1000*1000); + s = _rocksdb->compact_files(rocksdb::CompactionOptions(), _rocksdb->get_data_handle(), {"./rocks_db5/000030.sst"}, 5); + if (!s.ok()) { + DB_WARNING("compact_files error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + print_metadata_info(_rocksdb); + // bthread_usleep(10*1000*1000); + // s = _rocksdb->compact_files(rocksdb::CompactionOptions(), _rocksdb->get_data_handle(), {"./rocks_db5/000029.sst"}, 5); + // if (!s.ok()) { + // DB_WARNING("compact_files error: code=%d, msg=%s", + // s.code(), s.ToString().c_str()); + // } + // print_metadata_info(_rocksdb); + bthread_usleep(60 * 1000 * 1000); +} + +TEST_F(RocksdbTEST, ingest_test4) { + baikaldb::RocksWrapper* _rocksdb = RocksWrapper::get_instance(); + if (!_rocksdb) { + DB_FATAL("create rocksdb handler failed"); + return; + } + int ret = _rocksdb->init("./rocks_db4"); + if (ret != 0) { + DB_FATAL("rocksdb init failed: code:%d", ret); + return; + } + ON_SCOPE_EXIT([_rocksdb]() { + _rocksdb->close(); + }); + // auto dbimpl = static_cast(_rocksdb->get_db()->GetRootDB()); + for (int i = 0; i < 15; i++) { + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i).c_str()); + } + rocksdb::Slice start("0"); + rocksdb::Slice end("999999"); + rocksdb::CompactRangeOptions compact_options; + compact_options.exclusive_manual_compaction = false; + DB_WARNING("begin compact"); + auto s = _rocksdb->compact_range(compact_options, _rocksdb->get_data_handle(), &start, &end); + if (!s.ok()) { + DB_WARNING("compact_range error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + print_metadata_info(_rocksdb); + DB_WARNING("1111111111111"); + for (int i = 0; i < 10; i++) { + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i+1))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i+1).c_str()); + } + const rocksdb::Snapshot* snapshot = _rocksdb->get_snapshot(); + rocksdb::FlushOptions flush_options; + auto status = _rocksdb->flush(flush_options, _rocksdb->get_data_handle()); + if (!status.ok()) { + DB_WARNING("flush data to rocksdb fail, err_msg:%s", status.ToString().c_str()); + } + print_metadata_info(_rocksdb); + DB_WARNING("222222222222222"); + for (int i = 0; i < 7; i++) { + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i+2))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i+2).c_str()); + } + + status = _rocksdb->flush(flush_options, _rocksdb->get_data_handle()); + if (!status.ok()) { + DB_WARNING("flush data to rocksdb fail, err_msg:%s", status.ToString().c_str()); + } + + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice("9"), rocksdb::Slice("12")); + DB_WARNING("put key: %s, value: %s", "9", "12"); + _rocksdb->remove(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice("5")); + DB_WARNING("remove key: %s", "5"); + rocksdb::ReadOptions read_options; + std::unique_ptr iter2(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + for (iter2->SeekToFirst(); iter2->Valid(); iter2->Next()) { + rocksdb::Slice key_slice(iter2->key()); + rocksdb::Slice value_slice(iter2->value()); + DB_WARNING("scan key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } + + DB_WARNING("InternalTravel begin seq: %lu", snapshot->GetSequenceNumber()); + s = _rocksdb->get_db()->InternalTravel(read_options, _rocksdb->get_data_handle(), snapshot->GetSequenceNumber()); + if (!s.ok()) { + DB_WARNING("InternalTravel error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + + print_metadata_info(_rocksdb); + read_options.table_filter = [&](const rocksdb::TableProperties& props) { + if (props.orig_file_number == 20) { + DB_WARNING("RETRUN FALSE file_number: %lu TableProperties: %s", props.orig_file_number, props.ToString("; ", ": ").c_str()); + return false; + } else { + DB_WARNING("RETRUN TRUE file_number: %lu TableProperties: %s", props.orig_file_number, props.ToString("; ", ": ").c_str()); + return true; + } + }; + iter2.reset(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + for (iter2->SeekToFirst(); iter2->Valid(); iter2->Next()) { + rocksdb::Slice key_slice(iter2->key()); + rocksdb::Slice value_slice(iter2->value()); + DB_WARNING("scan key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } + read_options.snapshot = nullptr; + // kMaxSequenceNumber + s = _rocksdb->get_db()->InternalTravel(read_options, _rocksdb->get_data_handle(), rocksdb::kMaxSequenceNumber); + if (!s.ok()) { + DB_WARNING("InternalTravel error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + + // for (int i = 0 ; i < 5; i++) { + // DB_WARNING("BEGIN READ ROUND %d", i); + // iter2.reset(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + // for (iter2->SeekToFirst(); iter2->Valid(); iter2->Next()) { + // rocksdb::Slice key_slice(iter2->key()); + // rocksdb::Slice value_slice(iter2->value()); + // DB_WARNING("key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + // } + // DB_WARNING("END READ ROUND %d", i); + // } + + // rocksdb::Slice start("0"); + // rocksdb::Slice end("999999"); + // rocksdb::CompactRangeOptions compact_options; + // compact_options.exclusive_manual_compaction = false; + // DB_WARNING("begin compact"); + // s = _rocksdb->compact_range(compact_options, _rocksdb->get_data_handle(), &start, &end); + // if (!s.ok()) { + // DB_WARNING("compact_range error: code=%d, msg=%s", + // s.code(), s.ToString().c_str()); + // } + + // print_metadata_info(_rocksdb); + // print_properties_info(_rocksdb); + // for (int i = 0 ; i < 3; i++) { + // DB_WARNING("BEGIN READ ROUND %d", i); + // iter2.reset(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + // for (iter2->SeekToFirst(); iter2->Valid(); iter2->Next()) { + // rocksdb::Slice key_slice(iter2->key()); + // rocksdb::Slice value_slice(iter2->value()); + // DB_WARNING("key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + // } + // DB_WARNING("END READ ROUND %d", i); + // } + + bthread_usleep(60 * 1000 * 1000); +} +TEST_F(RocksdbTEST, ingest_test3) { + baikaldb::RocksWrapper* _rocksdb = RocksWrapper::get_instance(); + if (!_rocksdb) { + DB_FATAL("create rocksdb handler failed"); + return; + } + int ret = _rocksdb->init("./rocks_db3"); + if (ret != 0) { + DB_FATAL("rocksdb init failed: code:%d", ret); + return; + } + ON_SCOPE_EXIT([_rocksdb]() { + _rocksdb->close(); + }); + // auto dbimpl = static_cast(_rocksdb->get_db()->GetRootDB()); std::unique_ptr writer(new SstFileWriter(_rocksdb->get_options(_rocksdb->get_data_handle()))); std::unique_ptr writer2(new SstFileWriter(_rocksdb->get_options(_rocksdb->get_data_handle()))); std::string path1 = "./test_ingest1.sst"; @@ -65,41 +528,395 @@ TEST_F(RocksdbTEST, ingest_test) { writer2->put(rocksdb::Slice("99"), rocksdb::Slice("99")); writer2->finish(); + const rocksdb::Snapshot* snapshot = nullptr; rocksdb::IngestExternalFileOptions ifo; ifo.move_files = true; ifo.write_global_seqno = false; ifo.allow_blocking_flush = false; + rocksdb::ReadOptions read_options; s = _rocksdb->ingest_external_file(_rocksdb->get_data_handle(), {path1}, ifo); if (!s.ok()) { DB_WARNING("ingest fail: %s", path1.c_str()); } - rocksdb::ReadOptions read_options; - std::unique_ptr iter(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); - int idx = 0; - for (iter->SeekToFirst(); iter->Valid(); iter->Next()) { - rocksdb::Slice key_slice(iter->key()); - rocksdb::Slice value_slice(iter->value()); - DB_WARNING("key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); - idx++; - } - ASSERT_EQ(idx, 4); s = _rocksdb->ingest_external_file(_rocksdb->get_data_handle(), {path2}, ifo); if (!s.ok()) { DB_WARNING("ingest fail: %s", path1.c_str()); } + print_metadata_info(_rocksdb); + + + + s = _rocksdb->get_db()->InternalTravel(read_options, _rocksdb->get_data_handle(), rocksdb::kMaxSequenceNumber); + if (!s.ok()) { + DB_WARNING("InternalTravel error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + + print_properties_info(_rocksdb); std::unique_ptr iter2(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); - idx = 0; for (iter2->SeekToFirst(); iter2->Valid(); iter2->Next()) { rocksdb::Slice key_slice(iter2->key()); rocksdb::Slice value_slice(iter2->value()); + DB_WARNING("scan key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } + read_options.table_filter = [&](const rocksdb::TableProperties& props) { + if (butil::fast_rand() % 2 == 0) { + DB_WARNING("RETRUN TRUE file_number: %lu TableProperties: %s", props.orig_file_number, props.ToString("; ", ": ").c_str()); + return true; + } else { + DB_WARNING("RETRUN FALSE file_number: %lu TableProperties: %s", props.orig_file_number, props.ToString("; ", ": ").c_str()); + return false; + } + }; + for (int i = 0 ; i < 5; i++) { + DB_WARNING("BEGIN READ ROUND %d", i); + iter2.reset(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + for (iter2->SeekToFirst(); iter2->Valid(); iter2->Next()) { + rocksdb::Slice key_slice(iter2->key()); + rocksdb::Slice value_slice(iter2->value()); + DB_WARNING("key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } + DB_WARNING("END READ ROUND %d", i); + } + + rocksdb::Slice start("0"); + rocksdb::Slice end("999999"); + rocksdb::CompactRangeOptions compact_options; + compact_options.exclusive_manual_compaction = false; + DB_WARNING("begin compact"); + s = _rocksdb->compact_range(compact_options, _rocksdb->get_data_handle(), &start, &end); + if (!s.ok()) { + DB_WARNING("compact_range error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + + print_metadata_info(_rocksdb); + print_properties_info(_rocksdb); + for (int i = 0 ; i < 3; i++) { + DB_WARNING("BEGIN READ ROUND %d", i); + iter2.reset(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + for (iter2->SeekToFirst(); iter2->Valid(); iter2->Next()) { + rocksdb::Slice key_slice(iter2->key()); + rocksdb::Slice value_slice(iter2->value()); + DB_WARNING("key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } + DB_WARNING("END READ ROUND %d", i); + } + + +} + +TEST_F(RocksdbTEST, ingest_test2) { + baikaldb::RocksWrapper* _rocksdb = RocksWrapper::get_instance(); + if (!_rocksdb) { + DB_FATAL("create rocksdb handler failed"); + return; + } + int ret = _rocksdb->init("./rocks_db2"); + if (ret != 0) { + DB_FATAL("rocksdb init failed: code:%d", ret); + return; + } + ON_SCOPE_EXIT([_rocksdb]() { + _rocksdb->close(); + }); + + std::unique_ptr writer(new SstFileWriter(_rocksdb->get_options(_rocksdb->get_data_handle()))); + std::string path1 = "./test_ingest1.sst"; + auto s = writer->open(path1); + if (!s.ok()) { + DB_FATAL("open sst file path: %s failed, err: %s", path1.c_str(), s.ToString().c_str()); + return; + } + for (int i = 0; i < 5; i++) { + writer->put(rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i))); + } + writer->finish(); + + const rocksdb::Snapshot* snapshot = nullptr; + rocksdb::IngestExternalFileOptions ifo; + ifo.move_files = true; + ifo.write_global_seqno = false; + ifo.allow_blocking_flush = false; + rocksdb::ReadOptions read_options; + s = _rocksdb->ingest_external_file(_rocksdb->get_data_handle(), {path1}, ifo); + if (!s.ok()) { + DB_WARNING("ingest fail: %s", path1.c_str()); + } + + std::unique_ptr iter(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + for (iter->SeekToFirst(); iter->Valid(); iter->Next()) { + rocksdb::Slice key_slice(iter->key()); + rocksdb::Slice value_slice(iter->value()); + DB_WARNING("scan key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } + print_metadata_info(_rocksdb); + FLAGS_enable_remote_compaction = true; + rocksdb::Slice start("0"); + rocksdb::Slice end("999999"); + rocksdb::CompactRangeOptions compact_options; + compact_options.exclusive_manual_compaction = false; + DB_WARNING("begin compact"); + s = _rocksdb->compact_range(compact_options, _rocksdb->get_data_handle(), &start, &end); + if (!s.ok()) { + DB_WARNING("compact_range error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + print_metadata_info(_rocksdb); + iter.reset(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + for (iter->SeekToFirst(); iter->Valid(); iter->Next()) { + rocksdb::Slice key_slice(iter->key()); + rocksdb::Slice value_slice(iter->value()); + DB_WARNING("scan key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } + DB_WARNING("ingest_test2 ENDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDD"); + FLAGS_enable_remote_compaction = false; +} + +TEST_F(RocksdbTEST, ingest_test1) { + baikaldb::RocksWrapper* _rocksdb = RocksWrapper::get_instance(); + if (!_rocksdb) { + DB_FATAL("create rocksdb handler failed"); + return; + } + int ret = _rocksdb->init("./rocks_db1"); + if (ret != 0) { + DB_FATAL("rocksdb init failed: code:%d", ret); + return; + } + ON_SCOPE_EXIT([_rocksdb]() { + _rocksdb->close(); + }); + + rocksdb::Slice start("0"); + rocksdb::Slice end("999999"); + rocksdb::CompactRangeOptions compact_options; + compact_options.exclusive_manual_compaction = false; + + for (int i = 0; i < 7; i++) { + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i).c_str()); + } + DB_WARNING("begin compact"); + auto s = _rocksdb->compact_range(compact_options, _rocksdb->get_data_handle(), &start, &end); + if (!s.ok()) { + DB_WARNING("compact_range error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + print_metadata_info(_rocksdb); + rocksdb::ReadOptions read_options; + std::unique_ptr iter(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + DB_WARNING("begin iterator"); + for (iter->SeekToFirst(); iter->Valid(); iter->Next()) { + rocksdb::Slice key_slice(iter->key()); + rocksdb::Slice value_slice(iter->value()); + DB_WARNING("scan key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } + const rocksdb::Snapshot* snapshot = _rocksdb->get_snapshot(); + for (int i = 0; i < 7; i++) { + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i+1))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i+1).c_str()); + } + rocksdb::WriteOptions write_opt; + rocksdb::TransactionOptions txn_opt; + auto txn = _rocksdb->begin_transaction(write_opt, txn_opt); + for (int i = 5; i < 10; i++) { + txn->Put(_rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i+2))); + DB_WARNING("put key: %s, value: %s", std::to_string(i).c_str(), std::to_string(i+2).c_str()); + } + + + + + DB_WARNING("begin compact1"); + s = _rocksdb->compact_range(compact_options, _rocksdb->get_data_handle(), &start, &end); + if (!s.ok()) { + DB_WARNING("compact_range error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + bthread_usleep(60 * 1000 * 1000LL); + iter.reset(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + DB_WARNING("begin iterator1"); + for (iter->SeekToFirst(); iter->Valid(); iter->Next()) { + rocksdb::Slice key_slice(iter->key()); + rocksdb::Slice value_slice(iter->value()); + DB_WARNING("scan key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } + + print_metadata_info(_rocksdb); + + + txn->Prepare(); + txn->Commit(); + _rocksdb->release_snapshot(snapshot); + DB_WARNING("begin compact2"); + s = _rocksdb->compact_range(compact_options, _rocksdb->get_data_handle(), &start, &end); + if (!s.ok()) { + DB_WARNING("compact_range error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + bthread_usleep(60 * 1000 * 1000LL); + print_metadata_info(_rocksdb); + iter.reset(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + DB_WARNING("begin iterator2"); + for (iter->SeekToFirst(); iter->Valid(); iter->Next()) { + rocksdb::Slice key_slice(iter->key()); + rocksdb::Slice value_slice(iter->value()); + DB_WARNING("scan key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } +} + +TEST_F(RocksdbTEST, ingest_test) { + baikaldb::RocksWrapper* _rocksdb = RocksWrapper::get_instance(); + if (!_rocksdb) { + DB_FATAL("create rocksdb handler failed"); + return; + } + int ret = _rocksdb->init("./rocks_db"); + if (ret != 0) { + DB_FATAL("rocksdb init failed: code:%d", ret); + return; + } + ON_SCOPE_EXIT([_rocksdb]() { + _rocksdb->close(); + }); + + std::unique_ptr writer(new SstFileWriter(_rocksdb->get_options(_rocksdb->get_data_handle()))); + std::unique_ptr writer2(new SstFileWriter(_rocksdb->get_options(_rocksdb->get_data_handle()))); + std::string path1 = "./test_ingest1.sst"; + std::string path2 = "./test_ingest2.sst"; + auto s = writer->open(path1); + if (!s.ok()) { + DB_FATAL("open sst file path: %s failed, err: %s", path1.c_str(), s.ToString().c_str()); + return; + } + writer->put(rocksdb::Slice("1"), rocksdb::Slice("1")); + writer->put(rocksdb::Slice("4"), rocksdb::Slice("4")); + writer->put(rocksdb::Slice("7"), rocksdb::Slice("7")); + writer->put(rocksdb::Slice("9"), rocksdb::Slice("9")); + writer->finish(); + + s = writer2->open(path2); + if (!s.ok()) { + DB_FATAL("open sst file path: %s failed, err: %s", path2.c_str(), s.ToString().c_str()); + return; + } + writer2->put(rocksdb::Slice("11"), rocksdb::Slice("11")); + writer2->put(rocksdb::Slice("3"), rocksdb::Slice("3")); + writer2->put(rocksdb::Slice("6"), rocksdb::Slice("6")); + writer2->put(rocksdb::Slice("9"), rocksdb::Slice("9_2")); + writer2->put(rocksdb::Slice("99"), rocksdb::Slice("99")); + writer2->finish(); + + const rocksdb::Snapshot* snapshot = nullptr; + rocksdb::IngestExternalFileOptions ifo; + ifo.move_files = true; + ifo.write_global_seqno = false; + ifo.allow_blocking_flush = false; + rocksdb::ReadOptions read_options; + if (!FLAGS_only_ingest_external) { + s = _rocksdb->ingest_external_file(_rocksdb->get_data_handle(), {path1}, ifo); + if (!s.ok()) { + DB_WARNING("ingest fail: %s", path1.c_str()); + } + + std::unique_ptr iter(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + int idx = 0; + for (iter->SeekToFirst(); iter->Valid(); iter->Next()) { + rocksdb::Slice key_slice(iter->key()); + rocksdb::Slice value_slice(iter->value()); + DB_WARNING("key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + idx++; + } + // ASSERT_EQ(idx, 4); + + s = _rocksdb->ingest_external_file(_rocksdb->get_data_handle(), {path2}, ifo); + if (!s.ok()) { + DB_WARNING("ingest fail: %s", path1.c_str()); + } + std::unique_ptr iter2(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + idx = 0; + for (iter2->SeekToFirst(); iter2->Valid(); iter2->Next()) { + rocksdb::Slice key_slice(iter2->key()); + rocksdb::Slice value_slice(iter2->value()); + DB_WARNING("key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + idx++; + } + // ASSERT_EQ(idx, 8); + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice("5"), rocksdb::Slice("5")); + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice("8"), rocksdb::Slice("8")); + for (int i = 0; i < 10; i++) { + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i))); + } + snapshot = _rocksdb->get_snapshot(); + for (int i = 90; i < 110; i++) { + _rocksdb->put(rocksdb::WriteOptions(), _rocksdb->get_data_handle(), rocksdb::Slice(std::to_string(i)), rocksdb::Slice(std::to_string(i))); + } + std::unique_ptr iter3(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + for (iter3->SeekToFirst(); iter3->Valid(); iter3->Next()) { + rocksdb::Slice key_slice(iter3->key()); + rocksdb::Slice value_slice(iter3->value()); + DB_WARNING("key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } + // MutTableKey start_key; + // MutTableKey end_key; + // start_key.append_u64(0); + // end_key.append_u64(UINT64_MAX); + rocksdb::Slice start("0"); + rocksdb::Slice end("999999"); + rocksdb::CompactRangeOptions compact_options; + compact_options.exclusive_manual_compaction = false; + s = _rocksdb->compact_range(compact_options, _rocksdb->get_data_handle(), &start, &end); + if (!s.ok()) { + DB_WARNING("compact_range error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + } + + bthread_usleep(60 * 1000 * 1000LL); + + + if (FLAGS_ingest_file_name != "") { + s = _rocksdb->ingest_external_file(_rocksdb->get_data_handle(), {FLAGS_ingest_file_name}, ifo); + if (!s.ok()) { + DB_WARNING("ingest fail: %s error: code=%d, msg=%s", FLAGS_ingest_file_name.c_str(), s.code(), s.ToString().c_str()); + } else { + DB_WARNING("ingest succ: %s", FLAGS_ingest_file_name.c_str()); + } + } + + std::unique_ptr iter4(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + for (iter4->SeekToFirst(); iter4->Valid(); iter4->Next()) { + rocksdb::Slice key_slice(iter4->key()); + rocksdb::Slice value_slice(iter4->value()); DB_WARNING("key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); - idx++; } - ASSERT_EQ(idx, 8); + if (snapshot) { + _rocksdb->release_snapshot(snapshot); + } + rocksdb::Slice start("0"); + rocksdb::Slice end("999999"); + rocksdb::CompactRangeOptions compact_options; + compact_options.exclusive_manual_compaction = false; + s = _rocksdb->compact_range(compact_options, _rocksdb->get_data_handle(), &start, &end); + if (!s.ok()) { + DB_WARNING("compact_range error: code=%d, msg=%s", + s.code(), s.ToString().c_str()); + } + bthread_usleep(60 * 1000 * 1000LL); + DB_WARNING("after compact range"); + iter4.reset(_rocksdb->new_iterator(read_options, _rocksdb->get_data_handle())); + for (iter4->SeekToFirst(); iter4->Valid(); iter4->Next()) { + rocksdb::Slice key_slice(iter4->key()); + rocksdb::Slice value_slice(iter4->value()); + DB_WARNING("key: %s, value: %s", key_slice.ToString(false).c_str(), value_slice.ToString(false).c_str()); + } + } // TEST_F int main(int argc, char** argv) { + google::ParseCommandLineFlags(&argc, &argv, true); testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); } diff --git a/test/test_schema_manager.cpp b/test/test_schema_manager.cpp index 0f41b9ac0..0fafbb227 100644 --- a/test/test_schema_manager.cpp +++ b/test/test_schema_manager.cpp @@ -113,7 +113,7 @@ TEST_F(SchemaManagerTest, test_create_drop_modify) { //测试点：region 分离，分配一个新的region-id baikaldb::pb::MetaManagerRequest request_split_region_fc; request_split_region_fc.mutable_region_split()->set_region_id(1); - _schema_manager->split_region(request_split_region_fc, NULL); + _schema_manager->split_region(request_split_region_fc, 0, NULL); ASSERT_EQ(2, _schema_manager->_max_region_id); for (auto& table_id : _schema_manager->_table_id_map) { DB_WARNING("table_id:%ld, name:%s", table_id.second, table_id.first.c_str()); @@ -561,7 +561,7 @@ TEST_F(SchemaManagerTest, test_create_drop_modify) { } baikaldb::pb::MetaManagerRequest request_split_region_feed; request_split_region_feed.mutable_region_split()->set_region_id(3); - _schema_manager->split_region(request_split_region_feed, NULL); + _schema_manager->split_region(request_split_region_feed, 0, NULL); ASSERT_EQ(4, _schema_manager->_max_region_id); _schema_manager->load_snapshot(); ASSERT_EQ(4, _schema_manager->_max_region_id); diff --git a/test/test_table_manager.cpp b/test/test_table_manager.cpp index 525f4b10d..0c205f0d2 100644 --- a/test/test_table_manager.cpp +++ b/test/test_table_manager.cpp @@ -232,6 +232,8 @@ TEST_F(TableManagerTest, test_create_drop_modify) { request_create_table_fc.mutable_table_info()->set_table_name("userinfo"); request_create_table_fc.mutable_table_info()->set_database("FC_Word"); request_create_table_fc.mutable_table_info()->set_namespace_name("FengChao"); + request_create_table_fc.mutable_table_info()->set_engine(baikaldb::pb::ROCKSDB); + request_create_table_fc.mutable_table_info()->set_partition_num(1); // request_create_table_fc.mutable_table_info()->add_init_store("127.0.0.1:8010"); request_create_table_fc.mutable_table_info()->set_resource_tag("e0"); baikaldb::pb::FieldInfo* field = request_create_table_fc.mutable_table_info()->add_fields(); @@ -255,6 +257,8 @@ TEST_F(TableManagerTest, test_create_drop_modify) { index->set_index_type(baikaldb::pb::I_KEY); index->add_field_names("username"); index->add_field_names("type"); + baikaldb::pb::MetaManagerResponse response_create_table; + EXPECT_EQ(_schema_manager->pre_process_for_create_table(&request_create_table_fc, &response_create_table, 0), 0); _table_manager->create_table(request_create_table_fc, 1, NULL); ASSERT_EQ(2, _namespace_manager->_max_namespace_id); @@ -278,8 +282,8 @@ TEST_F(TableManagerTest, test_create_drop_modify) { ASSERT_EQ(1, _database_manager->_table_ids[1].size()); ASSERT_EQ(1, _database_manager->_database_info_map[1].version()); - std::unordered_map table_id_map; - std::unordered_map table_info_map; + std::unordered_map table_id_map; + std::unordered_map table_info_map; { baikaldb::DoubleBufferedTableMemMapping::ScopedPtr info; if (_table_manager->_table_mem_infos.Read(&info) != 0) { @@ -359,6 +363,8 @@ TEST_F(TableManagerTest, test_create_drop_modify) { request_create_table_fc_level.mutable_table_info()->set_table_name("planinfo"); request_create_table_fc_level.mutable_table_info()->set_database("FC_Word"); request_create_table_fc_level.mutable_table_info()->set_namespace_name("FengChao"); + request_create_table_fc_level.mutable_table_info()->set_engine(baikaldb::pb::ROCKSDB); + request_create_table_fc_level.mutable_table_info()->set_partition_num(1); // request_create_table_fc_level.mutable_table_info()->add_init_store("127.0.0.1:8010"); //request_create_table_fc_level.mutable_table_info()->set_upper_table_name("userinfo"); request_create_table_fc_level.mutable_table_info()->set_resource_tag("e0"); @@ -384,6 +390,7 @@ TEST_F(TableManagerTest, test_create_drop_modify) { index->set_index_type(baikaldb::pb::I_KEY); index->add_field_names("planname"); index->add_field_names("type"); + EXPECT_EQ(_schema_manager->pre_process_for_create_table(&request_create_table_fc_level, &response_create_table, 0), 0); _table_manager->create_table(request_create_table_fc_level, 2, NULL); ASSERT_EQ(2, _namespace_manager->_max_namespace_id); ASSERT_EQ(3, _database_manager->_max_database_id); From 174266495f0d5fa747457c4988961b0c534c90d0 Mon Sep 17 00:00:00 2001 From: lgqss Date: Thu, 12 Feb 2026 14:56:24 +0800 Subject: [PATCH 2/5] merge_2025_12_05 fix compile --- cmake/boost.cmake | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmake/boost.cmake b/cmake/boost.cmake index e348ede8e..18d9c1546 100644 --- a/cmake/boost.cmake +++ b/cmake/boost.cmake @@ -20,7 +20,7 @@ SET(Boost_VERSION "106300") SET(Boost_LIB_VERSION "1_63_0") SET(BOOST_VER "1.63.0") SET(BOOST_TAR "boost_1_63_0" CACHE STRING "" FORCE) -SET(BOOST_URL "https://jaist.dl.sourceforge.net/project/boost/boost/1.63.0/${BOOST_TAR}.tar.gz" CACHE STRING "" FORCE) +SET(BOOST_URL "https://sourceforge.net/project/boost/boost/1.63.0/${BOOST_TAR}.tar.gz" CACHE STRING "" FORCE) MESSAGE(STATUS "BOOST_TAR: ${BOOST_TAR}, BOOST_URL: ${BOOST_URL}") From fa8cdc4d171509905f8b911347297a230d8f6b72 Mon Sep 17 00:00:00 2001 From: lgqss Date: Thu, 12 Feb 2026 20:46:00 +0800 Subject: [PATCH 3/5] merge_2025_12_05 fix compile2 --- cmake/arrow.cmake | 6 +++--- cmake/boost.cmake | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/cmake/arrow.cmake b/cmake/arrow.cmake index df9ea8717..ec774ca68 100644 --- a/cmake/arrow.cmake +++ b/cmake/arrow.cmake @@ -23,7 +23,7 @@ SET(ARROW_ACERO_LIB "${ARROW_INSTALL_DIR}/lib/libarrow_acero.a" CACHE FILEPATH " SET(ARROW_BUNDLED_DEP_LIB "${ARROW_INSTALL_DIR}/lib/libarrow_bundled_dependencies.a" CACHE FILEPATH "arrow dependencies." FORCE) FILE(WRITE ${ARROW_SOURCES_DIR}/src/build.sh - "cd cpp && cmake -DCMAKE_BUILD_TYPE=release -DARROW_JEMALLOC=OFF -DARROW_BUILD_SHARED=OFF -DARROW_PARQUET=ON -DARROW_WITH_ZLIB=ON -DARROW_WITH_ZSTD=ON -DARROW_WITH_LZ4=ON -DARROW_WITH_SNAPPY=ON -DARROW_COMPUTE=ON -DARROW_ACERO=ON -DARROW_FILESYSTEM=ON -DARROW_JSON=ON -DARROW_PARQUET=ON -DARROW_BUILD_TESTS=OFF -DARROW_BUILD_STATIC=ON -DARROW_WITH_RE2=ON -DARROW_RE2_VENDORED=OFF -DRE2_HOME=${RE2_INSTALL_DIR} -DCMAKE_PREFIX_PATH=${RE2_INSTALL_DIR} && make -j4" + "#!/bin/bash\nset -ex\ncd cpp\ncmake -DCMAKE_BUILD_TYPE=release -DARROW_JEMALLOC=OFF -DARROW_BUILD_SHARED=OFF -DARROW_PARQUET=ON -DARROW_WITH_ZLIB=ON -DARROW_WITH_ZSTD=ON -DARROW_WITH_LZ4=ON -DARROW_WITH_SNAPPY=ON -DARROW_COMPUTE=ON -DARROW_ACERO=ON -DARROW_FILESYSTEM=ON -DARROW_JSON=ON -DARROW_PARQUET=ON -DARROW_BUILD_TESTS=OFF -DARROW_BUILD_STATIC=ON -DARROW_WITH_RE2=ON -DARROW_RE2_VENDORED=OFF -DRE2_HOME=${RE2_INSTALL_DIR} -DARROW_GFLAGS=OFF -DCMAKE_PREFIX_PATH=${RE2_INSTALL_DIR}:${GFLAGS_INSTALL_DIR} -DCMAKE_CXX_FLAGS=\"-I${GFLAGS_INSTALL_DIR}/include\" && make -j4\n" ) ExternalProject_Add( @@ -48,7 +48,7 @@ ExternalProject_Add( COMMAND cp -r ${ARROW_SOURCES_DIR}/src/extern_arrow/cpp/src/parquet ${ARROW_INCLUDE_DIR}/ ) -ADD_DEPENDENCIES(extern_arrow zlib snappy zstd lz4 re2 protobuf rapidjson) +ADD_DEPENDENCIES(extern_arrow zlib snappy zstd lz4 re2 protobuf rapidjson gflags) ADD_LIBRARY(arrow STATIC IMPORTED GLOBAL) SET_PROPERTY(TARGET arrow PROPERTY IMPORTED_LOCATION ${ARROW_LIBRARIES}) ADD_LIBRARY(parquet STATIC IMPORTED GLOBAL) @@ -57,4 +57,4 @@ ADD_LIBRARY(acero STATIC IMPORTED GLOBAL) SET_PROPERTY(TARGET acero PROPERTY IMPORTED_LOCATION ${ARROW_ACERO_LIB}) ADD_LIBRARY(arrow_deps STATIC IMPORTED GLOBAL) SET_PROPERTY(TARGET arrow_deps PROPERTY IMPORTED_LOCATION ${ARROW_BUNDLED_DEP_LIB}) -ADD_DEPENDENCIES(arrow parquet acero arrow_deps extern_arrow) +ADD_DEPENDENCIES(arrow parquet acero arrow_deps extern_arrow ) diff --git a/cmake/boost.cmake b/cmake/boost.cmake index 18d9c1546..2d9518ab5 100644 --- a/cmake/boost.cmake +++ b/cmake/boost.cmake @@ -20,7 +20,7 @@ SET(Boost_VERSION "106300") SET(Boost_LIB_VERSION "1_63_0") SET(BOOST_VER "1.63.0") SET(BOOST_TAR "boost_1_63_0" CACHE STRING "" FORCE) -SET(BOOST_URL "https://sourceforge.net/project/boost/boost/1.63.0/${BOOST_TAR}.tar.gz" CACHE STRING "" FORCE) +SET(BOOST_URL "https://sourceforge.net/projects/boost/files/boost/1.63.0/boost_1_63_0.tar.gz" CACHE STRING "" FORCE) MESSAGE(STATUS "BOOST_TAR: ${BOOST_TAR}, BOOST_URL: ${BOOST_URL}") From d75f98f716315674a21037db997b9e1cc65c98b5 Mon Sep 17 00:00:00 2001 From: lgqss Date: Fri, 13 Feb 2026 10:29:22 +0800 Subject: [PATCH 4/5] merge_2025_12_05 fix compile3 --- include/exec/file_scan_node.h | 199 +++++ include/exec/ttl_delete_node.h | 42 + include/physical_plan/condition_optimizer.h | 33 + src/common/statistics.cpp | 52 ++ src/exec/file_scan_node.cpp | 921 ++++++++++++++++++++ src/exec/ttl_delete_node.cpp | 85 ++ src/physical_plan/condition_optimizer.cpp | 154 ++++ 7 files changed, 1486 insertions(+) create mode 100644 include/exec/file_scan_node.h create mode 100644 include/exec/ttl_delete_node.h create mode 100644 include/physical_plan/condition_optimizer.h create mode 100644 src/common/statistics.cpp create mode 100644 src/exec/file_scan_node.cpp create mode 100644 src/exec/ttl_delete_node.cpp create mode 100644 src/physical_plan/condition_optimizer.cpp diff --git a/include/exec/file_scan_node.h b/include/exec/file_scan_node.h new file mode 100644 index 000000000..29956a05e --- /dev/null +++ b/include/exec/file_scan_node.h @@ -0,0 +1,199 @@ +// Copyright (c) 2018-present Baidu, Inc. All Rights Reserved. +// +// Licensed 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. + +#pragma once + +#include "arrow_io_excutor.h" +#include "file_system.h" +#include "scan_node.h" + +namespace baikaldb { + +DECLARE_int32(inner_file_scan_concurrency); + +class FileVectorizedReader; +class FileScanNode; + +class FileScanner { +public: + FileScanner(FileSystem* fs, const pb::PartitionFile& partition_file) : _fs(fs) { + _partition_vals.assign(partition_file.partition_vals().begin(), partition_file.partition_vals().end()); + _file_path = partition_file.file_path(); + } + virtual ~FileScanner() {} + virtual int init(FileVectorizedReader* vectorized_reader, FileScanNode* scan_node); + virtual int run() = 0; + + const std::vector& get_partition_vals() { + return _partition_vals; + } + FileVectorizedReader* get_vectorized_reader() { + return _vectorized_reader; + } + FileScanNode* get_scan_node() { + return _scan_node; + } + +protected: + FileSystem* _fs = nullptr; + std::vector _partition_vals; + std::string _file_path; + FileVectorizedReader* _vectorized_reader = nullptr; + FileScanNode* _scan_node = nullptr; +}; + +class CSVScanner : public FileScanner { + class BlockImpl; +public: + using FileScanner::FileScanner; + virtual ~CSVScanner() {} + int init(FileVectorizedReader* vectorized_reader, FileScanNode* scan_node) override; + int run() override; + +private: + FileInfo _file_info; +}; + +class ParquetScanner : public FileScanner { +public: + using FileScanner::FileScanner; + virtual ~ParquetScanner() {} + int init(FileVectorizedReader* vectorized_reader, FileScanNode* scan_node) override; + int run() override; + +private: + int process_record_batch(std::shared_ptr<::arrow::RecordBatch>& record_batch); + int get_file_reader(std::unique_ptr<::parquet::arrow::FileReader>& reader); + +private: + std::shared_ptr<::parquet::FileMetaData> _file_metadata; + // 本次需要获取，且parquet文件中存在的列 + std::vector _exist_column_indices; + // 本次需要获取，但parqeut文件中不存在的列 + std::vector _not_exist_columns; + // + std::unordered_map _column_name_map; + // key: baikaldb column name, value: baikaldb column type + std::unordered_map _column_type_map; + // + std::unordered_map _partition_id2val_map; +}; + +class FileVectorizedReader : public arrow::RecordBatchReader { +public: + FileVectorizedReader() + : _file_concurrency(FLAGS_inner_file_scan_concurrency) + , _file_concurrency_cond(-FLAGS_inner_file_scan_concurrency) + , _record_batches(FLAGS_inner_file_scan_concurrency * 2) {} + virtual ~FileVectorizedReader() {} + int init(FileScanNode* scan_node, RuntimeState* state); + arrow::Status ReadNext(std::shared_ptr* out) override; + std::shared_ptr schema() const override { + return _schema; + } + void close() { + _file_scan_bth.join(); + } + + BlockingQueue>& get_record_batches() { + return _record_batches; + } + BthreadCond& get_file_concurrency_cond() { + return _file_concurrency_cond; + } + std::atomic& get_is_succ() { + return _is_succ; + } + std::atomic& get_eos() { + return _eos; + } + +private: + void run_file_scan_thread(); + +private: + std::shared_ptr _schema; + FileScanNode* _scan_node = nullptr; + RuntimeState* _state = nullptr; + Bthread _file_scan_bth; + BlockingQueue> _record_batches; + // Block或RowGroup并发数 + int32_t _file_concurrency = -1; + BthreadCond _file_concurrency_cond; + // limit功能实现 + std::shared_ptr<::arrow::RecordBatch> _record_batch; + int64_t _row_idx_in_record_batch = 0; + int64_t _processed_row_cnt = 0; + // 读取状态 + std::atomic _is_succ = true; + std::atomic _eos = false; + // paralle模式使用 + bool _is_delay_fetch = false; + std::shared_ptr _index_cond; +}; + +class FileScanNode : public ScanNode { +public: + FileScanNode() { + _is_file_scan_node = true; + } + virtual ~FileScanNode() { + } + virtual int init(const pb::PlanNode& node) override; + virtual int open(RuntimeState* state) override; + virtual int get_next(RuntimeState* state, RowBatch* batch, bool* eos) override; + virtual void close(RuntimeState* state) override; + virtual void transfer_pb(int64_t region_id, pb::PlanNode* pb_node) override; + virtual int build_arrow_declaration(RuntimeState* state) override; + virtual bool can_use_arrow_vector(RuntimeState* state) override { + return true; + } + const pb::FileInfo& get_pb_file_info() { + return _pb_file_info; + } + const std::vector& get_field_id2slot() { + return _field_id2slot; + } + const std::vector& get_field_id2info() { + return _field_id2info; + } + const std::vector& get_partition_fields() { + return _partition_fields; + } + const std::vector& get_data_fields() { + return _data_fields; + } + const std::vector& get_files() { + return _files; + } + void set_files(const std::vector& files) { + _files = files; + } + +private: + pb::FileInfo _pb_file_info; + std::vector _files; // 需要访问的离线文件集合 + + SmartTable _table_info; + std::vector _field_id2slot; // + std::vector _field_id2info; // + + std::vector _partition_fields; // 分区字段 + std::vector _data_fields; // 数据字段，去掉分区字段后按表定义顺序排列 + + std::shared_ptr _vectorized_reader; + std::shared_ptr _arrow_io_executor; +}; + +} // namespace baikaldb \ No newline at end of file diff --git a/include/exec/ttl_delete_node.h b/include/exec/ttl_delete_node.h new file mode 100644 index 000000000..24d81456f --- /dev/null +++ b/include/exec/ttl_delete_node.h @@ -0,0 +1,42 @@ +// Copyright (c) 2018-present Baidu, Inc. All Rights Reserved. +// +// Licensed 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. + +#pragma once + +#include "exec_node.h" +#include "dml_node.h" + +namespace baikaldb { +class TTLDeleteNode : public DMLNode { +public: + TTLDeleteNode() { + } + virtual ~TTLDeleteNode() { + } + int init(const pb::PlanNode& node) override; + int open(RuntimeState* state) override; + + void add_delete_records(SmartRecord record) { + _delete_records.emplace_back(record); + } + void clean_delete_records() { + _delete_records.clear(); + } +private: + std::vector _delete_records; +}; + +} + +/* vim: set ts=4 sw=4 sts=4 tw=100 */ diff --git a/include/physical_plan/condition_optimizer.h b/include/physical_plan/condition_optimizer.h new file mode 100644 index 000000000..633cb59f0 --- /dev/null +++ b/include/physical_plan/condition_optimizer.h @@ -0,0 +1,33 @@ +// Copyright (c) 2018-present Baidu, Inc. All Rights Reserved. +// +// Licensed 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. + +#pragma once +#include "exec_node.h" +#include "query_context.h" +#include "scan_node.h" + +namespace baikaldb { +class ConditionOptimizer { +public: + ConditionOptimizer() {} + ~ConditionOptimizer() {} + ConditionOptimizer(QueryContext* ctx) : _ctx(ctx) {} + int analyze(QueryContext* ctx); + int adjust_huge_in_condition(ScanNode* scan_node, bool in_acero); +private: + QueryContext* _ctx = nullptr; +}; +} + +/* vim: set ts=4 sw=4 sts=4 tw=100 */ diff --git a/src/common/statistics.cpp b/src/common/statistics.cpp new file mode 100644 index 000000000..18f10a4c0 --- /dev/null +++ b/src/common/statistics.cpp @@ -0,0 +1,52 @@ +// Copyright (c) 2018-present Baidu, Inc. All Rights Reserved. +// +// Licensed 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. + +#pragma once +#include "schema_factory.h" + +namespace baikaldb { +std::map Statistics::get_field_names() const { + std::unordered_set field_ids; + for (auto& column : _field_histogram) { + field_ids.insert(column.first); + } + for (auto& column : _field_cmsketch) { + field_ids.insert(column.first); + } + for (auto& column : _field_hll) { + field_ids.insert(column.first); + } + auto schema_factory = SchemaFactory::get_instance(); + auto table_info = schema_factory->get_table_info_ptr(_table_id); + if (table_info == nullptr) { + DB_WARNING("statistics table not found, table id: %ld", _table_id); + return {}; + } + std::map field_names; + for (auto& field_id : field_ids) { + auto field_infos = table_info->fields; + auto iter = std::find_if(field_infos.begin(), field_infos.end(), + [&field_id](const FieldInfo& field_info) { + return field_id == field_info.id; + } + ); + if (iter == field_infos.end()) { + DB_WARNING("field not found, field id: %ld", field_id); + continue; + } + field_names[iter->short_name] = field_id; + } + return std::move(field_names); +} +} diff --git a/src/exec/file_scan_node.cpp b/src/exec/file_scan_node.cpp new file mode 100644 index 000000000..41ea903fe --- /dev/null +++ b/src/exec/file_scan_node.cpp @@ -0,0 +1,921 @@ +#include "file_scan_node.h" +#include "select_manager_node.h" +#include "vectorize_helpper.h" + +namespace baikaldb { + +DEFINE_int32(inner_file_scan_concurrency, 4, "inner_file_scan_concurrency"); +DECLARE_int32(file_block_size); +DECLARE_int32(chunk_size); + +/// FileScanNode +int FileScanNode::init(const pb::PlanNode& node) { + int ret = 0; + ret = ScanNode::init(node); + if (ret < 0) { + DB_WARNING("ExecNode::init fail, ret: %d", ret); + return ret; + } + _table_info = SchemaFactory::get_instance()->get_table_info_ptr(_table_id); + if (_table_info == nullptr) { + DB_WARNING("table info is null, table id: %ld", _table_id); + return -1; + } + _pb_file_info = _table_info->dblink_info.file_info(); + const auto& pb_files = node.derive_node().scan_node().files(); + _files.assign(pb_files.begin(), pb_files.end()); + return 0; +} + +int FileScanNode::open(RuntimeState* state) { + if (state == nullptr) { + DB_WARNING("state is nullptr"); + return -1; + } + START_LOCAL_TRACE(get_trace(), state->get_trace_cost(), OPEN_TRACE, nullptr); + int ret = 0; + ret = ScanNode::open(state); + if (ret < 0) { + DB_WARNING_STATE(state, "ExecNode::open fail:%d", ret); + return ret; + } + if (_table_info == nullptr) { + DB_WARNING("_table_info is nullptr"); + return -1; + } + // 离线文件查询走列式 + set_node_exec_type(pb::EXEC_ARROW_ACERO); + // 建立field_id到slot_id的映射，建立field_id到field_info的映射 + _field_id2slot = std::vector(_table_info->fields.back().id + 1, 0); + _field_id2info = std::vector(_table_info->fields.back().id + 1, nullptr); + for (const auto& slot : _tuple_desc->slots()) { + if (slot.field_id() >= _field_id2slot.size()) { + DB_WARNING("vector out of range, field_id: %d", slot.field_id()); + continue; + } + _field_id2slot[slot.field_id()] = slot.slot_id(); + FieldInfo* field_info = _table_info->get_field_ptr(slot.field_id()); + if (field_info == nullptr) { + DB_WARNING("field not found, field_id: %d", slot.field_id()); + continue; + } + _field_id2info[slot.field_id()] = field_info; + } + // 获取分区字段集合和数据字段集合 + std::unordered_set partition_field_set; + for (const auto& partition_field : _pb_file_info.partition_fields()) { + partition_field_set.insert(partition_field); + FieldInfo* field_info = _table_info->get_field_ptr(partition_field); + if (field_info == nullptr) { + DB_WARNING("field not found, field_name: %s", partition_field.c_str()); + return -1; + } + _partition_fields.emplace_back(field_info); + } + for (auto& field : _table_info->fields) { + if (partition_field_set.find(field.short_name) == partition_field_set.end()) { + _data_fields.emplace_back(&field); + } + } + _vectorized_reader = std::make_shared(); + if (_vectorized_reader == nullptr) { + DB_WARNING("_vectorized_reader is nullptr"); + return -1; + } + ret = _vectorized_reader->init(this, state); + if (ret != 0) { + DB_WARNING("Fail to init vectorized_reader"); + return -1; + } + return 0; +} + +int FileScanNode::get_next(RuntimeState* state, RowBatch* batch, bool* eos) { + // 离线文件查询只走列式 + *eos = true; + return 0; +} + +void FileScanNode::close(RuntimeState* state) { + ScanNode::close(state); + if (_vectorized_reader != nullptr) { + _vectorized_reader->close(); + } + _field_id2slot.clear(); + _field_id2info.clear(); + _partition_fields.clear(); + _data_fields.clear(); + _vectorized_reader = nullptr; + _arrow_io_executor = nullptr; +} + +void FileScanNode::transfer_pb(int64_t region_id, pb::PlanNode* pb_node) { + ExecNode::transfer_pb(region_id, pb_node); + auto* scan_pb = pb_node->mutable_derive_node()->mutable_scan_node(); + for (const auto& file : _files) { + scan_pb->add_files()->CopyFrom(file); + } +} + +int FileScanNode::build_arrow_declaration(RuntimeState* state) { + if (state == nullptr) { + DB_WARNING("state is nullptr"); + return -1; + } + START_LOCAL_TRACE_WITH_PARTITION_PROPERTY(get_trace(), state->get_trace_cost(), &_partition_property, OPEN_TRACE, nullptr); + std::function>()> iter_maker = [this] () { + arrow::Iterator> batch_it = arrow::MakeIteratorFromReader(_vectorized_reader); + return batch_it; + }; + bool is_delay_fetch = false; + if (_related_manager_node != nullptr) { + is_delay_fetch = _related_manager_node->is_delay_fetcher_store(); + } + if (state->vectorlized_parallel_execution == false) { + arrow::acero::Declaration dec{"record_batch_source", + arrow::acero::RecordBatchSourceNodeOptions{_vectorized_reader->schema(), std::move(iter_maker)}}; + LOCAL_TRACE_ARROW_PLAN_WITH_SCHEMA(dec, _vectorized_reader->schema(), &is_delay_fetch); + state->append_acero_declaration(dec); + } else { + auto executor = BthreadArrowExecutor::Make(1); + _arrow_io_executor = *executor; + arrow::acero::Declaration dec{"record_batch_source", + arrow::acero::RecordBatchSourceNodeOptions{_vectorized_reader->schema(), std::move(iter_maker), _arrow_io_executor.get()}}; + LOCAL_TRACE_ARROW_PLAN_WITH_SCHEMA(dec, _vectorized_reader->schema(), &is_delay_fetch); + state->append_acero_declaration(dec); + } + return 0; +} + +/// FileVectorizedReader +int FileVectorizedReader::init(FileScanNode* scan_node, RuntimeState* state) { + if (scan_node == nullptr) { + DB_WARNING("scan_node is nullptr"); + return -1; + } + _scan_node = scan_node; + if (state == nullptr) { + DB_WARNING("_state is nullptr"); + return -1; + } + _state = state; + // 生成schema + _schema = VectorizeHelpper::make_schema(_scan_node->get_tuple()); + if (_schema == nullptr) { + DB_WARNING("Fail to make_schema"); + return -1; + } + SelectManagerNode* related_manager_node = + static_cast(_scan_node->get_related_manager_node()); + if (related_manager_node != nullptr) { + _is_delay_fetch = related_manager_node->is_delay_fetcher_store(); + _index_cond = related_manager_node->get_index_collector_cond(); + } + if (_is_delay_fetch) { + return 0; + } + run_file_scan_thread(); + return 0; +} + +arrow::Status FileVectorizedReader::ReadNext(std::shared_ptr* out) { + int ret = 0; + if (_is_delay_fetch) { + if (_index_cond != nullptr) { + _index_cond->cond.wait(); + if (_index_cond->index_cnt == 0) { + // join驱动表没数据 + out->reset(); + return arrow::Status::OK(); + } + } + run_file_scan_thread(); + _is_delay_fetch = false; + } + if (_state->is_cancelled()) { + DB_WARNING_STATE(_state, "cancelled"); + _eos = true; + return arrow::Status::OK(); + } + const int64_t limit = _scan_node->get_limit(); + if (limit > 0 && _processed_row_cnt >= limit) { + _eos = true; + return arrow::Status::OK(); + } + if (_eos) { + out->reset(); + return arrow::Status::OK(); + } + if (_record_batch == nullptr) { + bool ret = _record_batches.blocking_get(&_record_batch); + if (!_is_succ) { + return arrow::Status::IOError("Fail to file scan"); + } + if (!ret) { + _eos = true; + return arrow::Status::OK(); + } + } + if (_record_batch != nullptr) { + // limit功能 + // 每次最多读取FLAGS_chunk_size行 + int64_t slice_cnt = FLAGS_chunk_size; + if (limit > 0 && _processed_row_cnt + slice_cnt >= limit) { + slice_cnt = limit - _processed_row_cnt; + } + *out = _record_batch->Slice(_row_idx_in_record_batch, slice_cnt); + _processed_row_cnt += (*out)->num_rows(); + _row_idx_in_record_batch += (*out)->num_rows(); + if (_record_batch->num_rows() <= _row_idx_in_record_batch) { + _record_batch.reset(); + _row_idx_in_record_batch = 0; + } + } + return arrow::Status::OK(); +} + +void FileVectorizedReader::run_file_scan_thread() { + auto file_scan_thread = [this] () { + const pb::FileInfo& file_info = _scan_node->get_pb_file_info(); + TimeCost tm; + int64_t create_filesystem_tm = -1; + int64_t total_file_scan_thread_tm = -1; + std::shared_ptr fs = create_filesystem(file_info.cluster(), + file_info.username(), + file_info.password(), + AFS_CLIENT_CONF_PATH); + ScopeGuard guard([this, &fs, &create_filesystem_tm, &total_file_scan_thread_tm] () { + destroy_filesystem(fs); + }); + if (fs == nullptr) { + DB_WARNING("Fail to create_filesystem"); + _is_succ = false; + return; + } + create_filesystem_tm = tm.get_time(); + // 生成并执行file scanner + const auto& files = _scan_node->get_files(); + std::vector> file_scanners; + for (const auto& file : files) { + std::shared_ptr file_scanner; + switch (file_info.format()) { + case pb::CSV: { + file_scanner = std::make_shared(fs.get(), file); + break; + } + case pb::PARQUET: { + file_scanner = std::make_shared(fs.get(), file); + break; + } + default: { + DB_WARNING("Invalid file_format: %d", file_info.format()); + _is_succ = false; + break; + } + } + if (!_is_succ) { + break; + } + if (file_scanner == nullptr) { + DB_WARNING("file_scanner is nullptr"); + _is_succ = false; + break; + } + int ret = file_scanner->init(this, _scan_node); + if (ret == -2) { // -2表示Parquet目录中非parquet文件，如_SUCCESS文件 + continue; + } + if (ret != 0) { + DB_WARNING("Fail to init file_scanner"); + _is_succ = false; + break; + } + ret = file_scanner->run(); + if (ret < 0) { + DB_WARNING("Fail to run file_scanner"); + _is_succ = false; + break; + } + file_scanners.emplace_back(file_scanner); + } + _file_concurrency_cond.wait(-_file_concurrency); + _record_batches.shutdown(); + total_file_scan_thread_tm = tm.get_time(); + DB_WARNING("create_filesystem_tm: %ld, file_scan_thread total tm: %ld", + create_filesystem_tm, total_file_scan_thread_tm); + }; + _file_scan_bth.run(file_scan_thread); +} + +/// FileScanner +int FileScanner::init(FileVectorizedReader* vectorized_reader, FileScanNode* scan_node) { + _vectorized_reader = vectorized_reader; + _scan_node = scan_node; + if (_vectorized_reader == nullptr) { + DB_WARNING("_vectorized_reader is nullptr"); + return -1; + } + if (_scan_node == nullptr) { + DB_WARNING("_scan_node is nullptr"); + return -1; + } + if (_fs == nullptr) { + DB_WARNING("_fs is nulpptr"); + return -1; + } + const auto& partition_fields = scan_node->get_partition_fields(); + if (partition_fields.size() != _partition_vals.size()) { + DB_WARNING("partition_fields.size[%lu] not equal to partition_vals.size[%lu]", + partition_fields.size(), _partition_vals.size()); + return -1; + } + return 0; +} + +/// CSVScanner::BlockImpl +class CSVScanner::BlockImpl { +public: + BlockImpl(CSVScanner* scanner, + FileSystem* fs, + const std::string& file_path, + size_t file_size, + size_t start_pos, + size_t end_pos, + std::atomic& eos, + std::atomic& is_succ) + : _scanner(scanner) + , _fs(fs) + , _file_path(file_path) + , _file_size(file_size) + , _start_pos(start_pos) + , _end_pos(end_pos) + , _cur_pos(start_pos) + , _eos(eos) + , _is_succ(is_succ) { + if (_start_pos == 0) { + _escape_first_line = false; + } + } + virtual ~BlockImpl() { + if (_read_buffer != nullptr) { + free(_read_buffer); + _read_buffer = nullptr; + } + if (_file != nullptr) { + if (_fs != nullptr) { + _fs->close_reader(_file); + } + _file = nullptr; + } + } + + int init(); + // @brief 读取一个块的数据，对每一行进行处理 + // @return 1: 已读完 / 0: 未读完 / others: 读失败 + int read_and_exec(); + // @brief 读取完一个块，处理未提交的_batch + int read_done(); + +private: + bool block_finish(); + int read_buffer_resize(int64_t size); + int process_line(const std::string& line); + +private: + CSVScanner* _scanner = nullptr; + FileSystem* _fs = nullptr; + std::shared_ptr _file; + std::string _file_path; + int64_t _file_size = -1; + int64_t _start_pos = -1; + int64_t _end_pos = -1; + int64_t _cur_pos = -1; + char* _read_buffer = nullptr; + int64_t _read_buffer_size = 0; + bool _escape_first_line = true; + RowBatch _batch; + std::atomic& _eos; + std::atomic& _is_succ; +}; + +int CSVScanner::BlockImpl::init() { + if (_fs == nullptr) { + DB_WARNING("_fs is nullptr"); + return -1; + } + _file = _fs->open_reader(_file_path); + if (_file == nullptr) { + DB_WARNING("Fail to open_reader, _file_path: %s", _file_path.c_str()); + return -1; + } + _read_buffer_size = FLAGS_file_buffer_size * 1024 * 1024ULL; + _read_buffer = static_cast(malloc(_read_buffer_size)); + if (_read_buffer == nullptr) { + DB_WARNING("Fail to malloc"); + return -1; + } + if (_scanner == nullptr) { + DB_WARNING("_scanner is nullptr"); + return -1; + } + if (_scanner->get_vectorized_reader() == nullptr) { + DB_WARNING("_scanner->get_vectorized_reader() is nullptr"); + return -1; + } + if (_scanner->get_scan_node() == nullptr) { + DB_WARNING("_scanner->get_scan_node() is nullptr"); + return -1; + } + if (_scanner->get_scan_node()->get_tuple() == nullptr) { + DB_WARNING("tuple_desc is nullptr"); + return -1; + } + _batch.init_chunk({_scanner->get_scan_node()->get_tuple()}, nullptr); + return 0; +} + +int CSVScanner::BlockImpl::read_buffer_resize(int64_t size) { + if (_read_buffer != nullptr) { + free(_read_buffer); + _read_buffer = nullptr; + } + _read_buffer_size = size; + _read_buffer = static_cast(malloc(_read_buffer_size)); + if (_read_buffer == nullptr) { + DB_WARNING("Fail to malloc"); + return -1; + } + return 0; +} + +bool CSVScanner::BlockImpl::block_finish() { + // 只有_cur_pos > _end_pos才能判断完成，相等说明_end_pos正好在行首，需要多读一行，因为下一个block跳过了首行 + if (_cur_pos > _end_pos) { + return true; + } + if (_cur_pos == _end_pos && _end_pos == _file_size) { + return true; + } + if (_start_pos == _end_pos) { + return true; + } + if (_eos) { + return true; + } + if (!_is_succ) { + return true; + } + return false; +} + +int CSVScanner::BlockImpl::read_and_exec() { + if (block_finish()) { + return 1; + } + const int64_t file_pos = _cur_pos; + int64_t buf_pos = 0; + int64_t buf_size = _file->read(file_pos, _read_buffer, _read_buffer_size); + if (buf_size < 0) { + DB_WARNING("Fail to read, file: %s, pos: %ld", _file_path.c_str(), file_pos); + return -1; + } + if (buf_size == 0) { + // 已读完 + return 1; + } + MemBuf mem_buf(_read_buffer, _read_buffer + buf_size); + std::istream f(&mem_buf); + // 跳过首行 + if (_escape_first_line) { + std::string line; + std::getline(f, line); + if (f.eof()) { + // 增大buffer重试 + if (read_buffer_resize(_read_buffer_size * 2) < 0) { + return -1; + } + return 0; + } + buf_pos += line.size() + 1; + _escape_first_line = false; + } + _cur_pos = file_pos + buf_pos; + if (block_finish()) { + return 1; + } + bool has_get_line = false; + while (!f.eof()) { + std::string line; + std::getline(f, line); + // eof直接退出不更新 _cur_pos, 下次从_cur_pos继续读 + if (f.eof()) { + buf_pos += line.size(); + // 最后一块特殊处理，不需要跳过 + if (file_pos + buf_pos == _end_pos && _end_pos == _file_size) { + if (!line.empty() && process_line(line) < 0) { + return -1; + } + return 1; + } + if (_cur_pos <= _end_pos && !has_get_line) { + // 增大buffer重试 + if (read_buffer_resize(_read_buffer_size * 2) < 0) { + return -1; + } + return 0; + } + return 0; + } + has_get_line = true; + buf_pos += line.size() + 1; + _cur_pos = file_pos + buf_pos; + if (!line.empty() && process_line(line) < 0) { + return -1; + } + if (block_finish()) { + return 1; + } + } + return 0; +} + +int CSVScanner::BlockImpl::read_done() { + if (!_batch.empty()) { + std::shared_ptr record_batch; + int ret = _batch.finish_and_make_record_batch(&record_batch); + if (ret < 0) { + DB_WARNING("arrow chunk finish and make record batch fail"); + return -1; + } + auto& record_batches = _scanner->get_vectorized_reader()->get_record_batches(); + record_batches.blocking_put(record_batch); + } + return 0; +} + +// 处理文件行 +// - 如果文件行字段数比表字段数多，则忽略多于字段； +// - 如果文件行字段数比表字段数少，则补充默认值； +int CSVScanner::BlockImpl::process_line(const std::string& line) { + FileScanNode* scan_node = _scanner->get_scan_node(); + const auto& tuple_id = scan_node->tuple_id(); + const auto& delimiter = scan_node->get_pb_file_info().delimiter(); + const auto& field_id2slot = scan_node->get_field_id2slot(); + const auto& partition_fields = scan_node->get_partition_fields(); + const auto& data_fields = scan_node->get_data_fields(); + std::vector line_field_vec; + boost::split(line_field_vec, line, boost::is_any_of(delimiter)); // 当前按单字符分割 + for (int i = 0; i < line_field_vec.size(); ++i) { + if (i >= data_fields.size()) { + break; + } + FieldInfo* field = data_fields[i]; + if (field == nullptr) { + DB_WARNING("field is nullptr"); + return -1; + } + if (field->id >= field_id2slot.size() || field_id2slot[field->id] <= 0) { + continue; + } + ExprValue value(field->type, line_field_vec[i]); + _batch.set_chunk_tmp_row_value(tuple_id, field_id2slot[field->id], value); + } + for (int i = line_field_vec.size(); i < data_fields.size(); ++i) { + FieldInfo* field = data_fields[i]; + if (field == nullptr) { + DB_WARNING("field is nullptr"); + return -1; + } + if (field->id >= field_id2slot.size() || field_id2slot[field->id] <= 0) { + continue; + } + _batch.set_chunk_tmp_row_value(tuple_id, field_id2slot[field->id], field->default_expr_value); + } + // 填充分区字段 + const auto& partition_vals = _scanner->get_partition_vals(); + for (int i = 0; i < partition_fields.size(); ++i) { + FieldInfo* field = partition_fields[i]; + if (field == nullptr) { + DB_WARNING("field is nullptr"); + return -1; + } + if (field->id >= field_id2slot.size() || field_id2slot[field->id] <= 0) { + continue; + } + ExprValue value(field->type, partition_vals[i]); + _batch.set_chunk_tmp_row_value(tuple_id, field_id2slot[field->id], value); + } + if (_batch.add_chunk_row() != 0) { + DB_FATAL("add chunk row fail"); + return -1; + } + if (_batch.is_full()) { + std::shared_ptr record_batch; + int ret = _batch.finish_and_make_record_batch(&record_batch); + if (ret < 0) { + DB_WARNING("arrow chunk finish and make record batch fail"); + return -1; + } + auto& record_batches = _scanner->get_vectorized_reader()->get_record_batches(); + record_batches.blocking_put(record_batch); + } + return 0; +} + +/// CSVScanner +int CSVScanner::init(FileVectorizedReader* vectorized_reader, FileScanNode* scan_node) { + int ret = FileScanner::init(vectorized_reader, scan_node); + if (ret != 0) { + return -1; + } + ret = _fs->get_file_info(_file_path, _file_info, nullptr); + if (ret < 0) { + DB_WARNING("Fil to get_file_info, ret: %d, file_path: %s", ret, _file_path.c_str()); + return -1; + } + if (_file_info.mode != FileMode::I_FILE) { + DB_WARNING("Invalid file_info mode: %d, file_path: %s", (int)_file_info.mode, _file_path.c_str()); + return -1; + } + return 0; +} + +int CSVScanner::run() { + BthreadCond& file_concurrency_cond = _vectorized_reader->get_file_concurrency_cond(); + std::atomic& is_succ = _vectorized_reader->get_is_succ(); + std::atomic& eos = _vectorized_reader->get_eos(); + const int64_t file_block_size = FLAGS_file_block_size * 1024 * 1024LL; + if (file_block_size <= 0) { + DB_WARNING("Invalid file_block_size: %ld", file_block_size); + return -1; + } + int64_t block_num = _file_info.size / file_block_size + 1; + for (int i = 0; i < block_num; ++i) { + if (!is_succ) { + break; + } + if (eos) { + break; + } + int64_t start_pos = i * file_block_size; + int64_t end_pos = (i + 1) * file_block_size; + end_pos = end_pos > _file_info.size ? _file_info.size : end_pos; + auto process_block = + [this, start_pos, end_pos, &file_concurrency_cond, &is_succ, &eos] () { + TimeCost tm; + bool is_cur_succ = false; + ON_SCOPE_EXIT(([&file_concurrency_cond, &is_succ, &is_cur_succ]() { + file_concurrency_cond.decrease_signal(); + if (!is_cur_succ) { + is_succ = false; + } + })); + BlockImpl block(this, _fs, _file_path, _file_info.size, start_pos, end_pos, eos, is_succ); + int ret = block.init(); + if (ret < 0) { + DB_WARNING("Fail to init block"); + return; + } + while (true) { + ret = block.read_and_exec(); + if (ret < 0) { + DB_WARNING("Fail to read_and_exec"); + return; + } + if (ret == 1) { + ret = block.read_done(); + if (ret < 0) { + return; + } + break; + } + } + is_cur_succ = true; + DB_WARNING("process_block tm: %ld", tm.get_time()); + }; + file_concurrency_cond.increase_wait(); + Bthread bth(&BTHREAD_ATTR_SMALL); + bth.run(process_block); + } + return 0; +} + +/// ParquetScanner +int ParquetScanner::init(FileVectorizedReader* vectorized_reader, FileScanNode* scan_node) { + int ret = FileScanner::init(vectorized_reader, scan_node); + if (ret != 0) { + return -1; + } + std::unique_ptr<::parquet::arrow::FileReader> reader; + ret = get_file_reader(reader); + if (ret != 0) { + DB_WARNING("Fail to get_file_reader, file: %s", _file_path.c_str()); + return ret; + } + _file_metadata = reader->parquet_reader()->metadata(); + if (_file_metadata == nullptr) { + DB_WARNING("FileMetaData is nullptr, file: %s", _file_path.c_str()); + return -1; + } + const ::parquet::SchemaDescriptor* schema = _file_metadata->schema(); + if (schema == nullptr) { + DB_WARNING("schema is nullptr, file: %s", _file_path.c_str()); + return -1; + } + // Schema处理 + std::unordered_map column_name2index_map; + for (int i = 0; i < schema->num_columns(); ++i) { + const ::parquet::ColumnDescriptor* column = schema->Column(i); + if (column == nullptr) { + DB_WARNING("column is nullptr, file: %s", _file_path.c_str()); + return -1; + } + column_name2index_map[column->name()] = i; + } + const auto& field_id2info = _scan_node->get_field_id2info(); + for (const auto& field_info : field_id2info) { + if (field_info == nullptr) { + continue; + } + const std::string& field_name = field_info->short_name; + if (column_name2index_map.find(field_name) != column_name2index_map.end()) { + _exist_column_indices.emplace_back(column_name2index_map.at(field_name)); + } else { + _not_exist_columns.emplace_back(field_info); + } + } + pb::TupleDescriptor* tuple_desc = _scan_node->get_tuple(); + for (const auto& slot : tuple_desc->slots()) { + if (slot.field_id() >= field_id2info.size() || field_id2info[slot.field_id()] == nullptr) { + continue; + } + const std::string& field_name = std::to_string(slot.tuple_id()) + "_" + std::to_string(slot.slot_id()); + _column_name_map[field_name] = field_id2info[slot.field_id()]->short_name; + _column_type_map[field_name] = field_id2info[slot.field_id()]->type; + } + const auto& partition_fields = scan_node->get_partition_fields(); + for (int i = 0; i < partition_fields.size(); ++i) { + _partition_id2val_map[partition_fields[i]->id] = _partition_vals[i]; + } + return 0; +} + +int ParquetScanner::run() { + BthreadCond& file_concurrency_cond = _vectorized_reader->get_file_concurrency_cond(); + std::atomic& is_succ = _vectorized_reader->get_is_succ(); + std::atomic& eos = _vectorized_reader->get_eos(); + auto& record_batches = _vectorized_reader->get_record_batches(); + for (size_t row_group_idx = 0; row_group_idx < _file_metadata->num_row_groups(); ++row_group_idx) { + if (!is_succ) { + break; + } + if (eos) { + break; + } + auto process_row_group = + [this, &file_concurrency_cond, &is_succ, &eos, &record_batches, row_group_idx] () { + TimeCost tm; + bool is_cur_succ = false; + ON_SCOPE_EXIT(([&file_concurrency_cond, &is_succ, &is_cur_succ]() { + file_concurrency_cond.decrease_signal(); + if (!is_cur_succ) { + is_succ = false; + } + })); + std::unique_ptr<::parquet::arrow::FileReader> reader; + if (get_file_reader(reader) != 0) { + DB_WARNING("Fail to get_file_reader, file_path: %s", _file_path.c_str()); + return; + } + std::unique_ptr<::arrow::RecordBatchReader> record_batch_reader; + auto status = reader->GetRecordBatchReader({ row_group_idx }, _exist_column_indices, &record_batch_reader); + if (!status.ok()) { + DB_WARNING("Fail to GetRecordBatchReader, row_group_idx: %ld, path: %s, reason: %s", + row_group_idx, _file_path.c_str(), status.message().c_str()); + return; + } + if (record_batch_reader == nullptr) { + DB_WARNING("RecordBatchReader is nullptr, row_group_idx: %ld, path: %s", + row_group_idx, _file_path.c_str()); + return; + } + while (true) { + if (!is_succ) { + break; + } + if (eos) { + break; + } + std::shared_ptr<::arrow::RecordBatch> batch; + auto status = record_batch_reader->ReadNext(&batch); + if (!status.ok()) { + DB_WARNING("RecordBatchReader fail to ReadNext, row_group_idx: %ld, path: %s, reason: %s", + row_group_idx, _file_path.c_str(), status.message().c_str()); + return; + } + if (batch == nullptr) { + break; + } + int ret = process_record_batch(batch); + if (ret != 0) { + DB_WARNING("Fail to process_record_batch"); + return; + } + record_batches.blocking_put(batch); + } + is_cur_succ = true; + DB_WARNING("process_row_group tm: %ld", tm.get_time()); + }; + file_concurrency_cond.increase_wait(); + Bthread bth(&BTHREAD_ATTR_SMALL); + bth.run(process_row_group); + } + return 0; +} + +// 填充文件中不存在的列，并转换schema +int ParquetScanner::process_record_batch(std::shared_ptr<::arrow::RecordBatch>& record_batch) { + const auto& schema = _vectorized_reader->schema(); + if (record_batch == nullptr) { + DB_WARNING("record_batch is nullptr"); + return -1; + } + // 填充文件中不存在的列 + for (auto* field_info : _not_exist_columns) { + std::shared_ptr arrow_field = VectorizeHelpper::make_field( + field_info->short_name, arrow::Type::type(primitive_to_arrow_type(field_info->type))); + if (arrow_field == nullptr) { + DB_WARNING("Fail to get arrow type, field_id: %d, field_type: %d", field_info->id, field_info->type); + return -1; + } + ExprValue value = field_info->default_expr_value; + if (_partition_id2val_map.find(field_info->id) != _partition_id2val_map.end()) { + // 分区字段值 + value = ExprValue(field_info->type, _partition_id2val_map[field_info->id]); + } + std::shared_ptr arrow_array = + VectorizeHelpper::make_array_from_exprvalue(field_info->type, value, record_batch->num_rows()); + if (arrow_array == nullptr) { + DB_WARNING("Fail to make array from expr value"); + return -1; + } + auto new_record_batch_ret = record_batch->AddColumn(record_batch->num_columns(), arrow_field, arrow_array); + if (!new_record_batch_ret.ok()) { + DB_WARNING("Fail to AddColumn"); + return -1; + } + record_batch = *new_record_batch_ret; + } + // 转换schema，列名转化为tupleid_slotid形式，类型也可能发生转化 + int ret = VectorizeHelpper::change_arrow_record_batch_schema( + _column_name_map, _column_type_map, schema, record_batch, &record_batch); + if (ret != 0) { + DB_WARNING("Fail to change arrow record batch schema"); + return -1; + } + return 0; +} + +int ParquetScanner::get_file_reader(std::unique_ptr<::parquet::arrow::FileReader>& reader) { + // Open parquet file + auto res = _fs->open_arrow_file(_file_path); + if (!res.ok()) { + DB_WARNING("Fail to open ParquetReader, reason: %s", res.status().message().c_str()); + return -1; + } + auto infile = std::move(res).ValueOrDie(); + // Avoid reading whole file data into memory at once + ::parquet::ReaderProperties read_properties = ::parquet::default_reader_properties(); + read_properties.enable_buffered_stream(); + read_properties.set_buffer_size(FLAGS_file_buffer_size * 1024 * 1024ULL); + // Set GetRecordBatchReader batch size, default: 65536 + ::parquet::ArrowReaderProperties arrow_properties = ::parquet::default_arrow_reader_properties(); + arrow_properties.set_batch_size(FLAGS_chunk_size); + arrow_properties.set_pre_buffer(false); + ::parquet::arrow::FileReaderBuilder builder; + builder.properties(arrow_properties); + auto status = builder.Open(infile, read_properties); + if (!status.ok()) { + // 非Parquet文件Open会失败 + DB_WARNING("FileBuilder fail to open file, file: %s, reason: %s", + _file_path.c_str(), status.message().c_str()); + return -2; + } + status = builder.Build(&reader); + if (!status.ok()) { + DB_WARNING("FileBuilder fail to build reader, file: %s reason: %s", + _file_path.c_str(), status.message().c_str()); + return -1; + } + if (reader == nullptr) { + DB_WARNING("FileReader is nullptr, file: %s", _file_path.c_str()); + return -1; + } + if (reader->parquet_reader() == nullptr) { + DB_WARNING("ParquetReader is nullptr, file: %s", _file_path.c_str()); + return -1; + } + return 0; +} + +} // namespace baikaldb \ No newline at end of file diff --git a/src/exec/ttl_delete_node.cpp b/src/exec/ttl_delete_node.cpp new file mode 100644 index 000000000..53fa5f7d2 --- /dev/null +++ b/src/exec/ttl_delete_node.cpp @@ -0,0 +1,85 @@ +// Copyright (c) 2018-present Baidu, Inc. All Rights Reserved. +// +// Licensed 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 "runtime_state.h" +#include "ttl_delete_node.h" + +namespace baikaldb { + +int TTLDeleteNode::init(const pb::PlanNode& node) { + _table_id = node.derive_node().delete_node().table_id(); + _global_index_id = _table_id; + if (nullptr == (_factory = SchemaFactory::get_instance())) { + DB_WARNING("get record encoder failed"); + return -1; + } + _node_type = pb::DELETE_NODE; + return 0; +} + +int TTLDeleteNode::open(RuntimeState* state) { + int num_affected_rows = 0; + ScopeGuard clear_guard([this](){ clean_delete_records(); }); + + _txn = state->txn(); + bool need_rollback = true; + ScopeGuard auto_rollback([&need_rollback, this]() { + if (need_rollback) { + _txn->rollback(); + } + }); + + int ret = 0; + ret = ExecNode::open(state); + if (ret < 0) { + DB_WARNING_STATE(state, "ExecNode::open fail:%d", ret); + return ret; + } + ret = init_schema_info(state); + if (ret == -1) { + DB_WARNING_STATE(state, "init schema failed fail:%d", ret); + return ret; + } + auto txn = state->txn(); + if (txn == nullptr) { + DB_WARNING_STATE(state, "txn is nullptr: region:%ld", _region_id); + return -1; + } + + for (auto& record: _delete_records) { + int ret = 0; + MutTableKey pk_key; + ret = record->encode_key(*_pri_info, pk_key, -1, false); + if (ret < 0) { + DB_WARNING_STATE(state, "encode key failed, ret:%d", ret); + return ret; + } + ret = remove_row(state, record, pk_key.data(), true); + if (ret < 0) { + DB_WARNING_STATE(state, "delete_row fail"); + return -1; + } + num_affected_rows += ret; + } + auto s = _txn->commit(); + if (!s.ok()) { + DB_FATAL("TTL delete rows commit failed, status: %s", s.getState()); + return -1; + } + need_rollback = false; + return num_affected_rows; +} +} + +/* vim: set ts=4 sw=4 sts=4 tw=100 */ diff --git a/src/physical_plan/condition_optimizer.cpp b/src/physical_plan/condition_optimizer.cpp new file mode 100644 index 000000000..7f422da69 --- /dev/null +++ b/src/physical_plan/condition_optimizer.cpp @@ -0,0 +1,154 @@ +// Copyright (c) 2018-present Baidu, Inc. All Rights Reserved. +// +// Licensed 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 "plan.pb.h" +#include "condition_optimizer.h" +#include "vectorize_helpper.h" +#include "filter_node.h" +#include "agg_node.h" +#include "limit_node.h" +#include "sort_node.h" +#include "rocksdb_scan_node.h" +#include "select_manager_node.h" +#include "schema_factory.h" + +namespace baikaldb { +int ConditionOptimizer::analyze(QueryContext* ctx) { + if (ctx->is_from_subquery + || ctx->is_union_subquery + || ctx->is_insert_select_subquery) { + return 0; + } + if (ctx->stmt_type != parser::NT_SELECT + && ctx->stmt_type != parser::NT_UNION) { + return 0; + } + _ctx = ctx; + // 当in filter很大时, 判断是否要调整成in filter不下推的计划, 避免fetcher_store oom + std::vector scan_nodes; + ctx->root->get_node_pass_subquery(pb::SCAN_NODE, scan_nodes); + for (auto scan_node : scan_nodes) { + adjust_huge_in_condition(static_cast(scan_node), false); + } + return 0; +} + +// 单拎出来为, join下推的in和普通in条件处理逻辑一致 +int ConditionOptimizer::adjust_huge_in_condition(ScanNode* scan_node, bool in_acero) { + if (_ctx == nullptr || scan_node == nullptr) { + return 0; + } + if (!scan_node->is_rocksdb_scan_node()) { + return 0; + } + ExecNode* parent_node = scan_node->get_parent(); + if (parent_node->node_type() != pb::WHERE_FILTER_NODE + && parent_node->node_type() != pb::TABLE_FILTER_NODE) { + return 0; + } + FilterNode* filter_node = static_cast(parent_node); + if (!filter_node->has_huge_in_condition()) { + return 0; + } + + ScanIndexInfo* scan_index_info = scan_node->main_scan_index(); + int64_t router_index_id = scan_index_info->router_index_id; + int64_t main_table_id = scan_node->table_id(); + if (router_index_id != main_table_id/* && !scan_index_info->covering_index*/) { + // 先不支持全局索引 + return 0; + } + + // sort agg需要调到selectmanager上面 + ExecNode* parent = scan_node->get_parent(); + SortNode* sort_node = nullptr; + AggNode* agg_node = nullptr; + SelectManagerNode* select_manager = nullptr; + bool has_limit = false; + while (parent != nullptr) { + if (parent->node_type() == pb::SORT_NODE) { + sort_node = static_cast(parent); + } else if (parent->node_type() == pb::AGG_NODE) { + agg_node = static_cast(parent); + } else if (parent->node_type() == pb::SELECT_MANAGER_NODE) { + select_manager = static_cast(parent); + break; + } + parent = parent->get_parent(); + } + if (select_manager == nullptr) { + return 0; + } + + if (agg_node != nullptr) { + if (in_acero) { + // join in条件runtime filter, 此时db acero计划已经定了并启动了 + return 0; + } + } else if (sort_node != nullptr && sort_node->need_projection()) { + // 遇到case,可以考虑在selectmanaeger那先做projection.行列都需要 + return 0; + } + + // filter干掉huge in, 并重新序列化filter_str + std::vector huge_in_exprs; + filter_node->cut_huge_in_condition(huge_in_exprs); + if (huge_in_exprs.empty()) { + return 0; + } + DB_WARNING("log_id: %lu, huge in condition size: %ld, cut huge in condition", + _ctx->stat_info.log_id, huge_in_exprs.size()); + if (select_manager->get_limit() != -1) { + has_limit = true; + } + if (agg_node != nullptr) { + // merge_agg_parent -> merge_agg -> selectmangaer -> agg -> filter ====> + // merge_agg_parent -> merge_agg transfer agg_node -> selectmanager(filter) -> filter (delete agg_node) + auto merge_agg_node = select_manager->get_parent(); + if (merge_agg_node == nullptr + || merge_agg_node->node_type() != pb::MERGE_AGG_NODE) { + return 0; + } + static_cast(merge_agg_node)->transfer_to_agg(); + auto agg_child = agg_node->children(0); + select_manager->replace_child(agg_node, agg_child); + agg_node->clear_children(); + delete agg_node; + } else if (sort_node != nullptr) { + // parent -> selectmanager(sort) -> sort -> filter ====> + // parent -> selectmanager(filter -> sort) -> filter (delete sort_node) + auto sort_child = sort_node->children(0); + select_manager->replace_child(sort_node, sort_child); + select_manager->steal_slot_order_exprs(sort_node); + sort_node->clear_children(); + delete sort_node; + } else if (has_limit) { + // limit -> selectmanager(limit) -> filter(limit) ====> + // limit -> selectmanager(filter -> limit) -> filter(limit:-1) + ExecNode* node = select_manager->children(0); + while (node != nullptr) { + node->set_limit(-1); + if (node->children_size() > 0) { + node = node->children(0); + } else { + break; + } + } + } + select_manager->add_conditions(huge_in_exprs); + return 0; +} +} + +/* vim: set ts=4 sw=4 sts=4 tw=100 */ From dadac3c30c7082134c9c505c0af266fd9bcdabbf Mon Sep 17 00:00:00 2001 From: lgqss Date: Fri, 13 Feb 2026 11:51:33 +0800 Subject: [PATCH 5/5] merge_2025_12_05 fix compile4 --- cmake/arrow.cmake | 2 +- src/store/region.cpp | 2 +- test/test_column.cpp | 630 +++++++++++++++++ test/test_column_storage.cpp | 31 + test/test_exchange.cpp | 1286 ++++++++++++++++++++++++++++++++++ test/test_faiss_sift1M.cpp | 607 ++++++++++++++++ test/test_parquet_writer.cpp | 77 ++ test/test_row2col.cpp | 27 + test/test_window.cpp | 441 ++++++++++++ 9 files changed, 3101 insertions(+), 2 deletions(-) create mode 100644 test/test_column.cpp create mode 100644 test/test_column_storage.cpp create mode 100644 test/test_exchange.cpp create mode 100644 test/test_faiss_sift1M.cpp create mode 100644 test/test_parquet_writer.cpp create mode 100644 test/test_row2col.cpp create mode 100644 test/test_window.cpp diff --git a/cmake/arrow.cmake b/cmake/arrow.cmake index ec774ca68..038a78582 100644 --- a/cmake/arrow.cmake +++ b/cmake/arrow.cmake @@ -23,7 +23,7 @@ SET(ARROW_ACERO_LIB "${ARROW_INSTALL_DIR}/lib/libarrow_acero.a" CACHE FILEPATH " SET(ARROW_BUNDLED_DEP_LIB "${ARROW_INSTALL_DIR}/lib/libarrow_bundled_dependencies.a" CACHE FILEPATH "arrow dependencies." FORCE) FILE(WRITE ${ARROW_SOURCES_DIR}/src/build.sh - "#!/bin/bash\nset -ex\ncd cpp\ncmake -DCMAKE_BUILD_TYPE=release -DARROW_JEMALLOC=OFF -DARROW_BUILD_SHARED=OFF -DARROW_PARQUET=ON -DARROW_WITH_ZLIB=ON -DARROW_WITH_ZSTD=ON -DARROW_WITH_LZ4=ON -DARROW_WITH_SNAPPY=ON -DARROW_COMPUTE=ON -DARROW_ACERO=ON -DARROW_FILESYSTEM=ON -DARROW_JSON=ON -DARROW_PARQUET=ON -DARROW_BUILD_TESTS=OFF -DARROW_BUILD_STATIC=ON -DARROW_WITH_RE2=ON -DARROW_RE2_VENDORED=OFF -DRE2_HOME=${RE2_INSTALL_DIR} -DARROW_GFLAGS=OFF -DCMAKE_PREFIX_PATH=${RE2_INSTALL_DIR}:${GFLAGS_INSTALL_DIR} -DCMAKE_CXX_FLAGS=\"-I${GFLAGS_INSTALL_DIR}/include\" && make -j4\n" + "#!/bin/bash\nset -ex\ncd cpp\ncmake -DCMAKE_BUILD_TYPE=release -DARROW_JEMALLOC=OFF -DARROW_BUILD_SHARED=OFF -DARROW_PARQUET=ON -DARROW_WITH_ZLIB=ON -DARROW_WITH_ZSTD=ON -DARROW_WITH_LZ4=ON -DARROW_WITH_SNAPPY=ON -DARROW_COMPUTE=ON -DARROW_ACERO=ON -DARROW_FILESYSTEM=ON -DARROW_JSON=ON -DARROW_PARQUET=ON -DARROW_BUILD_TESTS=OFF -DARROW_BUILD_STATIC=ON -DARROW_WITH_RE2=ON -DARROW_RE2_VENDORED=OFF -DRE2_HOME=${RE2_INSTALL_DIR} -DARROW_GFLAGS=OFF -DCMAKE_PREFIX_PATH=${RE2_INSTALL_DIR} -DCMAKE_CXX_FLAGS=\"-I${GFLAGS_INSTALL_DIR}/include\" && make -j4\n" ) ExternalProject_Add( diff --git a/src/store/region.cpp b/src/store/region.cpp index 1c6aad616..2059c1e5d 100644 --- a/src/store/region.cpp +++ b/src/store/region.cpp @@ -9115,7 +9115,7 @@ void Region::check_peer_latency() { } else { _node.get_status(&status); } - if (status.state == raft::STATE_ERROR) { + if (status.state == braft::STATE_ERROR) { _raft_status_error = true; } else { _raft_status_error = false; diff --git a/test/test_column.cpp b/test/test_column.cpp new file mode 100644 index 000000000..0931d724f --- /dev/null +++ b/test/test_column.cpp @@ -0,0 +1,630 @@ +// Copyright (c) 2018 Baidu, Inc. All Rights Reserved. +// +// Licensed 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 +#include +#include +#include +// #include "proto/meta.interface.pb.h" +// #include "proto/plan.pb.h" +#include "parquet_writer.h" +#include "sort_merge.h" +#include "file_manager.h" +#include "mut_table_key.h" +// #include "arrow/testing/gtest_util.h" + +int main(int argc, char* argv[]) +{ + testing::InitGoogleTest(&argc, argv); + google::ParseCommandLineFlags(&argc, &argv, true); + return RUN_ALL_TESTS(); +} + +namespace baikaldb { +DECLARE_string(db_path); +/* +struct ColumnSchemaInfo { + std::vector key_fields; + std::vector value_fields; + SmartIndex index_info = nullptr; + SmartTable table_info = nullptr; + int uniq_size = 0; // 主键或索引key的字段个数key_fields.size(); 在parquet中的idx为0~uniq_size-1 + int keytype_idx = -1; // key_fields.size() + value_fields.size(); + int raft_index_idx = -1; // keytype_idx + 1 + int batch_pos_idx = -1; // raft_index_idx + 1 + std::set need_sum_idx; + std::shared_ptr schema = nullptr; // 包含表中所有字段 + __key_type__; 顺序为key + value + __key_type__; key为主键序，value为非主键字段按field_idx顺序 + std::shared_ptr schema_with_order_info = nullptr; // 比schema多包含__raft_index__, __batch_pos__放在__key_type__后面 +}; +*/ + +std::shared_ptr test_make_column_schema(int64_t tableid) { + auto schema_ptr = std::make_shared(); + FieldInfo key; + key.id = 1; + key.type = pb::INT32; + key.lower_short_name = "k1"; + schema_ptr->key_fields.emplace_back(key); + + FieldInfo v1; + v1.id = 2; + v1.type = pb::INT32; + v1.lower_short_name = "v1"; + schema_ptr->value_fields.emplace_back(v1); + + FieldInfo v2; + v2.id = 3; + v2.type = pb::INT32; + v2.lower_short_name = "v2"; + schema_ptr->value_fields.emplace_back(v2); + + schema_ptr->uniq_size = schema_ptr->key_fields.size(); + schema_ptr->keytype_idx = schema_ptr->key_fields.size() + schema_ptr->value_fields.size(); + schema_ptr->raft_index_idx = schema_ptr->keytype_idx + 1; + schema_ptr->batch_pos_idx = schema_ptr->raft_index_idx + 1; + + schema_ptr->need_sum_idx.insert(1); + schema_ptr->need_sum_idx.insert(2); + std::vector> arrow_fields; + arrow_fields.reserve(schema_ptr->batch_pos_idx + 1); + for (const auto& field : schema_ptr->key_fields) { + auto arrow_type = primitive_to_arrow_type(field.type); + if (arrow_type < 0) { + DB_COLUMN_FATAL("field: %s primitive type:%d to arrow type failed", field.lower_short_name.c_str(), field.type); + return nullptr; + } + auto arrow_field = VectorizeHelpper::make_field(field.lower_short_name, arrow::Type::type(arrow_type)); + if (arrow_field == nullptr) { + DB_COLUMN_FATAL("field: %s make arrow schema failed", field.lower_short_name.c_str()); + return nullptr; + } + arrow_fields.emplace_back(arrow_field); + } + + for (const auto& field : schema_ptr->value_fields) { + auto arrow_type = primitive_to_arrow_type(field.type); + if (arrow_type < 0) { + DB_COLUMN_FATAL("field: %s primitive type:%d to arrow type failed", field.lower_short_name.c_str(), field.type); + return nullptr; + } + auto arrow_field = VectorizeHelpper::make_field(field.lower_short_name, arrow::Type::type(arrow_type)); + if (arrow_field == nullptr) { + DB_COLUMN_FATAL("field: %s make arrow schema failed", field.lower_short_name.c_str()); + return nullptr; + } + arrow_fields.emplace_back(arrow_field); + } + + arrow_fields.emplace_back(VectorizeHelpper::make_field(KEY_TYPE_NAME, arrow::Type::type::INT32)); + schema_ptr->schema = std::make_shared(arrow_fields); + arrow_fields.emplace_back(VectorizeHelpper::make_field(RAFT_INDEX_NAME, arrow::Type::type::INT64)); + arrow_fields.emplace_back(VectorizeHelpper::make_field(BATCH_POS_NAME, arrow::Type::type::INT32)); + schema_ptr->schema_with_order_info = std::make_shared(arrow_fields); + return schema_ptr; +} + +struct TESTRecordBatchReaderOptions { + std::shared_ptr schema = nullptr; + int batch = 0; +}; + +class TESTRecordBatchReader : public ::arrow::RecordBatchReader { +public: + TESTRecordBatchReader(std::shared_ptr schema) : _schema(schema) { + _column_record = std::make_shared(schema, 1024); + _column_record->init(); + _column_record->reserve(1024); + + } + virtual ~TESTRecordBatchReader() { } + + void append_row(const std::vector& row, int64_t raft_index = -1, int32_t batch_pos = -1) { + std::vector expr_row; + expr_row.reserve(row.size() + 2); + for (int i = 0; i < row.size(); ++i) { + ExprValue value; + value.type = pb::INT32; + value._u.int32_val = row[i]; + expr_row.push_back(value); + } + if (raft_index != -1 && batch_pos != -1) { + ExprValue raft_index_value; + raft_index_value.type = pb::INT64; + raft_index_value._u.int64_val = raft_index; + expr_row.push_back(raft_index_value); + ExprValue batch_pos_value; + batch_pos_value.type = pb::INT32; + batch_pos_value._u.int32_val = batch_pos; + expr_row.push_back(batch_pos_value); + } + + _column_record->append_row(expr_row); + } + + std::shared_ptr schema() const override { return _schema; } + + virtual ::arrow::Status ReadNext(std::shared_ptr<::arrow::RecordBatch>* batch) override { + if (_column_record->size() <= 0) { + batch->reset(); + return ::arrow::Status::OK(); + } + + int ret = _column_record->finish_and_make_record_batch(batch); + if (ret < 0) { + return arrow::Status::IOError("make record batch fail"); + } + DB_WARNING("record batch size:%ld", (*batch)->num_rows()); + return ::arrow::Status::OK(); + } +private: + std::shared_ptr _schema = nullptr; + std::shared_ptr _column_record = nullptr; +}; + +TEST(test_column_record, case_all) { + std::shared_ptr schema_ptr = test_make_column_schema(123); + auto reader = std::make_shared(schema_ptr->schema_with_order_info); + reader->append_row({1,1,1,0},0,0); + reader->append_row({2,2,2,0},0,0); + reader->append_row({3,3,3,0},0,0); + reader->append_row({4,4,4,0},0,0); + std::shared_ptr record_batch; + reader->ReadNext(&record_batch); + ColumnRecord::TEST_print_record_batch(record_batch); +} +TEST(test_sort_merge, unorder_sort) { + int ret = 0; + std::shared_ptr schema_ptr = test_make_column_schema(123); + auto reader1 = std::make_shared(schema_ptr->schema_with_order_info); + reader1->append_row({1,1,1,COLUMN_KEY_PUT},0,0); + reader1->append_row({3,3,3,COLUMN_KEY_PUT},0,0); + reader1->append_row({2,2,2,COLUMN_KEY_MERGE},1,0); + reader1->append_row({2,2,2,COLUMN_KEY_PUT},0,0); + reader1->append_row({1,1,1,COLUMN_KEY_DELETE},1,0); + reader1->append_row({1,1,1,COLUMN_KEY_MERGE},2,0); + + std::vector> readers; + readers.reserve(2); + auto usrr = std::make_shared(reader1, schema_ptr.get(), 5); + readers.emplace_back(usrr); + SortMergeOptions merge_options; + merge_options.batch_size = 10; + merge_options.is_base_compact = false; + merge_options.schema_info = schema_ptr; + auto sort_merge = std::make_shared(merge_options, readers); + std::shared_ptr record_batch; + auto s = sort_merge->ReadNext(&record_batch); + EXPECT_EQ(s.ok(), true); + ColumnRecord::TEST_print_record_batch(record_batch); + + auto result = std::make_shared(schema_ptr->schema); + result->append_row({1,1,1,COLUMN_KEY_PUT},-1,-1); + result->append_row({2,4,4,COLUMN_KEY_PUT},-1,-1); + result->append_row({3,3,3,COLUMN_KEY_PUT},-1,-1); + std::shared_ptr result_rb; + result->ReadNext(&result_rb); + + EXPECT_EQ(ColumnRecord::TEST_record_batch_diff(result_rb, record_batch), false); + +} + +TEST(test_sort_merge, minor_compact) { + int ret = 0; + std::shared_ptr schema_ptr = test_make_column_schema(123); + auto reader1 = std::make_shared(schema_ptr->schema_with_order_info); + reader1->append_row({1,1,1,2},2,0); + reader1->append_row({3,3,3,2},2,0); + reader1->append_row({2,2,2,2},2,0); + + auto reader2 = std::make_shared(schema_ptr->schema_with_order_info); + reader2->append_row({1,1,1,2},1,0); + reader2->append_row({2,2,2,2},1,0); + reader2->append_row({3,3,3,2},1,0); + reader2->append_row({4,4,4,2},1,0); + + auto reader3 = std::make_shared(schema_ptr->schema_with_order_info); + reader3->append_row({2,2,2,2},0,0); + reader3->append_row({3,3,3,2},0,0); + reader3->append_row({4,1,1,2},0,0); + reader3->append_row({5,5,5,2},0,0); + + std::vector> readers; + readers.reserve(2); + auto usrr = std::make_shared(reader1, schema_ptr.get(), 0); + auto osrr = std::make_shared(reader2, schema_ptr.get()); + auto osrr2 = std::make_shared(reader3, schema_ptr.get()); + + readers.emplace_back(usrr); + readers.emplace_back(osrr); + readers.emplace_back(osrr2); + + SortMergeOptions merge_options; + merge_options.batch_size = 10; + merge_options.is_base_compact = false; + merge_options.schema_info = schema_ptr; + auto sort_merge = std::make_shared(merge_options, readers); + std::shared_ptr record_batch; + auto s = sort_merge->ReadNext(&record_batch); + EXPECT_EQ(s.ok(), true); + ColumnRecord::TEST_print_record_batch(record_batch); + + auto result = std::make_shared(schema_ptr->schema); + result->append_row({1,2,2,COLUMN_KEY_MERGE},-1,-1); + result->append_row({2,6,6,COLUMN_KEY_MERGE},-1,-1); + result->append_row({3,9,9,COLUMN_KEY_MERGE},-1,-1); + result->append_row({4,5,5,COLUMN_KEY_MERGE},-1,-1); + result->append_row({5,5,5,COLUMN_KEY_MERGE},-1,-1); + std::shared_ptr result_rb; + result->ReadNext(&result_rb); + + EXPECT_EQ(ColumnRecord::TEST_record_batch_diff(result_rb, record_batch), false); +} + +TEST(test_acero_merge, acero_merge) { + int ret = 0; + std::shared_ptr schema_ptr = test_make_column_schema(123); + auto reader1 = std::make_shared(schema_ptr->schema_with_order_info); + reader1->append_row({1,1,1,COLUMN_KEY_MERGE},2,0); + reader1->append_row({3,3,3,COLUMN_KEY_MERGE},2,0); + reader1->append_row({2,2,2,COLUMN_KEY_MERGE},2,0); + reader1->append_row({1,1,1,COLUMN_KEY_MERGE},2,0); + reader1->append_row({2,2,2,COLUMN_KEY_MERGE},2,0); + + AceroMergeOptions acero_options; + acero_options.batch_size = 10; + acero_options.schema_info = schema_ptr; + std::shared_ptr acero_merge(new AceroMerge(acero_options, {reader1})); + + std::shared_ptr record_batch; + auto s = acero_merge->ReadNext(&record_batch); + EXPECT_EQ(s.ok(), true); + ColumnRecord::TEST_print_record_batch(record_batch); + auto result = std::make_shared(schema_ptr->schema); + result->append_row({1,2,2,COLUMN_KEY_MERGE},-1,-1); + result->append_row({2,4,4,COLUMN_KEY_MERGE},-1,-1); + result->append_row({3,3,3,COLUMN_KEY_MERGE},-1,-1); + std::shared_ptr result_rb; + result->ReadNext(&result_rb); + ColumnRecord::TEST_print_record_batch(result_rb); + // EXPECT_EQ(ColumnRecord::TEST_record_batch_diff(result_rb, record_batch), false); +} + +TEST(parquet_read_write, write_read) { + int64_t table_id = 123; + int64_t region_id = 456; + int ret = 0; + std::shared_ptr schema_info = test_make_column_schema(table_id); + auto reader1 = std::make_shared(schema_info->schema); + reader1->append_row({1,1,1,0}); + reader1->append_row({2,2,2,0}); + reader1->append_row({3,3,3,0}); + std::shared_ptr<::arrow::RecordBatch> record_batch; + auto s = reader1->ReadNext(&record_batch); + EXPECT_EQ(s.ok(), true); + EXPECT_EQ(record_batch->num_columns(), 4); + EXPECT_EQ(record_batch->num_rows(), 3); + FLAGS_db_path = "./test_data"; + boost::filesystem::path output_path(FLAGS_db_path + "_tmp"); + if (!boost::filesystem::exists(output_path)) { + // 创建目录 + if (!boost::filesystem::create_directories(output_path)) { + DB_COLUMN_FATAL("FATAL create output_path fail."); + return; + } + } + + boost::filesystem::path output_path2(FLAGS_db_path + "_column/" + std::to_string(table_id) + "/" + std::to_string(region_id)); + if (!boost::filesystem::exists(output_path2)) { + // 创建目录 + if (!boost::filesystem::create_directories(output_path2)) { + DB_COLUMN_FATAL("FATAL create output_path2 fail."); + return; + } + } + + ParquetWriteOptions parquet_options; + parquet_options.write_batch_length = 10; + parquet_options.max_file_rows = 1000; + parquet_options.max_row_group_length = 100; // 保证row group长度是write_batch_length的整数倍 + parquet_options.schema_info = schema_info; + auto writer = std::make_shared(parquet_options); + s = writer->init(); + EXPECT_EQ(s.ok(), true); + + RecordBatchInfo rb_info; + rb_info.row_count = record_batch->num_rows(); + s = writer->write_batch(record_batch, rb_info); + EXPECT_EQ(s.ok(), true); + + s = writer->finish(); + EXPECT_EQ(s.ok(), true); + + auto file_infos = writer->get_file_infos(); + EXPECT_EQ(file_infos.size(), 1); + + std::shared_ptr new_file = std::make_shared(table_id, region_id, 0, 123, 0, file_infos[0]); + + ret = ::link(file_infos[0].file_name.c_str(), new_file->full_path().c_str()); + EXPECT_EQ(ret, 0); + + auto file = std::make_shared(new_file); + ret = file->open(); + EXPECT_EQ(ret, 0); + if (ret < 0) { + DB_COLUMN_FATAL("open file:%s failed, ret:%d", new_file->full_path().c_str(), ret); + return; + } + ParquetFileReaderOptions options; + options.schema = schema_info->schema; + std::shared_ptr parquet_reader = std::make_shared(options, file); + ret = parquet_reader->init(); + EXPECT_EQ(ret, 0); + std::shared_ptr<::arrow::RecordBatch> record_batch2; + s = parquet_reader->ReadNext(&record_batch2); + EXPECT_EQ(s.ok(), true); + ColumnRecord::TEST_print_record_batch(record_batch2); + ColumnRecord::TEST_print_record_batch(record_batch); + EXPECT_EQ(ColumnRecord::TEST_record_batch_diff(record_batch2, record_batch), false); + +} +DEFINE_string(test_parquet_file, "", "rocks db path"); +struct TESTColumnFileInfo : public ColumnFileInfo { + TESTColumnFileInfo() : ColumnFileInfo(123, 245, 0, 123, 0, ColumnFileMeta()) { + + } + + std::string full_path() { + return FLAGS_test_parquet_file; + } +}; + +TEST(parquet_read_write, read) { + std::shared_ptr new_file(new TESTColumnFileInfo()); + if (new_file->full_path().empty()) { + DB_COLUMN_FATAL("Fail to get parquet file"); + return; + } + + auto parquet_file = std::make_shared(new_file); + + int ret = parquet_file->open(); + if (ret < 0) { + DB_COLUMN_FATAL("Fail to open parquet file"); + return; + } + std::unique_ptr<::arrow::RecordBatchReader> reader; + auto s = parquet_file->GetRecordBatchReader(&reader); + if (!s.ok()) { + DB_COLUMN_FATAL("Fail to get_record_batch_reader"); + return; + } + + std::shared_ptr<::arrow::RecordBatch> record_batch; + s = reader->ReadNext(&record_batch); + if (!s.ok()) { + DB_COLUMN_FATAL("Fail to ReadNext"); + return; + } + + ColumnRecord::TEST_print_record_batch(record_batch); +} +struct TestRow { + int64_t v = 0; + bool operator<(const TestRow& other) const { + DB_WARNING("v:%ld, other.v:%ld", v, other.v); + return v < other.v; + } +}; +TEST(heap_test, all) { + Heap heap; + std::vector v = {1, 2, 3, 5, 8, 13, 7, 6, 4, 10}; + for (auto& i : v) { + TestRow r; + r.v = i; + heap.push(r); + } + DB_WARNING("make heap size:%lu", heap.size()); + heap.make_heap(); + // auto& h = heap.heap(); + // for (auto& t : h) { + // DB_WARNING("v:%ld", t.v); + // } + + // TestRow r; + // r.v = 15; + // heap.replace_top(r); + // for (auto& t : h) { + // DB_WARNING("v:%ld", t.v); + // } + + +} + +TEST(check_interval_overlapped, all) { +{ + pb::PossibleIndex::Range range; + MutTableKey index_start_key; + index_start_key.append_u64(1).append_u64(2); + MutTableKey index_end_key; + index_end_key.append_u64(1).append_u64(2); + range.set_left_key(index_start_key.data()); + range.set_left_open(false); + range.set_right_key(index_end_key.data()); + range.set_right_open(false); + MutTableKey file_start_key; + file_start_key.append_u64(1).append_u64(2).append_u64(3); + MutTableKey file_end_key; + file_end_key.append_u64(1).append_u64(2).append_u64(3); + EXPECT_EQ(ParquetFile::check_interval_overlapped(range, false, false, false, file_start_key.data(), file_end_key.data()), true); +} + +{ + pb::PossibleIndex::Range range; + MutTableKey index_start_key; + index_start_key.append_u64(1).append_u64(1); + MutTableKey index_end_key; + index_end_key.append_u64(1).append_u64(2); + range.set_left_key(index_start_key.data()); + range.set_left_open(false); + range.set_right_key(index_end_key.data()); + range.set_right_open(false); + MutTableKey file_start_key; + file_start_key.append_u64(1).append_u64(2); + MutTableKey file_end_key; + file_end_key.append_u64(1).append_u64(2); + EXPECT_EQ(ParquetFile::check_interval_overlapped(range, false, false, false, file_start_key.data(), file_end_key.data()), true); +} + +{ + pb::PossibleIndex::Range range; + MutTableKey index_start_key; + index_start_key.append_u64(1).append_u64(1); + MutTableKey index_end_key; + index_end_key.append_u64(1).append_u64(2); + range.set_left_key(index_start_key.data()); + range.set_left_open(false); + range.set_right_key(index_end_key.data()); + range.set_right_open(true); + MutTableKey file_start_key; + file_start_key.append_u64(1).append_u64(2); + MutTableKey file_end_key; + file_end_key.append_u64(1).append_u64(2); + EXPECT_EQ(ParquetFile::check_interval_overlapped(range, false, false, true, file_start_key.data(), file_end_key.data()), false); +} + +{ + pb::PossibleIndex::Range range; + MutTableKey index_start_key; + index_start_key.append_u64(1).append_u64(1); + MutTableKey index_end_key; + index_end_key.append_u64(1).append_u64(3); + range.set_left_key(index_start_key.data()); + range.set_left_open(false); + range.set_right_key(index_end_key.data()); + range.set_right_open(true); + MutTableKey file_start_key; + file_start_key.append_u64(1).append_u64(2).append_u64(3); + MutTableKey file_end_key; + file_end_key.append_u64(1).append_u64(2).append_u64(3); + EXPECT_EQ(ParquetFile::check_interval_overlapped(range, false, false, true, file_start_key.data(), file_end_key.data()), true); +} + +} + +struct TestBlockContents { + int size = 0; + std::string value; + char* data = nullptr; + TestBlockContents(int size, const std::string& value) : size(size), value(value), data(new char[size]) { + DB_WARNING("size:%d, value:%s", size, value.c_str()); + } + ~TestBlockContents() { + DB_WARNING("delete data size:%d, value:%s", size, value.c_str()); + delete[] data; + } + static void delete_fn(void* value, rocksdb::MemoryAllocator* allocator) { + TestBlockContents* block = static_cast(value); + DB_WARNING("delete data size:%d, value:%s", block->size, block->value.c_str()); + delete block; + return; + } + + static rocksdb::Cache::CacheItemHelper kBasicHelper; +}; + +rocksdb::Cache::CacheItemHelper TestBlockContents::kBasicHelper{ + rocksdb::CacheEntryRole::kOtherBlock, &TestBlockContents::delete_fn +}; + +// // Returns the maximum configured capacity of the cache +// virtual size_t GetCapacity() const = 0; + +// // Returns the memory size for the entries residing in the cache. +// virtual size_t GetUsage() const = 0; + +// // Returns the number of entries currently tracked in the table. SIZE_MAX +// // means "not supported." This is used for inspecting the load factor, along +// // with GetTableAddressCount(). +// virtual size_t GetOccupancyCount() const { return SIZE_MAX; } + +// // Returns the number of ways the hash function is divided for addressing +// // entries. Zero means "not supported." This is used for inspecting the load +// // factor, along with GetOccupancyCount(). +// virtual size_t GetTableAddressCount() const { return 0; } + +// // Returns the memory size for a specific entry in the cache. +// virtual size_t GetUsage(Handle* handle) const = 0; + +// // Returns the memory size for the entries in use by the system +// virtual size_t GetPinnedUsage() const = 0; + +// // Returns the charge for the specific entry in the cache. +// virtual size_t GetCharge(Handle* handle) const = 0; + +TEST(lru_cache, all) { +{ + std::vector keys = {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"}; + std::shared_ptr block_cache = rocksdb::NewLRUCache(128 * 1024 * 1024LL, 8); + auto block = new TestBlockContents(1024, "abc"); + rocksdb::Cache::Handle* handle = nullptr; + rocksdb::Slice key = rocksdb::Slice("abc"); + auto s = block_cache->Insert(key, block, &TestBlockContents::kBasicHelper, 1024,&handle); + if (s.ok()) { + DB_WARNING("insert ok"); + EXPECT_NE(handle, nullptr); + auto value = block_cache->Value(handle); + EXPECT_EQ(value, block); + TestBlockContents* v = static_cast(value); + block_cache->Release(handle); + DB_WARNING("value:%s, GetUsage:%lu,GetCharge:%lu", v->value.c_str(), block_cache->GetUsage(handle), block_cache->GetCharge(handle)); + } + rocksdb::Cache::Handle* handle2 = nullptr; + auto block2 = new TestBlockContents(1024, "abcd"); + s = block_cache->Insert(key, block2, &TestBlockContents::kBasicHelper, 1024,&handle2); + if (s.ok()) { + block_cache->Release(handle2); + DB_WARNING("insert ok abcd"); + } else { + DB_WARNING("insert fail abcd %s", s.ToString().c_str()); + } + +} +{ + std::vector keys = {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"}; + std::shared_ptr block_cache = rocksdb::NewLRUCache(128 * 1024 * 1024LL, 8); + int i = 0; + for (auto& k : keys) { + auto block = new TestBlockContents(1024, k); + rocksdb::Slice key = rocksdb::Slice(k); + rocksdb::Cache::Handle* handle = nullptr; + auto s = block_cache->Insert(key, block, &TestBlockContents::kBasicHelper, 1024 + ++i, &handle); + EXPECT_EQ(s.ok(), true); + auto value = block_cache->Value(handle); + EXPECT_EQ(value, block); + TestBlockContents* v = static_cast(value); + block_cache->Release(handle); + DB_WARNING("value:%s, GetUsage:%lu,GetCharge:%lu", v->value.c_str(), block_cache->GetUsage(handle), block_cache->GetCharge(handle)); + } + DB_WARNING("GetCapacity:%lu, GetUsage:%lu", block_cache->GetCapacity(), block_cache->GetUsage()); + for (auto& k : keys) { + rocksdb::Slice key = rocksdb::Slice(k); + auto handle = block_cache->BasicLookup(key, nullptr); + EXPECT_NE(handle, nullptr); + auto value = block_cache->Value(handle); + TestBlockContents* v = static_cast(value); + EXPECT_EQ(k, v->value); + block_cache->Release(handle); + } +} +} +} // namespace baikal \ No newline at end of file diff --git a/test/test_column_storage.cpp b/test/test_column_storage.cpp new file mode 100644 index 000000000..a45c7a5b4 --- /dev/null +++ b/test/test_column_storage.cpp @@ -0,0 +1,31 @@ +// Copyright (c) 2022 Baidu, Inc. All Rights Reserved. +// +// Licensed 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 + +#include "parquet_scan_node.h" + +int main(int argc, char* argv[]) { + testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} + +namespace baikaldb { + +TEST(test_check_interval_overlapped, case_all) { + ParquetScanNode parquet_scan_node; + +} + +} // namespace baikaldb \ No newline at end of file diff --git a/test/test_exchange.cpp b/test/test_exchange.cpp new file mode 100644 index 000000000..c748607cf --- /dev/null +++ b/test/test_exchange.cpp @@ -0,0 +1,1286 @@ +// Copyright (c) 2022 Baidu, Inc. All Rights Reserved. +// +// Licensed 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 + +#include "arrow/array/concatenate.h" +#include "arrow_exec_node.h" +#include "data_stream_manager.h" +#include "exchange_sender_node.h" +#include "exchange_receiver_node.h" +#include "db_service.h" +#include "region.h" +#include +#include + +namespace baikaldb { + +using namespace arrow; + +//// 测试数据构造 +pb::TransmitDataParam make_invalid_transmit_data_param() { + pb::TransmitDataParam transmit_data_param; + transmit_data_param.set_exchange_state(pb::ES_DOING); + transmit_data_param.set_log_id(-1); + transmit_data_param.set_sender_fragment_instance_id(-1); + transmit_data_param.set_receiver_fragment_instance_id(-1); + transmit_data_param.set_receiver_node_id(-1); + return transmit_data_param; +} + +std::shared_ptr make_test_array(const std::vector& values) { + arrow::Int64Builder builder; + builder.AppendValues(values); + std::shared_ptr array; + builder.Finish(&array); + return array; +} +std::shared_ptr make_test_double_array(const std::vector& values) { + arrow::DoubleBuilder builder; + builder.AppendValues(values); + std::shared_ptr array; + builder.Finish(&array); + return array; +} +std::shared_ptr make_test_string_array(const std::vector& values) { + arrow::BinaryBuilder builder; + builder.AppendValues(values); + std::shared_ptr array; + builder.Finish(&array); + return array; +} +std::shared_ptr make_test_null_array() { + arrow::NullBuilder builder; + builder.AppendNull(); + builder.Append(nullptr); + std::shared_ptr array; + builder.Finish(&array); + return array; +} + +std::shared_ptr make_test_array_tmp(const std::vector& values) { + arrow::Int64Builder builder; + // builder.AppendValues(values); + builder.AppendNull(); + builder.AppendNull(); + std::shared_ptr array; + builder.Finish(&array); + return array; +} +std::shared_ptr make_test_schema() { + return arrow::schema({arrow::field("col1", arrow::int64()), + arrow::field("col2", arrow::int64()), + arrow::field("col3", arrow::int64())}); +} +std::shared_ptr make_test_schema2() { + return arrow::schema({arrow::field("col1", arrow::int64()), + arrow::field("col2", arrow::int64()), + arrow::field("col3", arrow::utf8())}); +} +std::shared_ptr make_test_schema3() { + return arrow::schema({arrow::field("col3", arrow::int64()), + arrow::field("col1", arrow::float64()), + arrow::field("col2", arrow::utf8())}); +} +std::shared_ptr make_test_null_schema() { + return arrow::schema({arrow::field("col1", arrow::int64()), + arrow::field("col2", arrow::int64()), + arrow::field("col3", arrow::int64())}); +} +std::shared_ptr make_test_empty_record_batch_schema1() { + auto out = arrow::RecordBatch::MakeEmpty(make_test_schema()); + return *out; +} +std::shared_ptr make_test_record_batch1_schema1() { + std::shared_ptr array1 = make_test_array({1,2,3,4,5,6,7,8,9,10}); + std::shared_ptr array2 = make_test_array({11,12,13,14,15,16,17,18,19,20}); + std::shared_ptr array3 = make_test_array({21,22,23,24,25,26,27,28,29,30}); + std::shared_ptr out = + arrow::RecordBatch::Make(make_test_schema(), 10, {array1, array2, array3}); + return out; +} +std::shared_ptr make_test_record_batch2_schema1() { + std::shared_ptr array1 = make_test_array({31,32,33,34,35,36,37,38}); + std::shared_ptr array2 = make_test_array({41,42,43,44,45,46,47,48}); + std::shared_ptr array3 = make_test_array({51,52,53,54,55,56,57,58}); + std::shared_ptr out = + arrow::RecordBatch::Make(make_test_schema(), 8, {array1, array2, array3}); + return out; +} +std::shared_ptr make_test_record_batch3_schema1() { + std::shared_ptr array1 = make_test_array({61,62,63,64,65,66,67,68,69}); + std::shared_ptr array2 = make_test_array({71,72,73,74,75,76,77,78,79}); + std::shared_ptr array3 = make_test_array({81,82,83,84,85,86,87,88,89}); + std::shared_ptr out = + arrow::RecordBatch::Make(make_test_schema(), 9, {array1, array2, array3}); + return out; +} +std::shared_ptr make_test_record_batch1_schema2() { + std::shared_ptr array1 = make_test_array({1,2,3,4,5,6,7,8,9,10}); + std::shared_ptr array2 = make_test_array({11,12,13,14,15,16,17,18,19,20}); + std::shared_ptr array3 = make_test_string_array({"21","22","23","24","25","26","27","28","29","20"}); + std::shared_ptr out = + arrow::RecordBatch::Make(make_test_schema2(), 10, {array1, array2, array3}); + return out; +} +std::shared_ptr make_test_record_batch_null_schema() { + std::shared_ptr array1 = make_test_null_array(); + std::shared_ptr array2 = make_test_null_array(); + std::shared_ptr array3 = make_test_null_array(); + std::shared_ptr out = + arrow::RecordBatch::Make(make_test_null_schema(), 2, {array1, array2, array3}); + return out; +} +pb::PlanNode make_exchange_sender_node() { + pb::PlanNode plan_node; + plan_node.set_node_type(pb::EXCHANGE_SENDER_NODE); + plan_node.set_num_children(1); + pb::DerivePlanNode* derive_node = plan_node.mutable_derive_node(); + pb::ExchangeSenderNode* exchange_sender_node = derive_node->mutable_exchange_sender_node(); + exchange_sender_node->set_log_id(100); + exchange_sender_node->set_fragment_id(7); + exchange_sender_node->set_receiver_fragment_id(8); + exchange_sender_node->set_fragment_instance_id(77); + exchange_sender_node->mutable_partition_property()->set_type(pb::HashPartitionType); + exchange_sender_node->mutable_partition_property()->add_hash_cols("col1"); + exchange_sender_node->mutable_partition_property()->add_hash_cols("col3"); + pb::ExchangeDestination* receiver_destination = exchange_sender_node->add_receiver_destinations(); + receiver_destination->set_fragment_instance_id(88); + receiver_destination->set_node_id(99); + receiver_destination->set_address("10.12.187.120:8787"); + return plan_node; +} +pb::PlanNode make_exchange_sender_node2() { + pb::PlanNode plan_node; + plan_node.set_node_type(pb::EXCHANGE_SENDER_NODE); + plan_node.set_num_children(1); + pb::DerivePlanNode* derive_node = plan_node.mutable_derive_node(); + pb::ExchangeSenderNode* exchange_sender_node = derive_node->mutable_exchange_sender_node(); + exchange_sender_node->set_log_id(100); + exchange_sender_node->set_fragment_id(7); + exchange_sender_node->set_receiver_fragment_id(8); + exchange_sender_node->set_fragment_instance_id(78); + exchange_sender_node->mutable_partition_property()->set_type(pb::HashPartitionType); + exchange_sender_node->mutable_partition_property()->add_hash_cols("col1"); + exchange_sender_node->mutable_partition_property()->add_hash_cols("col3"); + pb::ExchangeDestination* receiver_destination = exchange_sender_node->add_receiver_destinations(); + receiver_destination->set_fragment_instance_id(88); + receiver_destination->set_node_id(99); + receiver_destination->set_address("10.12.187.120:8787"); + return plan_node; +} +pb::PlanNode make_exchange_sender_node3() { + pb::PlanNode plan_node; + plan_node.set_node_type(pb::EXCHANGE_SENDER_NODE); + plan_node.set_num_children(1); + pb::DerivePlanNode* derive_node = plan_node.mutable_derive_node(); + pb::ExchangeSenderNode* exchange_sender_node = derive_node->mutable_exchange_sender_node(); + exchange_sender_node->set_log_id(100); + exchange_sender_node->set_fragment_id(7); + exchange_sender_node->set_receiver_fragment_id(8); + exchange_sender_node->set_fragment_instance_id(79); + exchange_sender_node->mutable_partition_property()->set_type(pb::HashPartitionType); + exchange_sender_node->mutable_partition_property()->add_hash_cols("col1"); + exchange_sender_node->mutable_partition_property()->add_hash_cols("col3"); + pb::ExchangeDestination* receiver_destination = exchange_sender_node->add_receiver_destinations(); + receiver_destination->set_fragment_instance_id(88); + receiver_destination->set_node_id(99); + receiver_destination->set_address("10.12.187.120:8787"); + return plan_node; +} +pb::PlanNode make_exchange_receiver_node() { + pb::PlanNode plan_node; + plan_node.set_node_type(pb::EXCHANGE_RECEIVER_NODE); + plan_node.set_num_children(1); + pb::DerivePlanNode* derive_node = plan_node.mutable_derive_node(); + pb::ExchangeReceiverNode* exchange_receiver_node = derive_node->mutable_exchange_receiver_node(); + exchange_receiver_node->set_log_id(100); + exchange_receiver_node->set_fragment_id(8); + exchange_receiver_node->set_sender_fragment_id(7); + exchange_receiver_node->set_fragment_instance_id(88); + exchange_receiver_node->set_node_id(99); + pb::ExchangeDestination* sender_destination = exchange_receiver_node->add_sender_destinations(); + sender_destination->set_fragment_instance_id(77); + sender_destination->set_node_id(99); + sender_destination->set_address("10.12.187.120:8788"); + + std::shared_ptr arrow_schema = make_test_schema(); + arrow::Result> schema_ret = + arrow::ipc::SerializeSchema(*arrow_schema, arrow::default_memory_pool()); + if (!schema_ret.ok()) { + DB_FATAL("arrow serialize schema fail, status: %s", schema_ret.status().ToString().c_str()); + return plan_node; + } + exchange_receiver_node->set_schema((*schema_ret)->data(), (*schema_ret)->size()); + + return plan_node; +} + +pb::PlanNode make_invalid_exchange_receiver_node() { + pb::PlanNode plan_node; + plan_node.set_node_type(pb::EXCHANGE_RECEIVER_NODE); + plan_node.set_num_children(0); + pb::DerivePlanNode* derive_node = plan_node.mutable_derive_node(); + pb::ExchangeReceiverNode* exchange_receiver_node = derive_node->mutable_exchange_receiver_node(); + exchange_receiver_node->set_log_id(-1); + exchange_receiver_node->set_fragment_id(-1); + exchange_receiver_node->set_sender_fragment_id(-1); + exchange_receiver_node->set_fragment_instance_id(-1); + exchange_receiver_node->set_node_id(-1); + pb::ExchangeDestination* sender_destination = exchange_receiver_node->add_sender_destinations(); + sender_destination->set_fragment_instance_id(-1); + sender_destination->set_node_id(-1); + sender_destination->set_address("10.12.187.120:8788"); + + std::shared_ptr arrow_schema = make_test_schema(); + arrow::Result> schema_ret = + arrow::ipc::SerializeSchema(*arrow_schema, arrow::default_memory_pool()); + if (!schema_ret.ok()) { + DB_FATAL("arrow serialize schema fail, status: %s", schema_ret.status().ToString().c_str()); + return plan_node; + } + exchange_receiver_node->set_schema((*schema_ret)->data(), (*schema_ret)->size()); + + return plan_node; +} + +pb::PlanNode make_exchange_receiver_node2() { + pb::PlanNode plan_node; + plan_node.set_node_type(pb::EXCHANGE_RECEIVER_NODE); + plan_node.set_num_children(1); + pb::DerivePlanNode* derive_node = plan_node.mutable_derive_node(); + pb::ExchangeReceiverNode* exchange_receiver_node = derive_node->mutable_exchange_receiver_node(); + exchange_receiver_node->set_log_id(100); + exchange_receiver_node->set_fragment_id(8); + exchange_receiver_node->set_sender_fragment_id(7); + exchange_receiver_node->set_fragment_instance_id(88); + exchange_receiver_node->set_node_id(99); + { + pb::ExchangeDestination* sender_destination = exchange_receiver_node->add_sender_destinations(); + sender_destination->set_fragment_instance_id(77); + sender_destination->set_node_id(99); + sender_destination->set_address("10.12.187.120:8788"); + } + { + pb::ExchangeDestination* sender_destination = exchange_receiver_node->add_sender_destinations(); + sender_destination->set_fragment_instance_id(78); + sender_destination->set_node_id(99); + sender_destination->set_address("10.12.187.120:8789"); + } + std::shared_ptr arrow_schema = make_test_schema(); + arrow::Result> schema_ret = + arrow::ipc::SerializeSchema(*arrow_schema, arrow::default_memory_pool()); + if (!schema_ret.ok()) { + DB_FATAL("arrow serialize schema fail, status: %s", schema_ret.status().ToString().c_str()); + return plan_node; + } + exchange_receiver_node->set_schema((*schema_ret)->data(), (*schema_ret)->size()); + return plan_node; +} +pb::PlanNode make_exchange_receiver_node3() { + pb::PlanNode plan_node; + plan_node.set_node_type(pb::EXCHANGE_RECEIVER_NODE); + plan_node.set_num_children(1); + pb::DerivePlanNode* derive_node = plan_node.mutable_derive_node(); + pb::ExchangeReceiverNode* exchange_receiver_node = derive_node->mutable_exchange_receiver_node(); + exchange_receiver_node->set_log_id(100); + exchange_receiver_node->set_fragment_id(8); + exchange_receiver_node->set_sender_fragment_id(7); + exchange_receiver_node->set_fragment_instance_id(88); + exchange_receiver_node->set_node_id(99); + { + pb::ExchangeDestination* sender_destination = exchange_receiver_node->add_sender_destinations(); + sender_destination->set_fragment_instance_id(77); + sender_destination->set_node_id(99); + sender_destination->set_address("10.12.187.120:8788"); + } + { + pb::ExchangeDestination* sender_destination = exchange_receiver_node->add_sender_destinations(); + sender_destination->set_fragment_instance_id(78); + sender_destination->set_node_id(99); + sender_destination->set_address("10.12.187.120:8789"); + } + { + pb::ExchangeDestination* sender_destination = exchange_receiver_node->add_sender_destinations(); + sender_destination->set_fragment_instance_id(79); + sender_destination->set_node_id(99); + sender_destination->set_address("10.12.187.120:8790"); + } + std::shared_ptr arrow_schema = make_test_schema(); + arrow::Result> schema_ret = + arrow::ipc::SerializeSchema(*arrow_schema, arrow::default_memory_pool()); + if (!schema_ret.ok()) { + DB_FATAL("arrow serialize schema fail, status: %s", schema_ret.status().ToString().c_str()); + return plan_node; + } + exchange_receiver_node->set_schema((*schema_ret)->data(), (*schema_ret)->size()); + return plan_node; +} + +// 接收db_77/db_78/db_79的数据 +pb::PlanNode make_exchange_receiver_node_db_schema3() { + pb::PlanNode plan_node; + plan_node.set_node_type(pb::EXCHANGE_RECEIVER_NODE); + plan_node.set_num_children(1); + pb::DerivePlanNode* derive_node = plan_node.mutable_derive_node(); + pb::ExchangeReceiverNode* exchange_receiver_node = derive_node->mutable_exchange_receiver_node(); + exchange_receiver_node->set_log_id(100); + exchange_receiver_node->set_fragment_id(8); + exchange_receiver_node->set_sender_fragment_id(7); + exchange_receiver_node->set_fragment_instance_id(88); + exchange_receiver_node->set_node_id(99); + { + pb::ExchangeDestination* sender_destination = exchange_receiver_node->add_sender_destinations(); + sender_destination->set_fragment_instance_id(77); + sender_destination->set_node_id(99); + sender_destination->set_address("10.12.187.120:8788"); + } + { + pb::ExchangeDestination* sender_destination = exchange_receiver_node->add_sender_destinations(); + sender_destination->set_fragment_instance_id(78); + sender_destination->set_node_id(99); + sender_destination->set_address("10.12.187.120:8789"); + } + { + pb::ExchangeDestination* sender_destination = exchange_receiver_node->add_sender_destinations(); + sender_destination->set_fragment_instance_id(79); + sender_destination->set_node_id(99); + sender_destination->set_address("10.12.187.120:8790"); + } + std::shared_ptr arrow_schema = make_test_schema3(); + arrow::Result> schema_ret = + arrow::ipc::SerializeSchema(*arrow_schema, arrow::default_memory_pool()); + if (!schema_ret.ok()) { + DB_FATAL("arrow serialize schema fail, status: %s", schema_ret.status().ToString().c_str()); + return plan_node; + } + exchange_receiver_node->set_schema((*schema_ret)->data(), (*schema_ret)->size()); + return plan_node; +} +// 接收store region1/region2/region3的数据 +pb::PlanNode make_exchange_receiver_node_region_schema3() { + pb::PlanNode plan_node; + plan_node.set_node_type(pb::EXCHANGE_RECEIVER_NODE); + plan_node.set_num_children(1); + pb::DerivePlanNode* derive_node = plan_node.mutable_derive_node(); + pb::ExchangeReceiverNode* exchange_receiver_node = derive_node->mutable_exchange_receiver_node(); + exchange_receiver_node->set_log_id(100); + exchange_receiver_node->set_fragment_id(8); + exchange_receiver_node->set_sender_fragment_id(7); + exchange_receiver_node->set_fragment_instance_id(88); + exchange_receiver_node->set_node_id(99); + { + pb::RegionInfo* region = exchange_receiver_node->add_regions(); + region->set_region_id(1); + region->set_version(1); + region->set_table_id(1); + region->set_partition_id(1); + region->set_replica_num(3); + region->set_conf_version(1); + } + { + pb::RegionInfo* region = exchange_receiver_node->add_regions(); + region->set_region_id(2); + region->set_version(1); + region->set_table_id(1); + region->set_partition_id(1); + region->set_replica_num(3); + region->set_conf_version(1); + } + { + pb::RegionInfo* region = exchange_receiver_node->add_regions(); + region->set_region_id(3); + region->set_version(1); + region->set_table_id(1); + region->set_partition_id(1); + region->set_replica_num(3); + region->set_conf_version(1); + } + std::shared_ptr arrow_schema = make_test_schema3(); + arrow::Result> schema_ret = + arrow::ipc::SerializeSchema(*arrow_schema, arrow::default_memory_pool()); + if (!schema_ret.ok()) { + DB_FATAL("arrow serialize schema fail, status: %s", schema_ret.status().ToString().c_str()); + return plan_node; + } + exchange_receiver_node->set_schema((*schema_ret)->data(), (*schema_ret)->size()); + return plan_node; +} + +// 接收store region1/region3/region4的数据 +pb::PlanNode make_exchange_receiver_node_region_new_schema3() { + pb::PlanNode plan_node; + plan_node.set_node_type(pb::EXCHANGE_RECEIVER_NODE); + plan_node.set_num_children(1); + pb::DerivePlanNode* derive_node = plan_node.mutable_derive_node(); + pb::ExchangeReceiverNode* exchange_receiver_node = derive_node->mutable_exchange_receiver_node(); + exchange_receiver_node->set_log_id(100); + exchange_receiver_node->set_fragment_id(8); + exchange_receiver_node->set_sender_fragment_id(7); + exchange_receiver_node->set_fragment_instance_id(88); + exchange_receiver_node->set_node_id(99); + { + pb::RegionInfo* region = exchange_receiver_node->add_regions(); + region->set_region_id(1); + region->set_version(1); + region->set_table_id(1); + region->set_partition_id(1); + region->set_replica_num(3); + region->set_conf_version(1); + } + { + pb::RegionInfo* region = exchange_receiver_node->add_regions(); + region->set_region_id(3); + region->set_version(1); + region->set_table_id(1); + region->set_partition_id(1); + region->set_replica_num(3); + region->set_conf_version(1); + } + { + pb::RegionInfo* region = exchange_receiver_node->add_regions(); + region->set_region_id(4); + region->set_version(1); + region->set_table_id(1); + region->set_partition_id(1); + region->set_replica_num(3); + region->set_conf_version(1); + } + std::shared_ptr arrow_schema = make_test_schema3(); + arrow::Result> schema_ret = + arrow::ipc::SerializeSchema(*arrow_schema, arrow::default_memory_pool()); + if (!schema_ret.ok()) { + DB_FATAL("arrow serialize schema fail, status: %s", schema_ret.status().ToString().c_str()); + return plan_node; + } + exchange_receiver_node->set_schema((*schema_ret)->data(), (*schema_ret)->size()); + return plan_node; +} + +/// 分裂合并场景测试数据构造 +pb::RegionInfo make_region_info(const int64_t region_id, const int64_t version) { + pb::RegionInfo region_info; + region_info.set_region_id(region_id); + region_info.set_partition_id(0); + region_info.set_replica_num(3); + region_info.set_conf_version(0); + region_info.set_version(version); + return region_info; +} +// region1分裂成region1/region2 +pb::RegionInfo make_region1_old() { + return make_region_info(1,1); +} +pb::RegionInfo make_region1_new() { + return make_region_info(1,2); +} +pb::RegionInfo make_region2() { + return make_region_info(2,1); +} +pb::TransmitDataParam make_transmit_data_param_split() { + pb::TransmitDataParam transmit_data_param; + transmit_data_param.set_exchange_state(pb::ES_VERSION_OLD); + transmit_data_param.set_log_id(100); + transmit_data_param.set_sender_fragment_instance_id(77); + transmit_data_param.set_receiver_fragment_instance_id(88); + transmit_data_param.set_receiver_node_id(99); + transmit_data_param.add_region_infos()->CopyFrom(make_region1_new()); + transmit_data_param.add_region_infos()->CopyFrom(make_region2()); + return transmit_data_param; +} +pb::PlanNode make_exchange_receiver_node_split() { + pb::PlanNode plan_node; + plan_node.set_node_type(pb::EXCHANGE_RECEIVER_NODE); + plan_node.set_num_children(1); + pb::DerivePlanNode* derive_node = plan_node.mutable_derive_node(); + pb::ExchangeReceiverNode* exchange_receiver_node = derive_node->mutable_exchange_receiver_node(); + exchange_receiver_node->set_log_id(100); + exchange_receiver_node->set_fragment_id(8); + exchange_receiver_node->set_sender_fragment_id(7); + exchange_receiver_node->set_fragment_instance_id(88); + exchange_receiver_node->set_node_id(99); + exchange_receiver_node->add_regions()->CopyFrom(make_region1_old()); + return plan_node; +} +// region3/region4合并成region3 +pb::RegionInfo make_region3_old() { + return make_region_info(3,1); +} +pb::RegionInfo make_region3_new() { + return make_region_info(3,2); +} +pb::RegionInfo make_region4_old() { + return make_region_info(4,1); +} +pb::RegionInfo make_region4_new() { + return make_region_info(4,2); +} +pb::TransmitDataParam make_transmit_data_param_merge3() { + pb::TransmitDataParam transmit_data_param; + transmit_data_param.set_exchange_state(pb::ES_VERSION_OLD); + transmit_data_param.set_log_id(100); + transmit_data_param.set_sender_fragment_instance_id(77); + transmit_data_param.set_receiver_fragment_instance_id(88); + transmit_data_param.set_receiver_node_id(99); + transmit_data_param.set_is_merge(true); + transmit_data_param.add_region_infos()->CopyFrom(make_region3_new()); + return transmit_data_param; +} +pb::TransmitDataParam make_transmit_data_param_merge4() { + pb::TransmitDataParam transmit_data_param; + transmit_data_param.set_exchange_state(pb::ES_VERSION_OLD); + transmit_data_param.set_log_id(100); + transmit_data_param.set_sender_fragment_instance_id(77); + transmit_data_param.set_receiver_fragment_instance_id(88); + transmit_data_param.set_receiver_node_id(99); + transmit_data_param.set_is_merge(true); + transmit_data_param.add_region_infos()->CopyFrom(make_region4_new()); + transmit_data_param.add_region_infos()->CopyFrom(make_region3_new()); + return transmit_data_param; +} +pb::PlanNode make_exchange_receiver_node_merge3() { + pb::PlanNode plan_node; + plan_node.set_node_type(pb::EXCHANGE_RECEIVER_NODE); + plan_node.set_num_children(1); + pb::DerivePlanNode* derive_node = plan_node.mutable_derive_node(); + pb::ExchangeReceiverNode* exchange_receiver_node = derive_node->mutable_exchange_receiver_node(); + exchange_receiver_node->set_log_id(100); + exchange_receiver_node->set_fragment_id(8); + exchange_receiver_node->set_sender_fragment_id(7); + exchange_receiver_node->set_fragment_instance_id(88); + exchange_receiver_node->set_node_id(99); + exchange_receiver_node->add_regions()->CopyFrom(make_region3_old()); + return plan_node; +} +pb::PlanNode make_exchange_receiver_node_merge4() { + pb::PlanNode plan_node; + plan_node.set_node_type(pb::EXCHANGE_RECEIVER_NODE); + plan_node.set_num_children(1); + pb::DerivePlanNode* derive_node = plan_node.mutable_derive_node(); + pb::ExchangeReceiverNode* exchange_receiver_node = derive_node->mutable_exchange_receiver_node(); + exchange_receiver_node->set_log_id(100); + exchange_receiver_node->set_fragment_id(8); + exchange_receiver_node->set_sender_fragment_id(7); + exchange_receiver_node->set_fragment_instance_id(88); + exchange_receiver_node->set_node_id(99); + exchange_receiver_node->add_regions()->CopyFrom(make_region4_old()); + return plan_node; +} +pb::PlanNode make_exchange_receiver_node_merge34() { + pb::PlanNode plan_node; + plan_node.set_node_type(pb::EXCHANGE_RECEIVER_NODE); + plan_node.set_num_children(1); + pb::DerivePlanNode* derive_node = plan_node.mutable_derive_node(); + pb::ExchangeReceiverNode* exchange_receiver_node = derive_node->mutable_exchange_receiver_node(); + exchange_receiver_node->set_log_id(100); + exchange_receiver_node->set_fragment_id(8); + exchange_receiver_node->set_sender_fragment_id(7); + exchange_receiver_node->set_fragment_instance_id(88); + exchange_receiver_node->set_node_id(99); + exchange_receiver_node->add_regions()->CopyFrom(make_region3_old()); + exchange_receiver_node->add_regions()->CopyFrom(make_region4_old()); + return plan_node; +} + +//// ER/ER交互 +TEST(test_rpc_store, case_all) { + DB_WARNING("-------------------------------test_rpc_store-------------------------------"); + // RecordBatch + std::shared_ptr batch = make_test_record_batch1_schema1(); + std::shared_ptr batch2 = make_test_record_batch2_schema1(); + std::shared_ptr batch3 = make_test_record_batch3_schema1(); + // Receiver + RuntimeState recv_state; + pb::PlanNode pb_exchange_receiver_node = make_exchange_receiver_node_region_schema3(); + ExchangeReceiverNode exchange_receiver; + exchange_receiver.init(pb_exchange_receiver_node); + // exchange_receiver.set_limit(24); + exchange_receiver.open(&recv_state); + exchange_receiver.build_arrow_declaration(&recv_state); + Bthread bth; + bth.run([&recv_state] () { + arrow::Result> final_table = arrow::acero::DeclarationToTable( + arrow::acero::Declaration::Sequence(std::move(recv_state.acero_declarations)), false); + if (!final_table.ok()) { + DB_WARNING("Fail to DeclarationToTable"); + return; + } + std::shared_ptr table = *final_table; + DB_WARNING("breakpoint table: %s", table->ToString().c_str()); + }); + + // Sender + RuntimeState send_state; + send_state._region_id = 1; + send_state._region_version = 1; + pb::PlanNode pb_exchange_sender_node = make_exchange_sender_node(); + ExchangeSenderNode exchange_sender; + exchange_sender.init(pb_exchange_sender_node); + exchange_sender.open(&send_state); + exchange_sender.send_curr_record_batch(&send_state, batch); + exchange_sender.send_eof_record_batch(&send_state, make_test_empty_record_batch_schema1()); + + RuntimeState send_state2; + send_state2._region_id = 2; + send_state2._region_version = 1; + pb::PlanNode pb_exchange_sender_node2 = make_exchange_sender_node2(); + ExchangeSenderNode exchange_sender2; + exchange_sender2.init(pb_exchange_sender_node2); + exchange_sender2.open(&send_state2); + exchange_sender2.send_curr_record_batch(&send_state2, batch2); + exchange_sender2.send_eof_record_batch(&send_state2, make_test_empty_record_batch_schema1()); + + RuntimeState send_state3; + send_state3._region_id = 3; + send_state3._region_version = 1; + pb::PlanNode pb_exchange_sender_node3 = make_exchange_sender_node3(); + ExchangeSenderNode exchange_sender3; + exchange_sender3.init(pb_exchange_sender_node3); + exchange_sender3.open(&send_state3); + exchange_sender3.send_curr_record_batch(&send_state3, batch3); + exchange_sender3.send_eof_record_batch(&send_state3, make_test_empty_record_batch_schema1()); + + bth.join(); + exchange_receiver.close(&recv_state); + exchange_sender.close(&send_state); + exchange_sender2.close(&send_state2); + exchange_sender3.close(&send_state3); +} + +TEST(test_rpc_db, case_all) { + DB_WARNING("-------------------------------test_rpc_db-------------------------------"); + // RecordBatch + std::shared_ptr batch = make_test_record_batch1_schema1(); + std::shared_ptr batch2 = make_test_record_batch2_schema1(); + std::shared_ptr batch3 = make_test_record_batch3_schema1(); + + // Receiver + RuntimeState recv_state; + pb::PlanNode pb_exchange_receiver_node = make_exchange_receiver_node_db_schema3(); + ExchangeReceiverNode exchange_receiver; + exchange_receiver.init(pb_exchange_receiver_node); + // exchange_receiver.set_limit(24); + exchange_receiver.open(&recv_state); + exchange_receiver.build_arrow_declaration(&recv_state); + Bthread bth; + bth.run([&recv_state] () { + arrow::Result> final_table = arrow::acero::DeclarationToTable( + arrow::acero::Declaration::Sequence(std::move(recv_state.acero_declarations)), false); + if (!final_table.ok()) { + DB_WARNING("Fail to DeclarationToTable"); + return; + } + std::shared_ptr table = *final_table; + DB_WARNING("breakpoint table: %s", table->ToString().c_str()); + }); + + // Sender + { + pb::PlanNode pb_exchange_receiver_node1 = make_invalid_exchange_receiver_node(); + ExchangeReceiverNode exchange_receiver1; + exchange_receiver1.init(pb_exchange_receiver_node1); + ExchangeSenderNode exchange_sender; + pb::PlanNode pb_exchange_sender_node = make_exchange_sender_node(); + exchange_sender.init(pb_exchange_sender_node); + exchange_sender.add_child(&exchange_receiver1); + + RuntimeState send_state; + exchange_sender.open(&send_state); + exchange_sender.build_arrow_declaration(&send_state); + + pb::ExtraRes extra_res; + auto param = make_invalid_transmit_data_param(); + param.set_record_batch_seq(0); + exchange_receiver1.get_data_stream_receiver()->add_record_batch(param, batch, extra_res); + param.set_record_batch_seq(0); + exchange_receiver1.get_data_stream_receiver()->add_record_batch(param, batch, extra_res); + param.set_record_batch_seq(1); + param.set_exchange_state(pb::ES_EOF); + exchange_receiver1.get_data_stream_receiver()->add_record_batch(param, batch, extra_res); + + arrow::Result> final_table = arrow::acero::DeclarationToTable( + arrow::acero::Declaration::Sequence(std::move(send_state.acero_declarations)), false); + if (!final_table.ok()) { + DB_WARNING("Fail to DeclarationToTable, %s", final_table.status().ToString().c_str()); + exchange_sender.clear_children(); + return; + } + std::shared_ptr table = *final_table; + DB_WARNING("breakpoint table: %s", table->ToString().c_str()); + exchange_sender.close(&send_state); + exchange_sender.clear_children(); + } + { + pb::PlanNode pb_exchange_receiver_node1 = make_invalid_exchange_receiver_node(); + ExchangeReceiverNode exchange_receiver1; + exchange_receiver1.init(pb_exchange_receiver_node1); + ExchangeSenderNode exchange_sender; + pb::PlanNode pb_exchange_sender_node = make_exchange_sender_node2(); + exchange_sender.init(pb_exchange_sender_node); + exchange_sender.add_child(&exchange_receiver1); + + RuntimeState send_state; + exchange_sender.open(&send_state); + exchange_sender.build_arrow_declaration(&send_state); + + pb::ExtraRes extra_res; + auto param = make_invalid_transmit_data_param(); + param.set_record_batch_seq(0); + exchange_receiver1.get_data_stream_receiver()->add_record_batch(param, batch2, extra_res); + param.set_record_batch_seq(0); + exchange_receiver1.get_data_stream_receiver()->add_record_batch(param, batch2, extra_res); + param.set_record_batch_seq(1); + param.set_exchange_state(pb::ES_EOF); + exchange_receiver1.get_data_stream_receiver()->add_record_batch(param, batch2, extra_res); + + arrow::Result> final_table = arrow::acero::DeclarationToTable( + arrow::acero::Declaration::Sequence(std::move(send_state.acero_declarations)), false); + if (!final_table.ok()) { + DB_WARNING("Fail to DeclarationToTable, %s", final_table.status().ToString().c_str()); + exchange_sender.clear_children(); + return; + } + std::shared_ptr table = *final_table; + DB_WARNING("breakpoint table: %s", table->ToString().c_str()); + exchange_sender.close(&send_state); + exchange_sender.clear_children(); + } + { + pb::PlanNode pb_exchange_receiver_node1 = make_invalid_exchange_receiver_node(); + ExchangeReceiverNode exchange_receiver1; + exchange_receiver1.init(pb_exchange_receiver_node1); + ExchangeSenderNode exchange_sender; + pb::PlanNode pb_exchange_sender_node = make_exchange_sender_node3(); + exchange_sender.init(pb_exchange_sender_node); + exchange_sender.add_child(&exchange_receiver1); + + RuntimeState send_state; + exchange_sender.open(&send_state); + exchange_sender.build_arrow_declaration(&send_state); + + pb::ExtraRes extra_res; + auto param = make_invalid_transmit_data_param(); + param.set_record_batch_seq(0); + exchange_receiver1.get_data_stream_receiver()->add_record_batch(param, batch, extra_res); + param.set_record_batch_seq(1); + exchange_receiver1.get_data_stream_receiver()->add_record_batch(param, batch2, extra_res); + param.set_record_batch_seq(2); + param.set_exchange_state(pb::ES_EOF); + exchange_receiver1.get_data_stream_receiver()->add_record_batch(param, batch3, extra_res); + + arrow::Result> final_table = arrow::acero::DeclarationToTable( + arrow::acero::Declaration::Sequence(std::move(send_state.acero_declarations)), false); + if (!final_table.ok()) { + DB_WARNING("Fail to DeclarationToTable, %s", final_table.status().ToString().c_str()); + exchange_sender.clear_children(); + return; + } + std::shared_ptr table = *final_table; + DB_WARNING("breakpoint table: %s", table->ToString().c_str()); + exchange_sender.close(&send_state); + exchange_sender.clear_children(); + } + + bth.join(); + exchange_receiver.close(&recv_state); +} + +TEST(test_rpc_error_new, case_all) { + DB_WARNING("-------------------------------test_rpc_error_new-------------------------------"); + // RecordBatch + std::shared_ptr batch = make_test_record_batch1_schema1(); + std::shared_ptr batch2 = make_test_record_batch2_schema1(); + std::shared_ptr batch3 = make_test_record_batch3_schema1(); + + // Receiver + RuntimeState recv_state; + pb::PlanNode pb_exchange_receiver_node = make_exchange_receiver_node_region_schema3(); + ExchangeReceiverNode exchange_receiver; + exchange_receiver.init(pb_exchange_receiver_node); + // exchange_receiver.set_limit(24); + exchange_receiver.open(&recv_state); + exchange_receiver.build_arrow_declaration(&recv_state); + Bthread bth; + bth.run([&recv_state] () { + arrow::Result> final_table = arrow::acero::DeclarationToTable( + arrow::acero::Declaration::Sequence(std::move(recv_state.acero_declarations)), false); + if (!final_table.ok()) { + DB_WARNING("Fail to DeclarationToTable"); + return; + } + std::shared_ptr table = *final_table; + DB_WARNING("breakpoint table: %s", table->ToString().c_str()); + }); + + // Sender + RuntimeState send_state; + send_state._region_id = 1; + send_state._region_version = 1; + pb::PlanNode pb_exchange_sender_node = make_exchange_sender_node(); + ExchangeSenderNode exchange_sender; + exchange_sender.init(pb_exchange_sender_node); + exchange_sender.open(&send_state); + exchange_sender.send_exec_fail(); + + RuntimeState send_state2; + send_state2._region_id = 2; + send_state2._region_version = 1; + pb::PlanNode pb_exchange_sender_node2 = make_exchange_sender_node2(); + ExchangeSenderNode exchange_sender2; + exchange_sender2.init(pb_exchange_sender_node2); + exchange_sender2.open(&send_state2); + exchange_sender2.send_curr_record_batch(&send_state2, batch2); + exchange_sender2.send_eof_record_batch(&send_state2, make_test_empty_record_batch_schema1()); + + RuntimeState send_state3; + send_state3._region_id = 3; + send_state3._region_version = 1; + pb::PlanNode pb_exchange_sender_node3 = make_exchange_sender_node3(); + ExchangeSenderNode exchange_sender3; + exchange_sender3.init(pb_exchange_sender_node3); + exchange_sender3.open(&send_state3); + exchange_sender3.send_curr_record_batch(&send_state3, batch3); + exchange_sender3.send_eof_record_batch(&send_state3, make_test_empty_record_batch_schema1()); + + bth.join(); + exchange_receiver.close(&recv_state); + exchange_sender.close(&send_state); + exchange_sender2.close(&send_state2); + exchange_sender3.close(&send_state3); +} + +// //// ES测试 +// repartition +TEST(test_sender, case_all) { + DB_WARNING("-------------------------------test_sender-------------------------------"); + { + // RecordBatch + std::shared_ptr record_batch = make_test_record_batch1_schema1(); + std::unordered_map> hash_batch_map; + ExchangeSenderNode exchange_sender_node; + auto status = exchange_sender_node.repartition(record_batch, {"col1", "col3"}, {}, 1, hash_batch_map); + if (!status.ok()) { + DB_WARNING("Fail to repartition"); + } + ASSERT_EQ(status.ok(), true); + for (auto& [bucket, batch] : hash_batch_map) { + DB_WARNING("bucket: %d, batch: %s", bucket, batch->ToString().c_str()); + } + } + { + // RecordBatch + std::shared_ptr record_batch = make_test_record_batch1_schema1(); + std::unordered_map> hash_batch_map; + ExchangeSenderNode exchange_sender_node; + auto status = exchange_sender_node.repartition(record_batch, {"col1", "col3"}, {}, 3, hash_batch_map); + if (!status.ok()) { + DB_WARNING("Fail to repartition"); + } + ASSERT_EQ(status.ok(), true); + for (auto& [bucket, batch] : hash_batch_map) { + DB_WARNING("bucket: %d, batch: %s", bucket, batch->ToString().c_str()); + } + } + { + // RecordBatch + std::shared_ptr record_batch = make_test_record_batch1_schema1(); + std::unordered_map> hash_batch_map; + ExchangeSenderNode exchange_sender_node; + auto status = exchange_sender_node.repartition(record_batch, {"col1", "col3"}, {"col3"}, 3, hash_batch_map); + if (!status.ok()) { + DB_WARNING("Fail to repartition"); + } + ASSERT_EQ(status.ok(), true); + for (auto& [bucket, batch] : hash_batch_map) { + DB_WARNING("bucket: %d, batch: %s", bucket, batch->ToString().c_str()); + } + } + { + // RecordBatch + std::shared_ptr record_batch = make_test_record_batch1_schema2(); + DB_WARNING("schema: %s, record_batch: %s", record_batch->schema()->ToString().c_str(), record_batch->ToString().c_str()); + std::unordered_map> hash_batch_map; + ExchangeSenderNode exchange_sender_node; + auto status = exchange_sender_node.repartition(record_batch, {"col1", "col3"}, {"col3"}, 3, hash_batch_map); + if (!status.ok()) { + DB_WARNING("Fail to repartition"); + } + ASSERT_EQ(status.ok(), true); + for (auto& [bucket, batch] : hash_batch_map) { + DB_WARNING("bucket: %d, batch: %s", bucket, batch->ToString().c_str()); + } + } + { + // RecordBatch + std::shared_ptr record_batch = make_test_record_batch1_schema1(); + std::unordered_map> hash_batch_map; + ExchangeSenderNode exchange_sender_node; + auto status = exchange_sender_node.repartition(record_batch, {"col1", "col5"}, {}, 3, hash_batch_map); + if (!status.ok()) { + DB_WARNING("Fail to repartition"); + return; + } + ASSERT_EQ(status.ok(), true); + for (auto& [bucket, batch] : hash_batch_map) { + DB_WARNING("bucket: %d, batch: %s", bucket, batch->ToString().c_str()); + } + } +} + +//// ER测试 +/// handle_version_old() +// 分裂场景，region1分裂成region1/region2 +// ExchangeReceiverNode原本是region1_old，现在收到region1_new +TEST(test_split, case_all) { + DB_WARNING("-------------------------------test_split-------------------------------"); + RuntimeState state; + ExchangeReceiverNode exchange_receiver; + pb::PlanNode pb_exchange_receiver_node = make_exchange_receiver_node_split(); + exchange_receiver.init(pb_exchange_receiver_node); + exchange_receiver.open(&state); + pb::TransmitDataParam transmit_data_param = make_transmit_data_param_split(); + std::shared_ptr receiver = + DataStreamManager::get_instance()->get_receiver(transmit_data_param.log_id(), + transmit_data_param.receiver_fragment_instance_id(), + transmit_data_param.receiver_node_id()); + if (receiver == nullptr) { + DB_WARNING("receiver is nullptr"); + return; + } + receiver->to_string(); + receiver->handle_version_old(transmit_data_param); + receiver->to_string(); +} +// 合并场景，region3/region4合并成region3 +// ExchangeReceiverNode原本是region3_old，现在收到region3_new +TEST(test_merge, case_all) { + DB_WARNING("-------------------------------test_merge-------------------------------"); + RuntimeState state; + ExchangeReceiverNode exchange_receiver; + pb::PlanNode pb_exchange_receiver_node = make_exchange_receiver_node_merge3(); + exchange_receiver.init(pb_exchange_receiver_node); + exchange_receiver.open(&state); + pb::TransmitDataParam transmit_data_param = make_transmit_data_param_merge3(); + std::shared_ptr receiver = + DataStreamManager::get_instance()->get_receiver(transmit_data_param.log_id(), + transmit_data_param.receiver_fragment_instance_id(), + transmit_data_param.receiver_node_id()); + if (receiver == nullptr) { + DB_WARNING("receiver is nullptr"); + return; + } + receiver->to_string(); + receiver->handle_version_old(transmit_data_param); + receiver->to_string(); +} +// ExchangeReceiverNode原本是region4 +TEST(test_merge4, case_all) { + DB_WARNING("-------------------------------test_merge4-------------------------------"); + RuntimeState state; + ExchangeReceiverNode exchange_receiver; + pb::PlanNode pb_exchange_receiver_node = make_exchange_receiver_node_merge4(); + exchange_receiver.init(pb_exchange_receiver_node); + exchange_receiver.open(&state); + pb::TransmitDataParam transmit_data_param = make_transmit_data_param_merge4(); + std::shared_ptr receiver = + DataStreamManager::get_instance()->get_receiver(transmit_data_param.log_id(), + transmit_data_param.receiver_fragment_instance_id(), + transmit_data_param.receiver_node_id()); + if (receiver == nullptr) { + DB_WARNING("receiver is nullptr"); + return; + } + receiver->to_string(); + receiver->handle_version_old(transmit_data_param); + receiver->to_string(); +} +// ExchangeReceiverNode原本是region3_old/region4，现在收到region3_new +TEST(test_merge34, case_all) { + DB_WARNING("-------------------------------test_merge34-------------------------------"); + RuntimeState state; + ExchangeReceiverNode exchange_receiver; + pb::PlanNode pb_exchange_receiver_node = make_exchange_receiver_node_merge34(); + exchange_receiver.init(pb_exchange_receiver_node); + exchange_receiver.open(&state); + { + pb::TransmitDataParam transmit_data_param = make_transmit_data_param_merge4(); + std::shared_ptr receiver = + DataStreamManager::get_instance()->get_receiver(transmit_data_param.log_id(), + transmit_data_param.receiver_fragment_instance_id(), + transmit_data_param.receiver_node_id()); + if (receiver == nullptr) { + DB_WARNING("receiver is nullptr"); + return; + } + receiver->to_string(); + receiver->handle_version_old(transmit_data_param); + receiver->to_string(); + } + { + pb::TransmitDataParam transmit_data_param = make_transmit_data_param_merge3(); + std::shared_ptr receiver = + DataStreamManager::get_instance()->get_receiver(transmit_data_param.log_id(), + transmit_data_param.receiver_fragment_instance_id(), + transmit_data_param.receiver_node_id()); + if (receiver == nullptr) { + DB_WARNING("receiver is nullptr"); + return; + } + receiver->to_string(); + receiver->handle_version_old(transmit_data_param); + receiver->to_string(); + } +} + +TEST(test_merge34_new, case_all) { + DB_WARNING("-------------------------------test_merge34_new-------------------------------"); + RuntimeState state; + ExchangeReceiverNode exchange_receiver; + pb::PlanNode pb_exchange_receiver_node = make_exchange_receiver_node_merge34(); + exchange_receiver.init(pb_exchange_receiver_node); + exchange_receiver.open(&state); + { + pb::TransmitDataParam transmit_data_param = make_transmit_data_param_merge3(); + std::shared_ptr receiver = + DataStreamManager::get_instance()->get_receiver(transmit_data_param.log_id(), + transmit_data_param.receiver_fragment_instance_id(), + transmit_data_param.receiver_node_id()); + if (receiver == nullptr) { + DB_WARNING("receiver is nullptr"); + return; + } + receiver->to_string(); + receiver->handle_version_old(transmit_data_param); + receiver->to_string(); + } + { + pb::TransmitDataParam transmit_data_param = make_transmit_data_param_merge4(); + std::shared_ptr receiver = + DataStreamManager::get_instance()->get_receiver(transmit_data_param.log_id(), + transmit_data_param.receiver_fragment_instance_id(), + transmit_data_param.receiver_node_id()); + if (receiver == nullptr) { + DB_WARNING("receiver is nullptr"); + return; + } + receiver->to_string(); + receiver->handle_version_old(transmit_data_param); + receiver->to_string(); + } +} + +// CAST +TEST(test_cast, case_all) { + DB_WARNING("-------------------------------test_cast-------------------------------"); + std::shared_ptr record_batch = make_test_record_batch1_schema1(); + auto array = record_batch->column(0); + auto array_utf8_ret = arrow::compute::Cast(*array, arrow::utf8()); + auto array_utf8 = *array_utf8_ret; + DB_WARNING("array: %s", array->ToString().c_str()); + DB_WARNING("array_utf8: %s", array_utf8->ToString().c_str()); + + auto array_lb1_ret = arrow::compute::Cast(*array, arrow::large_binary()); + auto array_lb1 = *array_lb1_ret; + auto array_lb2_ret = arrow::compute::Cast(*array_utf8, arrow::large_binary()); + auto array_lb2 = *array_lb2_ret; + DB_WARNING("array_lb1: %s", array_lb1->ToString().c_str()); + DB_WARNING("array_lb2: %s", array_lb2->ToString().c_str()); +} + +// mpp_send_version_old +TEST(test_mpp_send_version_old_split, case_all) { + DB_WARNING("-------------------------------test_mpp_send_version_old_split-------------------------------"); + RuntimeState recv_state; + pb::PlanNode pb_exchange_receiver_node = make_exchange_receiver_node_region_new_schema3(); + ExchangeReceiverNode exchange_receiver; + exchange_receiver.init(pb_exchange_receiver_node); + exchange_receiver.open(&recv_state); + exchange_receiver.get_data_stream_receiver()->to_string(); + pb::PlanNode pb_exchange_sender_node = make_exchange_sender_node(); + ::google::protobuf::RepeatedPtrField region_infos; + { + auto region_info = region_infos.Add(); + region_info->set_region_id(1); + region_info->set_version(2); + region_info->set_table_id(1); + region_info->set_partition_id(1); + region_info->set_replica_num(3); + region_info->set_conf_version(1); + } + { + auto region_info = region_infos.Add(); + region_info->set_region_id(2); + region_info->set_version(1); + region_info->set_table_id(1); + region_info->set_partition_id(1); + region_info->set_replica_num(3); + region_info->set_conf_version(1); + } + Region::mpp_send_version_old(pb_exchange_sender_node.derive_node().exchange_sender_node(), false, region_infos); + exchange_receiver.get_data_stream_receiver()->to_string(); + exchange_receiver.close(&recv_state); +} + +TEST(test_mpp_send_version_old_merge1, case_all) { + DB_WARNING("-------------------------------test_mpp_send_version_old_merge1-------------------------------"); + RuntimeState recv_state; + pb::PlanNode pb_exchange_receiver_node = make_exchange_receiver_node_region_new_schema3(); + ExchangeReceiverNode exchange_receiver; + exchange_receiver.init(pb_exchange_receiver_node); + exchange_receiver.open(&recv_state); + exchange_receiver.get_data_stream_receiver()->to_string(); + pb::PlanNode pb_exchange_sender_node = make_exchange_sender_node(); + ::google::protobuf::RepeatedPtrField region_infos; + { + auto region_info = region_infos.Add(); + region_info->set_region_id(4); + region_info->set_version(2); + region_info->set_table_id(1); + region_info->set_partition_id(1); + region_info->set_replica_num(3); + region_info->set_conf_version(1); + } + { + auto region_info = region_infos.Add(); + region_info->set_region_id(3); + region_info->set_version(2); + region_info->set_table_id(1); + region_info->set_partition_id(1); + region_info->set_replica_num(3); + region_info->set_conf_version(1); + } + Region::mpp_send_version_old(pb_exchange_sender_node.derive_node().exchange_sender_node(), true, region_infos); + exchange_receiver.get_data_stream_receiver()->to_string(); + exchange_receiver.close(&recv_state); +} + +TEST(test_mpp_send_version_old_merge2, case_all) { + DB_WARNING("-------------------------------test_mpp_send_version_old_merge2-------------------------------"); + RuntimeState recv_state; + pb::PlanNode pb_exchange_receiver_node = make_exchange_receiver_node_region_new_schema3(); + ExchangeReceiverNode exchange_receiver; + exchange_receiver.init(pb_exchange_receiver_node); + exchange_receiver.open(&recv_state); + exchange_receiver.get_data_stream_receiver()->to_string(); + pb::PlanNode pb_exchange_sender_node = make_exchange_sender_node(); + ::google::protobuf::RepeatedPtrField region_infos; + { + auto region_info = region_infos.Add(); + region_info->set_region_id(3); + region_info->set_version(2); + region_info->set_table_id(1); + region_info->set_partition_id(1); + region_info->set_replica_num(3); + region_info->set_conf_version(1); + } + Region::mpp_send_version_old(pb_exchange_sender_node.derive_node().exchange_sender_node(), false, region_infos); + exchange_receiver.get_data_stream_receiver()->to_string(); + exchange_receiver.close(&recv_state); +} + +TEST(test_null_cast, case_all) { + DB_WARNING("-------------------------------test_null_cast-------------------------------"); + std::shared_ptr record_batch = make_test_record_batch_null_schema(); + { + auto array = record_batch->column(0); + auto array_utf8_ret = arrow::compute::Cast(*array, arrow::utf8()); + auto array_utf8 = *array_utf8_ret; + DB_WARNING("array: %s", array->ToString().c_str()); + DB_WARNING("array_utf8: %s", array_utf8->ToString().c_str()); + + auto array_lb1_ret = arrow::compute::Cast(*array, arrow::large_binary()); + auto array_lb1 = *array_lb1_ret; + auto array_lb2_ret = arrow::compute::Cast(*array_utf8, arrow::large_binary()); + auto array_lb2 = *array_lb2_ret; + DB_WARNING("array_lb1: %s", array_lb1->ToString().c_str()); + DB_WARNING("array_lb2: %s", array_lb2->ToString().c_str()); + } + { + auto array = record_batch->column(1); + auto array_utf8_ret = arrow::compute::Cast(*array, arrow::utf8()); + auto array_utf8 = *array_utf8_ret; + DB_WARNING("array: %s", array->ToString().c_str()); + DB_WARNING("array_utf8: %s", array_utf8->ToString().c_str()); + + auto array_lb1_ret = arrow::compute::Cast(*array, arrow::large_binary()); + auto array_lb1 = *array_lb1_ret; + auto array_lb2_ret = arrow::compute::Cast(*array_utf8, arrow::large_binary()); + auto array_lb2 = *array_lb2_ret; + DB_WARNING("array_lb1: %s", array_lb1->ToString().c_str()); + DB_WARNING("array_lb2: %s", array_lb2->ToString().c_str()); + } +} + +TEST(test_null_hash, case_all) { + DB_WARNING("-------------------------------test_null_hash-------------------------------"); + std::shared_ptr array1 = make_test_array_tmp({1,2,3,4,5,6,7,8,9,10}); + std::shared_ptr array2 = make_test_array_tmp({11,12,13,14,15,16,17,18,19,20}); + std::shared_ptr array3 = make_test_array_tmp({21,22,23,24,25,26,27,28,29,30}); + DB_WARNING("array length: %ld", array1->length()); + std::shared_ptr record_batch = + arrow::RecordBatch::Make(make_test_schema(), 2, {array1, array2, array3}); + DB_WARNING("record_batch: %s", record_batch->ToString().c_str()); + std::unordered_map> hash_batch_map; + ExchangeSenderNode exchange_sender_node; + auto status = exchange_sender_node.repartition(record_batch, {"col1", "col2"}, {"col1"}, 3, hash_batch_map); + if (!status.ok()) { + DB_WARNING("Fail to repartition"); + } + ASSERT_EQ(status.ok(), true); + for (auto& [bucket, batch] : hash_batch_map) { + DB_WARNING("bucket: %d, batch: %s", bucket, batch->ToString().c_str()); + } +} + +} // namespace baikaldb + +using namespace baikaldb; + +int main(int argc, char* argv[]) { + google::SetCommandLineOption("flagfile", "conf/gflags.conf"); + google::ParseCommandLineFlags(&argc, &argv, true); + + // 初始化日志 + if (baikaldb::init_log(argv[0]) != 0) { + fprintf(stderr, "log init failed."); + return -1; + } + DB_WARNING("breakpoint main begin"); + + if (ArrowExecNodeManager::RegisterAllArrowExecNode() != 0) { + DB_FATAL("RegisterAllArrowExecNode failed"); + return -1; + } + + // 启动Server + brpc::Server server; + baikaldb::DbService* db_service = baikaldb::DbService::get_instance(); + if (0 != server.AddService(db_service, brpc::SERVER_DOESNT_OWN_SERVICE)) { + DB_FATAL("Fail to Add idonlyeService"); + return -1; + } + if (server.Start("10.12.187.120:8787", NULL) != 0) { + DB_FATAL("Fail to start server"); + return -1; + } + // UT + testing::InitGoogleTest(&argc, argv); + int ret = RUN_ALL_TESTS(); + DB_WARNING("breakpoint UT end"); + // 停止Server + server.Stop(0); + server.Join(); + DB_WARNING("breakpoint main end"); + return 0; +} \ No newline at end of file diff --git a/test/test_faiss_sift1M.cpp b/test/test_faiss_sift1M.cpp new file mode 100644 index 000000000..121afd6a9 --- /dev/null +++ b/test/test_faiss_sift1M.cpp @@ -0,0 +1,607 @@ +// Copyright (c) 2018 Baidu, Inc. All Rights Reserved. +// +// Licensed 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "mysql_interact.h" +#include + +DEFINE_int32(nb, 100000, "number of base vectors for testing"); +DEFINE_int32(nq, 1000, "number of query vectors for testing"); +DEFINE_int32(m, 16, "HNSW parameter M"); +DEFINE_int32(tp, 0, "thread pool size for parallel processing"); +DEFINE_int32(nlist, 0, "number of IVF clusters (0 for auto-calculate)"); +DEFINE_bool(read, false, "read existing index from file"); +DEFINE_int32(batch, 0, "batch size for vector insertion (0 for single batch)"); + +DEFINE_bool(test_recall, false, "run recall test scenarios"); +DEFINE_int32(d, 384, "vector dimension"); +DEFINE_int32(vector_num, 0, "number of vectors (0 for auto-detect)"); +DEFINE_int32(recall_s, 10, "sample size for recall calculation"); +DEFINE_int32(refine_factor, 5, "refinement factor for RFlat indexes"); +DEFINE_int32(ivf_nlist_factor, 1, "IVF nlist scaling factor"); +DEFINE_int32(ivf_nprobe_factor, 10, "IVF nprobe scaling factor"); +DEFINE_int32(pq_factor, 4, "PQ sub-quantizer factor"); +DEFINE_int32(pqfs_factor, 2, "PQFS sub-quantizer factor"); +DEFINE_string(target, "", "target index type to test"); +DEFINE_int32(search_num, 10, "number of nearest neighbors to search"); +DEFINE_int32(query_times, 100, "number of query iterations"); +DEFINE_string(dataset, "sift", "dataset name (sift) - http://corpus-texmex.irisa.fr/"); +DEFINE_int32(efsearch, 16, "HNSW efSearch parameter"); +DEFINE_int32(efconstruction, 40, "HNSW efConstruction parameter"); +DEFINE_int32(learn_vec_cnt, 100000, "number of vectors for index training"); + +DEFINE_bool(insert_to_baikaldb, false, "insert test data to BaikalDB"); +DEFINE_int32(insert_row_cnt, 100000, "insert test data to BaikalDB"); +DEFINE_bool(query_baikaldb, false, "read test data from BaikalDB"); +DEFINE_string(table, "vec.test_ivf", "BaikalDB table name for testing"); +DEFINE_int32(sql_print_time_limit, 100000, "SQL execution time limit for printing"); +DEFINE_string(order, "asc", "order by use"); +DEFINE_string(id2_value, "333", "id2 value"); + + +// UTApplication('test_faiss_sift1m', Sources('test/test_faiss_sift1M.cpp', CxxFlags('-g -DBAIDU_INTERNAL -fno-access-control -Wno-sign-compare -Wno-unused-variable -std=c++17 -fopenmp')), +// Libraries( +// '$OUT/lib/libproto.a', +// '$OUT/lib/libcapture.a', +// '$OUT/lib/libwatt_proto.a', +// '$OUT/lib/libcommon.a')) + +int cnt = 0; +using idx_t = int64_t; + +int main(int argc, char* argv[]) +{ + testing::InitGoogleTest(&argc, argv); + google::ParseCommandLineFlags(&argc, &argv, true); + //cnt = std::stoi(argv[1]); + return RUN_ALL_TESTS(); +} +class TimeCost { +public: + TimeCost() { + _start = base::gettimeofday_us(); + } + + ~TimeCost() {} + + void reset() { + _start = base::gettimeofday_us(); + } + + int64_t get_time() const { + return base::gettimeofday_us() - _start; + } + +private: + int64_t _start; +}; +void query_and_calc_recall(faiss::Index* idx, int d, int nb, int nq, int nl, float* xb, float* xq, float* xl, idx_t* ids, int sn, std::string desc, std::vector>& flat_res, int ivf_nprobe) { + faiss::Index& index = *idx; + bool is_ivf = (desc.find("IVF") != std::string::npos); + bool is_refine = (desc.find("RFlat") != std::string::npos); + bool is_hnsw = (desc.find("HNSW") != std::string::npos); + if (is_ivf) { + faiss::ParameterSpace().set_index_parameter(idx, "nprobe", ivf_nprobe); + } + if (is_refine) { + faiss::ParameterSpace().set_index_parameter(idx, "k_factor_rf", FLAGS_refine_factor); + } + if (is_hnsw) { + faiss::ParameterSpace().set_index_parameter(idx, "efSearch", FLAGS_efsearch); + faiss::ParameterSpace().set_index_parameter(idx, "efConstruction", FLAGS_efconstruction); + } + + int k = sn; + bool is_flat = flat_res.empty(); + std::vector> res_vec; + TimeCost cost; + { // sanity check: search 5 first vectors of xb + idx_t* I = new idx_t[k * FLAGS_recall_s]; + float* D = new float[k * FLAGS_recall_s]; + cost.reset(); + index.search(FLAGS_recall_s, xb, k, D, I); + printf("search %d in %zd, cost=:%ld\n", FLAGS_recall_s, index.ntotal, cost.get_time()); + cost.reset(); + + // print results + printf("I=\n"); + for (int i = 0; i < FLAGS_recall_s; i++) { + for (int j = 0; j < k; j++) { + printf("%5zd ", I[i * k + j]); + } + printf("\n"); + } + + printf("D=\n"); + for (int i = 0; i < FLAGS_recall_s; i++) { + for (int j = 0; j < k; j++) + printf("%7g ", D[i * k + j]); + printf("\n"); + } + + delete[] I; + delete[] D; + } + + { // search xq + idx_t* I = new idx_t[k * nq]; + float* D = new float[k * nq]; + TimeCost cost2; + for (int i = 0; i < nq; i++) { + cost.reset(); + index.search(1, xq + i * d, k, D + i * k, I + i * k); + //printf("search i = %d, cost=:%ld\n", i, cost.get_time()); + } + printf("search avg, cost=:%ld\n", cost2.get_time()/nq); + + // print results + printf("I (5 first results)=\n"); + for (int i = 0; i < 5; i++) { + for (int j = 0; j < k; j++) + printf("%5zd ", I[i * k + j]); + printf("\n"); + } + + printf("I (5 last results)=\n"); + for (int i = nq - 5; i < nq; i++) { + for (int j = 0; j < k; j++) + printf("%5zd ", I[i * k + j]); + printf("\n"); + } + + bool is_flat_res = (flat_res.empty()); + float recall_rate = 0; + for (int i = 0; i < nq; i++) { + std::set one_vec_res; + for (int j = 0; j < k; j++) { + one_vec_res.insert(I[i * k + j]); + } + if (is_flat_res) { + flat_res.emplace_back(one_vec_res); + } else { + int recall_cnt = 0; + for (int j = 0; j < k; j++) { + if (flat_res[i].count(I[i * k + j]) > 0) { + recall_cnt++; + } + } + recall_rate += (float)recall_cnt / k; + } + } + std::cout << "recall_rate:" << recall_rate / nq << std::endl; + } +} + +void run_one_test(int d, int nb, int nq, int nl, float* xb, float* xq, float* xl, idx_t* ids, int sn, std::string desc, std::vector>& flat_res, int ivf_nprobe) { + faiss::IndexFlatL2 quantizer(d); // call constructor + printf("RUN TEST: dimension: %d, vector_cnt: %d, query_cnt: %d, ivf_nprobe: %d, desc: %s=======================\n", d, nb, nq, ivf_nprobe, desc.c_str()); + faiss::Index* idx = faiss::index_factory(d, desc.c_str()); + TimeCost cost; + + faiss::Index& index = *idx; + index.train(nl, xl); // add vectors to the index + printf("ntotal = %zd, nb:%d,nq:%d, train cost=:%ld\n", index.ntotal, nb, nq, cost.get_time()); + cost.reset(); + if (FLAGS_batch == 0) { + index.add_with_ids(nb, xb, ids); // add vectors to the index + } else { + for (int i = 0; i < nb; i+=FLAGS_batch) { + TimeCost cost2; + index.add_with_ids(FLAGS_batch, xb + i * d, ids + i); // add vectors to the index + if (i % 1000 == 0) { + printf(" batch= %d, cost=:%ld\n", FLAGS_batch, cost2.get_time()); + } + } + } + printf("ntotal = %zd, nb:%d,nq:%d, cost=:%ld\n", index.ntotal, nb, nq, cost.get_time()); + cost.reset(); + std::string file_name = desc + ".index"; + faiss::write_index(&index, file_name.c_str()); + printf("ntotal = %zd, nb:%d,nq:%d, write cost=:%ld\n", index.ntotal, nb, nq, cost.get_time()); + + query_and_calc_recall(idx, d, nb, nq, nl, xb, xq, xl, ids, sn, desc, flat_res, ivf_nprobe); + printf("done=====================================================\n\n\n"); + sleep(1); +} + +void load_data(const std::string filename, float*& data, unsigned& num, unsigned& dim) { + std::ifstream in(filename.c_str(), std::ios::binary); //以二进制的方式打开文件 + if (!in.is_open()) { + std::cout << "open file error" << std::endl; + exit(-1); + } + in.read((char*)&dim, 4); //读取向量维度 + in.seekg(0, std::ios::end); //光标定位到文件末尾 + std::ios::pos_type ss = in.tellg(); //获取文件大小（多少字节） + size_t fsize = (size_t)ss; + num = (unsigned)(fsize / (dim + 1) / 4); //数据的个数 + data = new float[(size_t)num * (size_t)dim]; + + in.seekg(0, std::ios::beg); //光标定位到起始处 + for (size_t i = 0; i < num; i++) { + in.seekg(4, std::ios::cur); //光标向右移动4个字节 + in.read((char*)(data + i * dim), dim * 4); //读取数据到一维数据data中 + } + for(size_t i = 0; i < num * dim; i++) { //输出数据 + std::cout << (float)data[i]; + if(!i) { + std::cout << " "; + continue; + } + if(i % (dim - 1) != 0) { + std::cout << " "; + } + else{ + std::cout << std::endl; + break; + } + } + in.close(); +} + +void load_shif1M(unsigned& d, unsigned& nb, unsigned& nq, unsigned& nl, float*& xb, float*& xq, float*& xl, idx_t*& ids) { + unsigned base_dim = 0; + unsigned learn_dim = 0; + unsigned query_dim = 0; + unsigned base_num = 0; + unsigned learn_num = 0; + unsigned query_num = 0; + load_data("./sift_base.fvecs", xb, base_num, base_dim); + load_data("./sift_learn.fvecs", xl, learn_num, learn_dim); + load_data("./sift_query.fvecs", xq, query_num, query_dim); + if (base_dim != learn_dim || base_dim != query_dim) { + std::cout << "dim error: base_dim: " << base_dim << " learn_dim: " << learn_dim << " query_dim: " << query_dim << std::endl; + exit(-1); + } + std::cout << "load finish: \n" << "base_num: " << base_num << " learn_num: " << learn_num << " query_num: " << query_num << std::endl; + d = base_dim; + nb = base_num; + nq = query_num; + nl = learn_num; + + ids = new idx_t[nb]; + for (int i = 0; i < nb; i++) { + ids[i] = i * 2; + } + return; +} + +TEST(test_sift_1m, case_all) { + if (!FLAGS_test_recall) { + return; + } + unsigned dim, nb, nq, nl, nrl; + float* xb; // base + float* xq; // query + float* xl; // learn + float* xrl; // real_learn + idx_t* ids;// ids + + load_shif1M(dim, nb, nq, nl, xb, xq, xl, ids); + + nq = std::min(nq, (unsigned)FLAGS_query_times); + std::vector> flat_res; + + int ivf_nlist = FLAGS_ivf_nlist_factor * sqrt(nb); + int ivf_nprobe = ivf_nlist / FLAGS_ivf_nprobe_factor; + int pq_nsub = dim / FLAGS_pq_factor; + int pqfs_nsub = dim / FLAGS_pqfs_factor; + int sn = FLAGS_search_num; + if (FLAGS_learn_vec_cnt <= 100000) { + xrl = xl; + } else { + xrl = xb; + } + nrl = FLAGS_learn_vec_cnt; + std::cout << "ivf_nlist: " << ivf_nlist << " ivf_nprobe: " << ivf_nprobe << "\npq_nsub: " << pq_nsub << " pqfs_nsub: " << pqfs_nsub << std::endl; + std::cout << "search_num: " << sn << std::endl; + std::cout << "learn_vec_cnt: " << nrl << std::endl; + std::vector targets; + targets.emplace_back("IDMap2,Flat"); + if (FLAGS_target != "") { + targets.emplace_back(FLAGS_target); + } else { + targets.emplace_back("IDMap2,HNSW32"); + { + char desc[48]; + sprintf(desc, "IDMap2,HNSW32,PQ%d", pq_nsub); + targets.emplace_back(desc); + } + + { + char desc[48]; + sprintf(desc, "IDMap2,PQ%d", pq_nsub); + targets.emplace_back(desc); + } + + { + char desc[48]; + sprintf(desc, "IDMap2,PQ%d,RFlat", pq_nsub); + targets.emplace_back(desc); + } + + { + char desc[48]; + sprintf(desc, "IDMap2,PQ%dx4fs", pqfs_nsub); + targets.emplace_back(desc); + } + + { + char desc[48]; + sprintf(desc, "IDMap2,PQ%dx4fs,RFlat", pqfs_nsub); + targets.emplace_back(desc); + } + + { + char desc[48]; + sprintf(desc, "IDMap2,IVF%d,Flat", ivf_nlist); + targets.emplace_back(desc); + } + + { + char desc[48]; + sprintf(desc, "IDMap2,IVF%d,PQ%d", ivf_nlist, pq_nsub); + targets.emplace_back(desc); + } + + { + char desc[48]; + sprintf(desc, "IDMap2,IVF%d,PQ%d,RFlat", ivf_nlist, pq_nsub); + targets.emplace_back(desc); + } + + { + char desc[48]; + sprintf(desc, "IDMap2,IVF%d,PQ%dx4fs", ivf_nlist, pqfs_nsub); + targets.emplace_back(desc); + } + + { + char desc[48]; + sprintf(desc, "IDMap2,IVF%d,PQ%dx4fs,RFlat", ivf_nlist, pqfs_nsub); + targets.emplace_back(desc); + } + } + + for (auto target : targets) { + run_one_test(dim, nb, nq, nrl, xb, xq, xrl, ids, sn, target, flat_res, ivf_nprobe); + } + + delete[] xb; + delete[] xq; + delete[] ids; +} + +TEST(test_load, case_all) { + if (!FLAGS_read || FLAGS_target == "") { + return; + } + unsigned dim, nb, nq, nl, nrl; + float* xb; // base + float* xq; // query + float* xl; // learn + float* xrl; // real_learn + idx_t* ids;// ids + + load_shif1M(dim, nb, nq, nl, xb, xq, xl, ids); + + nq = std::min(nq, (unsigned)FLAGS_query_times); + std::vector> flat_res; + + int ivf_nlist = FLAGS_ivf_nlist_factor * sqrt(nb); + int ivf_nprobe = ivf_nlist / FLAGS_ivf_nprobe_factor; + int pq_nsub = dim / FLAGS_pq_factor; + int pqfs_nsub = dim / FLAGS_pqfs_factor; + int sn = FLAGS_search_num; + if (FLAGS_learn_vec_cnt <= 100000) { + xrl = xl; + } else { + xrl = xb; + } + nrl = FLAGS_learn_vec_cnt; + std::cout << "ivf_nlist: " << ivf_nlist << " ivf_nprobe: " << ivf_nprobe << "\npq_nsub: " << pq_nsub << " pqfs_nsub: " << pqfs_nsub << std::endl; + std::cout << "search_num: " << sn << std::endl; + std::cout << "learn_vec_cnt: " << nrl << std::endl; + + std::vector targets = {"IDMap2,Flat", FLAGS_target}; + for (auto target : targets) { + run_one_test(dim, nb, nq, nrl, xb, xq, xrl, ids, sn, target, flat_res, ivf_nprobe); + } + + // read + faiss::Index* load_idx; + TimeCost cost; + std::string file_name = FLAGS_target + ".index"; + load_idx = faiss::read_index(file_name.c_str()); + printf("nb:%d,nq:%d, read cost=:%ld, is_trained: %d\n" ,nb, nq, cost.get_time(), load_idx->is_trained); + cost.reset(); + + if (!load_idx->is_trained) { + load_idx->train(nrl, xrl); + printf("train cost=:%ld\n", cost.get_time()); + cost.reset(); + } + query_and_calc_recall(load_idx, dim, nb, nq, nl, xb, xq, xl, ids, sn, FLAGS_target, flat_res, ivf_nprobe); + // load_idx->add_with_ids(nrl, xrl, ids); + // printf("add cost=:%ld\n", cost.get_time()); + delete[] xb; + delete[] xq; + delete[] ids; +} + +int query(std::string sql, baikal::client::Service* _baikaldb, baikal::client::SmartConnection& connection) { + int ret = 0; + int retry = 0; + int affected_rows = 0; + do { + if (connection == nullptr) { + connection = _baikaldb->fetch_connection(); + } + baikal::client::ResultSet result_set; + ret = connection->execute(sql, true, &result_set); + if (ret == 0) { + affected_rows = result_set.get_affected_rows(); + break; + } + bthread_usleep(1000000); + if (connection != nullptr) { + connection->close(); + } + connection = _baikaldb->fetch_connection(); + while (connection == nullptr) { + if (retry > 3) { + printf("service fetch_connection() failed"); + break; + } + bthread_usleep(1000000); + connection = _baikaldb->fetch_connection(); + ++retry; + } + } while (++retry < 3); + if (ret != 0) { + printf("sql_len:%lu query fail finally: %s", sql.size(), sql.c_str()); + return -1; + } + return affected_rows; +} + +TEST(test_query_baikaldb, case_all) { + if (!FLAGS_query_baikaldb) { + return; + } + baikal::client::Manager _manager; + baikal::client::Service* _baikaldb; + int rc = 0; + rc = _manager.init("conf", "baikal_client.conf"); + if (rc != 0) { + printf("baikal client init fail:%d", rc); + return; + } + _baikaldb = _manager.get_service("baikaldb"); + if (_baikaldb == NULL) { + printf("baikaldb is null"); + return; + } + + unsigned dim, nb, nq, nl, nrl; + float* xb; // base + float* xq; // query + float* xl; // learn + float* xrl; // real_learn + idx_t* ids;// ids + + load_shif1M(dim, nb, nq, nl, xb, xq, xl, ids); + + std::string sql; + std::string vec_string; + std::vector sqls; + for (int i = 0; i < nq; ++i) { + sql = ""; + vec_string = ""; + for (int j = 0; j < dim; ++j) { + vec_string += std::to_string(xq[i * dim + j]) + ","; + } + vec_string.pop_back(); + sql = "select id from " + FLAGS_table + " where match(vec) against ('" + vec_string + "' in vector mode) order by __weight " + FLAGS_order + " limit 10;"; + sqls.emplace_back(sql); + } + + TimeCost total_time; + TimeCost one_query_time; + auto connection = _baikaldb->fetch_connection(); + for (int i = 0; i < nq; ++i) { + one_query_time.reset(); + query(sqls[i], _baikaldb, connection); + if (one_query_time.get_time() > FLAGS_sql_print_time_limit) { + std::cout << "query " << i << " cost: " << one_query_time.get_time() << std::endl; + } + } + printf("select %d vector, cost=:%ld\n", nq, total_time.get_time()); +} + + +TEST(test_insert_baikaldb, case_all) { + if (!FLAGS_insert_to_baikaldb) { + return; + } + baikal::client::Manager _manager; + baikal::client::Service* _baikaldb; + int rc = 0; + rc = _manager.init("conf", "baikal_client.conf"); + if (rc != 0) { + printf("baikal client init fail:%d", rc); + return; + } + _baikaldb = _manager.get_service("baikaldb"); + if (_baikaldb == NULL) { + printf("baikaldb is null"); + return; + } + + unsigned dim, nb, nq, nl, nrl; + float* xb; // base + float* xq; // query + float* xl; // learn + float* xrl; // real_learn + idx_t* ids;// ids + + load_shif1M(dim, nb, nq, nl, xb, xq, xl, ids); + + std::string sql; + std::string vec_string; + std::vector vec_strings; + std::vector tables; + boost::split(tables, FLAGS_table, boost::is_any_of(",")); + for (int i = 0; i < nb && i < FLAGS_insert_row_cnt; ++i) { + vec_string = ""; + for (int j = 0; j < dim; ++j) { + vec_string += std::to_string(xb[i * dim + j]) + ","; + } + vec_string.pop_back(); + vec_strings.emplace_back(vec_string); + } + + TimeCost total_time; + TimeCost one_query_time; + auto connection = _baikaldb->fetch_connection(); + for (auto table : tables) { + total_time.reset(); + for (int i = 0; i < vec_strings.size(); ++i) { + one_query_time.reset(); + sql = "replace into " + table + "(id, vec, id2) values (" + std::to_string(i) + ", '" + vec_strings[i] + "', '" + FLAGS_id2_value + "')"; + query(sql, _baikaldb, connection); + if (one_query_time.get_time() > FLAGS_sql_print_time_limit) { + std::cout << "insert " << i << " cost: " << one_query_time.get_time() << std::endl; + } + } + printf("INSERT table: %s, %d vector, cost=:%ld\n", table.c_str(), nb, total_time.get_time()); + } +} + diff --git a/test/test_parquet_writer.cpp b/test/test_parquet_writer.cpp new file mode 100644 index 000000000..f50e7b47b --- /dev/null +++ b/test/test_parquet_writer.cpp @@ -0,0 +1,77 @@ +// Copyright (c) 2018 Baidu, Inc. All Rights Reserved. +// +// Licensed 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 +#include "parquet_writer.h" + +int main(int argc, char* argv[]) +{ + testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} + +namespace baikaldb { +void build_arrow_record_batch(std::shared_ptr record_batch) { + // 定义Schema + auto schema = arrow::schema({ + arrow::field("id", arrow::int32()), + arrow::field("score", arrow::float32()) + }); + + // 创建ArrayBuilder + arrow::Int32Builder id_builder; + arrow::FloatBuilder score_builder; + + // 添加数据到ArrayBuilder + id_builder.Append(1); + id_builder.Append(2); + id_builder.Append(3); + + score_builder.Append(4.5); + score_builder.Append(8.9); + score_builder.Append(7.2); + + // 从ArrayBuilder创建Array + std::shared_ptr id_array; + std::shared_ptr score_array; + id_builder.Finish(&id_array); + score_builder.Finish(&score_array); + + // 将Array包装成ArrayData + std::vector> arrays = { + id_array->data(), + score_array->data() + }; + + // 创建RecordBatch + record_batch = arrow::RecordBatch::Make(schema, 3, arrays); +} + + +TEST(test_arrow_vector_execute, case_all) { + // std::shared_ptr record_batch; + // build_arrow_record_batch(record_batch); + + // ParquetWriteOptions options; + // ParquetWriter writer; + // writer.init(options); + // writer.write_batch(record_batch); + + // std::vector > file_infos = writer.get_file_infos(); + // ASSERT_EQ(file_infos.size(), 1); + // writer.finish(); + + +} +} // namespace baikal \ No newline at end of file diff --git a/test/test_row2col.cpp b/test/test_row2col.cpp new file mode 100644 index 000000000..4e1d8f599 --- /dev/null +++ b/test/test_row2col.cpp @@ -0,0 +1,27 @@ +// Copyright (c) 2018 Baidu, Inc. All Rights Reserved. +// +// Licensed 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 +#include "row2column.h" + +int main(int argc, char* argv[]) +{ + testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} + +namespace baikaldb { + + +} // namespace baikal \ No newline at end of file diff --git a/test/test_window.cpp b/test/test_window.cpp new file mode 100644 index 000000000..4fa7bcff5 --- /dev/null +++ b/test/test_window.cpp @@ -0,0 +1,441 @@ +// Copyright (c) 2022 Baidu, Inc. All Rights Reserved. +// +// Licensed 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 + +#include "logical_planner.h" +#include "parser.h" +#include "scan_node.h" +#include "expr_optimizer.h" +#include "window_node.h" + +int main(int argc, char* argv[]) { + testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} + +namespace baikaldb { + +void update_table() { + auto schema_factory = SchemaFactory::get_instance(); + pb::SchemaInfo info; + info.set_namespace_name("test_namespace"); + info.set_database("test_database"); + info.set_table_name("test_window"); + info.set_partition_num(1); + info.set_namespace_id(111); + info.set_database_id(222); + for (int idx = 1; idx < 5; idx++) { + baikaldb::pb::FieldInfo *field_string = info.add_fields(); + field_string->set_field_name("col" + std::to_string(idx)); + field_string->set_field_id(idx); + field_string->set_mysql_type(baikaldb::pb::INT32); + } + baikaldb::pb::IndexInfo *index_pk = info.add_indexs(); + index_pk->set_index_type(baikaldb::pb::I_PRIMARY); + index_pk->set_index_name("pk_index"); + index_pk->add_field_ids(1); + index_pk->set_index_id(1); + info.set_table_id(1); + info.set_version(2); + schema_factory->init(); + schema_factory->update_table(info); +} + +void construct_row_batch(RuntimeState* state, RowBatch& row_batch) { + for (int i = 0; i <= 5; ++i) { + for (int j = 0; j <= i; ++j) { + int loop_cnt = 1; + if (i == 4) { + loop_cnt = 2; + } + for (int k = 0; k < loop_cnt; ++k) { + std::unique_ptr row = state->mem_row_desc()->fetch_mem_row(); + { + ExprValue tmp(pb::INT32); + tmp._u.int32_val = i; + row->set_value(0, 1, tmp); + } + { + ExprValue tmp(pb::INT32); + tmp._u.int32_val = i; + row->set_value(0, 2, tmp); + } + { + ExprValue tmp(pb::INT32); + tmp._u.int32_val = i*10 + j; + row->set_value(0, 3, tmp); + } + row_batch.move_row(std::move(row)); + } + } + } +} + +void construct_row_batch(std::shared_ptr mem_row_desc, RowBatch& row_batch, int start) { + for (int i = start; i <= start + 5; ++i) { + for (int j = 0; j <= i; ++j) { + int loop_cnt = 1; + if (i == 4) { + loop_cnt = 2; + } + for (int k = 0; k < loop_cnt; ++k) { + std::unique_ptr row = mem_row_desc->fetch_mem_row(); + { + ExprValue tmp(pb::INT32); + tmp._u.int32_val = 1; + row->set_value(0, 1, tmp); + } + { + ExprValue tmp(pb::INT32); + tmp._u.int32_val = i; + row->set_value(0, 2, tmp); + } + { + ExprValue tmp(pb::INT32); + tmp._u.int32_val = i*10 + j; + row->set_value(0, 3, tmp); + } + row_batch.move_row(std::move(row)); + } + } + } +} + + +class MockScanNode : public ExecNode { +public: + virtual int get_next(RuntimeState* state, RowBatch* batch, bool* eos) override { + construct_row_batch(state, *batch); + *eos = true; + std::cout << "MockScanNode batch size: " << batch->size() << std::endl; + return 0; + } +}; + +TEST(test_window, case_sql_parser) { + parser::SqlParser parser; + const std::string& sql = + "SELECT " + "id, " + "COUNT(*) OVER (PARTITION BY id ORDER BY id ROWS 1 PRECEDING), " + "COUNT(col1) OVER (PARTITION BY id ORDER BY id ROWS BETWEEN 1 PRECEDING AND CURRENT ROW), " + "SUM(col1) OVER (PARTITION BY id ORDER BY id ROWS BETWEEN CURRENT ROW AND 1 FOLLOWING), " + "AVG(col1) OVER (PARTITION BY id ORDER BY id ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "MIN(col1) OVER (PARTITION BY id ORDER BY id ROWS BETWEEN 3 PRECEDING AND 2 PRECEDING), " + "MAX(col1) OVER (PARTITION BY id ORDER BY id ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING), " + "ROW_NUMBER() OVER (PARTITION BY id ORDER BY id), " + "RANK() OVER (PARTITION BY id ORDER BY id), " + "DENSE_RANK() OVER (PARTITION BY id ORDER BY id), " + "PERCENT_RANK() OVER (PARTITION BY id ORDER BY id), " + "CUME_DIST() OVER (PARTITION BY id ORDER BY id), " + "NTILE(5) OVER (PARTITION BY id ORDER BY id), " + "LEAD(id) OVER (PARTITION BY id ORDER BY id), " + "LEAD(id, 2) OVER (PARTITION BY id ORDER BY id), " + "LEAD(id, 2, id + col1 + 1) OVER (PARTITION BY id ORDER BY id), " + "LEAD(id, 2, id + col1 + 1) IGNORE NULLS OVER (PARTITION BY id ORDER BY id), " + "LEAD(id, 2, id + col1 + 1) RESPECT NULLS OVER (PARTITION BY id ORDER BY id), " + "LAG(id) OVER (PARTITION BY id ORDER BY id), " + "LAG(id, 2) OVER (PARTITION BY id ORDER BY id), " + "LAG(id, 2, id + col1 + 1) OVER (PARTITION BY id ORDER BY id), " + "FIRST_VALUE(id + 2) OVER (PARTITION BY id ORDER BY id ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "LAST_VALUE(id) RESPECT NULLS OVER (PARTITION BY id ORDER BY id ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "NTH_VALUE(id * 3, 2) IGNORE NULLS OVER (PARTITION BY id ORDER BY id ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "NTH_VALUE(id * 3, 2) IGNORE NULLS OVER (PARTITION BY id ORDER BY id), " + "NTH_VALUE(id * 3, 2) FROM LAST OVER ()" + "FROM " + "test_window;"; + parser.parse(sql); + ASSERT_EQ(parser.result.size(), 1); + ASSERT_EQ(parser.result[0]->node_type, parser::NT_SELECT); + parser::SelectStmt* stmt = (parser::SelectStmt*)parser.result[0]; + ASSERT_NE(stmt, nullptr); + + const std::string& stmt_str = stmt->to_string(); + const std::string& expect_stmt_str = + "SELECT " + "id, " + "(COUNT(*)) OVER ( PARTITION BY id ASC ORDER BY id ASC ROWS 1 PRECEDING), " + "(COUNT(col1)) OVER ( PARTITION BY id ASC ORDER BY id ASC ROWS BETWEEN 1 PRECEDING AND CURRENT ROW), " + "(SUM(col1)) OVER ( PARTITION BY id ASC ORDER BY id ASC ROWS BETWEEN CURRENT ROW AND 1 FOLLOWING), " + "(AVG(col1)) OVER ( PARTITION BY id ASC ORDER BY id ASC ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "(MIN(col1)) OVER ( PARTITION BY id ASC ORDER BY id ASC ROWS BETWEEN 3 PRECEDING AND 2 PRECEDING), " + "(MAX(col1)) OVER ( PARTITION BY id ASC ORDER BY id ASC ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING), " + "(ROW_NUMBER()) OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(RANK()) OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(DENSE_RANK()) OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(PERCENT_RANK()) OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(CUME_DIST()) OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(NTILE(5)) OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(LEAD(id)) RESPECT NULLS OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(LEAD(id,2)) RESPECT NULLS OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(LEAD(id,2,((id + col1) + 1))) RESPECT NULLS OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(LEAD(id,2,((id + col1) + 1))) IGNORE NULLS OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(LEAD(id,2,((id + col1) + 1))) RESPECT NULLS OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(LAG(id)) RESPECT NULLS OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(LAG(id,2)) RESPECT NULLS OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(LAG(id,2,((id + col1) + 1))) RESPECT NULLS OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(FIRST_VALUE((id + 2))) RESPECT NULLS OVER ( PARTITION BY id ASC ORDER BY id ASC ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "(LAST_VALUE(id)) RESPECT NULLS OVER ( PARTITION BY id ASC ORDER BY id ASC ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "(NTH_VALUE((id * 3),2)) FROM FIRST IGNORE NULLS OVER ( PARTITION BY id ASC ORDER BY id ASC ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "(NTH_VALUE((id * 3),2)) FROM FIRST IGNORE NULLS OVER ( PARTITION BY id ASC ORDER BY id ASC), " + "(NTH_VALUE((id * 3),2)) FROM LAST RESPECT NULLS OVER () " + "FROM " + "test_window"; + EXPECT_EQ(stmt_str, expect_stmt_str); +} + +TEST(test_window, case_plan) { + // 更新schema信息 + update_table(); + // logical plan测试 + std::shared_ptr user_info(new (std::nothrow)UserInfo); + user_info->namespace_ = "test_namespace"; + user_info->is_super = true; + SmartSocket client = std::make_shared(); + client->user_info = user_info; + client->send_buf = new (std::nothrow) DataBuffer; + client->reset_query_ctx(new (std::nothrow) QueryContext(client->user_info, "test_database")); + client->query_ctx->client_conn = client.get(); + const std::string& sql = + "SELECT " + "col1, " + "COUNT(*) OVER (PARTITION BY col1 ORDER BY col1 ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "AVG(col2) OVER (PARTITION BY col1 ORDER BY col1 ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "SUM(col3) OVER (PARTITION BY col1 ORDER BY col1 ROWS BETWEEN UNBOUNDED PRECEDING AND 2 FOLLOWING), " + "MIN(col3) OVER (PARTITION BY col1 ORDER BY col1 ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "MAX(col3) OVER (PARTITION BY col1 ORDER BY col1 ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "ROW_NUMBER() OVER (PARTITION BY col1 ORDER BY col3), " + "RANK() OVER (PARTITION BY col1 ORDER BY col3), " + "DENSE_RANK() OVER (PARTITION BY col1 ORDER BY col3), " + "PERCENT_RANK() OVER (PARTITION BY col1 ORDER BY col3), " + "CUME_DIST() OVER (PARTITION BY col1 ORDER BY col1), " + "NTILE(5) OVER (PARTITION BY col1 ORDER BY col1), " + "LEAD(col1) OVER (PARTITION BY col1 ORDER BY col1), " + "LEAD(col1, 2) OVER (PARTITION BY col1 ORDER BY col1), " + "LEAD(col1, 2, col1 + col1 + 1) OVER (PARTITION BY col1 ORDER BY col1), " + "LEAD(col1, 2, col1 + col1 + 1) OVER (PARTITION BY col1 ORDER BY col1), " + "LEAD(col1, 2, col1 + col1 + 1) RESPECT NULLS OVER (PARTITION BY col1 ORDER BY col1), " + "LAG(col1) OVER (PARTITION BY col1 ORDER BY col1), " + "LAG(col1, 2) OVER (PARTITION BY col1 ORDER BY col1), " + "LAG(col1, 2, col1 + col1 + 1) OVER (PARTITION BY col1 ORDER BY col1), " + "FIRST_VALUE(col1 + 2) OVER (PARTITION BY col1 ORDER BY col1 ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "LAST_VALUE(col1) RESPECT NULLS OVER (PARTITION BY col1 ORDER BY col1 ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "NTH_VALUE(col1 * 3, 2) OVER (PARTITION BY col1 ORDER BY col1 ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING), " + "NTH_VALUE(col1 * 3, 2) OVER (PARTITION BY col1 ORDER BY col1), " + "NTH_VALUE(col1 * 3, 2) OVER () " + "FROM " + "test_window"; + client->query_ctx->sql = sql; + int ret = LogicalPlanner::analyze(client->query_ctx.get()); + ASSERT_EQ(ret, 0); + // 生成逻辑计划树测试 + ret = client->query_ctx->create_plan_tree(); + ASSERT_EQ(ret, 0); + // 物理计划优化 + // 执行ExprOptimizer analyze就行 + ret = ExprOptimize().analyze(client->query_ctx.get()); + ASSERT_EQ(ret, 0); + pb::Plan plan; + ExecNode::create_pb_plan(0, &plan, client->query_ctx->root); + std::cout << plan.ShortDebugString() << std::endl; + // 执行物理计划 + RuntimeState& state = *client->query_ctx->get_runtime_state(); + ret = state.init(client->query_ctx.get(), client->send_buf); + ASSERT_EQ(ret, 0); + ScanNode* scan_node = static_cast(client->query_ctx->root->get_node(pb::SCAN_NODE)); + ASSERT_NE(scan_node, nullptr); + ExecNode* parent = scan_node->get_parent(); + ASSERT_NE(scan_node, nullptr); + std::unique_ptr mock_scan_node(new (std::nothrow) MockScanNode); + parent->replace_child(scan_node, mock_scan_node.release()); + state.set_from_subquery(true); + ret = client->query_ctx->root->open(&state); + ASSERT_EQ(ret, 0); + bool eos = false; + while (!eos) { + RowBatch row_batch; + ret = client->query_ctx->root->get_next(&state, &row_batch, &eos); + ASSERT_EQ(ret, 0); + for (int i = 0; i < row_batch.size(); ++i) { + std::cout << "tuple 0: " << row_batch.get_row(i)->debug_string(0) << std::endl; + std::cout << "tuple 1: " << row_batch.get_row(i)->debug_string(1) << std::endl; + } + } +} + +TEST(test_window, case_split_into_partitions) { + // 构造Tuple + pb::TupleDescriptor tuple_desc; + tuple_desc.set_tuple_id(0); + pb::SlotDescriptor* slot1 = tuple_desc.add_slots(); + slot1->set_tuple_id(0); + slot1->set_slot_id(1); + slot1->set_slot_type(pb::INT32); + pb::SlotDescriptor* slot2 = tuple_desc.add_slots(); + slot2->set_tuple_id(0); + slot2->set_slot_id(2); + slot2->set_slot_type(pb::INT32); + pb::SlotDescriptor* slot3 = tuple_desc.add_slots(); + slot3->set_tuple_id(0); + slot3->set_slot_id(3); + slot3->set_slot_type(pb::INT32); + std::vector tuple_descs {tuple_desc}; + + std::shared_ptr mem_row_desc = std::make_shared(); + mem_row_desc->init(tuple_descs); + + // 构造WindowNode + { + // OVER (PARTITION BY col1, col2) + pb::PlanNode pb_plan_node; + pb::DerivePlanNode* derive = pb_plan_node.mutable_derive_node(); + pb::WindowNode* window = derive->mutable_window_node(); + pb::WindowSpec* window_spec = window->mutable_window_spec(); + pb::Expr pb_expr1; + pb::ExprNode* pb_expr_node1 = pb_expr1.add_nodes(); + pb_expr_node1->set_node_type(pb::SLOT_REF); + pb_expr_node1->set_col_type(pb::INT32); + pb_expr_node1->mutable_derive_node()->set_tuple_id(0); + pb_expr_node1->mutable_derive_node()->set_slot_id(1); + window_spec->add_partition_exprs()->CopyFrom(pb_expr1); + pb::Expr pb_expr2; + pb::ExprNode* pb_expr_node2 = pb_expr2.add_nodes(); + pb_expr_node2->set_node_type(pb::SLOT_REF); + pb_expr_node2->set_col_type(pb::INT32); + pb_expr_node2->mutable_derive_node()->set_tuple_id(0); + pb_expr_node2->mutable_derive_node()->set_slot_id(2); + window_spec->add_partition_exprs()->CopyFrom(pb_expr2); + + WindowNode window_node; + window_node.init(pb_plan_node); + + RowBatch row_batch1; + construct_row_batch(mem_row_desc, row_batch1, 0); + bool is_first_partition_belong_to_prev = false; + window_node.split_into_partitions(&row_batch1, is_first_partition_belong_to_prev); + EXPECT_EQ(is_first_partition_belong_to_prev, false); + std::string expected_str = "[0,1),[1,3),[3,6),[6,10),[10,20),[20,26)"; + std::string str; + while (window_node.has_cache_partition()) { + int32_t start = -1; + int32_t end = -1; + window_node.get_next_partition(start, end); + str += "[" + std::to_string(start) + "," + std::to_string(end) + "),"; + } + str.pop_back(); + EXPECT_EQ(str, expected_str); + + RowBatch row_batch2; + construct_row_batch(mem_row_desc, row_batch2, 5); + window_node.split_into_partitions(&row_batch2, is_first_partition_belong_to_prev); + EXPECT_EQ(is_first_partition_belong_to_prev, true); + expected_str = "[0,6),[6,13),[13,21),[21,30),[30,40),[40,51)"; + str.clear(); + while (window_node.has_cache_partition()) { + int32_t start = -1; + int32_t end = -1; + window_node.get_next_partition(start, end); + str += "[" + std::to_string(start) + "," + std::to_string(end) + "),"; + } + str.pop_back(); + EXPECT_EQ(str, expected_str); + } + { + // OVER (PARTITION BY col1) + pb::PlanNode pb_plan_node; + pb::DerivePlanNode* derive = pb_plan_node.mutable_derive_node(); + pb::WindowNode* window = derive->mutable_window_node(); + pb::WindowSpec* window_spec = window->mutable_window_spec(); + pb::Expr pb_expr1; + pb::ExprNode* pb_expr_node1 = pb_expr1.add_nodes(); + pb_expr_node1->set_node_type(pb::SLOT_REF); + pb_expr_node1->set_col_type(pb::INT32); + pb_expr_node1->mutable_derive_node()->set_tuple_id(0); + pb_expr_node1->mutable_derive_node()->set_slot_id(1); + window_spec->add_partition_exprs()->CopyFrom(pb_expr1); + WindowNode window_node; + window_node.init(pb_plan_node); + + RowBatch row_batch1; + construct_row_batch(mem_row_desc, row_batch1, 0); + bool is_first_partition_belong_to_prev = false; + window_node.split_into_partitions(&row_batch1, is_first_partition_belong_to_prev); + EXPECT_EQ(is_first_partition_belong_to_prev, false); + std::string expected_str = "[0,26)"; + std::string str; + while (window_node.has_cache_partition()) { + int32_t start = -1; + int32_t end = -1; + window_node.get_next_partition(start, end); + str += "[" + std::to_string(start) + "," + std::to_string(end) + "),"; + } + str.pop_back(); + EXPECT_EQ(str, expected_str); + + RowBatch row_batch2; + construct_row_batch(mem_row_desc, row_batch2, 5); + window_node.split_into_partitions(&row_batch2, is_first_partition_belong_to_prev); + expected_str = "[0,51)"; + str.clear(); + EXPECT_EQ(is_first_partition_belong_to_prev, true); + while (window_node.has_cache_partition()) { + int32_t start = -1; + int32_t end = -1; + window_node.get_next_partition(start, end); + str += "[" + std::to_string(start) + "," + std::to_string(end) + "),"; + } + str.pop_back(); + EXPECT_EQ(str, expected_str); + } + { + // OVER () + pb::PlanNode pb_plan_node; + WindowNode window_node; + window_node.init(pb_plan_node); + + RowBatch row_batch1; + construct_row_batch(mem_row_desc, row_batch1, 0); + bool is_first_partition_belong_to_prev = false; + window_node.split_into_partitions(&row_batch1, is_first_partition_belong_to_prev); + EXPECT_EQ(is_first_partition_belong_to_prev, true); + std::string expected_str = "[0,26)"; + std::string str; + while (window_node.has_cache_partition()) { + int32_t start = -1; + int32_t end = -1; + window_node.get_next_partition(start, end); + str += "[" + std::to_string(start) + "," + std::to_string(end) + "),"; + } + str.pop_back(); + EXPECT_EQ(str, expected_str); + + RowBatch row_batch2; + construct_row_batch(mem_row_desc, row_batch2, 5); + window_node.split_into_partitions(&row_batch2, is_first_partition_belong_to_prev); + EXPECT_EQ(is_first_partition_belong_to_prev, true); + expected_str = "[0,51)"; + str.clear(); + while (window_node.has_cache_partition()) { + int32_t start = -1; + int32_t end = -1; + window_node.get_next_partition(start, end); + str += "[" + std::to_string(start) + "," + std::to_string(end) + "),"; + } + str.pop_back(); + EXPECT_EQ(str, expected_str); + } +} + +} // namespace baikaldb \ No newline at end of file