Add multilinestring_segment_manager for segment related methods in multilinestring ranges

cpp/include/cuspatial/detail/distance/linestring_polygon_distance.cuh

@@ -22,8 +22,6 @@
 #include <cuspatial/detail/algorithm/is_point_in_polygon.cuh>
 #include <cuspatial/detail/kernel/pairwise_distance.cuh>
 
-#include <thrust/fill.h>
-
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/exec_policy.hpp>
 
@@ -53,17 +51,23 @@ OutputIt pairwise_linestring_polygon_distance(MultiLinestringRange multilinestri
 
   auto polygons_as_linestrings = multipolygons.as_multilinestring_range();
 
-  thrust::fill(rmm::exec_policy(stream),
-               distances_first,
-               distances_first + size,
-               std::numeric_limits<T>::max());
+  thrust::transform(rmm::exec_policy(stream),
+                    multilinestrings.begin(),
+                    multilinestrings.end(),
+                    multipolygons.begin(),
+                    distances_first,
+                    [] __device__(auto multilinestring, auto multipolygon) {
+                      return (multilinestring.is_empty() || multipolygon.is_empty())
+                               ? std::numeric_limits<T>::quiet_NaN()
+                               : std::numeric_limits<T>::max();
+                    });
 
   auto [threads_per_block, num_blocks] = grid_1d(multilinestrings.num_points());
 
   detail::linestring_distance<<<num_blocks, threads_per_block, 0, stream.value()>>>(
     multilinestrings, polygons_as_linestrings, intersects.begin(), distances_first);
 
-  return distances_first + multilinestrings.num_multilinestrings();
+  return distances_first + size;
 }
 
 }  // namespace cuspatial

cpp/include/cuspatial/detail/functors.cuh

@@ -46,97 +46,5 @@ struct offset_pair_to_count_functor {
   }
 };
 
-/**
- * @brief Convert counts of points to counts of segments in a linestring.
- *
- * A Multilinestring is composed of a series of Linestrings. Each Linestring is composed of a
- * segments. The number of segments in a multilinestring is the number of points in the
- * multilinestring minus the number of linestrings.
- *
- * Caveats: This has a strong assumption that the Multilinestring does not contain empty
- * linestrings. While each non-empty linestring in the multilinestring represents 1 extra segment,
- * an empty multilinestring does not introduce any extra segments since it does not contain any
- * points.
- *
- * Used to create segment count iterators, such as `multi*_segment_count_begin`.
- *
- * @tparam IndexPair Must be iterator to a pair of counts
- * @param n_point_linestring_pair A pair of counts, the first is the number of points, the second is
- * the number of linestrings.
- */
-struct point_count_to_segment_count_functor {
-  template <typename IndexPair>
-  CUSPATIAL_HOST_DEVICE auto operator()(IndexPair n_point_linestring_pair)
-  {
-    auto nPoints      = thrust::get<0>(n_point_linestring_pair);
-    auto nLinestrings = thrust::get<1>(n_point_linestring_pair);
-    return nPoints - nLinestrings;
-  }
-};
-
-/**
- * @brief Given an offset iterator it, returns an iterator of the distance between it and an input
- * index i
- *
- * @tparam OffsetIterator Iterator type to the offset
- *
- * Caveats: This has a strong assumption that the Multilinestring does not contain empty
- * linestrings. While each non-empty linestring in the multilinestring represents 1 extra segment,
- * an empty multilinestring does not introduce any extra segments since it does not contain any
- * points.
- *
- * Used to create iterator to segment offsets, such as `segment_offset_begin`.
- */
-template <typename OffsetIterator>
-struct to_distance_iterator {
-  OffsetIterator begin;
-
-  template <typename IndexType>
-  CUSPATIAL_HOST_DEVICE auto operator()(IndexType i)
-  {
-    return begin[i] - i;
-  }
-};
-
-/// Deduction guide for to_distance_iterator
-template <typename OffsetIterator>
-to_distance_iterator(OffsetIterator) -> to_distance_iterator<OffsetIterator>;
-
-/**
- * @brief Return a segment from the a partitioned range of points
- *
- * Used in a counting transform iterator. Given an index of the segment, offset it by the number of
- * skipped segments preceding i in the partitioned range of points. Dereference the corresponding
- * point and the point following to make a segment.
- *
- * Used to create iterator to segments, such as `segment_begin`.
- *
- * @tparam OffsetIterator the iterator type indicating partitions of the point range.
- * @tparam CoordinateIterator the iterator type to the point range.
- */
-template <typename OffsetIterator, typename CoordinateIterator>
-struct to_valid_segment_functor {
-  using element_t = iterator_vec_base_type<CoordinateIterator>;
-
-  OffsetIterator begin;
-  OffsetIterator end;
-  CoordinateIterator point_begin;
-
-  template <typename IndexType>
-  CUSPATIAL_HOST_DEVICE segment<element_t> operator()(IndexType i)
-  {
-    auto kit = thrust::upper_bound(thrust::seq, begin, end, i);
-    auto k   = thrust::distance(begin, kit);
-    auto pid = i + k - 1;
-
-    return segment<element_t>{point_begin[pid], point_begin[pid + 1]};
-  }
-};
-
-/// Deduction guide for to_valid_segment_functor
-template <typename OffsetIterator, typename CoordinateIterator>
-to_valid_segment_functor(OffsetIterator, OffsetIterator, CoordinateIterator)
-  -> to_valid_segment_functor<OffsetIterator, CoordinateIterator>;
-
 }  // namespace detail
 }  // namespace cuspatial

cpp/include/cuspatial/detail/kernel/pairwise_distance.cuh

@@ -19,6 +19,7 @@
 #include <cuspatial/detail/utility/device_atomics.cuh>
 #include <cuspatial/detail/utility/linestring.cuh>
 
+#include <limits>
 #include <rmm/device_uvector.hpp>
 
 #include <thrust/optional.h>
@@ -40,6 +41,8 @@ namespace detail {
  * `intersects` is an optional pointer to a boolean range where the `i`th element indicates the
  * `i`th output should be set to 0 and bypass distance computation. This argument is optional, if
  * set to nullopt, no distance computation will be bypassed.
+ *
+ * @note This kernel does not compute pairs that contains empty geometry.
  */
 template <class MultiLinestringRange1, class MultiLinestringRange2, class OutputIt>
 __global__ void linestring_distance(MultiLinestringRange1 multilinestrings1,
@@ -55,11 +58,14 @@ __global__ void linestring_distance(MultiLinestringRange1 multilinestrings1,
     if (!multilinestrings1.is_valid_segment_id(idx, part_idx)) continue;
     auto const geometry_idx = multilinestrings1.geometry_idx_from_part_idx(part_idx);
 
+    if (multilinestrings1[geometry_idx].is_empty() || multilinestrings2[geometry_idx].is_empty()) {
+      continue;
+    }
+
     if (intersects.has_value() && intersects.value()[geometry_idx]) {
       distances_first[geometry_idx] = 0;
       continue;
     }
-
     auto [a, b]            = multilinestrings1.segment(idx);
     T min_distance_squared = std::numeric_limits<T>::max();
 

cpp/include/cuspatial/detail/multilinestring_segment.cuh

@@ -0,0 +1,128 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * 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 <cuspatial/detail/functors.cuh>
+#include <cuspatial/detail/range/multilinestring_segment_range.cuh>
+#include <cuspatial/detail/utility/zero_data.cuh>
+#include <cuspatial/iterator_factory.cuh>
+#include <cuspatial/range/range.cuh>
+#include <cuspatial/traits.hpp>
+
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/device_uvector.hpp>
+#include <rmm/exec_policy.hpp>
+#include <rmm/mr/device/per_device_resource.hpp>
+
+#include <thrust/device_vector.h>
+#include <thrust/iterator/discard_iterator.h>
+#include <thrust/iterator/transform_iterator.h>
+#include <thrust/logical.h>
+
+namespace cuspatial {
+namespace detail {
+
+/**
+ * @internal
+ * @brief Functor that returns true if current value is greater than 0.
+ */
+template <typename IndexType>
+struct greater_than_zero_functor {
+  __device__ IndexType operator()(IndexType x) const { return x > 0; }
+};
+
+/**
+ * @internal
+ * @brief Owning class to provide iterators to segments in a multilinestring range
+ *
+ * The owned memory in this struct is the vector `_non_empty_linestring_prefix_sum` of size equal to
+ * the number of linestrings in the multilinestring range plus 1. This vector holds the number
+ * of non empty linestrings that precedes the current linestring.
+ *
+ * This class is only meant for tracking the life time of the owned memory. To access the
+ * segment iterators, call `segment_range()` function to create a non-owning object of this class.
+ *
+ * For detailed explanation on the implementation of the segment iterators, see documentation
+ * of `multilinestring_segment_range`.
+ *
+ * @note To use this class with a multipolygon range, cast the multipolygon range as a
+ * multilinestring range.
+ *
+ * TODO: Optimization: for ranges that do not contain any empty linestrings,
+ * `_non_empty_linestring_prefix_sum` can replaced by a `counting_iterator`.
+ *
+ * @tparam MultilinestringRange The multilinestring range to initialize this class with.
+ */
+template <typename MultilinestringRange>
+class multilinestring_segment_manager {
+  using index_t = iterator_value_type<typename MultilinestringRange::part_it_t>;
+
+ public:
+  /**
+   * @brief Construct a new multilinestring segment object
+   *
+   * @note multilinestring_segment is always internal use, thus memory consumed is always
+   * temporary, therefore always use default device memory resource.
+   *
+   * @param parent The parent multilinestring object to construct from
+   * @param stream The stream to perform computation on
+   */
+  multilinestring_segment_manager(MultilinestringRange parent, rmm::cuda_stream_view stream)
+    : _parent(parent), _non_empty_linestring_prefix_sum(parent.num_linestrings() + 1, stream)
+  {
+    auto offset_range = ::cuspatial::range{_parent.part_offset_begin(), _parent.part_offset_end()};
+    auto count_begin  = thrust::make_transform_iterator(
+      thrust::make_zip_iterator(offset_range.begin(), thrust::next(offset_range.begin())),
+      offset_pair_to_count_functor{});
+
+    auto count_greater_than_zero =
+      thrust::make_transform_iterator(count_begin, greater_than_zero_functor<index_t>{});
+
+    zero_data_async(
+      _non_empty_linestring_prefix_sum.begin(), _non_empty_linestring_prefix_sum.end(), stream);
+
+    thrust::inclusive_scan(rmm::exec_policy(stream),
+                           count_greater_than_zero,
+                           count_greater_than_zero + _parent.num_linestrings(),
+                           thrust::next(_non_empty_linestring_prefix_sum.begin()));
+
+    _num_segments = _parent.num_points() - _non_empty_linestring_prefix_sum.element(
+                                             _non_empty_linestring_prefix_sum.size() - 1, stream);
+  }
+
+  /**
+   * @brief Return a non-owning `multilinestring_segment_range` object from this class
+   *
+   * @return multilinestring_segment_range
+   */
+  auto segment_range()
+  {
+    auto index_range = ::cuspatial::range{_non_empty_linestring_prefix_sum.begin(),
+                                          _non_empty_linestring_prefix_sum.end()};
+    return multilinestring_segment_range<MultilinestringRange, decltype(index_range)>{
+      _parent, index_range, _num_segments};
+  }
+
+ private:
+  MultilinestringRange _parent;
+  index_t _num_segments;
+  rmm::device_uvector<index_t> _non_empty_linestring_prefix_sum;
+};
+
+}  // namespace detail
+
+}  // namespace cuspatial