tweak spatial join

Author: Maxxen

src/spatial/operators/spatial_join_physical.cpp

@@ -21,6 +21,8 @@ namespace duckdb {
 // Flat RTree
 //======================================================================================================================
 
+// #define DUCKDB_SPATIAL_DEBUG_JOIN
+
 namespace {
 
 template <class T>
@@ -132,11 +134,12 @@ class FlatRTree {
 		return current_position++;
 	}
 
-	void Sort(vector<uint32_t> &curve) {
-		Sort(curve, 0, curve.size() - 1);
+	static void Sort(vector<uint32_t> &curve, typed_view<Box> &box_array, typed_view<uint32_t> &idx_array) {
+		Sort(curve, box_array, idx_array, 0, curve.size() - 1);
 	}
 
-	void Sort(vector<uint32_t> &curve, size_t l_idx, size_t r_idx) {
+	static void Sort(vector<uint32_t> &curve, typed_view<Box> &box_array, typed_view<uint32_t> &idx_array, size_t l_idx,
+	                 size_t r_idx) {
 		if (l_idx < r_idx) {
 			const auto pivot = curve[(l_idx + r_idx) >> 1];
 			auto pivot_l = l_idx - 1;
@@ -161,8 +164,55 @@ class FlatRTree {
 				std::swap(idx_array[pivot_l], idx_array[pivot_r]);
 			}
 
-			Sort(curve, l_idx, pivot_r);
-			Sort(curve, pivot_r + 1, r_idx);
+			Sort(curve, box_array, idx_array, l_idx, pivot_r);
+			Sort(curve, box_array, idx_array, pivot_r + 1, r_idx);
+		}
+	}
+
+	void STRSort(typed_view<Box> &box_array, typed_view<uint32_t> idx_array) {
+		// Perform Sort-tile-recursive (STR) packing
+
+		const auto num_leaf_nodes = (item_count + node_size - 1) / node_size;
+		const auto num_vertical_slices = static_cast<uint32_t>(std::ceil(std::sqrt(num_leaf_nodes)));
+		const auto slice_size = (item_count + num_vertical_slices - 1) / num_vertical_slices;
+
+		vector<uint32_t> indexes;
+		for (uint32_t i = 0; i < item_count; i++) {
+			indexes.push_back(i);
+		}
+
+		// Sort by x-axis into vertical slices
+		std::sort(indexes.begin(), indexes.end(),
+		          [&](uint32_t a, uint32_t b) { return box_array[a].Center().x < box_array[b].Center().x; });
+
+		// Then sort each vertical slice by y-axis
+		for (uint32_t slice_idx = 0; slice_idx < num_vertical_slices; slice_idx++) {
+			const auto slice_beg = slice_idx * slice_size;
+			const auto slice_end = MinValue<uint32_t>(slice_beg + slice_size, item_count);
+			std::sort(indexes.begin() + slice_beg, indexes.begin() + slice_end,
+			          [&](uint32_t a, uint32_t b) { return box_array[a].Center().y < box_array[b].Center().y; });
+		}
+
+		// Reorder the box_array and idx_array based on the sorted indexes
+		// DO this in-place. There cannot be any cycles since all indexes are unique
+		for (int i = 0; i < item_count; i++) {
+			if (indexes[i] == -1)
+				continue; // index `i` has been processed, skip
+			auto box = box_array[i];
+			auto idx = idx_array[i];
+
+			int x = i, y = indexes[i]; // `x` is the current index, `y` is the "target" index
+			while (y != i) {
+				indexes[x] = -1; // mark index as processed
+				box_array[x] = box_array[y];
+				idx_array[x] = idx_array[y];
+				x = y;
+				y = indexes[x];
+			}
+			// Now `x` is the index that satisfies `indices[x] == i`.
+			box_array[x] = box;
+			idx_array[x] = idx;
+			indexes[x] = -1;
 		}
 	}
 
@@ -176,6 +226,7 @@ class FlatRTree {
 
 		// Generate hilbert curve values
 		// TODO: Parallelize this with tasks when the number of items is large?
+
 		constexpr auto max_hilbert = std::numeric_limits<uint16_t>::max();
 		const auto hw = max_hilbert / (tree_box.max.x - tree_box.min.x);
 		const auto hh = max_hilbert / (tree_box.max.y - tree_box.min.y);
@@ -191,7 +242,8 @@ class FlatRTree {
 		}
 
 		// Now, sort the indices based on their curve value
-		Sort(curve);
+		Sort(curve, box_array, idx_array);
+		//STRSort(box_array, idx_array);
 
 		size_t layer_idx = 0;
 		size_t entry_idx = 0;
@@ -723,11 +775,17 @@ class SpatialJoinGlobalOperatorState final : public GlobalOperatorState {
 	// TODO: Move this into proper profiling metrics later
 	// Statistics
 	atomic<idx_t> total_rtree_probes = {0};
+	atomic<idx_t> total_rtree_successfull_probes = {0};
 	atomic<idx_t> total_rtree_candidates = {0};
+	atomic<idx_t> max_candidates = {0};
+	atomic<idx_t> min_candidates = {std::numeric_limits<idx_t>::max()};
 
 	~SpatialJoinGlobalOperatorState() override {
-		Printer::PrintF("Spatial Join Stats: RTree Probes: %llu, RTree Candidates: %llu\n", total_rtree_probes.load(),
-		                total_rtree_candidates.load());
+		Printer::PrintF("Spatial Join RTree Probes: %llu\n", total_rtree_probes.load());
+		Printer::PrintF("Spatial Join RTree Successful Probes: %llu\n", total_rtree_successfull_probes.load());
+		Printer::PrintF("Spatial Join RTree Candidates: %llu\n", total_rtree_candidates.load());
+		Printer::PrintF("Spatial Join RTree Max Candidates per Probe: %llu\n", max_candidates.load());
+		Printer::PrintF("Spatial Join RTree Min Candidates per Probe: %llu\n", min_candidates.load());
 	}
 #endif
 };
@@ -843,14 +901,20 @@ OperatorResultType PhysicalSpatialJoin::ExecuteInternal(ExecutionContext &contex
 
 			gstate.rtree->InitScan(lstate.scan, bbox);
 
+#ifdef DUCKDB_SPATIAL_DEBUG_JOIN
+			gstate.total_rtree_probes += 1;
+#endif
+
 			if (!gstate.rtree->Scan(lstate.scan)) {
 				lstate.input_index++;
 				continue;
 			}
 
 #ifdef DUCKDB_SPATIAL_DEBUG_JOIN
+			gstate.total_rtree_successfull_probes += 1;
 			gstate.total_rtree_candidates += lstate.scan.matches_count;
-			gstate.total_rtree_probes += 1;
+			gstate.max_candidates = MaxValue(gstate.max_candidates.load(), lstate.scan.matches_count);
+			gstate.min_candidates = MinValue(gstate.min_candidates.load(), lstate.scan.matches_count);
 #endif
 
 			lstate.state = SpatialJoinState::SCAN;
@@ -1042,25 +1106,34 @@ class SpatialJoinGlobalSourceState final : public GlobalSourceState {
 		column_ids.push_back(op.build_side_key_types.size() + op.build_side_payload_types.size());
 
 		// We dont need to keep the tuples aroun after scanning
-		state.collection->InitializeScan(scan_state, std::move(column_ids),
-		                                 TupleDataPinProperties::KEEP_EVERYTHING_PINNED);
+		state.collection->InitializeScan(scan_state, std::move(column_ids), TupleDataPinProperties::UNPIN_AFTER_DONE);
 
 		tuples_maximum = state.collection->Count();
 	}
 
 	const PhysicalSpatialJoin &op;
+
+	mutex scan_lock;
 	TupleDataParallelScanState scan_state;
+
 	// How many tuples we have scanned so far
 	idx_t tuples_maximum = 0;
 	atomic<idx_t> tuples_scanned = {0};
 
 public:
 	idx_t MaxThreads() override {
 		const auto &state = op.op_state->Cast<SpatialJoinGlobalOperatorState>();
-		const auto count = state.collection->Count();
+		return state.collection->ChunkCount();
+	}
+
+	bool Scan(TupleDataLocalScanState &local_scan, DataChunk &chunk) {
+		const auto &collection = op.op_state->Cast<SpatialJoinGlobalOperatorState>().collection;
 
-		// Rough approximation of the number of threads to use
-		return count / (STANDARD_VECTOR_SIZE * 10ULL);
+		lock_guard<mutex> guard(scan_lock);
+		const auto not_empty = collection->Scan(scan_state, local_scan, chunk);
+		tuples_scanned += chunk.size();
+
+		return not_empty;
 	}
 };
 
@@ -1106,45 +1179,41 @@ SourceResultType PhysicalSpatialJoin::GetDataInternal(ExecutionContext &context,
 	auto &gstate = input.global_state.Cast<SpatialJoinGlobalSourceState>();
 	auto &lstate = input.local_state.Cast<SpatialJoinLocalSourceState>();
 
-	const auto &tuples = gstate.op.op_state->Cast<SpatialJoinGlobalOperatorState>().collection;
-
-	while (tuples->Scan(gstate.scan_state, lstate.scan_state, lstate.scan_chunk)) {
-		gstate.tuples_scanned += lstate.scan_chunk.size();
+	if (!gstate.Scan(lstate.scan_state, lstate.scan_chunk)) {
+		return SourceResultType::FINISHED;
+	}
 
-		const auto matches = FlatVector::GetData<bool>(lstate.scan_chunk.data.back());
+	const auto matches = FlatVector::GetData<bool>(lstate.scan_chunk.data.back());
 
-		idx_t result_count = 0;
-		for (idx_t i = 0; i < lstate.scan_chunk.size(); i++) {
-			if (!matches[i]) {
-				lstate.match_sel.set_index(result_count++, i);
-			}
+	idx_t result_count = 0;
+	for (idx_t i = 0; i < lstate.scan_chunk.size(); i++) {
+		if (!matches[i]) {
+			lstate.match_sel.set_index(result_count++, i);
 		}
+	}
 
-		if (result_count > 0) {
+	if (result_count > 0) {
 
-			const auto lhs_col_count = probe_side_output_columns.size();
-			const auto rhs_col_count = build_side_output_columns.size();
+		const auto lhs_col_count = probe_side_output_columns.size();
+		const auto rhs_col_count = build_side_output_columns.size();
 
-			// Null the LHS columns
-			for (idx_t i = 0; i < lhs_col_count; i++) {
-				auto &target = chunk.data[i];
-				target.SetVectorType(VectorType::CONSTANT_VECTOR);
-				ConstantVector::SetNull(target, true);
-			}
-
-			// Set the RHS columns
-			for (idx_t i = 0; i < rhs_col_count; i++) {
-				auto &target = chunk.data[lhs_col_count + i];
-				// Offset by one here to skip the match column
-				target.Slice(lstate.scan_chunk.data[i], lstate.match_sel, result_count);
-			}
+		// Null the LHS columns
+		for (idx_t i = 0; i < lhs_col_count; i++) {
+			auto &target = chunk.data[i];
+			target.SetVectorType(VectorType::CONSTANT_VECTOR);
+			ConstantVector::SetNull(target, true);
+		}
 
-			chunk.SetCardinality(result_count);
-			return SourceResultType::HAVE_MORE_OUTPUT;
+		// Set the RHS columns
+		for (idx_t i = 0; i < rhs_col_count; i++) {
+			auto &target = chunk.data[lhs_col_count + i];
+			// Offset by one here to skip the match column
+			target.Slice(lstate.scan_chunk.data[i], lstate.match_sel, result_count);
 		}
 	}
 
-	return SourceResultType::FINISHED;
+	chunk.SetCardinality(result_count);
+	return SourceResultType::HAVE_MORE_OUTPUT;
 }
 
 //----------------------------------------------------------------------------------------------------------------------