Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/include/index/rtree_module.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
88 changes: 8 additions & 80 deletions src/index/rtree_index_create_physical.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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<row_t>(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<LogicalType> 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<row_t> 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<string_t>(vec_format);
auto blob_data = stbox_data_ptr[vec_idx];
const uint8_t *stbox_bytes = reinterpret_cast<const uint8_t*>(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<mutex> 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) {
Expand Down
119 changes: 66 additions & 53 deletions src/index/rtree_module.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand All @@ -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;
Expand All @@ -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_) {
Expand Down Expand Up @@ -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<BoundReferenceExpression>(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<BoundReferenceExpression>(row_type, create_index.info->scan_types.size() - 1)
);

auto &projection = planner.Make<PhysicalProjection>(new_column_types, std::move(select_list),
auto &projection = planner.Make<PhysicalProjection>(new_column_types, std::move(select_list),
create_index.estimated_cardinality);
projection.children.push_back(input.table_scan);


auto &physical_create_index = planner.Make<PhysicalCreateTRTreeIndex>(
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;
}

//------------------------------------------------------------------------------
Expand Down Expand Up @@ -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<string_t>(vector)[i];
const uint8_t *data = reinterpret_cast<const uint8_t*>(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<string_t>(vector)[i];
const uint8_t *data = reinterpret_cast<const uint8_t *>(blob_data.GetData());
size_t data_size = blob_data.GetSize();

if (indexes_temporal) {
const Temporal *temp = reinterpret_cast<const Temporal *>(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<int>(row_data[i]));
free(box);
} else {
rtree_insert_temporal(rtree_, temp, static_cast<int>(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<int64_t>(row_data[i]));
rtree_insert(rtree_, target, static_cast<int>(row_data[i]));
free(box);
}
free(boxes);

if (boxes) free(boxes);
}


Expand Down Expand Up @@ -476,18 +499,8 @@ unique_ptr<ExpressionMatcher> TRTreeIndex::MakeFunctionMatcher() const {
matcher->expr_type = make_uniq<SpecificExpressionTypeMatcher>(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<ExpressionMatcher>();
lhs_matcher->type = make_uniq<SpecificTypeMatcher>(index_type);
lhs_matcher->type = make_uniq<SpecificTypeMatcher>(column_type_);
matcher->matchers.push_back(std::move(lhs_matcher));

// Right operand
Expand Down
22 changes: 22 additions & 0 deletions test/sql/parity/077_index_types.test
Original file line number Diff line number Diff line change
@@ -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
Loading