From c91ea02789a58b4d353a0cf56fb5029e87dbbb45 Mon Sep 17 00:00:00 2001 From: Ngoc Hoang Date: Wed, 15 Jul 2026 11:05:52 +0000 Subject: [PATCH] Added rtree support for tgeompoint with overlaps predicate --- src/include/index/rtree_module.hpp | 1 + src/index/rtree_index_create_physical.cpp | 88 ++-------------- src/index/rtree_module.cpp | 119 ++++++++++++---------- test/sql/parity/077_index_types.test | 22 ++++ 4 files changed, 97 insertions(+), 133 deletions(-) create mode 100644 test/sql/parity/077_index_types.test diff --git a/src/include/index/rtree_module.hpp b/src/include/index/rtree_module.hpp index 5cb45811..9edb258e 100644 --- a/src/include/index/rtree_module.hpp +++ b/src/include/index/rtree_module.hpp @@ -86,6 +86,7 @@ class TRTreeIndex : public BoundIndex { meosType bbox_type_; size_t bbox_size_; + LogicalType column_type_; size_t current_size_ = 0; size_t current_capacity_ = 0; diff --git a/src/index/rtree_index_create_physical.cpp b/src/index/rtree_index_create_physical.cpp index 8ea0cb61..190c8533 100644 --- a/src/index/rtree_index_create_physical.cpp +++ b/src/index/rtree_index_create_physical.cpp @@ -153,8 +153,6 @@ class TRTreeIndexConstructTask final : public ExecutorTask { auto &vec_vec = scan_chunk.data[0]; auto &rowid_vec = scan_chunk.data[1]; - - auto vector_type = vec_vec.GetType(); if (vec_vec.GetVectorType() != VectorType::FLAT_VECTOR) { vec_vec.Flatten(count); @@ -163,87 +161,17 @@ class TRTreeIndexConstructTask final : public ExecutorTask { rowid_vec.Flatten(count); } - UnifiedVectorFormat vec_format; - UnifiedVectorFormat rowid_format; - - vec_vec.ToUnifiedFormat(count, vec_format); - rowid_vec.ToUnifiedFormat(count, rowid_format); - - const auto row_ptr = UnifiedVectorFormat::GetData(rowid_format); - STBox* boxes = (STBox*)malloc(sizeof(STBox) * count); - if (!boxes) { - executor.PushError(ErrorData("Failed to allocate memory for STBox array")); - return TaskExecutionResult::TASK_ERROR; - } + DataChunk col_chunk; + vector col_types = {vec_vec.GetType()}; + col_chunk.Initialize(Allocator::DefaultAllocator(), col_types); + col_chunk.SetCardinality(count); + col_chunk.data[0].Reference(vec_vec); - idx_t valid_count = 0; - vector valid_row_ids; - - for (idx_t i = 0; i < count; i++) { - const auto vec_idx = vec_format.sel->get_index(i); - const auto row_idx = rowid_format.sel->get_index(i); - - const auto vec_valid = vec_format.validity.RowIsValid(vec_idx); - const auto rowid_valid = rowid_format.validity.RowIsValid(row_idx); - - if (!vec_valid || !rowid_valid) { - continue; - } - - fprintf(stderr, "Processing row %zu (vec_idx=%zu, row_idx=%zu)\n", i, vec_idx, row_idx); - - STBox *box = nullptr; - - if (vector_type.id() == LogicalTypeId::BLOB) { - - const auto stbox_data_ptr = UnifiedVectorFormat::GetData(vec_format); - auto blob_data = stbox_data_ptr[vec_idx]; - const uint8_t *stbox_bytes = reinterpret_cast(blob_data.GetData()); - size_t stbox_size = blob_data.GetSize(); - box = (STBox*)malloc(stbox_size); - memcpy(box, stbox_bytes, stbox_size); - - int32_t box_srid = stbox_srid(box); - - if (box_srid != 0) { - STBox *normalized_box = stbox_set_srid(box, 0); - if (normalized_box) { - free(box); - box = normalized_box; - } - } - - // Copy to our batch array - memcpy(&boxes[valid_count], box, sizeof(STBox)); - valid_row_ids.push_back(row_ptr[row_idx]); - valid_count++; - - free(box); - - - } - else { - free(boxes); - executor.PushError(ErrorData("Unsupported data type for RTree index: " + vector_type.ToString())); - return TaskExecutionResult::TASK_ERROR; - } + { + lock_guard l(gstate.glock); + gstate.global_index->Construct(col_chunk, rowid_vec); } - // Now batch insert the valid STBoxes into the index - if (valid_count > 0) { - auto &rtree_index = gstate.global_index; - - auto result = rtree_index->BulkConstruct(boxes, valid_row_ids.data(), valid_count); - if (result.HasError()) { - free(boxes); - executor.PushError(result); - return TaskExecutionResult::TASK_ERROR; - } - - } - - free(boxes); - gstate.built_count += count; if (mode == TaskExecutionMode::PROCESS_PARTIAL) { diff --git a/src/index/rtree_module.cpp b/src/index/rtree_module.cpp index d05e4be8..ed6349c7 100644 --- a/src/index/rtree_module.cpp +++ b/src/index/rtree_module.cpp @@ -27,6 +27,7 @@ #include "duckdb/optimizer/matcher/expression_matcher.hpp" #include "index/rtree_module.hpp" #include "geo/stbox.hpp" +#include "geo/tgeompoint.hpp" #include "index/rtree_index_create_physical.hpp" #include "time_util.hpp" @@ -48,6 +49,9 @@ TRTreeIndex::TRTreeIndex(const string &name, IndexConstraintType constraint_type auto &type = unbound_expressions[0]->return_type; + column_type_ = type; + + // R-tree's bbox type is determined by the type of the indexed column if (type == StboxType::STBOX()) { bbox_type_ = T_STBOX; @@ -57,8 +61,12 @@ TRTreeIndex::TRTreeIndex(const string &name, IndexConstraintType constraint_type bbox_type_ = T_TSTZSPAN; bbox_size_ = sizeof(Span); rtree_ = rtree_create_tstzspan(); + } else if (type == TgeompointType::TGEOMPOINT()) { + bbox_type_ = T_STBOX; + bbox_size_ = sizeof(STBox); + rtree_ = rtree_create_stbox(); } else { - throw InternalException("RTree index only supports STBOX and TSTZSPAN types, got: " + type.ToString()); + throw InternalException("RTree index only supports STBOX, TSTZSPAN, and TGEOMPOINT types, got: " + type.ToString()); } if (!rtree_) { @@ -103,23 +111,23 @@ PhysicalOperator &TRTreeIndex::CreatePlan(PlanIndexInput &input) { select_list.push_back(std::move(expression)); } - // new_column_types.emplace_back(LogicalType::ROW_TYPE); - // select_list.push_back( - // make_uniq(LogicalType::ROW_TYPE, create_index.info->scan_types.size() - 1)); + LogicalType row_type = LogicalType::ROW_TYPE; + new_column_types.push_back(row_type); + select_list.push_back( + make_uniq(row_type, create_index.info->scan_types.size() - 1) + ); - auto &projection = planner.Make(new_column_types, std::move(select_list), + auto &projection = planner.Make(new_column_types, std::move(select_list), create_index.estimated_cardinality); projection.children.push_back(input.table_scan); - auto &physical_create_index = planner.Make( - create_index.types, create_index.table, create_index.info->column_ids, - std::move(create_index.info), std::move(create_index.unbound_expressions), + create_index.types, create_index.table, create_index.info->column_ids, + std::move(create_index.info), std::move(create_index.unbound_expressions), create_index.estimated_cardinality); - + physical_create_index.children.push_back(projection); return physical_create_index; - return input.table_scan; } //------------------------------------------------------------------------------ @@ -224,58 +232,73 @@ void TRTreeIndex::Construct(DataChunk &expression_result, Vector &row_identifier if (vector.GetVectorType() != VectorType::FLAT_VECTOR) { vector.Flatten(expression_result.size()); } - - auto vector_type = vector.GetType(); - - void* boxes = malloc(bbox_size_ * expression_result.size()); + const bool indexes_temporal = column_type_ == TgeompointType::TGEOMPOINT(); + + void *boxes = indexes_temporal ? nullptr : malloc(bbox_size_ * expression_result.size()); for (idx_t i = 0; i < expression_result.size(); i++) { if (FlatVector::IsNull(vector, i)) { - continue; + continue; } - void *box = nullptr; - - if (vector_type.id() == LogicalTypeId::BLOB) { - auto blob_data = FlatVector::GetData(vector)[i]; - const uint8_t *data = reinterpret_cast(blob_data.GetData()); - size_t data_size = blob_data.GetSize(); - - - if (data_size != bbox_size_) { - continue; - } - - box = malloc(data_size); - memcpy(box, data, data_size); + if (vector.GetType().id() != LogicalTypeId::BLOB) { + continue; + } + + auto blob_data = FlatVector::GetData(vector)[i]; + const uint8_t *data = reinterpret_cast(blob_data.GetData()); + size_t data_size = blob_data.GetSize(); + if (indexes_temporal) { + const Temporal *temp = reinterpret_cast(data); if (bbox_type_ == T_STBOX) { - STBox *stbox = (STBox*)box; - int32_t box_srid = stbox_srid(stbox); - if (box_srid != 0) { - STBox *normalized_box = stbox_set_srid(stbox, 0); - if (normalized_box) { + STBox *box = tspatial_to_stbox(temp); + if (!box) { + continue; + } + + if (stbox_srid(box) != 0) { + STBox *normalized = stbox_set_srid(box, 0); + if (normalized) { free(box); - box = normalized_box; + box = normalized; } } + rtree_insert(rtree_, box, static_cast(row_data[i])); + free(box); + } else { + rtree_insert_temporal(rtree_, temp, static_cast(row_data[i])); } - } else { continue; } - if (box == nullptr) { + if (data_size != bbox_size_) { continue; } - - void* target = (char*)boxes + (i * bbox_size_); + + void *box = malloc(data_size); + memcpy(box, data, data_size); + + if (bbox_type_ == T_STBOX) { + STBox *stbox = (STBox *) box; + int32_t box_srid = stbox_srid(stbox); + if (box_srid != 0) { + STBox *normalized_box = stbox_set_srid(stbox, 0); + if (normalized_box) { + free(box); + box = normalized_box; + } + } + } + + void *target = (char *) boxes + (i * bbox_size_); memcpy(target, box, bbox_size_); - rtree_insert(rtree_, target, static_cast(row_data[i])); + rtree_insert(rtree_, target, static_cast(row_data[i])); free(box); } - - free(boxes); + + if (boxes) free(boxes); } @@ -476,18 +499,8 @@ unique_ptr TRTreeIndex::MakeFunctionMatcher() const { matcher->expr_type = make_uniq(ExpressionType::BOUND_FUNCTION); matcher->policy = SetMatcher::Policy::UNORDERED; - LogicalType index_type; - if (bbox_type_ == T_STBOX) { - index_type = StboxType::STBOX(); - } else if (bbox_type_ == T_TSTZSPAN) { - index_type = SpanTypes::TSTZSPAN(); - } else { - index_type = LogicalType::BLOB; - } - - // Left operand auto lhs_matcher = make_uniq(); - lhs_matcher->type = make_uniq(index_type); + lhs_matcher->type = make_uniq(column_type_); matcher->matchers.push_back(std::move(lhs_matcher)); // Right operand diff --git a/test/sql/parity/077_index_types.test b/test/sql/parity/077_index_types.test new file mode 100644 index 00000000..d8a7aacc --- /dev/null +++ b/test/sql/parity/077_index_types.test @@ -0,0 +1,22 @@ +# name: test/sql/parity/077_index_types.test +# description: TRTREE index coverage for tgeompoint +# group: [sql] + +require mobilityduck + +statement ok +CREATE TABLE idx_tgeompoint(t tgeompoint); + +statement ok +INSERT INTO idx_tgeompoint VALUES + ('Point(0 0)@2000-01-01'::tgeompoint), + ('Point(10 10)@2001-01-01'::tgeompoint); + +statement ok +CREATE INDEX i_tgeompoint ON idx_tgeompoint USING TRTREE (t); + +query I +SELECT count(*) FROM idx_tgeompoint + WHERE t && 'STBOX XT(((-1,-1),(1,1)),[2000-01-01, 2000-01-02])'::stbox; +---- +1 \ No newline at end of file