row_operators.cuh

/*
 * Copyright (c) 2019-2024, 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
 *
 * https://apache.ac.cn/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 <cudf/column/column_device_view.cuh>
#include <cudf/detail/utilities/assert.cuh>
#include <cudf/hashing/detail/hash_functions.cuh>
#include <cudf/hashing/detail/hashing.hpp>
#include <cudf/table/table_device_view.cuh>
#include <cudf/utilities/traits.hpp>
#include <cudf/utilities/type_dispatcher.hpp>

#include <cuda/std/limits>
#include <thrust/equal.h>
#include <thrust/execution_policy.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/transform_reduce.h>

namespace CUDF_EXPORT cudf {

enum class weak_ordering {
 LESS,
 EQUIVALENT,
 GREATER
};

namespace detail {
template <typename Element>
__device__ weak_ordering compare_elements(Element lhs, Element rhs)
{
 if (lhs < rhs) {
 return weak_ordering::LESS;
 } else if (rhs < lhs) {
 return weak_ordering::GREATER;
 }
 return weak_ordering::EQUIVALENT;
}
} // namespace detail

template <typename Element, std::enable_if_t<std::is_floating_point_v<Element>>* = nullptr>
__device__ weak_ordering relational_compare(Element lhs, Element rhs)
{
 if (isnan(lhs) and isnan(rhs)) {
 return weak_ordering::EQUIVALENT;
 } else if (isnan(rhs)) {
 return weak_ordering::LESS;
 } else if (isnan(lhs)) {
 return weak_ordering::GREATER;
 }

 return detail::compare_elements(lhs, rhs);
}

inline __device__ auto null_compare(bool lhs_is_null, bool rhs_is_null, null_order null_precedence)
{
 if (lhs_is_null and rhs_is_null) { // null <? null
 return weak_ordering::EQUIVALENT;
 } else if (lhs_is_null) { // null <? x
 return (null_precedence == null_order::BEFORE) ? weak_ordering::LESS : weak_ordering::GREATER;
 } else if (rhs_is_null) { // x <? null
 return (null_precedence == null_order::AFTER) ? weak_ordering::LESS : weak_ordering::GREATER;
 }
 return weak_ordering::EQUIVALENT;
}

template <typename Element, std::enable_if_t<not std::is_floating_point_v<Element>>* = nullptr>
__device__ weak_ordering relational_compare(Element lhs, Element rhs)
{
 return detail::compare_elements(lhs, rhs);
}

template <typename Element, std::enable_if_t<std::is_floating_point_v<Element>>* = nullptr>
__device__ bool equality_compare(Element lhs, Element rhs)
{
 if (isnan(lhs) and isnan(rhs)) { return true; }
 return lhs == rhs;
}

template <typename Element, std::enable_if_t<not std::is_floating_point_v<Element>>* = nullptr>
__device__ bool equality_compare(Element const lhs, Element const rhs)
{
 return lhs == rhs;
}

template <typename Nullate>
class element_equality_comparator {
 public
 __host__ __device__
 element_equality_comparator(Nullate has_nulls,
 column_device_view lhs,
 column_device_view rhs,
 null_equality nulls_are_equal = null_equality::EQUAL)
 : lhs{lhs}, rhs{rhs}, nulls{has_nulls}, nulls_are_equal{nulls_are_equal}
 {
 }

 template <typename Element,
 std::enable_if_t<cudf::is_equality_comparable<Element, Element>()>* = nullptr>
 __device__ bool operator()(size_type lhs_element_index,
 size_type rhs_element_index) const noexcept
 {
 if (nulls) {
 bool const lhs_is_null{lhs.is_null(lhs_element_index)};
 bool const rhs_is_null{rhs.is_null(rhs_element_index)};
 if (lhs_is_null and rhs_is_null) {
 return nulls_are_equal == null_equality::EQUAL;
 } else if (lhs_is_null != rhs_is_null) {
 return false;
 }
 }

 return equality_compare(lhs.element<Element>(lhs_element_index),
 rhs.element<Element>(rhs_element_index));
 }

 // @cond
 template <typename Element,
 std::enable_if_t<not cudf::is_equality_comparable<Element, Element>()>* = nullptr>
 __device__ bool operator()(size_type lhs_element_index, size_type rhs_element_index)
 {
 CUDF_UNREACHABLE("Attempted to compare elements of uncomparable types.");
 }
 // @endcond

 private
 column_device_view lhs;
 column_device_view rhs;
 Nullate nulls;
 null_equality nulls_are_equal;
};

template <typename Nullate>
class row_equality_comparator {
 public
 row_equality_comparator(Nullate has_nulls,
 table_device_view lhs,
 table_device_view rhs,
 null_equality nulls_are_equal = null_equality::EQUAL)
 : lhs{lhs}, rhs{rhs}, nulls{has_nulls}, nulls_are_equal{nulls_are_equal}
 {
 CUDF_EXPECTS(lhs.num_columns() == rhs.num_columns(), "Mismatched number of columns.");
 }

 __device__ bool operator()(size_type lhs_row_index, size_type rhs_row_index) const noexcept
 {
 auto equal_elements = [=](column_device_view l, column_device_view r) {
 return cudf::type_dispatcher(l.type(),
 element_equality_comparator{nulls, l, r, nulls_are_equal},
 lhs_row_index,
 rhs_row_index);
 };

 return thrust::equal(thrust::seq, lhs.begin(), lhs.end(), rhs.begin(), equal_elements);
 }

 private
 table_device_view lhs;
 table_device_view rhs;
 Nullate nulls;
 null_equality nulls_are_equal;
};

template <typename Nullate>
class element_relational_comparator {
 public
 __host__ __device__ element_relational_comparator(Nullate has_nulls,
 column_device_view lhs,
 column_device_view rhs,
 null_order null_precedence)
 : lhs{lhs}, rhs{rhs}, nulls{has_nulls}, null_precedence{null_precedence}
 {
 }

 __host__ __device__ element_relational_comparator(Nullate has_nulls,
 column_device_view lhs,
 column_device_view rhs)
 : lhs{lhs}, rhs{rhs}, nulls{has_nulls}
 {
 }

 template <typename Element,
 std::enable_if_t<cudf::is_relationally_comparable<Element, Element>()>* = nullptr>
 __device__ weak_ordering operator()(size_type lhs_element_index,
 size_type rhs_element_index) const noexcept
 {
 if (nulls) {
 bool const lhs_is_null{lhs.is_null(lhs_element_index)};
 bool const rhs_is_null{rhs.is_null(rhs_element_index)};

 if (lhs_is_null or rhs_is_null) { // at least one is null
 return null_compare(lhs_is_null, rhs_is_null, null_precedence);
 }
 }

 return relational_compare(lhs.element<Element>(lhs_element_index),
 rhs.element<Element>(rhs_element_index));
 }

 // @cond
 template <typename Element,
 std::enable_if_t<not cudf::is_relationally_comparable<Element, Element>()>* = nullptr>
 __device__ weak_ordering operator()(size_type lhs_element_index, size_type rhs_element_index)
 {
 CUDF_UNREACHABLE("Attempted to compare elements of uncomparable types.");
 }
 // @endcond

 private
 column_device_view lhs;
 column_device_view rhs;
 Nullate nulls;
 null_order null_precedence{};
};

template <typename Nullate>
class row_lexicographic_comparator {
 public
 row_lexicographic_comparator(Nullate has_nulls,
 table_device_view lhs,
 table_device_view rhs,
 order const* column_order = nullptr,
 null_order const* null_precedence = nullptr)
 : _lhs{lhs},
 _rhs{rhs},
 _nulls{has_nulls},
 _column_order{column_order},
 _null_precedence{null_precedence}
 {
 CUDF_EXPECTS(_lhs.num_columns() == _rhs.num_columns(), "Mismatched number of columns.");
 }

 __device__ bool operator()(size_type lhs_index, size_type rhs_index) const noexcept
 {
 for (size_type i = 0; i < _lhs.num_columns(); ++i) {
 bool ascending = (_column_order == nullptr) or (_column_order[i] == order::ASCENDING);

 null_order null_precedence =
 _null_precedence == nullptr ? null_order::BEFORE : _null_precedence[i];

 auto comparator =
 element_relational_comparator{_nulls, _lhs.column(i), _rhs.column(i), null_precedence};

 weak_ordering state =
 cudf::type_dispatcher(_lhs.column(i).type(), comparator, lhs_index, rhs_index);

 if (state == weak_ordering::EQUIVALENT) { continue; }

 return state == (ascending ? weak_ordering::LESS : weak_ordering::GREATER);
 }
 return false;
 }

 private
 table_device_view _lhs;
 table_device_view _rhs;
 Nullate _nulls{};
 null_order const* _null_precedence{};
 order const* _column_order{};
}; // class row_lexicographic_comparator

template <template <typename> class hash_function, typename Nullate>
class element_hasher {
 public
 template <typename T, CUDF_ENABLE_IF(column_device_view::has_element_accessor<T>())>
 __device__ hash_value_type operator()(column_device_view col, size_type row_index) const
 {
 if (has_nulls && col.is_null(row_index)) {
 return cuda::std::numeric_limits<hash_value_type>::max();
 }
 return hash_function<T>{}(col.element<T>(row_index));
 }

 template <typename T, CUDF_ENABLE_IF(not column_device_view::has_element_accessor<T>())>
 __device__ hash_value_type operator()(column_device_view col, size_type row_index) const
 {
 CUDF_UNREACHABLE("Unsupported type in hash.");
 }

 Nullate has_nulls;
};

template <template <typename> class hash_function, typename Nullate>
class element_hasher_with_seed {
 public
 __device__ element_hasher_with_seed(Nullate has_nulls, uint32_t seed)
 : _seed{seed}, _has_nulls{has_nulls}
 {
 }

 __device__ element_hasher_with_seed(Nullate has_nulls, uint32_t seed, hash_value_type null_hash)
 : _seed{seed}, _null_hash{null_hash}, _has_nulls{has_nulls}
 {
 }

 template <typename T, CUDF_ENABLE_IF(column_device_view::has_element_accessor<T>())>
 __device__ hash_value_type operator()(column_device_view col, size_type row_index) const
 {
 if (_has_nulls && col.is_null(row_index)) { return _null_hash; }
 return hash_function<T>{_seed}(col.element<T>(row_index));
 }

 template <typename T, CUDF_ENABLE_IF(not column_device_view::has_element_accessor<T>())>
 __device__ hash_value_type operator()(column_device_view col, size_type row_index) const
 {
 CUDF_UNREACHABLE("Unsupported type in hash.");
 }

 private
 uint32_t _seed{DEFAULT_HASH_SEED};
 hash_value_type _null_hash{cuda::std::numeric_limits<hash_value_type>::max()};
 Nullate _has_nulls;
};

template <template <typename> class hash_function, typename Nullate>
class row_hasher {
 public
 row_hasher() = delete;

 CUDF_HOST_DEVICE row_hasher(Nullate has_nulls, table_device_view t)
 : _table{t}, _has_nulls{has_nulls}
 {
 }
 CUDF_HOST_DEVICE row_hasher(Nullate has_nulls, table_device_view t, uint32_t seed)
 : _table{t}, _seed(seed), _has_nulls{has_nulls}
 {
 }

 __device__ auto operator()(size_type row_index) const
 {
 // Hash the first column w/ the seed
 auto const initial_hash = cudf::hashing::detail::hash_combine(
 hash_value_type{0},
 type_dispatcher<dispatch_storage_type>(
 _table.column(0).type(),
 element_hasher_with_seed<hash_function, Nullate>{_has_nulls, _seed},
 _table.column(0),
 row_index));

 // Hashes an element in a column
 auto hasher = [=](size_type column_index) {
 return cudf::type_dispatcher<dispatch_storage_type>(
 _table.column(column_index).type(),
 element_hasher<hash_function, Nullate>{_has_nulls},
 _table.column(column_index),
 row_index);
 };

 // Hash each element and combine all the hash values together
 return thrust::transform_reduce(
 thrust::seq,
 // note that this starts at 1 and not 0 now since we already hashed the first column
 thrust::make_counting_iterator(1),
 thrust::make_counting_iterator(_table.num_columns()),
 hasher,
 initial_hash,
 [](hash_value_type lhs, hash_value_type rhs) {
 return cudf::hashing::detail::hash_combine(lhs, rhs);
 });
 }

 private
 table_device_view _table;
 Nullate _has_nulls;
 uint32_t _seed{DEFAULT_HASH_SEED};
};

} // namespace CUDF_EXPORT cudf