Skip to content

Commit

Permalink
Implement interleave_columns for lists with arbitrary nested type (#…
Browse files Browse the repository at this point in the history 
…9130)

This PR adds more support for lists column in the `interleave_column` API. In particular, it adds nested types support for the list entries. As such, now we can call `interleave_column` on a lists column of any type such as lists of structs, lists of lists, lists of lists of lists and so on. In addition, this PR also does a simple refactor of the existing overload functions with a new style of SFINAE implementation.

Closes #9106.

Authors:
  - Nghia Truong (https://github.com/ttnghia)

Approvers:
  - MithunR (https://github.com/mythrocks)
  - https://github.com/nvdbaranec

URL: #9130
  • Loading branch information
@ttnghia
ttnghia authored Sep 2, 2021
1 parent 58467fa commit 4556b5b
Show file tree
Hide file tree
Showing 2 changed files with 420 additions and 59 deletions.
129 changes: 94 additions & 35 deletions cpp/src/lists/interleave_columns.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,11 @@
*/

#include <cudf/column/column_factories.hpp>
#include <cudf/detail/concatenate.hpp>
#include <cudf/detail/copy.hpp>
#include <cudf/detail/gather.cuh>
#include <cudf/detail/get_value.cuh>
#include <cudf/detail/iterator.cuh>
#include <cudf/detail/valid_if.cuh>
#include <cudf/lists/lists_column_view.hpp>
#include <cudf/strings/detail/utilities.cuh>
Expand Down Expand Up @@ -86,6 +89,40 @@ generate_list_offsets_and_validities(table_view const& input,
return {std::move(list_offsets), std::move(validities)};
}

/**
* @brief Concatenate all input columns into one column and gather its rows to generate an output
* column that is the result of interleaving the input columns.
*/
std::unique_ptr<column> concatenate_and_gather_lists(host_span<column_view const> columns_to_concat,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
{
// Concatenate all columns into a single (temporary) column.
auto const concatenated_col =
cudf::detail::concatenate(columns_to_concat, stream, rmm::mr::get_current_device_resource());

// The number of input columns is known to be non-zero thus it's safe to call `front()` here.
auto const num_cols = columns_to_concat.size();
auto const num_input_rows = columns_to_concat.front().size();

// Generate the gather map that interleaves the input columns.
auto const iter_gather = cudf::detail::make_counting_transform_iterator(
0, [num_cols, num_input_rows] __device__(auto const idx) {
auto const source_col_idx = idx % num_cols;
auto const source_row_idx = idx / num_cols;
return source_col_idx * num_input_rows + source_row_idx;
});

// The gather API should be able to handle any data type for the input columns.
auto result = cudf::detail::gather(table_view{{concatenated_col->view()}},
iter_gather,
iter_gather + concatenated_col->size(),
out_of_bounds_policy::DONT_CHECK,
stream,
mr);
return std::move(result->release()[0]);
}

/**
* @brief Compute string sizes, string validities, and interleave string lists functor.
*
Expand Down Expand Up @@ -156,20 +193,25 @@ struct compute_string_sizes_and_interleave_lists_fn {
}
};

/**
* @brief Struct used in type_dispatcher to interleave list entries of the input lists columns and
* output the results into a destination column.
*/
struct interleave_list_entries_fn {
template <class T>
std::enable_if_t<std::is_same_v<T, cudf::string_view>, std::unique_ptr<column>> operator()(
table_view const& input,
column_view const& output_list_offsets,
size_type num_output_lists,
size_type num_output_entries,
bool data_has_null_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr) const noexcept
// Error case when no other overload or specialization is available
template <typename T, typename Enable = void>
struct interleave_list_entries_impl {
template <typename... Args>
std::unique_ptr<column> operator()(Args&&...)
{
CUDF_FAIL("Called `interleave_list_entries_fn()` on non-supported types.");
}
};

template <typename T>
struct interleave_list_entries_impl<T, std::enable_if_t<std::is_same_v<T, cudf::string_view>>> {
std::unique_ptr<column> operator()(table_view const& input,
column_view const& output_list_offsets,
size_type num_output_lists,
size_type num_output_entries,
bool data_has_null_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr) const noexcept
{
auto const table_dv_ptr = table_device_view::create(input);
auto comp_fn = compute_string_sizes_and_interleave_lists_fn{
Expand All @@ -191,16 +233,17 @@ struct interleave_list_entries_fn {
null_count,
std::move(null_mask));
}
};

template <class T>
std::enable_if_t<cudf::is_fixed_width<T>(), std::unique_ptr<column>> operator()(
table_view const& input,
column_view const& output_list_offsets,
size_type num_output_lists,
size_type num_output_entries,
bool data_has_null_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr) const noexcept
template <typename T>
struct interleave_list_entries_impl<T, std::enable_if_t<cudf::is_fixed_width<T>()>> {
std::unique_ptr<column> operator()(table_view const& input,
column_view const& output_list_offsets,
size_type num_output_lists,
size_type num_output_entries,
bool data_has_null_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr) const noexcept
{
auto const num_cols = input.num_columns();
auto const table_dv_ptr = table_device_view::create(input);
Expand Down Expand Up @@ -265,20 +308,29 @@ struct interleave_list_entries_fn {

return output;
}
};

/**
* @brief Struct used in type_dispatcher to interleave list entries of the input lists columns and
* output the results into a destination column.
*/
struct interleave_list_entries_fn {
template <class T>
std::enable_if_t<not std::is_same_v<T, cudf::string_view> and not cudf::is_fixed_width<T>(),
std::unique_ptr<column>>
operator()(table_view const&,
column_view const&,
size_type,
size_type,
bool,
rmm::cuda_stream_view,
rmm::mr::device_memory_resource*) const
std::unique_ptr<column> operator()(table_view const& input,
column_view const& output_list_offsets,
size_type num_output_lists,
size_type num_output_entries,
bool data_has_null_mask,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr) const
{
// Currently, only support string_view and fixed-width types
CUDF_FAIL("Called `interleave_list_entries_fn()` on non-supported types.");
return interleave_list_entries_impl<T>{}(input,
output_list_offsets,
num_output_lists,
num_output_entries,
data_has_null_mask,
stream,
mr);
}
};

Expand All @@ -299,14 +351,21 @@ std::unique_ptr<column> interleave_columns(table_view const& input,
"All columns of the input table must be of lists column type.");

auto const child_col = lists_column_view(col).child();
CUDF_EXPECTS(not cudf::is_nested(child_col.type()), "Nested types are not supported.");
CUDF_EXPECTS(entry_type == child_col.type(),
"The types of entries in the input columns must be the same.");
}

if (input.num_rows() == 0) { return cudf::empty_like(input.column(0)); }
if (input.num_columns() == 1) { return std::make_unique<column>(*(input.begin()), stream, mr); }

// For nested types, we rely on the `concatenate_and_gather` method, which costs more memory due
// to concatenation of the input columns into a temporary column. For non-nested types, we can
// directly interleave the input columns into the output column for better efficiency.
if (cudf::is_nested(entry_type)) {
auto const input_columns = std::vector<column_view>(input.begin(), input.end());
return concatenate_and_gather_lists(host_span<column_view const>{input_columns}, stream, mr);
}

// Generate offsets of the output lists column.
auto [list_offsets, list_validities] =
generate_list_offsets_and_validities(input, has_null_mask, stream, mr);
Expand Down
Loading

0 comments on commit 4556b5b

Please sign in to comment.