[FEA] Optimizations for groupby::scan

[jrhemstad]

The groupby object has a parameter allowing a caller to indicate that their keys (and likewise values to aggregate) are already in sorted order such that all values belonging to a particular group are already contiguous.

However, this information is not fully leveraged in the groupby::scan implementation.

Internally, the groupby::scan implementation calls get_grouped_values() to rearrange the values such that all values in the same group are contiguous

cudf/cpp/src/groupby/sort/scan.cpp

Lines 74 to 75 in 6728c75

	detail::sum_scan(
	get_grouped_values(), helper.num_groups(stream), helper.group_labels(stream), stream, mr));

This in turn calls the sort_groupby_helper::grouped_values to perform the grouping:

cudf/cpp/src/groupby/sort/sort_helper.cu

Lines 289 to 302 in 6728c75

	sort_groupby_helper::column_ptr sort_groupby_helper::grouped_values(
	column_view const& values, rmm::cuda_stream_view stream, rmm::mr::device_memory_resource* mr)
	{
	auto gather_map = key_sort_order(stream);
	auto grouped_values_table = cudf::detail::gather(table_view({values}),
	gather_map,
	cudf::out_of_bounds_policy::DONT_CHECK,
	cudf::detail::negative_index_policy::NOT_ALLOWED,
	stream,
	mr);
	return std::move(grouped_values_table->release()[0]);
	}

Which in turn calls key_sort_order to get the indices of the sorted order of the keys:

cudf/cpp/src/groupby/sort/sort_helper.cu

Lines 107 to 157 in 6728c75

	column_view sort_groupby_helper::key_sort_order(rmm::cuda_stream_view stream)
	{
	auto sliced_key_sorted_order = [stream, this]() {
	return cudf::detail::slice(this->_key_sorted_order->view(), 0, this->num_keys(stream));
	};
	if (_key_sorted_order) { return sliced_key_sorted_order(); }
	// TODO (dm): optimization. When keys are pre sorted but ignore nulls is true,
	// we still want all rows with nulls in the end. Sort is costly, so
	// do a copy_if(counting, sorted_order, {bitmask.is_valid(i)})
	if (_keys_pre_sorted == sorted::YES) {
	_key_sorted_order = make_numeric_column(
	data_type(type_to_id<size_type>()), _keys.num_rows(), mask_state::UNALLOCATED, stream);
	auto d_key_sorted_order = _key_sorted_order->mutable_view().data<size_type>();
	thrust::sequence(rmm::exec_policy(stream),
	d_key_sorted_order,
	d_key_sorted_order + _key_sorted_order->size(),
	0);
	return sliced_key_sorted_order();
	}
	if (_include_null_keys == null_policy::INCLUDE || !cudf::has_nulls(_keys)) { // SQL style
	_key_sorted_order = cudf::detail::stable_sorted_order(
	_keys,
	{},
	std::vector<null_order>(_keys.num_columns(), null_order::AFTER),
	stream,
	rmm::mr::get_current_device_resource());
	} else { // Pandas style
	// Temporarily prepend the keys table with a column that indicates the
	// presence of a null value within a row. This allows moving all rows that
	// contain a null value to the end of the sorted order.
	auto augmented_keys = table_view({table_view({keys_bitmask_column(stream)}), _keys});
	_key_sorted_order = cudf::detail::stable_sorted_order(
	augmented_keys,
	{},
	std::vector<null_order>(_keys.num_columns() + 1, null_order::AFTER),
	stream,
	rmm::mr::get_current_device_resource());
	// All rows with one or more null values are at the end of the resulting sorted order.
	}
	return sliced_key_sorted_order();
	}

Then the values are gathered based on the sorted key order:

cudf/cpp/src/groupby/sort/sort_helper.cu

Lines 294 to 299 in 6728c75

	auto grouped_values_table = cudf::detail::gather(table_view({values}),
	gather_map,
	cudf::out_of_bounds_policy::DONT_CHECK,
	cudf::detail::negative_index_policy::NOT_ALLOWED,
	stream,
	mr);

However, the steps of materializing the key_sort_order and performing the gather are completely redundant if the user has already specified that the inputs are already sorted.

I haven't full thought through the solution, but my intuition is that sort_groupby_helper should short circuit when they keys are already sorted and just return a view to the input values.

[jrhemstad]

This impacts more than just groupby::scan. This optimization will improve any time grouped_values is called when the input keys are already sorted.

[github-actions]

This issue has been labeled inactive-90d due to no recent activity in the past 90 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed.

[revans2]

Not sure if this is still an issue or not, but I would love to have groupby::scan go faster for already sorted data. As this is a use case we commonly have in Spark.

[PointKernel]

Oops, I completely forgot about this issue. Will work on this shortly.

[PointKernel]

Still desired.