Sort the mergeable segments before computing merged segments

cpp/include/cuspatial/detail/find/find_and_combine_segment.cuh

@@ -19,18 +19,20 @@
 #include <cuspatial/cuda_utils.hpp>
 #include <cuspatial/detail/utility/linestring.cuh>
 #include <cuspatial/error.hpp>
+#include <cuspatial/iterator_factory.cuh>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/exec_policy.hpp>
 
-#include <thrust/uninitialized_fill.h>
+#include <thrust/sort.h>
 
 namespace cuspatial {
 namespace detail {
 
 /**
  * @internal
- * @brief Kernel to merge segments, naive n^2 algorithm.
+ * @brief Kernel to merge segments. Each thread works on one segment space,
+ * with a naive n^2 algorithm.
  */
 template <typename OffsetRange, typename SegmentRange, typename OutputIt>
 void __global__ simple_find_and_combine_segments_kernel(OffsetRange offsets,
@@ -44,6 +46,9 @@ void __global__ simple_find_and_combine_segments_kernel(OffsetRange offsets,
       merged_flag[i] = 0;
     }
 
+    // For each of the segment, loop over the rest of the segment in the space and see
+    // if it is mergeable with the current segment.
+    // Note that if the current segment is already merged. Skip checking.
     for (auto i = offsets[pair_idx]; i < offsets[pair_idx + 1] && merged_flag[i] != 1; i++) {
       for (auto j = i + 1; j < offsets[pair_idx + 1]; j++) {
         auto res = maybe_merge_segments(segments[i], segments[j]);
@@ -57,20 +62,66 @@ void __global__ simple_find_and_combine_segments_kernel(OffsetRange offsets,
   }
 }
 
+/**
+ * @brief Comparator for sorting the segment range.
+ *
+ * This comparator makes sure that the segment range are sorted by the following keys:
+ * 1. Segments with the same space id are grouped together.
+ * 2. Segments within the same space are grouped by their slope.
+ * 3. Within each slope group, segments are sorted by their lower left point.
+ */
+template <typename index_t, typename T>
+struct segment_comparator {
+  bool __device__ operator()(thrust::tuple<index_t, segment<T>> const& lhs,
+                             thrust::tuple<index_t, segment<T>> const& rhs) const
+  {
+    auto lhs_index   = thrust::get<0>(lhs);
+    auto rhs_index   = thrust::get<0>(rhs);
+    auto lhs_segment = thrust::get<1>(lhs);
+    auto rhs_segment = thrust::get<1>(rhs);
+
+    // Compare space id
+    if (lhs_index == rhs_index) {
+      // Compare slope
+      if (lhs_segment.collinear(rhs_segment)) {
+        // Sort by the lower left point of the segment.
+        return lhs_segment.lower_left() < rhs_segment.lower_left();
+      }
+      return lhs_segment.slope() < rhs_segment.slope();
+    }
+    return lhs_index < rhs_index;
+  }
+};
+
 /**
  * @internal
  * @brief For each pair of mergeable segment, overwrites the first segment with merged result,
  * sets the flag for the second segment as 1.
+ *
+ * @note This function will alter the input segment range by rearranging the order of the segments
+ * within each space so that merging kernel can take place.
  */
 template <typename OffsetRange, typename SegmentRange, typename OutputIt>
 void find_and_combine_segment(OffsetRange offsets,
                               SegmentRange segments,
                               OutputIt merged_flag,
                               rmm::cuda_stream_view stream)
 {
+  using index_t   = typename OffsetRange::value_type;
+  using T         = typename SegmentRange::value_type::value_type;
   auto num_spaces = offsets.size() - 1;
   if (num_spaces == 0) return;
 
+  // Construct a key iterator using the offsets of the segment and the segment itself.
+  auto space_id_iter         = make_geometry_id_iterator<index_t>(offsets.begin(), offsets.end());
+  auto space_id_segment_iter = thrust::make_zip_iterator(space_id_iter, segments.begin());
+
+  thrust::sort_by_key(rmm::exec_policy(stream),
+                      space_id_segment_iter,
+                      space_id_segment_iter + segments.size(),
+                      segments.begin(),
+                      segment_comparator<index_t, T>{});
+
   auto [threads_per_block, num_blocks] = grid_1d(num_spaces);
   simple_find_and_combine_segments_kernel<<<num_blocks, threads_per_block, 0, stream.value()>>>(
     offsets, segments, merged_flag);

cpp/include/cuspatial/geometry/segment.cuh

@@ -55,6 +55,17 @@ class alignas(sizeof(Vertex)) segment {
   /// Return the length squared of segment.
   T CUSPATIAL_HOST_DEVICE length2() const { return dot(v2 - v1, v2 - v1); }
 
+  /// Return slope of segment.
+  T CUSPATIAL_HOST_DEVICE slope() { return (v2.y - v1.y) / (v2.x - v1.x); }
+
+  /// Return the lower left vertex of segment.
+  Vertex CUSPATIAL_HOST_DEVICE lower_left() { return v1 < v2 ? v1 : v2; }
+
+  bool CUSPATIAL_HOST_DEVICE collinear(segment<T> const& other)
+  {
+    return (v1.x - v1.y) * (other.v2.x - other.v2.y) == (v2.x - v2.y) * (other.v1.x - other.v1.y);
+  }
+
  private:
   friend std::ostream& operator<<(std::ostream& os, segment<T> const& seg)
   {

cpp/tests/find/find_and_combine_segments_test.cu

@@ -281,3 +281,65 @@ TYPED_TEST(FindAndCombineSegmentsTest, nooverlap3)
                      {0, 0},
                      {S{P{0.0, 0.0}, P{1.0, 1.0}}, S{P{0.0, 1.0}, P{1.0, 0.0}}});
 }
+
+TYPED_TEST(FindAndCombineSegmentsTest, twospaces)
+{
+  using T       = TypeParam;
+  using index_t = std::size_t;
+  using P       = vec_2d<T>;
+  using S       = segment<T>;
+
+  auto segments = make_segment_array<index_t, T>({0, 2, 4},
+                                                 {S{P{0.0, 0.0}, P{1.0, 1.0}},
+                                                  S{P{1.0, 1.0}, P{2.0, 2.0}},
+                                                  S{P{1.0, 1.0}, P{0.0, 0.0}},
+                                                  S{P{2.0, 2.0}, P{1.0, 1.0}}});
+
+  CUSPATIAL_RUN_TEST(this->run_single_test,
+                     segments,
+                     {0, 1, 0, 1},
+                     {S{P{0.0, 0.0}, P{2.0, 2.0}},
+                      S{P{1.0, 1.0}, P{2.0, 2.0}},
+                      S{P{0.0, 0.0}, P{2.0, 2.0}},
+                      S{P{2.0, 2.0}, P{1.0, 1.0}}});
+}
+
+TYPED_TEST(FindAndCombineSegmentsTest, twospaces_non_contiguous_segments_with_empty)
+{
+  using T       = TypeParam;
+  using index_t = std::size_t;
+  using P       = vec_2d<T>;
+  using S       = segment<T>;
+
+  auto segments = make_segment_array<index_t, T>({0, 4, 4},
+                                                 {S{P{1.0, 1.0}, P{2.0, 2.0}},
+                                                  S{P{3.0, 3.0}, P{4.0, 4.0}},
+                                                  S{P{0.0, 0.0}, P{1.0, 1.0}},
+                                                  S{P{2.0, 2.0}, P{3.0, 3.0}}});
+
+  CUSPATIAL_RUN_TEST(this->run_single_test,
+                     segments,
+                     {0, 1, 1, 1},
+                     {S{P{0.0, 0.0}, P{4.0, 4.0}},
+                      S{P{1.0, 1.0}, P{2.0, 2.0}},
+                      S{P{2.0, 2.0}, P{3.0, 3.0}},
+                      S{P{3.0, 3.0}, P{4.0, 4.0}}});
+}
+
+TYPED_TEST(FindAndCombineSegmentsTest, onespace_non_contiguous_segments_overlaps)
+{
+  using T       = TypeParam;
+  using index_t = std::size_t;
+  using P       = vec_2d<T>;
+  using S       = segment<T>;
+
+  auto segments = make_segment_array<index_t, T>(
+    {0, 3},
+    {S{P{1.0, 1.0}, P{2.0, 2.0}}, S{P{4.0, 4.0}, P{5.0, 5.0}}, S{P{-1.0, -1.0}, P{4.0, 4.0}}});
+
+  CUSPATIAL_RUN_TEST(
+    this->run_single_test,
+    segments,
+    {0, 1, 1},
+    {S{P{-1.0, -1.0}, P{5.0, 5.0}}, S{P{1.0, 1.0}, P{2.0, 2.0}}, S{P{4.0, 4.0}, P{5.0, 5.0}}});
+}

cpp/tests/intersection/linestring_intersection_large_test.cu

@@ -27,6 +27,7 @@
 #include <cuspatial/geometry/vec_2d.hpp>
 #include <cuspatial/intersection.cuh>
 
+#include <rmm/exec_policy.hpp>
 #include <thrust/host_vector.h>
 
 template <typename T>