How to handle null values in a UDAF?

[rspears74]

I am trying to handle null values in the update_batch method of an Accumulator when defining a UDAF. I've been trying something like the following:

fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
    let vals: Vec<String> = (&values[0])
        .as_string::<i32>()
        .iter()
        .flatten()
        .map(|s| String::from(s))
        .collect();

    self.merge_values(vals);
    Ok(())
}

... but I get

thread 'tokio-runtime-worker' panicked at /.../.cargo/registry/src/index.crates.io-6f17d22bba15001f/arrow-data-49.0.0/src/transform/list.rs:36:69:
range end index 2 out of range for slice of length 1

I was able to make it work for numeric types by sort of mimicking what some of the internal aggregators do:

fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
    let vals = (&values[0])
        .as_primitive::<Float64Type>()
        .values()
        .to_vec();

    self.merge_values(vals);
    Ok(())
}

... but this doesn't seem possible with Strings since calling values on a StringArray returns a Vec<u8> instead of a Vec<&str> (or String). Is there something I'm doing wrong in the first example where I'm using flatten? That should flatten a list of [Some("string"), Some("other string"), None] to ["string", "other string"], but I suppose there's something going on internally that's not allowing that to work.

[haohuaijin]

maybe you can use the code in string_agg.
https://github.com/apache/arrow-datafusion/blob/ee7ab0b83ff5bf2934519cd6718ccea620edea1e/datafusion/physical-expr/src/aggregate/string_agg.rs#L134-L139

[rspears74]

Hmm, I still the same error with this method. Maybe I have a problem elsewhere.

[rspears74]

According to my stack trace, this is the last spot where I'm in datafusion https://github.com/apache/arrow-datafusion/blob/26933842e48d69f510f9461a1f2c87af587d5986/datafusion/physical-expr/src/aggregate/groups_accumulator/adapter.rs#L305. Here's the relevant section of the stack trace:

16:        0x104378218 - core::slice::index::slice_end_index_len_fail_rt::h0681d576c470bd14
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/slice/index.rs:76:5
  17:        0x104378218 - core::slice::index::slice_end_index_len_fail::h7c5d9dfb02ae2622
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/slice/index.rs:68:9
  18:        0x10418337c - <core::ops::range::Range<usize> as core::slice::index::SliceIndex<[T]>>::index::h8c86056ae34ef931
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/slice/index.rs:411:13
  19:        0x1041cbea0 - core::slice::index::<impl core::ops::index::Index<I> for [T]>::index::ha2d31dd7721987c2
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/slice/index.rs:18:9
  20:        0x1041beac8 - arrow_data::transform::list::build_extend::{{closure}}::hf5338b7ffd8b0e21
                               at /Users/PHUSS6E/.cargo/registry/src/index.crates.io-6f17d22bba15001f/arrow-data-49.0.0/src/transform/list.rs:36:69
  21:        0x1041a1cb0 - <alloc::boxed::Box<F,A> as core::ops::function::Fn<Args>>::call::h65bfd06117252cbd
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/alloc/src/boxed.rs:2021:9
  22:        0x1041b3268 - arrow_data::transform::MutableArrayData::extend::h7fa915339d2267ec
                               at /Users/PHUSS6E/.cargo/registry/src/index.crates.io-6f17d22bba15001f/arrow-data-49.0.0/src/transform/mod.rs:611:9
  23:        0x1025eb7a0 - datafusion_common::scalar::ScalarValue::iter_to_array::build_list_array::hf994cbec5f0c2c50
                               at /Users/PHUSS6E/.cargo/registry/src/index.crates.io-6f17d22bba15001f/datafusion-common-34.0.0/src/scalar.rs:1387:21
  24:        0x1025e59d4 - datafusion_common::scalar::ScalarValue::iter_to_array::hf8faadf1214b6f17
                               at /Users/PHUSS6E/.cargo/registry/src/index.crates.io-6f17d22bba15001f/datafusion-common-34.0.0/src/scalar.rs:1468:59
  25:        0x1022cca20 - core::ops::function::FnMut::call_mut::ha849797f138c850c
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/ops/function.rs:166:5
  26:        0x101d6fe5c - core::iter::adapters::map::map_try_fold::{{closure}}::h88058b2ae9583aee
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/iter/adapters/map.rs:91:28
  27:        0x101ffad88 - core::iter::traits::iterator::Iterator::try_fold::h4214caaf7a7a1fd2
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/iter/traits/iterator.rs:2461:21
  28:        0x101d50d80 - <core::iter::adapters::map::Map<I,F> as core::iter::traits::iterator::Iterator>::try_fold::hd6e1b6d53c2cbb33
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/iter/adapters/map.rs:117:9
  29:        0x101cc769c - <core::iter::adapters::GenericShunt<I,R> as core::iter::traits::iterator::Iterator>::try_fold::h46d2196d1594abed
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/iter/adapters/mod.rs:199:9
  30:        0x101cbff08 - core::iter::traits::iterator::Iterator::try_for_each::h7aac0fe70898c6d8
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/iter/traits/iterator.rs:2523:9
  31:        0x101cbff08 - <core::iter::adapters::GenericShunt<I,R> as core::iter::traits::iterator::Iterator>::next::hbfb96ea280234f10
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/iter/adapters/mod.rs:182:14
  32:        0x101c0f8d8 - <alloc::vec::Vec<T> as alloc::vec::spec_from_iter_nested::SpecFromIterNested<T,I>>::from_iter::h03711e7132d5c742
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/alloc/src/vec/spec_from_iter_nested.rs:26:32
  33:        0x101c673cc - alloc::vec::in_place_collect::<impl alloc::vec::spec_from_iter::SpecFromIter<T,I> for alloc::vec::Vec<T>>::from_iter::hf6003222470dd4d7
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/alloc/src/vec/in_place_collect.rs:167:20
  34:        0x101c6f384 - <alloc::vec::Vec<T> as core::iter::traits::collect::FromIterator<T>>::from_iter::h31faf12aa816df9d
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/alloc/src/vec/mod.rs:2749:9
  35:        0x1022c521c - core::iter::traits::iterator::Iterator::collect::ha7b09e1c4eff033a
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/iter/traits/iterator.rs:2053:9
  36:        0x1022c521c - <core::result::Result<V,E> as core::iter::traits::collect::FromIterator<core::result::Result<A,E>>>::from_iter::{{closure}}::he676bc3b91ef6b73
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/result.rs:1933:51
  37:        0x101d1f144 - core::iter::adapters::try_process::hc1ae3e2152e234d0
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/iter/adapters/mod.rs:168:17
  38:        0x1022c24dc - <core::result::Result<V,E> as core::iter::traits::collect::FromIterator<core::result::Result<A,E>>>::from_iter::hd2848351d632b3a8
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/result.rs:1933:9
  39:        0x101d5f73c - core::iter::traits::iterator::Iterator::collect::h17a2951b68f72865
                               at /rustc/a28077b28a02b92985b3a3faecf92813155f1ea1/library/core/src/iter/traits/iterator.rs:2053:9
  40:        0x101facaf4 - <datafusion_physical_expr::aggregate::groups_accumulator::adapter::GroupsAccumulatorAdapter as datafusion_physical_expr::aggregate::groups_accumulator::GroupsAccumulator>::state::hcccd4ac9cf03e89c
                               at /Users/PHUSS6E/.cargo/registry/src/index.crates.io-6f17d22bba15001f/datafusion-physical-expr-34.0.0/src/aggregate/groups_accumulator/adapter.rs:305:22
  41:        0x101944834 - datafusion_physical_plan::aggregates::row_hash::GroupedHashAggregateStream::emit::h61c42a3c5293430a
                               at /Users/PHUSS6E/.cargo/registry/src/index.crates.io-6f17d22bba15001f/datafusion-physical-plan-34.0.0/src/aggregates/row_hash.rs:644:57
  42:        0x101946440 - datafusion_physical_plan::aggregates::row_hash::GroupedHashAggregateStream::set_input_done_and_produce_output::h515ba0821c105cc1
                               at /Users/PHUSS6E/.cargo/registry/src/index.crates.io-6f17d22bba15001f/datafusion-physical-plan-34.0.0/src/aggregates/row_hash.rs:789:25
  43:        0x101942600 - <datafusion_physical_plan::aggregates::row_hash::GroupedHashAggregateStream as futures_core::stream::Stream>::poll_next::hdbc42528a9c45788
                               at /Users/PHUSS6E/.cargo/registry/src/index.crates.io-6f17d22bba15001f/datafusion-physical-plan-34.0.0/src/aggregates/row_hash.rs:473:41
  44:        0x101b42ea4 - <core::pin::Pin<P> as futures_core::stream::Stream>::poll_next::h2748215f59e8b8d5
                               at /Users/PHUSS6E/.cargo/registry/src/index.crates.io-6f17d22bba15001f/futures-core-0.3.30/src/stream.rs:120:9
  45:        0x101b3d204 - futures_util::stream::stream::StreamExt::poll_next_unpin::ha7a9da89bdf720c5
                               at /Users/PHUSS6E/.cargo/registry/src/index.crates.io-6f17d22bba15001f/futures-util-0.3.30/src/stream/stream/mod.rs:1638:9
  46:        0x1019e7450 - <futures_util::stream::stream::next::Next<St> as core::future::future::Future>::poll::hfa814403d6acf5ea
                               at /Users/PHUSS6E/.cargo/registry/src/index.crates.io-6f17d22bba15001f/futures-util-0.3.30/src/stream/stream/next.rs:32:9
  47:        0x101a7b86c - datafusion_physical_plan::repartition::RepartitionExec::pull_from_input::{{closure}}::h0f1d8c88f93e8310
                               at /Users/PHUSS6E/.cargo/registry/src/index.crates.io-6f17d22bba15001f/datafusion-physical-plan-34.0.0/src/repartition/mod.rs:702:40

[rspears74]

I'm now thinking that my original .flatten strategy likely works correctly, but that there's a bug in terms of writing out the empty lists that I'm aggregating. I think my strategy and this are both valid strategies of only aggregating non-null values.