feat(rust, python): in rolling aggregation functions, sort under-the-hood for user if data is unsorted (with warning) #11134
Conversation
github-actions
bot
added
enhancement
MarcoGorelli
force-pushed
the
sort-unsort
branch
6 times, most recently
from
e3ccb77
to
435a2e9
Compare
MarcoGorelli
force-pushed
the
sort-unsort
branch
2 times, most recently
from
711fd19
to
6ed0188
Compare
…hood for user if data is unsorted (with warning)
MarcoGorelli
force-pushed
the
sort-unsort
branch
from
6ed0188
to
9f8d828
Compare
MarcoGorelli
requested review from
ritchie46,
stinodego,
alexander-beedie and
orlp
as code owners
|
Was there some discussion on this? I am curious if we should do this implicitly or require more strictness/awareness. |
|
yup the discussion's all in this thread #10991 (comment), sorry for not having linked it right away |
crates/polars-plan/src/dsl/mod.rs
Outdated
| }, | ||
| IsSorted::Not => { | ||
| if options.warn_if_unsorted { | ||
| eprintln!( |
There was a problem hiding this comment.
Can you use polars_error's polars_warn function?
| unsafe { by = by.take_unchecked(&sorting_indices)? }; | ||
| unsafe { series = s[0].take_unchecked(&sorting_indices)? }; | ||
| let int_range = | ||
| UInt32Chunked::from_iter_values("", 0..s[0].len() as u32) |
There was a problem hiding this comment.
This should be IdxCa and IdxSize instead of u32, no? EDIT: see below, you shouldn't save this at all and instead save sorting_indices.
| }, | ||
| IsSorted::Not => { | ||
| let res = rolling_fn(&series, options)?; | ||
| let indices = &original_indices.unwrap().arg_sort(Default::default()); |
There was a problem hiding this comment.
This should not be another sort and take, instead this should be a scatter: out[sorting_indices[i]] = res[i]. We currently lack a good scatter kernel though, as far as I'm aware. So this may be blocked until I get to that.
There was a problem hiding this comment.
thanks - OK marking as blocked then
There was a problem hiding this comment.
hi @orlp - just checking whether a scatter kernel might be on the horizon, and if not whether you have suggestions on what to do instead
There was a problem hiding this comment.
@MarcoGorelli We don't have a scatter kernel yet, but we do have PolarsDataType::ZeroablePhysical now, which allows you to write a reasonably efficient scatter in a couple lines of code.
You can look in gather_skip_nulls.rs for an example, but basically you do
let mut out: Vec<T::ZeroablePhysical<'a>> = zeroed_vec(len);
unsafe {
for i in 0..n {
let out_idx = *sorted_indices.get_unchecked(i);
*out.get_unchecked_mut(out_idx) = res.get(i).into();
}
}
let out_arr = T::Array::from_zeroable_vec(out, dtype);There was a problem hiding this comment.
Depending on the behavior of nulls you may also have to construct a null bitmap.
There was a problem hiding this comment.
Thanks
Noob question but where would I get T from in this function? Do I first need to refactor and make finish_rolling generic?
MarcoGorelli
added
the
blocked
|
closing for now as I think doing this correctly is a bit above me at the moment, but @orlp feel free to take over, would be very interested to see how you'd do it |
got a bit of time in the end, so trying to address this