Skip to content

GH-36252: [Python] Add non decomposable hash aggregate UDF - #36253

Merged
icexelloss merged 11 commits into
apache:mainfrom
icexelloss:compute-hash-aggregate-UDF
Jun 29, 2023
Merged

GH-36252: [Python] Add non decomposable hash aggregate UDF #36253
icexelloss merged 11 commits into
apache:mainfrom
icexelloss:compute-hash-aggregate-UDF

Conversation

@icexelloss

@icexellossicexelloss commented Jun 22, 2023•
edited
Loading

Copy link
Copy Markdown
Contributor

Rationale for this change

In #35515,

I have implemented a Scalar version of the non decomposable UDF (Scalar as in SCALAR_AGGREGATE). I would like to support the Hash version of it (Hash as in HASH_AGGREGATE)

With this PR, user can register an aggregate UDF once with pc.register_aggregate_function and it can be used as both scalar aggregate function and hash aggregate function.

Example:

def median(x):
return pa.scalar(np.nanmedian(x))
pc.register_aggregate_function(func=median, func_name='median_udf', ...)
table = ...
table.groupby("id").aggregate([("v", 'median_udf')])

What changes are included in this PR?

The main changes are:

  • In ResigterAggregateFunction (udf.cc), we now register the function both as a scalar aggregate function and a hash aggregate function (with signature adjustment for hash aggregate kernel because we need to append the grouping key)
  • Implemented PythonUdfHashAggregateImpl, similar to the PythonUdfScalarAggregateImpl. In Consume, it will accumulate both the input batches and the group id array. In Merge, it will merge the input batches and group id array (with the group_id_mapping). In Finalize, it will apply groupings to the accumulated batches to create one record batch per group, then apply the UDF over each group.
  • Some code clean up - UdfWrapperCallback objects are named cb (previously, agg_cb or wrapper) now and the user defined python function is now just called function (previously agg_function)

For table.groupby().aggregate(...), the space complexity is O(n) where n is the size of the table (and therefore, is not very useful). However, this is more useful in the segmented aggregation case, where the space complexity of O(s), where s the size of the segments.

Are these changes tested?

Added new test in test_udf.py (with table.group_by().aggregate() and test_substrait.py (with segmented aggregation)

Are there any user-facing changes?

Yes with this change, user can call use registered aggregate UDF with table.group_by().aggregate() or Acero's segmented aggregation.

Checklist

  • Self Review
  • API Documentation

@github-actionsgithub-actionsBot added the awaiting committer review Awaiting committer review label Jun 22, 2023
icexelloss
icexelloss commented Jun 22, 2023
View reviewed changes
Comment threadcpp/src/arrow/compute/row/grouper.h Outdated
@github-actionsgithub-actionsBot added awaiting changes Awaiting changes and removed awaiting committer review Awaiting committer review labels Jun 22, 2023
icexelloss
icexelloss commented Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/conftest.py Outdated
icexelloss
icexelloss commented Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/src/arrow/python/udf.cc
wrapper, options, registry);
}

Status AddAggKernel(std::shared_ptr<compute::KernelSignature> sig,

@icexellossicexellossJun 22, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is inlined now.

@github-actionsgithub-actionsBot added Component: Python awaiting change review Awaiting change review and removed awaiting changes Awaiting changes labels Jun 22, 2023
@icexelloss

icexelloss commented Jun 22, 2023•
edited
Loading

Copy link
Copy Markdown
ContributorAuthor

@westonpace I would like a request a review on this PR. The code should be relatively straight forward and similar to #35514 so hopefully no confusion/surprises here.

For the implementation of the grouping, I used an approach similar to GroupedListImpl aggregator and partition.cc

For the registration, I decided to register both scalar/hash kernel with one API register_aggregate_function because I think scalar/hash difference is not really something user should be worry about (from user's point of view, it is just "aggregation", whether it is hash/scalar is an compute implementation detail).

More details in the PR description.

Let me know if those sounds OK to you.

@icexellossicexelloss changed the title GH-36252: [Python] Compute hash aggregate udfGH-36252: [Python] Add non decomposable hash aggregate UDF Jun 22, 2023
@westonpace

westonpace commented Jun 23, 2023

Copy link
Copy Markdown
Member

@icexelloss I should have some time to take a look tomorrow

westonpace
westonpace requested changes Jun 24, 2023
View reviewed changes

@westonpacewestonpace left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a good set of tests. It's nice and convenient that the same python implementation can work for both. I have a few minor suggestions / questions. I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?

Comment threadpython/pyarrow/src/arrow/python/udf.cc Outdated
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Comment threadpython/pyarrow/src/arrow/python/udf.cc Outdated
Comment threadpython/pyarrow/src/arrow/python/udf.cc Outdated
const ArraySpan& groups_array_data = batch[batch.num_values() - 1].array;
DCHECK_EQ(groups_array_data.offset, 0);
int64_t batch_num_values = groups_array_data.length;
const auto* batch_groups = groups_array_data.GetValues<uint32_t>(1, 0);

@westonpacewestonpaceJun 24, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not just groups_array_data.GetValues<uint32_t>(1);?

@icexellossicexellossJun 26, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - updated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
}

num_values += other.num_values;
return Status::OK();

@westonpacewestonpaceJun 24, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does num_groups need to be updated here?

@icexellossicexellossJun 26, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think num_groups need to be updated here. Reasoning:

From the code in https://github.com/apache/arrow/blob/main/cpp/src/arrow/compute/kernels/hash_aggregate.cc#L233 and https://github.com/apache/arrow/blob/main/cpp/src/arrow/acero/groupby_aggregate_node.cc#L248
(1) Other hash kernel implementation updates the num_groups in Resize
(2) resize is always called before consume and merge

Comment threadpython/pyarrow/src/arrow/python/udf.cc Outdated
UdfContext udf_context{ctx->memory_pool(), table->num_rows()};

if (rb->num_rows() == 0) {
return Status::Invalid("Finalized is called with empty inputs");

@westonpacewestonpaceJun 24, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this a problem?

@icexellossicexellossJun 26, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - I was being lazy here and didn't want to bother with empty aggregation. But now I look at this I can just return empty result here. Will update.

@icexellossicexellossJun 26, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated

Comment threadpython/pyarrow/src/arrow/python/udf.cc

ARROW_ASSIGN_OR_RAISE(auto table,
arrow::Table::FromRecordBatches(input_schema, values));
ARROW_ASSIGN_OR_RAISE(auto rb, table->CombineChunksToBatch(ctx->memory_pool()));

@westonpacewestonpaceJun 24, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some cases where this won't be possible. For example, if you have an array of strings then the array may only have 2GB of string data (regardless of how many elements it has). So any single group can't have more than 2GB of string data. I don't know this is fatal but you may want to mention in the user docs somewhere or wrap this failure with extra context.

@icexellossicexellossJun 26, 2023•
edited
Loading

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added note in the user doc: https://github.com/apache/arrow/pull/36253/files#diff-439f91c435cc8136d40eaba8c168aaa1cec00f08b10ba92ce376e47f29f81814R2770

Comment threadpython/pyarrow/src/arrow/python/udf.cc Outdated
@github-actionsgithub-actionsBot added awaiting changes Awaiting changes and removed awaiting change review Awaiting change review labels Jun 24, 2023
@icexelloss@westonpace
Apply suggestions from code review
007221a
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@github-actionsgithub-actionsBot added awaiting change review Awaiting change review and removed awaiting changes Awaiting changes labels Jun 26, 2023
@icexelloss

icexelloss commented Jun 26, 2023

Copy link
Copy Markdown
ContributorAuthor

I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?

Thanks @westonpace. Currently we only plan to use this with segmented aggregation so each group is not going to be very large. (grouping inside a segment), so I don't think it would be a problem.

@github-actionsgithub-actionsBot removed the awaiting change review Awaiting change review label Jun 26, 2023
@github-actionsgithub-actionsBot added the awaiting changes Awaiting changes label Jun 26, 2023
@github-actionsgithub-actionsBot added awaiting change review Awaiting change review awaiting changes Awaiting changes and removed awaiting changes Awaiting changes awaiting change review Awaiting change review labels Jun 26, 2023
@icexelloss

icexelloss commented Jun 26, 2023•
edited
Loading

Copy link
Copy Markdown
ContributorAuthor

@westonpace This should be clean now (all comments addressed, CI green) - another look?

@icexelloss

icexelloss commented Jun 27, 2023

Copy link
Copy Markdown
ContributorAuthor

Gentle ping @westonpace anything else you want me to change here?

westonpace
westonpace approved these changes Jun 28, 2023
View reviewed changes

@westonpacewestonpace left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor wording suggestion. Otherwise this looks good.

Comment threadpython/pyarrow/_compute.pyx Outdated
Comment threadpython/pyarrow/src/arrow/python/udf.cc
std::vector<std::shared_ptr<DataType>> input_types,
std::shared_ptr<DataType> output_type)
: function(function), cb(std::move(cb)), output_type(std::move(output_type)) {
Py_INCREF(function->obj());

@westonpacewestonpaceJun 28, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These INCREF's still seem superfluous to me but I don't think it's critical. We could test in a follow-up using temporary function registries to see if we are preventing UDF functions from being garbage collected.

@icexellossicexellossJun 28, 2023

Copy link
Copy Markdown
ContributorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with you. I plan to address this in #36000 but haven't got to it.

@github-actionsgithub-actionsBot added awaiting merge Awaiting merge and removed awaiting changes Awaiting changes labels Jun 28, 2023
@icexelloss@westonpace
Apply suggestions from code review
c29ae72
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@icexelloss

icexelloss commented Jun 28, 2023

Copy link
Copy Markdown
ContributorAuthor

Thanks @westonpace. I applied your suggestion and will merge once CI passes.

@github-actionsgithub-actionsBot added awaiting changes Awaiting changes and removed awaiting merge Awaiting merge labels Jun 28, 2023
@icexelloss

icexelloss commented Jun 29, 2023

Copy link
Copy Markdown
ContributorAuthor

CI failure is unrelated. Merging.

@icexelloss
icexelloss merged commit baf17a2 into apache:mainJun 29, 2023
@conbench-apache-arrow

conbench-apache-arrowBot commented Jul 4, 2023

Copy link
Copy Markdown

Conbench analyzed the 6 benchmark runs on commit baf17a20.

There was 1 benchmark result with an error:

There were no benchmark performance regressions. 🎉

The full Conbench report has more details.

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@westonpacewestonpacewestonpace approved these changes
@AlenkaFAlenkaFAwaiting requested review from AlenkaFAlenkaF is a code owner

Assignees

No one assigned

Labels

awaiting changesAwaiting changesComponent: Python

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[C++][Python] Non decomposable aggregation UDF (Hash version)

2 participants

@icexelloss@westonpace
, 'i'); if (__m === '*' || __re.test(location.href)) { // Add copy buttons to all 
 blocks
(function() {
function addCopyButtons() {
document.querySelectorAll('pre code').forEach(function(codeBlock) {
if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;
codeBlock.parentElement.setAttribute('data-copy-added', 'true');
var btn = document.createElement('button');
btn.textContent = 'Copy';
btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';
btn.onmouseover = function() { this.style.opacity = '1'; };
btn.onmouseout = function() { this.style.opacity = '0.7'; };
btn.onclick = function() {
navigator.clipboard.writeText(codeBlock.textContent).then(function() {
btn.textContent = 'Copied!';
setTimeout(function() { btn.textContent = 'Copy'; }, 1500);
});
};
codeBlock.parentElement.style.position = 'relative';
codeBlock.parentElement.appendChild(btn);
});
}
addCopyButtons();
// Re-run on dynamic content
var observer = new MutationObserver(addCopyButtons);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
GH-36252: [Python] Add non decomposable hash aggregate UDF by icexelloss · Pull Request #36253 · apache/arrow · GitHub
Skip to content
GH-36252: [Python] Add non decomposable hash aggregate UDF  - #36253
Merged
icexelloss merged 11 commits into
apache:mainfrom 
icexelloss:compute-hash-aggregate-UDF
Jun 29, 2023
Merged
GH-36252: [Python] Add non decomposable hash aggregate UDF #36253
icexelloss merged 11 commits into
apache:mainfrom 
icexelloss:compute-hash-aggregate-UDF
Conversation
@icexelloss
@icexellossicexelloss
commented
Jun 22, 2023•

edited

Loading

Copy link


Copy Markdown

Contributor
Rationale for this change
In #35515,
I have implemented a Scalar version of the non decomposable UDF (Scalar as in SCALAR_AGGREGATE). I would like to support the Hash version of it (Hash as in HASH_AGGREGATE)
With this PR, user can register an aggregate UDF once with pc.register_aggregate_function and it can be used as both scalar aggregate function and hash aggregate function.
Example:
def median(x):
return pa.scalar(np.nanmedian(x))
pc.register_aggregate_function(func=median, func_name='median_udf', ...)
table = ...
table.groupby("id").aggregate([("v", 'median_udf')])

What changes are included in this PR?
The main changes are:
  • In ResigterAggregateFunction (udf.cc), we now register the function both as a scalar aggregate function and a hash aggregate function (with signature adjustment for hash aggregate kernel because we need to append the grouping key)
  • Implemented PythonUdfHashAggregateImpl, similar to the PythonUdfScalarAggregateImpl. In Consume, it will accumulate both the input batches and the group id array. In Merge, it will merge the input batches and group id array (with the group_id_mapping). In Finalize, it will apply groupings to the accumulated batches to create one record batch per group, then apply the UDF over each group.
  • Some code clean up - UdfWrapperCallback objects are named cb (previously, agg_cb or wrapper) now and the user defined python function is now just called function (previously agg_function)
For table.groupby().aggregate(...), the space complexity is O(n) where n is the size of the table (and therefore, is not very useful). However, this is more useful in the segmented aggregation case, where the space complexity of O(s), where s the size of the segments.
Are these changes tested?
Added new test in test_udf.py (with table.group_by().aggregate() and test_substrait.py (with segmented aggregation)
Are there any user-facing changes?
Yes with this change, user can call use registered aggregate UDF with table.group_by().aggregate()  or Acero's segmented aggregation.
Checklist
  •  Self Review
  •  API Documentation
@github-actionsgithub-actionsBot
added
the 
awaiting committer review
Awaiting committer review
label
Jun 22, 2023
icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadcpp/src/arrow/compute/row/grouper.h
Outdated

@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting committer review
Awaiting committer review
labels
Jun 22, 2023
icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/conftest.py
Outdated

icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/src/arrow/python/udf.cc
 wrapper, options, registry);
}

Status AddAggKernel(std::shared_ptr<compute::KernelSignature> sig,
@icexellossicexellossJun 22, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is inlined now.
@github-actionsgithub-actionsBot
added

Component: Python

awaiting change review
Awaiting change review
and removed

awaiting changes
Awaiting changes
labels
Jun 22, 2023
@icexelloss
icexelloss
commented
Jun 22, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
@westonpace I would like a request a review on this PR. The code should be relatively straight forward and similar to #35514 so hopefully no confusion/surprises here.
For the implementation of the grouping, I used an approach similar to GroupedListImpl aggregator and partition.cc
For the registration, I decided to register both scalar/hash kernel with one API register_aggregate_function because I think scalar/hash difference is not really something user should be worry about (from user's point of view, it is just "aggregation", whether it is hash/scalar is an compute implementation detail).
More details in the PR description.
Let me know if those sounds OK to you.
@icexellossicexelloss
changed the title
GH-36252: [Python] Compute hash aggregate udfGH-36252: [Python] Add non decomposable hash aggregate UDF Jun 22, 2023
@westonpace
westonpace
commented
Jun 23, 2023

Copy link


Copy Markdown

Member
@icexelloss I should have some time to take a look tomorrow
westonpace
westonpace
requested changes
Jun 24, 2023
View reviewed changes
@westonpacewestonpace
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a good set of tests. It's nice and convenient that the same python implementation can work for both. I have a few minor suggestions / questions. I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

 const ArraySpan& groups_array_data = batch[batch.num_values() - 1].array;
 DCHECK_EQ(groups_array_data.offset, 0);
 int64_t batch_num_values = groups_array_data.length;
 const auto* batch_groups = groups_array_data.GetValues<uint32_t>(1, 0);
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not just groups_array_data.GetValues<uint32_t>(1);?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - updated
Comment threadpython/pyarrow/src/arrow/python/udf.cc
 }

 num_values += other.num_values;
 return Status::OK();
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does num_groups need to be updated here?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think num_groups need to be updated here. Reasoning:
From the code in https://github.com/apache/arrow/blob/main/cpp/src/arrow/compute/kernels/hash_aggregate.cc#L233 and https://github.com/apache/arrow/blob/main/cpp/src/arrow/acero/groupby_aggregate_node.cc#L248

(1) Other hash kernel implementation updates the num_groups in Resize

(2) resize is always called before consume and merge
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

 UdfContext udf_context{ctx->memory_pool(), table->num_rows()};

 if (rb->num_rows() == 0) {
 return Status::Invalid("Finalized is called with empty inputs");
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this a problem?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - I was being lazy here and didn't want to bother with empty aggregation. But now I look at this I can just return empty result here. Will update.
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated
Comment threadpython/pyarrow/src/arrow/python/udf.cc

 ARROW_ASSIGN_OR_RAISE(auto table,
 arrow::Table::FromRecordBatches(input_schema, values));
 ARROW_ASSIGN_OR_RAISE(auto rb, table->CombineChunksToBatch(ctx->memory_pool()));
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some cases where this won't be possible. For example, if you have an array of strings then the array may only have 2GB of string data (regardless of how many elements it has). So any single group can't have more than 2GB of string data. I don't know this is fatal but you may want to mention in the user docs somewhere or wrap this failure with extra context.
@icexellossicexellossJun 26, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added note in the user doc: https://github.com/apache/arrow/pull/36253/files#diff-439f91c435cc8136d40eaba8c168aaa1cec00f08b10ba92ce376e47f29f81814R2770
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting change review
Awaiting change review
labels
Jun 24, 2023
@icexelloss@westonpace
Apply suggestions from code review
007221a
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@github-actionsgithub-actionsBot
added

awaiting change review
Awaiting change review
and removed

awaiting changes
Awaiting changes
labels
Jun 26, 2023
@icexelloss
icexelloss
commented
Jun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?
Thanks @westonpace. Currently we only plan to use this with segmented aggregation so each group is not going to be very large. (grouping inside a segment), so I don't think it would be a problem.
@github-actionsgithub-actionsBot
removed
the 
awaiting change review
Awaiting change review
label
Jun 26, 2023
@github-actionsgithub-actionsBot
added
the 
awaiting changes
Awaiting changes
label
Jun 26, 2023
@github-actionsgithub-actionsBot
added

awaiting change review
Awaiting change review
awaiting changes
Awaiting changes
and removed

awaiting changes
Awaiting changes
awaiting change review
Awaiting change review
labels
Jun 26, 2023
@icexelloss
icexelloss
commented
Jun 26, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
@westonpace This should be clean now (all comments addressed, CI green) - another look?
@icexelloss
icexelloss
commented
Jun 27, 2023

Copy link


Copy Markdown

ContributorAuthor
Gentle ping @westonpace anything else you want me to change here?
westonpace
westonpace
approved these changes
Jun 28, 2023
View reviewed changes
@westonpacewestonpace
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor wording suggestion. Otherwise this looks good.
Comment threadpython/pyarrow/_compute.pyx
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
 std::vector<std::shared_ptr<DataType>> input_types,
 std::shared_ptr<DataType> output_type)
 : function(function), cb(std::move(cb)), output_type(std::move(output_type)) {
 Py_INCREF(function->obj());
@westonpacewestonpaceJun 28, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These INCREF's still seem superfluous to me but I don't think it's critical. We could test in a follow-up using temporary function registries to see if we are preventing UDF functions from being garbage collected.
@icexellossicexellossJun 28, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with you. I plan to address this in #36000 but haven't got to it.
@github-actionsgithub-actionsBot
added

awaiting merge
Awaiting merge
and removed

awaiting changes
Awaiting changes
labels
Jun 28, 2023
@icexelloss@westonpace
Apply suggestions from code review
c29ae72
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@icexelloss
icexelloss
commented
Jun 28, 2023

Copy link


Copy Markdown

ContributorAuthor
Thanks @westonpace. I applied your suggestion and will merge once CI passes.
@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting merge
Awaiting merge
labels
Jun 28, 2023
@icexelloss
icexelloss
commented
Jun 29, 2023

Copy link


Copy Markdown

ContributorAuthor
CI failure is unrelated. Merging.
@icexelloss
icexelloss
merged commit baf17a2
into
apache:mainJun 29, 2023
@conbench-apache-arrow
conbench-apache-arrowBot
commented
Jul 4, 2023

Copy link


Copy Markdown

Conbench analyzed the 6 benchmark runs on commit baf17a20.
There was 1 benchmark result with an error:
There were no benchmark performance regressions. 🎉
The full Conbench report has more details.
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@westonpacewestonpacewestonpace approved these changes
@AlenkaFAlenkaFAwaiting requested review from AlenkaFAlenkaF is a code owner

Assignees


No one assigned


Labels

awaiting changesAwaiting changesComponent: Python

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

[C++][Python] Non decomposable aggregation UDF (Hash version)


2 participants

@icexelloss@westonpace

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Force GitHub README to respect dark mode
(function() {
var style = document.createElement('style');
style.textContent = '
.markdown-body {
color-scheme: dark light;
}
.markdown-body pre { background: #161b22 !important; }
.markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }
.markdown-body table th, .markdown-body table td { border-color: #30363d !important; }
.markdown-body img { background: #0d1117; }
.markdown-body blockquote { border-left-color: #8b949e; }
.markdown-body hr { border-color: #30363d; }
';
document.head.appendChild(style);
})();
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
GH-36252: [Python] Add non decomposable hash aggregate UDF by icexelloss · Pull Request #36253 · apache/arrow · GitHub
Skip to content
GH-36252: [Python] Add non decomposable hash aggregate UDF  - #36253
Merged
icexelloss merged 11 commits into
apache:mainfrom 
icexelloss:compute-hash-aggregate-UDF
Jun 29, 2023
Merged
GH-36252: [Python] Add non decomposable hash aggregate UDF #36253
icexelloss merged 11 commits into
apache:mainfrom 
icexelloss:compute-hash-aggregate-UDF
Conversation
@icexelloss
@icexellossicexelloss
commented
Jun 22, 2023•

edited

Loading

Copy link


Copy Markdown

Contributor
Rationale for this change
In #35515,
I have implemented a Scalar version of the non decomposable UDF (Scalar as in SCALAR_AGGREGATE). I would like to support the Hash version of it (Hash as in HASH_AGGREGATE)
With this PR, user can register an aggregate UDF once with pc.register_aggregate_function and it can be used as both scalar aggregate function and hash aggregate function.
Example:
def median(x):
return pa.scalar(np.nanmedian(x))
pc.register_aggregate_function(func=median, func_name='median_udf', ...)
table = ...
table.groupby("id").aggregate([("v", 'median_udf')])

What changes are included in this PR?
The main changes are:
  • In ResigterAggregateFunction (udf.cc), we now register the function both as a scalar aggregate function and a hash aggregate function (with signature adjustment for hash aggregate kernel because we need to append the grouping key)
  • Implemented PythonUdfHashAggregateImpl, similar to the PythonUdfScalarAggregateImpl. In Consume, it will accumulate both the input batches and the group id array. In Merge, it will merge the input batches and group id array (with the group_id_mapping). In Finalize, it will apply groupings to the accumulated batches to create one record batch per group, then apply the UDF over each group.
  • Some code clean up - UdfWrapperCallback objects are named cb (previously, agg_cb or wrapper) now and the user defined python function is now just called function (previously agg_function)
For table.groupby().aggregate(...), the space complexity is O(n) where n is the size of the table (and therefore, is not very useful). However, this is more useful in the segmented aggregation case, where the space complexity of O(s), where s the size of the segments.
Are these changes tested?
Added new test in test_udf.py (with table.group_by().aggregate() and test_substrait.py (with segmented aggregation)
Are there any user-facing changes?
Yes with this change, user can call use registered aggregate UDF with table.group_by().aggregate()  or Acero's segmented aggregation.
Checklist
  •  Self Review
  •  API Documentation
@github-actionsgithub-actionsBot
added
the 
awaiting committer review
Awaiting committer review
label
Jun 22, 2023
icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadcpp/src/arrow/compute/row/grouper.h
Outdated

@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting committer review
Awaiting committer review
labels
Jun 22, 2023
icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/conftest.py
Outdated

icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/src/arrow/python/udf.cc
 wrapper, options, registry);
}

Status AddAggKernel(std::shared_ptr<compute::KernelSignature> sig,
@icexellossicexellossJun 22, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is inlined now.
@github-actionsgithub-actionsBot
added

Component: Python

awaiting change review
Awaiting change review
and removed

awaiting changes
Awaiting changes
labels
Jun 22, 2023
@icexelloss
icexelloss
commented
Jun 22, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
@westonpace I would like a request a review on this PR. The code should be relatively straight forward and similar to #35514 so hopefully no confusion/surprises here.
For the implementation of the grouping, I used an approach similar to GroupedListImpl aggregator and partition.cc
For the registration, I decided to register both scalar/hash kernel with one API register_aggregate_function because I think scalar/hash difference is not really something user should be worry about (from user's point of view, it is just "aggregation", whether it is hash/scalar is an compute implementation detail).
More details in the PR description.
Let me know if those sounds OK to you.
@icexellossicexelloss
changed the title
GH-36252: [Python] Compute hash aggregate udfGH-36252: [Python] Add non decomposable hash aggregate UDF Jun 22, 2023
@westonpace
westonpace
commented
Jun 23, 2023

Copy link


Copy Markdown

Member
@icexelloss I should have some time to take a look tomorrow
westonpace
westonpace
requested changes
Jun 24, 2023
View reviewed changes
@westonpacewestonpace
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a good set of tests. It's nice and convenient that the same python implementation can work for both. I have a few minor suggestions / questions. I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

 const ArraySpan& groups_array_data = batch[batch.num_values() - 1].array;
 DCHECK_EQ(groups_array_data.offset, 0);
 int64_t batch_num_values = groups_array_data.length;
 const auto* batch_groups = groups_array_data.GetValues<uint32_t>(1, 0);
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not just groups_array_data.GetValues<uint32_t>(1);?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - updated
Comment threadpython/pyarrow/src/arrow/python/udf.cc
 }

 num_values += other.num_values;
 return Status::OK();
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does num_groups need to be updated here?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think num_groups need to be updated here. Reasoning:
From the code in https://github.com/apache/arrow/blob/main/cpp/src/arrow/compute/kernels/hash_aggregate.cc#L233 and https://github.com/apache/arrow/blob/main/cpp/src/arrow/acero/groupby_aggregate_node.cc#L248

(1) Other hash kernel implementation updates the num_groups in Resize

(2) resize is always called before consume and merge
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

 UdfContext udf_context{ctx->memory_pool(), table->num_rows()};

 if (rb->num_rows() == 0) {
 return Status::Invalid("Finalized is called with empty inputs");
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this a problem?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - I was being lazy here and didn't want to bother with empty aggregation. But now I look at this I can just return empty result here. Will update.
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated
Comment threadpython/pyarrow/src/arrow/python/udf.cc

 ARROW_ASSIGN_OR_RAISE(auto table,
 arrow::Table::FromRecordBatches(input_schema, values));
 ARROW_ASSIGN_OR_RAISE(auto rb, table->CombineChunksToBatch(ctx->memory_pool()));
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some cases where this won't be possible. For example, if you have an array of strings then the array may only have 2GB of string data (regardless of how many elements it has). So any single group can't have more than 2GB of string data. I don't know this is fatal but you may want to mention in the user docs somewhere or wrap this failure with extra context.
@icexellossicexellossJun 26, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added note in the user doc: https://github.com/apache/arrow/pull/36253/files#diff-439f91c435cc8136d40eaba8c168aaa1cec00f08b10ba92ce376e47f29f81814R2770
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting change review
Awaiting change review
labels
Jun 24, 2023
@icexelloss@westonpace
Apply suggestions from code review
007221a
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@github-actionsgithub-actionsBot
added

awaiting change review
Awaiting change review
and removed

awaiting changes
Awaiting changes
labels
Jun 26, 2023
@icexelloss
icexelloss
commented
Jun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?
Thanks @westonpace. Currently we only plan to use this with segmented aggregation so each group is not going to be very large. (grouping inside a segment), so I don't think it would be a problem.
@github-actionsgithub-actionsBot
removed
the 
awaiting change review
Awaiting change review
label
Jun 26, 2023
@github-actionsgithub-actionsBot
added
the 
awaiting changes
Awaiting changes
label
Jun 26, 2023
@github-actionsgithub-actionsBot
added

awaiting change review
Awaiting change review
awaiting changes
Awaiting changes
and removed

awaiting changes
Awaiting changes
awaiting change review
Awaiting change review
labels
Jun 26, 2023
@icexelloss
icexelloss
commented
Jun 26, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
@westonpace This should be clean now (all comments addressed, CI green) - another look?
@icexelloss
icexelloss
commented
Jun 27, 2023

Copy link


Copy Markdown

ContributorAuthor
Gentle ping @westonpace anything else you want me to change here?
westonpace
westonpace
approved these changes
Jun 28, 2023
View reviewed changes
@westonpacewestonpace
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor wording suggestion. Otherwise this looks good.
Comment threadpython/pyarrow/_compute.pyx
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
 std::vector<std::shared_ptr<DataType>> input_types,
 std::shared_ptr<DataType> output_type)
 : function(function), cb(std::move(cb)), output_type(std::move(output_type)) {
 Py_INCREF(function->obj());
@westonpacewestonpaceJun 28, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These INCREF's still seem superfluous to me but I don't think it's critical. We could test in a follow-up using temporary function registries to see if we are preventing UDF functions from being garbage collected.
@icexellossicexellossJun 28, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with you. I plan to address this in #36000 but haven't got to it.
@github-actionsgithub-actionsBot
added

awaiting merge
Awaiting merge
and removed

awaiting changes
Awaiting changes
labels
Jun 28, 2023
@icexelloss@westonpace
Apply suggestions from code review
c29ae72
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@icexelloss
icexelloss
commented
Jun 28, 2023

Copy link


Copy Markdown

ContributorAuthor
Thanks @westonpace. I applied your suggestion and will merge once CI passes.
@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting merge
Awaiting merge
labels
Jun 28, 2023
@icexelloss
icexelloss
commented
Jun 29, 2023

Copy link


Copy Markdown

ContributorAuthor
CI failure is unrelated. Merging.
@icexelloss
icexelloss
merged commit baf17a2
into
apache:mainJun 29, 2023
@conbench-apache-arrow
conbench-apache-arrowBot
commented
Jul 4, 2023

Copy link


Copy Markdown

Conbench analyzed the 6 benchmark runs on commit baf17a20.
There was 1 benchmark result with an error:
There were no benchmark performance regressions. 🎉
The full Conbench report has more details.
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@westonpacewestonpacewestonpace approved these changes
@AlenkaFAlenkaFAwaiting requested review from AlenkaFAlenkaF is a code owner

Assignees


No one assigned


Labels

awaiting changesAwaiting changesComponent: Python

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

[C++][Python] Non decomposable aggregation UDF (Hash version)


2 participants

@icexelloss@westonpace

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Highlight search terms from Google/DuckDuckGo/Bing referrer
(function() {
var ref = document.referrer;
var terms = [];
if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {
var url = new URL(ref);
var q = url.searchParams.get('q') || url.searchParams.get('p');
if (q) {
terms = q.split(/\s+/).filter(function(t) { return t.length > 2; });
}
}
if (terms.length === 0) return;
var style = document.createElement('style');
style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';
document.head.appendChild(style);
function highlight(node) {
if (node.nodeType === 3) { // text node
var text = node.textContent;
var found = false;
terms.forEach(function(term) {
var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\]\\]/g, '\\') + ')', 'gi');
if (regex.test(text)) {
found = true;
var frag = document.createDocumentFragment();
var parts = text.split(regex);
parts.forEach(function(part, i) {
if (i % 2 === 0) {
frag.appendChild(document.createTextNode(part));
} else {
var span = document.createElement('span');
span.className = 'userscript-highlight';
span.textContent = part;
frag.appendChild(span);
}
});
node.parentNode.replaceChild(frag, node);
}
});
} else if (node.nodeType === 1 && node.childNodes) { // element
var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];
if (!skipTags.includes(node.tagName)) {
Array.from(node.childNodes).forEach(highlight);
}
}
}
highlight(document.body);
// Re-highlight on dynamic content
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1 || node.nodeType === 3) highlight(node);
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
GH-36252: [Python] Add non decomposable hash aggregate UDF by icexelloss · Pull Request #36253 · apache/arrow · GitHub
Skip to content
GH-36252: [Python] Add non decomposable hash aggregate UDF  - #36253
Merged
icexelloss merged 11 commits into
apache:mainfrom 
icexelloss:compute-hash-aggregate-UDF
Jun 29, 2023
Merged
GH-36252: [Python] Add non decomposable hash aggregate UDF #36253
icexelloss merged 11 commits into
apache:mainfrom 
icexelloss:compute-hash-aggregate-UDF
Conversation
@icexelloss
@icexellossicexelloss
commented
Jun 22, 2023•

edited

Loading

Copy link


Copy Markdown

Contributor
Rationale for this change
In #35515,
I have implemented a Scalar version of the non decomposable UDF (Scalar as in SCALAR_AGGREGATE). I would like to support the Hash version of it (Hash as in HASH_AGGREGATE)
With this PR, user can register an aggregate UDF once with pc.register_aggregate_function and it can be used as both scalar aggregate function and hash aggregate function.
Example:
def median(x):
return pa.scalar(np.nanmedian(x))
pc.register_aggregate_function(func=median, func_name='median_udf', ...)
table = ...
table.groupby("id").aggregate([("v", 'median_udf')])

What changes are included in this PR?
The main changes are:
  • In ResigterAggregateFunction (udf.cc), we now register the function both as a scalar aggregate function and a hash aggregate function (with signature adjustment for hash aggregate kernel because we need to append the grouping key)
  • Implemented PythonUdfHashAggregateImpl, similar to the PythonUdfScalarAggregateImpl. In Consume, it will accumulate both the input batches and the group id array. In Merge, it will merge the input batches and group id array (with the group_id_mapping). In Finalize, it will apply groupings to the accumulated batches to create one record batch per group, then apply the UDF over each group.
  • Some code clean up - UdfWrapperCallback objects are named cb (previously, agg_cb or wrapper) now and the user defined python function is now just called function (previously agg_function)
For table.groupby().aggregate(...), the space complexity is O(n) where n is the size of the table (and therefore, is not very useful). However, this is more useful in the segmented aggregation case, where the space complexity of O(s), where s the size of the segments.
Are these changes tested?
Added new test in test_udf.py (with table.group_by().aggregate() and test_substrait.py (with segmented aggregation)
Are there any user-facing changes?
Yes with this change, user can call use registered aggregate UDF with table.group_by().aggregate()  or Acero's segmented aggregation.
Checklist
  •  Self Review
  •  API Documentation
@github-actionsgithub-actionsBot
added
the 
awaiting committer review
Awaiting committer review
label
Jun 22, 2023
icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadcpp/src/arrow/compute/row/grouper.h
Outdated

@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting committer review
Awaiting committer review
labels
Jun 22, 2023
icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/conftest.py
Outdated

icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/src/arrow/python/udf.cc
 wrapper, options, registry);
}

Status AddAggKernel(std::shared_ptr<compute::KernelSignature> sig,
@icexellossicexellossJun 22, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is inlined now.
@github-actionsgithub-actionsBot
added

Component: Python

awaiting change review
Awaiting change review
and removed

awaiting changes
Awaiting changes
labels
Jun 22, 2023
@icexelloss
icexelloss
commented
Jun 22, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
@westonpace I would like a request a review on this PR. The code should be relatively straight forward and similar to #35514 so hopefully no confusion/surprises here.
For the implementation of the grouping, I used an approach similar to GroupedListImpl aggregator and partition.cc
For the registration, I decided to register both scalar/hash kernel with one API register_aggregate_function because I think scalar/hash difference is not really something user should be worry about (from user's point of view, it is just "aggregation", whether it is hash/scalar is an compute implementation detail).
More details in the PR description.
Let me know if those sounds OK to you.
@icexellossicexelloss
changed the title
GH-36252: [Python] Compute hash aggregate udfGH-36252: [Python] Add non decomposable hash aggregate UDF Jun 22, 2023
@westonpace
westonpace
commented
Jun 23, 2023

Copy link


Copy Markdown

Member
@icexelloss I should have some time to take a look tomorrow
westonpace
westonpace
requested changes
Jun 24, 2023
View reviewed changes
@westonpacewestonpace
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a good set of tests. It's nice and convenient that the same python implementation can work for both. I have a few minor suggestions / questions. I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

 const ArraySpan& groups_array_data = batch[batch.num_values() - 1].array;
 DCHECK_EQ(groups_array_data.offset, 0);
 int64_t batch_num_values = groups_array_data.length;
 const auto* batch_groups = groups_array_data.GetValues<uint32_t>(1, 0);
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not just groups_array_data.GetValues<uint32_t>(1);?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - updated
Comment threadpython/pyarrow/src/arrow/python/udf.cc
 }

 num_values += other.num_values;
 return Status::OK();
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does num_groups need to be updated here?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think num_groups need to be updated here. Reasoning:
From the code in https://github.com/apache/arrow/blob/main/cpp/src/arrow/compute/kernels/hash_aggregate.cc#L233 and https://github.com/apache/arrow/blob/main/cpp/src/arrow/acero/groupby_aggregate_node.cc#L248

(1) Other hash kernel implementation updates the num_groups in Resize

(2) resize is always called before consume and merge
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

 UdfContext udf_context{ctx->memory_pool(), table->num_rows()};

 if (rb->num_rows() == 0) {
 return Status::Invalid("Finalized is called with empty inputs");
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this a problem?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - I was being lazy here and didn't want to bother with empty aggregation. But now I look at this I can just return empty result here. Will update.
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated
Comment threadpython/pyarrow/src/arrow/python/udf.cc

 ARROW_ASSIGN_OR_RAISE(auto table,
 arrow::Table::FromRecordBatches(input_schema, values));
 ARROW_ASSIGN_OR_RAISE(auto rb, table->CombineChunksToBatch(ctx->memory_pool()));
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some cases where this won't be possible. For example, if you have an array of strings then the array may only have 2GB of string data (regardless of how many elements it has). So any single group can't have more than 2GB of string data. I don't know this is fatal but you may want to mention in the user docs somewhere or wrap this failure with extra context.
@icexellossicexellossJun 26, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added note in the user doc: https://github.com/apache/arrow/pull/36253/files#diff-439f91c435cc8136d40eaba8c168aaa1cec00f08b10ba92ce376e47f29f81814R2770
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting change review
Awaiting change review
labels
Jun 24, 2023
@icexelloss@westonpace
Apply suggestions from code review
007221a
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@github-actionsgithub-actionsBot
added

awaiting change review
Awaiting change review
and removed

awaiting changes
Awaiting changes
labels
Jun 26, 2023
@icexelloss
icexelloss
commented
Jun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?
Thanks @westonpace. Currently we only plan to use this with segmented aggregation so each group is not going to be very large. (grouping inside a segment), so I don't think it would be a problem.
@github-actionsgithub-actionsBot
removed
the 
awaiting change review
Awaiting change review
label
Jun 26, 2023
@github-actionsgithub-actionsBot
added
the 
awaiting changes
Awaiting changes
label
Jun 26, 2023
@github-actionsgithub-actionsBot
added

awaiting change review
Awaiting change review
awaiting changes
Awaiting changes
and removed

awaiting changes
Awaiting changes
awaiting change review
Awaiting change review
labels
Jun 26, 2023
@icexelloss
icexelloss
commented
Jun 26, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
@westonpace This should be clean now (all comments addressed, CI green) - another look?
@icexelloss
icexelloss
commented
Jun 27, 2023

Copy link


Copy Markdown

ContributorAuthor
Gentle ping @westonpace anything else you want me to change here?
westonpace
westonpace
approved these changes
Jun 28, 2023
View reviewed changes
@westonpacewestonpace
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor wording suggestion. Otherwise this looks good.
Comment threadpython/pyarrow/_compute.pyx
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
 std::vector<std::shared_ptr<DataType>> input_types,
 std::shared_ptr<DataType> output_type)
 : function(function), cb(std::move(cb)), output_type(std::move(output_type)) {
 Py_INCREF(function->obj());
@westonpacewestonpaceJun 28, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These INCREF's still seem superfluous to me but I don't think it's critical. We could test in a follow-up using temporary function registries to see if we are preventing UDF functions from being garbage collected.
@icexellossicexellossJun 28, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with you. I plan to address this in #36000 but haven't got to it.
@github-actionsgithub-actionsBot
added

awaiting merge
Awaiting merge
and removed

awaiting changes
Awaiting changes
labels
Jun 28, 2023
@icexelloss@westonpace
Apply suggestions from code review
c29ae72
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@icexelloss
icexelloss
commented
Jun 28, 2023

Copy link


Copy Markdown

ContributorAuthor
Thanks @westonpace. I applied your suggestion and will merge once CI passes.
@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting merge
Awaiting merge
labels
Jun 28, 2023
@icexelloss
icexelloss
commented
Jun 29, 2023

Copy link


Copy Markdown

ContributorAuthor
CI failure is unrelated. Merging.
@icexelloss
icexelloss
merged commit baf17a2
into
apache:mainJun 29, 2023
@conbench-apache-arrow
conbench-apache-arrowBot
commented
Jul 4, 2023

Copy link


Copy Markdown

Conbench analyzed the 6 benchmark runs on commit baf17a20.
There was 1 benchmark result with an error:
There were no benchmark performance regressions. 🎉
The full Conbench report has more details.
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@westonpacewestonpacewestonpace approved these changes
@AlenkaFAlenkaFAwaiting requested review from AlenkaFAlenkaF is a code owner

Assignees


No one assigned


Labels

awaiting changesAwaiting changesComponent: Python

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

[C++][Python] Non decomposable aggregation UDF (Hash version)


2 participants

@icexelloss@westonpace

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Strip utm_, fbclid, gclid, etc. from all links on page
(function() {
var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',
'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',
'ref', 'ref_src', 'source', 'medium', 'campaign'];
function cleanUrl(url) {
try {
var u = new URL(url, window.location.origin);
var changed = false;
trackingParams.forEach(function(p) {
if (u.searchParams.has(p)) {
u.searchParams.delete(p);
changed = true;
}
});
return changed ? u.toString() : url;
} catch (e) {
return url;
}
}
function cleanLinks() {
document.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
cleanLinks();
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1) {
if (node.tagName === 'A') cleanLinks();
node.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '
GH-36252: [Python] Add non decomposable hash aggregate UDF by icexelloss · Pull Request #36253 · apache/arrow · GitHub
Skip to content
GH-36252: [Python] Add non decomposable hash aggregate UDF  - #36253
Merged
icexelloss merged 11 commits into
apache:mainfrom 
icexelloss:compute-hash-aggregate-UDF
Jun 29, 2023
Merged
GH-36252: [Python] Add non decomposable hash aggregate UDF #36253
icexelloss merged 11 commits into
apache:mainfrom 
icexelloss:compute-hash-aggregate-UDF
Conversation
@icexelloss
@icexellossicexelloss
commented
Jun 22, 2023•

edited

Loading

Copy link


Copy Markdown

Contributor
Rationale for this change
In #35515,
I have implemented a Scalar version of the non decomposable UDF (Scalar as in SCALAR_AGGREGATE). I would like to support the Hash version of it (Hash as in HASH_AGGREGATE)
With this PR, user can register an aggregate UDF once with pc.register_aggregate_function and it can be used as both scalar aggregate function and hash aggregate function.
Example:
def median(x):
return pa.scalar(np.nanmedian(x))
pc.register_aggregate_function(func=median, func_name='median_udf', ...)
table = ...
table.groupby("id").aggregate([("v", 'median_udf')])

What changes are included in this PR?
The main changes are:
  • In ResigterAggregateFunction (udf.cc), we now register the function both as a scalar aggregate function and a hash aggregate function (with signature adjustment for hash aggregate kernel because we need to append the grouping key)
  • Implemented PythonUdfHashAggregateImpl, similar to the PythonUdfScalarAggregateImpl. In Consume, it will accumulate both the input batches and the group id array. In Merge, it will merge the input batches and group id array (with the group_id_mapping). In Finalize, it will apply groupings to the accumulated batches to create one record batch per group, then apply the UDF over each group.
  • Some code clean up - UdfWrapperCallback objects are named cb (previously, agg_cb or wrapper) now and the user defined python function is now just called function (previously agg_function)
For table.groupby().aggregate(...), the space complexity is O(n) where n is the size of the table (and therefore, is not very useful). However, this is more useful in the segmented aggregation case, where the space complexity of O(s), where s the size of the segments.
Are these changes tested?
Added new test in test_udf.py (with table.group_by().aggregate() and test_substrait.py (with segmented aggregation)
Are there any user-facing changes?
Yes with this change, user can call use registered aggregate UDF with table.group_by().aggregate()  or Acero's segmented aggregation.
Checklist
  •  Self Review
  •  API Documentation
@github-actionsgithub-actionsBot
added
the 
awaiting committer review
Awaiting committer review
label
Jun 22, 2023
icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadcpp/src/arrow/compute/row/grouper.h
Outdated

@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting committer review
Awaiting committer review
labels
Jun 22, 2023
icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/conftest.py
Outdated

icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/src/arrow/python/udf.cc
 wrapper, options, registry);
}

Status AddAggKernel(std::shared_ptr<compute::KernelSignature> sig,
@icexellossicexellossJun 22, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is inlined now.
@github-actionsgithub-actionsBot
added

Component: Python

awaiting change review
Awaiting change review
and removed

awaiting changes
Awaiting changes
labels
Jun 22, 2023
@icexelloss
icexelloss
commented
Jun 22, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
@westonpace I would like a request a review on this PR. The code should be relatively straight forward and similar to #35514 so hopefully no confusion/surprises here.
For the implementation of the grouping, I used an approach similar to GroupedListImpl aggregator and partition.cc
For the registration, I decided to register both scalar/hash kernel with one API register_aggregate_function because I think scalar/hash difference is not really something user should be worry about (from user's point of view, it is just "aggregation", whether it is hash/scalar is an compute implementation detail).
More details in the PR description.
Let me know if those sounds OK to you.
@icexellossicexelloss
changed the title
GH-36252: [Python] Compute hash aggregate udfGH-36252: [Python] Add non decomposable hash aggregate UDF Jun 22, 2023
@westonpace
westonpace
commented
Jun 23, 2023

Copy link


Copy Markdown

Member
@icexelloss I should have some time to take a look tomorrow
westonpace
westonpace
requested changes
Jun 24, 2023
View reviewed changes
@westonpacewestonpace
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a good set of tests. It's nice and convenient that the same python implementation can work for both. I have a few minor suggestions / questions. I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

 const ArraySpan& groups_array_data = batch[batch.num_values() - 1].array;
 DCHECK_EQ(groups_array_data.offset, 0);
 int64_t batch_num_values = groups_array_data.length;
 const auto* batch_groups = groups_array_data.GetValues<uint32_t>(1, 0);
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not just groups_array_data.GetValues<uint32_t>(1);?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - updated
Comment threadpython/pyarrow/src/arrow/python/udf.cc
 }

 num_values += other.num_values;
 return Status::OK();
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does num_groups need to be updated here?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think num_groups need to be updated here. Reasoning:
From the code in https://github.com/apache/arrow/blob/main/cpp/src/arrow/compute/kernels/hash_aggregate.cc#L233 and https://github.com/apache/arrow/blob/main/cpp/src/arrow/acero/groupby_aggregate_node.cc#L248

(1) Other hash kernel implementation updates the num_groups in Resize

(2) resize is always called before consume and merge
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

 UdfContext udf_context{ctx->memory_pool(), table->num_rows()};

 if (rb->num_rows() == 0) {
 return Status::Invalid("Finalized is called with empty inputs");
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this a problem?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - I was being lazy here and didn't want to bother with empty aggregation. But now I look at this I can just return empty result here. Will update.
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated
Comment threadpython/pyarrow/src/arrow/python/udf.cc

 ARROW_ASSIGN_OR_RAISE(auto table,
 arrow::Table::FromRecordBatches(input_schema, values));
 ARROW_ASSIGN_OR_RAISE(auto rb, table->CombineChunksToBatch(ctx->memory_pool()));
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some cases where this won't be possible. For example, if you have an array of strings then the array may only have 2GB of string data (regardless of how many elements it has). So any single group can't have more than 2GB of string data. I don't know this is fatal but you may want to mention in the user docs somewhere or wrap this failure with extra context.
@icexellossicexellossJun 26, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added note in the user doc: https://github.com/apache/arrow/pull/36253/files#diff-439f91c435cc8136d40eaba8c168aaa1cec00f08b10ba92ce376e47f29f81814R2770
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting change review
Awaiting change review
labels
Jun 24, 2023
@icexelloss@westonpace
Apply suggestions from code review
007221a
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@github-actionsgithub-actionsBot
added

awaiting change review
Awaiting change review
and removed

awaiting changes
Awaiting changes
labels
Jun 26, 2023
@icexelloss
icexelloss
commented
Jun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?
Thanks @westonpace. Currently we only plan to use this with segmented aggregation so each group is not going to be very large. (grouping inside a segment), so I don't think it would be a problem.
@github-actionsgithub-actionsBot
removed
the 
awaiting change review
Awaiting change review
label
Jun 26, 2023
@github-actionsgithub-actionsBot
added
the 
awaiting changes
Awaiting changes
label
Jun 26, 2023
@github-actionsgithub-actionsBot
added

awaiting change review
Awaiting change review
awaiting changes
Awaiting changes
and removed

awaiting changes
Awaiting changes
awaiting change review
Awaiting change review
labels
Jun 26, 2023
@icexelloss
icexelloss
commented
Jun 26, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
@westonpace This should be clean now (all comments addressed, CI green) - another look?
@icexelloss
icexelloss
commented
Jun 27, 2023

Copy link


Copy Markdown

ContributorAuthor
Gentle ping @westonpace anything else you want me to change here?
westonpace
westonpace
approved these changes
Jun 28, 2023
View reviewed changes
@westonpacewestonpace
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor wording suggestion. Otherwise this looks good.
Comment threadpython/pyarrow/_compute.pyx
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
 std::vector<std::shared_ptr<DataType>> input_types,
 std::shared_ptr<DataType> output_type)
 : function(function), cb(std::move(cb)), output_type(std::move(output_type)) {
 Py_INCREF(function->obj());
@westonpacewestonpaceJun 28, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These INCREF's still seem superfluous to me but I don't think it's critical. We could test in a follow-up using temporary function registries to see if we are preventing UDF functions from being garbage collected.
@icexellossicexellossJun 28, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with you. I plan to address this in #36000 but haven't got to it.
@github-actionsgithub-actionsBot
added

awaiting merge
Awaiting merge
and removed

awaiting changes
Awaiting changes
labels
Jun 28, 2023
@icexelloss@westonpace
Apply suggestions from code review
c29ae72
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@icexelloss
icexelloss
commented
Jun 28, 2023

Copy link


Copy Markdown

ContributorAuthor
Thanks @westonpace. I applied your suggestion and will merge once CI passes.
@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting merge
Awaiting merge
labels
Jun 28, 2023
@icexelloss
icexelloss
commented
Jun 29, 2023

Copy link


Copy Markdown

ContributorAuthor
CI failure is unrelated. Merging.
@icexelloss
icexelloss
merged commit baf17a2
into
apache:mainJun 29, 2023
@conbench-apache-arrow
conbench-apache-arrowBot
commented
Jul 4, 2023

Copy link


Copy Markdown

Conbench analyzed the 6 benchmark runs on commit baf17a20.
There was 1 benchmark result with an error:
There were no benchmark performance regressions. 🎉
The full Conbench report has more details.
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@westonpacewestonpacewestonpace approved these changes
@AlenkaFAlenkaFAwaiting requested review from AlenkaFAlenkaF is a code owner

Assignees


No one assigned


Labels

awaiting changesAwaiting changesComponent: Python

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

[C++][Python] Non decomposable aggregation UDF (Hash version)


2 participants

@icexelloss@westonpace

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Auto-enable theater mode on YouTube
(function() {
function tryTheater() {
var btn = document.querySelector('button[aria-label="Theater mode"], ytd-player #player button[title="Theater mode"]');
if (btn && !btn.classList.contains('activated')) {
btn.click();
}
}
// Try immediately
tryTheater();
// Try after navigation (SPA)
var lastUrl = location.href;
setInterval(function() {
if (location.href !== lastUrl) {
lastUrl = location.href;
setTimeout(tryTheater, 500);
}
}, 1000);
// Also try on player load
var observer = new MutationObserver(tryTheater);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
GH-36252: [Python] Add non decomposable hash aggregate UDF by icexelloss · Pull Request #36253 · apache/arrow · GitHub
Skip to content
GH-36252: [Python] Add non decomposable hash aggregate UDF  - #36253
Merged
icexelloss merged 11 commits into
apache:mainfrom 
icexelloss:compute-hash-aggregate-UDF
Jun 29, 2023
Merged
GH-36252: [Python] Add non decomposable hash aggregate UDF #36253
icexelloss merged 11 commits into
apache:mainfrom 
icexelloss:compute-hash-aggregate-UDF
Conversation
@icexelloss
@icexellossicexelloss
commented
Jun 22, 2023•

edited

Loading

Copy link


Copy Markdown

Contributor
Rationale for this change
In #35515,
I have implemented a Scalar version of the non decomposable UDF (Scalar as in SCALAR_AGGREGATE). I would like to support the Hash version of it (Hash as in HASH_AGGREGATE)
With this PR, user can register an aggregate UDF once with pc.register_aggregate_function and it can be used as both scalar aggregate function and hash aggregate function.
Example:
def median(x):
return pa.scalar(np.nanmedian(x))
pc.register_aggregate_function(func=median, func_name='median_udf', ...)
table = ...
table.groupby("id").aggregate([("v", 'median_udf')])

What changes are included in this PR?
The main changes are:
  • In ResigterAggregateFunction (udf.cc), we now register the function both as a scalar aggregate function and a hash aggregate function (with signature adjustment for hash aggregate kernel because we need to append the grouping key)
  • Implemented PythonUdfHashAggregateImpl, similar to the PythonUdfScalarAggregateImpl. In Consume, it will accumulate both the input batches and the group id array. In Merge, it will merge the input batches and group id array (with the group_id_mapping). In Finalize, it will apply groupings to the accumulated batches to create one record batch per group, then apply the UDF over each group.
  • Some code clean up - UdfWrapperCallback objects are named cb (previously, agg_cb or wrapper) now and the user defined python function is now just called function (previously agg_function)
For table.groupby().aggregate(...), the space complexity is O(n) where n is the size of the table (and therefore, is not very useful). However, this is more useful in the segmented aggregation case, where the space complexity of O(s), where s the size of the segments.
Are these changes tested?
Added new test in test_udf.py (with table.group_by().aggregate() and test_substrait.py (with segmented aggregation)
Are there any user-facing changes?
Yes with this change, user can call use registered aggregate UDF with table.group_by().aggregate()  or Acero's segmented aggregation.
Checklist
  •  Self Review
  •  API Documentation
@github-actionsgithub-actionsBot
added
the 
awaiting committer review
Awaiting committer review
label
Jun 22, 2023
icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadcpp/src/arrow/compute/row/grouper.h
Outdated

@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting committer review
Awaiting committer review
labels
Jun 22, 2023
icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/conftest.py
Outdated

icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/src/arrow/python/udf.cc
 wrapper, options, registry);
}

Status AddAggKernel(std::shared_ptr<compute::KernelSignature> sig,
@icexellossicexellossJun 22, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is inlined now.
@github-actionsgithub-actionsBot
added

Component: Python

awaiting change review
Awaiting change review
and removed

awaiting changes
Awaiting changes
labels
Jun 22, 2023
@icexelloss
icexelloss
commented
Jun 22, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
@westonpace I would like a request a review on this PR. The code should be relatively straight forward and similar to #35514 so hopefully no confusion/surprises here.
For the implementation of the grouping, I used an approach similar to GroupedListImpl aggregator and partition.cc
For the registration, I decided to register both scalar/hash kernel with one API register_aggregate_function because I think scalar/hash difference is not really something user should be worry about (from user's point of view, it is just "aggregation", whether it is hash/scalar is an compute implementation detail).
More details in the PR description.
Let me know if those sounds OK to you.
@icexellossicexelloss
changed the title
GH-36252: [Python] Compute hash aggregate udfGH-36252: [Python] Add non decomposable hash aggregate UDF Jun 22, 2023
@westonpace
westonpace
commented
Jun 23, 2023

Copy link


Copy Markdown

Member
@icexelloss I should have some time to take a look tomorrow
westonpace
westonpace
requested changes
Jun 24, 2023
View reviewed changes
@westonpacewestonpace
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a good set of tests. It's nice and convenient that the same python implementation can work for both. I have a few minor suggestions / questions. I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

 const ArraySpan& groups_array_data = batch[batch.num_values() - 1].array;
 DCHECK_EQ(groups_array_data.offset, 0);
 int64_t batch_num_values = groups_array_data.length;
 const auto* batch_groups = groups_array_data.GetValues<uint32_t>(1, 0);
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not just groups_array_data.GetValues<uint32_t>(1);?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - updated
Comment threadpython/pyarrow/src/arrow/python/udf.cc
 }

 num_values += other.num_values;
 return Status::OK();
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does num_groups need to be updated here?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think num_groups need to be updated here. Reasoning:
From the code in https://github.com/apache/arrow/blob/main/cpp/src/arrow/compute/kernels/hash_aggregate.cc#L233 and https://github.com/apache/arrow/blob/main/cpp/src/arrow/acero/groupby_aggregate_node.cc#L248

(1) Other hash kernel implementation updates the num_groups in Resize

(2) resize is always called before consume and merge
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

 UdfContext udf_context{ctx->memory_pool(), table->num_rows()};

 if (rb->num_rows() == 0) {
 return Status::Invalid("Finalized is called with empty inputs");
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this a problem?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - I was being lazy here and didn't want to bother with empty aggregation. But now I look at this I can just return empty result here. Will update.
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated
Comment threadpython/pyarrow/src/arrow/python/udf.cc

 ARROW_ASSIGN_OR_RAISE(auto table,
 arrow::Table::FromRecordBatches(input_schema, values));
 ARROW_ASSIGN_OR_RAISE(auto rb, table->CombineChunksToBatch(ctx->memory_pool()));
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some cases where this won't be possible. For example, if you have an array of strings then the array may only have 2GB of string data (regardless of how many elements it has). So any single group can't have more than 2GB of string data. I don't know this is fatal but you may want to mention in the user docs somewhere or wrap this failure with extra context.
@icexellossicexellossJun 26, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added note in the user doc: https://github.com/apache/arrow/pull/36253/files#diff-439f91c435cc8136d40eaba8c168aaa1cec00f08b10ba92ce376e47f29f81814R2770
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting change review
Awaiting change review
labels
Jun 24, 2023
@icexelloss@westonpace
Apply suggestions from code review
007221a
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@github-actionsgithub-actionsBot
added

awaiting change review
Awaiting change review
and removed

awaiting changes
Awaiting changes
labels
Jun 26, 2023
@icexelloss
icexelloss
commented
Jun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?
Thanks @westonpace. Currently we only plan to use this with segmented aggregation so each group is not going to be very large. (grouping inside a segment), so I don't think it would be a problem.
@github-actionsgithub-actionsBot
removed
the 
awaiting change review
Awaiting change review
label
Jun 26, 2023
@github-actionsgithub-actionsBot
added
the 
awaiting changes
Awaiting changes
label
Jun 26, 2023
@github-actionsgithub-actionsBot
added

awaiting change review
Awaiting change review
awaiting changes
Awaiting changes
and removed

awaiting changes
Awaiting changes
awaiting change review
Awaiting change review
labels
Jun 26, 2023
@icexelloss
icexelloss
commented
Jun 26, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
@westonpace This should be clean now (all comments addressed, CI green) - another look?
@icexelloss
icexelloss
commented
Jun 27, 2023

Copy link


Copy Markdown

ContributorAuthor
Gentle ping @westonpace anything else you want me to change here?
westonpace
westonpace
approved these changes
Jun 28, 2023
View reviewed changes
@westonpacewestonpace
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor wording suggestion. Otherwise this looks good.
Comment threadpython/pyarrow/_compute.pyx
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
 std::vector<std::shared_ptr<DataType>> input_types,
 std::shared_ptr<DataType> output_type)
 : function(function), cb(std::move(cb)), output_type(std::move(output_type)) {
 Py_INCREF(function->obj());
@westonpacewestonpaceJun 28, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These INCREF's still seem superfluous to me but I don't think it's critical. We could test in a follow-up using temporary function registries to see if we are preventing UDF functions from being garbage collected.
@icexellossicexellossJun 28, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with you. I plan to address this in #36000 but haven't got to it.
@github-actionsgithub-actionsBot
added

awaiting merge
Awaiting merge
and removed

awaiting changes
Awaiting changes
labels
Jun 28, 2023
@icexelloss@westonpace
Apply suggestions from code review
c29ae72
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@icexelloss
icexelloss
commented
Jun 28, 2023

Copy link


Copy Markdown

ContributorAuthor
Thanks @westonpace. I applied your suggestion and will merge once CI passes.
@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting merge
Awaiting merge
labels
Jun 28, 2023
@icexelloss
icexelloss
commented
Jun 29, 2023

Copy link


Copy Markdown

ContributorAuthor
CI failure is unrelated. Merging.
@icexelloss
icexelloss
merged commit baf17a2
into
apache:mainJun 29, 2023
@conbench-apache-arrow
conbench-apache-arrowBot
commented
Jul 4, 2023

Copy link


Copy Markdown

Conbench analyzed the 6 benchmark runs on commit baf17a20.
There was 1 benchmark result with an error:
There were no benchmark performance regressions. 🎉
The full Conbench report has more details.
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@westonpacewestonpacewestonpace approved these changes
@AlenkaFAlenkaFAwaiting requested review from AlenkaFAlenkaF is a code owner

Assignees


No one assigned


Labels

awaiting changesAwaiting changesComponent: Python

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

[C++][Python] Non decomposable aggregation UDF (Hash version)


2 participants

@icexelloss@westonpace

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Remove or un-stick sticky/fixed headers that block content
(function() {
function unstick() {
document.querySelectorAll('header, nav, [role="banner"], .header, .navbar, .sticky, .fixed-top, [style*="position: fixed"], [style*="position:sticky"]').forEach(function(el) {
if (el.style.position === 'fixed' || el.style.position === 'sticky' || getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {
el.style.position = 'static';
el.style.top = 'auto';
el.style.zIndex = 'auto';
}
});
}
unstick();
var observer = new MutationObserver(unstick);
observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });
})();
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
GH-36252: [Python] Add non decomposable hash aggregate UDF by icexelloss · Pull Request #36253 · apache/arrow · GitHub
Skip to content
GH-36252: [Python] Add non decomposable hash aggregate UDF  - #36253
Merged
icexelloss merged 11 commits into
apache:mainfrom 
icexelloss:compute-hash-aggregate-UDF
Jun 29, 2023
Merged
GH-36252: [Python] Add non decomposable hash aggregate UDF #36253
icexelloss merged 11 commits into
apache:mainfrom 
icexelloss:compute-hash-aggregate-UDF
Conversation
@icexelloss
@icexellossicexelloss
commented
Jun 22, 2023•

edited

Loading

Copy link


Copy Markdown

Contributor
Rationale for this change
In #35515,
I have implemented a Scalar version of the non decomposable UDF (Scalar as in SCALAR_AGGREGATE). I would like to support the Hash version of it (Hash as in HASH_AGGREGATE)
With this PR, user can register an aggregate UDF once with pc.register_aggregate_function and it can be used as both scalar aggregate function and hash aggregate function.
Example:
def median(x):
return pa.scalar(np.nanmedian(x))
pc.register_aggregate_function(func=median, func_name='median_udf', ...)
table = ...
table.groupby("id").aggregate([("v", 'median_udf')])

What changes are included in this PR?
The main changes are:
  • In ResigterAggregateFunction (udf.cc), we now register the function both as a scalar aggregate function and a hash aggregate function (with signature adjustment for hash aggregate kernel because we need to append the grouping key)
  • Implemented PythonUdfHashAggregateImpl, similar to the PythonUdfScalarAggregateImpl. In Consume, it will accumulate both the input batches and the group id array. In Merge, it will merge the input batches and group id array (with the group_id_mapping). In Finalize, it will apply groupings to the accumulated batches to create one record batch per group, then apply the UDF over each group.
  • Some code clean up - UdfWrapperCallback objects are named cb (previously, agg_cb or wrapper) now and the user defined python function is now just called function (previously agg_function)
For table.groupby().aggregate(...), the space complexity is O(n) where n is the size of the table (and therefore, is not very useful). However, this is more useful in the segmented aggregation case, where the space complexity of O(s), where s the size of the segments.
Are these changes tested?
Added new test in test_udf.py (with table.group_by().aggregate() and test_substrait.py (with segmented aggregation)
Are there any user-facing changes?
Yes with this change, user can call use registered aggregate UDF with table.group_by().aggregate()  or Acero's segmented aggregation.
Checklist
  •  Self Review
  •  API Documentation
@github-actionsgithub-actionsBot
added
the 
awaiting committer review
Awaiting committer review
label
Jun 22, 2023
icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadcpp/src/arrow/compute/row/grouper.h
Outdated

@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting committer review
Awaiting committer review
labels
Jun 22, 2023
icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/conftest.py
Outdated

icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/src/arrow/python/udf.cc
 wrapper, options, registry);
}

Status AddAggKernel(std::shared_ptr<compute::KernelSignature> sig,
@icexellossicexellossJun 22, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is inlined now.
@github-actionsgithub-actionsBot
added

Component: Python

awaiting change review
Awaiting change review
and removed

awaiting changes
Awaiting changes
labels
Jun 22, 2023
@icexelloss
icexelloss
commented
Jun 22, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
@westonpace I would like a request a review on this PR. The code should be relatively straight forward and similar to #35514 so hopefully no confusion/surprises here.
For the implementation of the grouping, I used an approach similar to GroupedListImpl aggregator and partition.cc
For the registration, I decided to register both scalar/hash kernel with one API register_aggregate_function because I think scalar/hash difference is not really something user should be worry about (from user's point of view, it is just "aggregation", whether it is hash/scalar is an compute implementation detail).
More details in the PR description.
Let me know if those sounds OK to you.
@icexellossicexelloss
changed the title
GH-36252: [Python] Compute hash aggregate udfGH-36252: [Python] Add non decomposable hash aggregate UDF Jun 22, 2023
@westonpace
westonpace
commented
Jun 23, 2023

Copy link


Copy Markdown

Member
@icexelloss I should have some time to take a look tomorrow
westonpace
westonpace
requested changes
Jun 24, 2023
View reviewed changes
@westonpacewestonpace
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a good set of tests. It's nice and convenient that the same python implementation can work for both. I have a few minor suggestions / questions. I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

 const ArraySpan& groups_array_data = batch[batch.num_values() - 1].array;
 DCHECK_EQ(groups_array_data.offset, 0);
 int64_t batch_num_values = groups_array_data.length;
 const auto* batch_groups = groups_array_data.GetValues<uint32_t>(1, 0);
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not just groups_array_data.GetValues<uint32_t>(1);?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - updated
Comment threadpython/pyarrow/src/arrow/python/udf.cc
 }

 num_values += other.num_values;
 return Status::OK();
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does num_groups need to be updated here?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think num_groups need to be updated here. Reasoning:
From the code in https://github.com/apache/arrow/blob/main/cpp/src/arrow/compute/kernels/hash_aggregate.cc#L233 and https://github.com/apache/arrow/blob/main/cpp/src/arrow/acero/groupby_aggregate_node.cc#L248

(1) Other hash kernel implementation updates the num_groups in Resize

(2) resize is always called before consume and merge
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

 UdfContext udf_context{ctx->memory_pool(), table->num_rows()};

 if (rb->num_rows() == 0) {
 return Status::Invalid("Finalized is called with empty inputs");
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this a problem?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - I was being lazy here and didn't want to bother with empty aggregation. But now I look at this I can just return empty result here. Will update.
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated
Comment threadpython/pyarrow/src/arrow/python/udf.cc

 ARROW_ASSIGN_OR_RAISE(auto table,
 arrow::Table::FromRecordBatches(input_schema, values));
 ARROW_ASSIGN_OR_RAISE(auto rb, table->CombineChunksToBatch(ctx->memory_pool()));
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some cases where this won't be possible. For example, if you have an array of strings then the array may only have 2GB of string data (regardless of how many elements it has). So any single group can't have more than 2GB of string data. I don't know this is fatal but you may want to mention in the user docs somewhere or wrap this failure with extra context.
@icexellossicexellossJun 26, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added note in the user doc: https://github.com/apache/arrow/pull/36253/files#diff-439f91c435cc8136d40eaba8c168aaa1cec00f08b10ba92ce376e47f29f81814R2770
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting change review
Awaiting change review
labels
Jun 24, 2023
@icexelloss@westonpace
Apply suggestions from code review
007221a
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@github-actionsgithub-actionsBot
added

awaiting change review
Awaiting change review
and removed

awaiting changes
Awaiting changes
labels
Jun 26, 2023
@icexelloss
icexelloss
commented
Jun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?
Thanks @westonpace. Currently we only plan to use this with segmented aggregation so each group is not going to be very large. (grouping inside a segment), so I don't think it would be a problem.
@github-actionsgithub-actionsBot
removed
the 
awaiting change review
Awaiting change review
label
Jun 26, 2023
@github-actionsgithub-actionsBot
added
the 
awaiting changes
Awaiting changes
label
Jun 26, 2023
@github-actionsgithub-actionsBot
added

awaiting change review
Awaiting change review
awaiting changes
Awaiting changes
and removed

awaiting changes
Awaiting changes
awaiting change review
Awaiting change review
labels
Jun 26, 2023
@icexelloss
icexelloss
commented
Jun 26, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
@westonpace This should be clean now (all comments addressed, CI green) - another look?
@icexelloss
icexelloss
commented
Jun 27, 2023

Copy link


Copy Markdown

ContributorAuthor
Gentle ping @westonpace anything else you want me to change here?
westonpace
westonpace
approved these changes
Jun 28, 2023
View reviewed changes
@westonpacewestonpace
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor wording suggestion. Otherwise this looks good.
Comment threadpython/pyarrow/_compute.pyx
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
 std::vector<std::shared_ptr<DataType>> input_types,
 std::shared_ptr<DataType> output_type)
 : function(function), cb(std::move(cb)), output_type(std::move(output_type)) {
 Py_INCREF(function->obj());
@westonpacewestonpaceJun 28, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These INCREF's still seem superfluous to me but I don't think it's critical. We could test in a follow-up using temporary function registries to see if we are preventing UDF functions from being garbage collected.
@icexellossicexellossJun 28, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with you. I plan to address this in #36000 but haven't got to it.
@github-actionsgithub-actionsBot
added

awaiting merge
Awaiting merge
and removed

awaiting changes
Awaiting changes
labels
Jun 28, 2023
@icexelloss@westonpace
Apply suggestions from code review
c29ae72
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@icexelloss
icexelloss
commented
Jun 28, 2023

Copy link


Copy Markdown

ContributorAuthor
Thanks @westonpace. I applied your suggestion and will merge once CI passes.
@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting merge
Awaiting merge
labels
Jun 28, 2023
@icexelloss
icexelloss
commented
Jun 29, 2023

Copy link


Copy Markdown

ContributorAuthor
CI failure is unrelated. Merging.
@icexelloss
icexelloss
merged commit baf17a2
into
apache:mainJun 29, 2023
@conbench-apache-arrow
conbench-apache-arrowBot
commented
Jul 4, 2023

Copy link


Copy Markdown

Conbench analyzed the 6 benchmark runs on commit baf17a20.
There was 1 benchmark result with an error:
There were no benchmark performance regressions. 🎉
The full Conbench report has more details.
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@westonpacewestonpacewestonpace approved these changes
@AlenkaFAlenkaFAwaiting requested review from AlenkaFAlenkaF is a code owner

Assignees


No one assigned


Labels

awaiting changesAwaiting changesComponent: Python

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

[C++][Python] Non decomposable aggregation UDF (Hash version)


2 participants

@icexelloss@westonpace

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Universal Dark Mode - works on any site
(function() {
var enabled = true;
function applyDarkMode() {
if (!enabled) return;
// Create style element if it doesn't exist
var style = document.getElementById('universal-dark-mode-style');
if (!style) {
style = document.createElement('style');
style.id = 'universal-dark-mode-style';
document.head.appendChild(style);
}
// Dark mode CSS - inverts colors but preserves images/video
style.textContent = '
/* Invert everything except media */
html {
filter: invert(1) hue-rotate(180deg) !important;
background: #1a1a2e !important;
}
/* Restore images, videos, iframes, canvas */
img, video, iframe, canvas, svg, picture, [style*="background-image"] {
filter: invert(1) hue-rotate(180deg) !important;
}
/* Preserve specific elements that should not be inverted */
.no-dark-mode, .no-dark-mode *,
[data-theme="light"], [data-theme="light"],
.ace_editor, .ace_editor *,
.CodeMirror, .CodeMirror *,
.monaco-editor, .monaco-editor *,
.markdown-body pre, .markdown-body pre *,
.highlight, .highlight *,
pre code, pre code * {
filter: none !important;
}
/* Fix common UI elements */
.modal, .popup, .dropdown-menu, .tooltip, .popover {
filter: invert(1) hue-rotate(180deg) !important;
background: #2d2d44 !important;
border-color: #444 !important;
}
/* Scrollbars */
::-webkit-scrollbar { background: #1a1a2e !important; }
::-webkit-scrollbar-thumb { background: #444 !important; }
::-webkit-scrollbar-thumb:hover { background: #555 !important; }
/* Selection */
::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
';
}
function removeDarkMode() {
var style = document.getElementById('universal-dark-mode-style');
if (style) style.remove();
}
// Toggle with Alt+Shift+D
document.addEventListener('keydown', function(e) {
if (e.altKey && e.shiftKey && e.key === 'D') {
e.preventDefault();
enabled = !enabled;
if (enabled) {
applyDarkMode();
console.log('[Universal Dark Mode] Enabled');
} else {
removeDarkMode();
console.log('[Universal Dark Mode] Disabled');
}
}
});
// Apply on load
applyDarkMode();
// Re-apply on dynamic content
var observer = new MutationObserver(function(mutations) {
if (enabled && !document.getElementById('universal-dark-mode-style')) {
applyDarkMode();
}
});
observer.observe(document.head, { childList: true });
console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');
})();
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();
GH-36252: [Python] Add non decomposable hash aggregate UDF by icexelloss · Pull Request #36253 · apache/arrow · GitHub
Skip to content
GH-36252: [Python] Add non decomposable hash aggregate UDF  - #36253
Merged
icexelloss merged 11 commits into
apache:mainfrom 
icexelloss:compute-hash-aggregate-UDF
Jun 29, 2023
Merged
GH-36252: [Python] Add non decomposable hash aggregate UDF #36253
icexelloss merged 11 commits into
apache:mainfrom 
icexelloss:compute-hash-aggregate-UDF
Conversation
@icexelloss
@icexellossicexelloss
commented
Jun 22, 2023•

edited

Loading

Copy link


Copy Markdown

Contributor
Rationale for this change
In #35515,
I have implemented a Scalar version of the non decomposable UDF (Scalar as in SCALAR_AGGREGATE). I would like to support the Hash version of it (Hash as in HASH_AGGREGATE)
With this PR, user can register an aggregate UDF once with pc.register_aggregate_function and it can be used as both scalar aggregate function and hash aggregate function.
Example:
def median(x):
return pa.scalar(np.nanmedian(x))
pc.register_aggregate_function(func=median, func_name='median_udf', ...)
table = ...
table.groupby("id").aggregate([("v", 'median_udf')])

What changes are included in this PR?
The main changes are:
  • In ResigterAggregateFunction (udf.cc), we now register the function both as a scalar aggregate function and a hash aggregate function (with signature adjustment for hash aggregate kernel because we need to append the grouping key)
  • Implemented PythonUdfHashAggregateImpl, similar to the PythonUdfScalarAggregateImpl. In Consume, it will accumulate both the input batches and the group id array. In Merge, it will merge the input batches and group id array (with the group_id_mapping). In Finalize, it will apply groupings to the accumulated batches to create one record batch per group, then apply the UDF over each group.
  • Some code clean up - UdfWrapperCallback objects are named cb (previously, agg_cb or wrapper) now and the user defined python function is now just called function (previously agg_function)
For table.groupby().aggregate(...), the space complexity is O(n) where n is the size of the table (and therefore, is not very useful). However, this is more useful in the segmented aggregation case, where the space complexity of O(s), where s the size of the segments.
Are these changes tested?
Added new test in test_udf.py (with table.group_by().aggregate() and test_substrait.py (with segmented aggregation)
Are there any user-facing changes?
Yes with this change, user can call use registered aggregate UDF with table.group_by().aggregate()  or Acero's segmented aggregation.
Checklist
  •  Self Review
  •  API Documentation
@github-actionsgithub-actionsBot
added
the 
awaiting committer review
Awaiting committer review
label
Jun 22, 2023
icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadcpp/src/arrow/compute/row/grouper.h
Outdated

@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting committer review
Awaiting committer review
labels
Jun 22, 2023
icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/conftest.py
Outdated

icexelloss
icexelloss
commented
Jun 22, 2023
View reviewed changes
Comment threadpython/pyarrow/src/arrow/python/udf.cc
 wrapper, options, registry);
}

Status AddAggKernel(std::shared_ptr<compute::KernelSignature> sig,
@icexellossicexellossJun 22, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is inlined now.
@github-actionsgithub-actionsBot
added

Component: Python

awaiting change review
Awaiting change review
and removed

awaiting changes
Awaiting changes
labels
Jun 22, 2023
@icexelloss
icexelloss
commented
Jun 22, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
@westonpace I would like a request a review on this PR. The code should be relatively straight forward and similar to #35514 so hopefully no confusion/surprises here.
For the implementation of the grouping, I used an approach similar to GroupedListImpl aggregator and partition.cc
For the registration, I decided to register both scalar/hash kernel with one API register_aggregate_function because I think scalar/hash difference is not really something user should be worry about (from user's point of view, it is just "aggregation", whether it is hash/scalar is an compute implementation detail).
More details in the PR description.
Let me know if those sounds OK to you.
@icexellossicexelloss
changed the title
GH-36252: [Python] Compute hash aggregate udfGH-36252: [Python] Add non decomposable hash aggregate UDF Jun 22, 2023
@westonpace
westonpace
commented
Jun 23, 2023

Copy link


Copy Markdown

Member
@icexelloss I should have some time to take a look tomorrow
westonpace
westonpace
requested changes
Jun 24, 2023
View reviewed changes
@westonpacewestonpace
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a good set of tests. It's nice and convenient that the same python implementation can work for both. I have a few minor suggestions / questions. I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

 const ArraySpan& groups_array_data = batch[batch.num_values() - 1].array;
 DCHECK_EQ(groups_array_data.offset, 0);
 int64_t batch_num_values = groups_array_data.length;
 const auto* batch_groups = groups_array_data.GetValues<uint32_t>(1, 0);
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not just groups_array_data.GetValues<uint32_t>(1);?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - updated
Comment threadpython/pyarrow/src/arrow/python/udf.cc
 }

 num_values += other.num_values;
 return Status::OK();
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does num_groups need to be updated here?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think num_groups need to be updated here. Reasoning:
From the code in https://github.com/apache/arrow/blob/main/cpp/src/arrow/compute/kernels/hash_aggregate.cc#L233 and https://github.com/apache/arrow/blob/main/cpp/src/arrow/acero/groupby_aggregate_node.cc#L248

(1) Other hash kernel implementation updates the num_groups in Resize

(2) resize is always called before consume and merge
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

 UdfContext udf_context{ctx->memory_pool(), table->num_rows()};

 if (rb->num_rows() == 0) {
 return Status::Invalid("Finalized is called with empty inputs");
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this a problem?
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch - I was being lazy here and didn't want to bother with empty aggregation. But now I look at this I can just return empty result here. Will update.
@icexellossicexellossJun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated
Comment threadpython/pyarrow/src/arrow/python/udf.cc

 ARROW_ASSIGN_OR_RAISE(auto table,
 arrow::Table::FromRecordBatches(input_schema, values));
 ARROW_ASSIGN_OR_RAISE(auto rb, table->CombineChunksToBatch(ctx->memory_pool()));
@westonpacewestonpaceJun 24, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some cases where this won't be possible. For example, if you have an array of strings then the array may only have 2GB of string data (regardless of how many elements it has). So any single group can't have more than 2GB of string data. I don't know this is fatal but you may want to mention in the user docs somewhere or wrap this failure with extra context.
@icexellossicexellossJun 26, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added note in the user doc: https://github.com/apache/arrow/pull/36253/files#diff-439f91c435cc8136d40eaba8c168aaa1cec00f08b10ba92ce376e47f29f81814R2770
Comment threadpython/pyarrow/src/arrow/python/udf.cc
Outdated

@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting change review
Awaiting change review
labels
Jun 24, 2023
@icexelloss@westonpace
Apply suggestions from code review
007221a
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@github-actionsgithub-actionsBot
added

awaiting change review
Awaiting change review
and removed

awaiting changes
Awaiting changes
labels
Jun 26, 2023
@icexelloss
icexelloss
commented
Jun 26, 2023

Copy link


Copy Markdown

ContributorAuthor
I think the only concerning thing is that we need to restrict group sizes to things that will fit in a single batch. Do you think this will be a problem for your use cases?
Thanks @westonpace. Currently we only plan to use this with segmented aggregation so each group is not going to be very large. (grouping inside a segment), so I don't think it would be a problem.
@github-actionsgithub-actionsBot
removed
the 
awaiting change review
Awaiting change review
label
Jun 26, 2023
@github-actionsgithub-actionsBot
added
the 
awaiting changes
Awaiting changes
label
Jun 26, 2023
@github-actionsgithub-actionsBot
added

awaiting change review
Awaiting change review
awaiting changes
Awaiting changes
and removed

awaiting changes
Awaiting changes
awaiting change review
Awaiting change review
labels
Jun 26, 2023
@icexelloss
icexelloss
commented
Jun 26, 2023•

edited

Loading

Copy link


Copy Markdown

ContributorAuthor
@westonpace This should be clean now (all comments addressed, CI green) - another look?
@icexelloss
icexelloss
commented
Jun 27, 2023

Copy link


Copy Markdown

ContributorAuthor
Gentle ping @westonpace anything else you want me to change here?
westonpace
westonpace
approved these changes
Jun 28, 2023
View reviewed changes
@westonpacewestonpace
left a comment


Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor wording suggestion. Otherwise this looks good.
Comment threadpython/pyarrow/_compute.pyx
Outdated

Comment threadpython/pyarrow/src/arrow/python/udf.cc
 std::vector<std::shared_ptr<DataType>> input_types,
 std::shared_ptr<DataType> output_type)
 : function(function), cb(std::move(cb)), output_type(std::move(output_type)) {
 Py_INCREF(function->obj());
@westonpacewestonpaceJun 28, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These INCREF's still seem superfluous to me but I don't think it's critical. We could test in a follow-up using temporary function registries to see if we are preventing UDF functions from being garbage collected.
@icexellossicexellossJun 28, 2023

Copy link


Copy Markdown

ContributorAuthor
Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with you. I plan to address this in #36000 but haven't got to it.
@github-actionsgithub-actionsBot
added

awaiting merge
Awaiting merge
and removed

awaiting changes
Awaiting changes
labels
Jun 28, 2023
@icexelloss@westonpace
Apply suggestions from code review
c29ae72
Co-authored-by: Weston Pace <weston.pace@gmail.com>
@icexelloss
icexelloss
commented
Jun 28, 2023

Copy link


Copy Markdown

ContributorAuthor
Thanks @westonpace. I applied your suggestion and will merge once CI passes.
@github-actionsgithub-actionsBot
added

awaiting changes
Awaiting changes
and removed

awaiting merge
Awaiting merge
labels
Jun 28, 2023
@icexelloss
icexelloss
commented
Jun 29, 2023

Copy link


Copy Markdown

ContributorAuthor
CI failure is unrelated. Merging.
@icexelloss
icexelloss
merged commit baf17a2
into
apache:mainJun 29, 2023
@conbench-apache-arrow
conbench-apache-arrowBot
commented
Jul 4, 2023

Copy link


Copy Markdown

Conbench analyzed the 6 benchmark runs on commit baf17a20.
There was 1 benchmark result with an error:
There were no benchmark performance regressions. 🎉
The full Conbench report has more details.
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment

Reviewers

@westonpacewestonpacewestonpace approved these changes
@AlenkaFAlenkaFAwaiting requested review from AlenkaFAlenkaF is a code owner

Assignees


No one assigned


Labels

awaiting changesAwaiting changesComponent: Python

Projects


None yet


Milestone


No milestone


Development

Successfully merging this pull request may close these issues.

[C++][Python] Non decomposable aggregation UDF (Hash version)


2 participants

@icexelloss@westonpace