jayzhan211 commented on code in PR #26357:
URL: https://github.com/apache/datafusion/pull/26357#discussion_r4294697539
##########
datafusion/functions-aggregate/src/array_agg.rs:
##########
@@ -832,6 +834,418 @@ impl GroupsAccumulator for ArrayAggGroupsAccumulator {
}
}
+/// Converts each input row into a one-element list, the state shape of
+/// `ARRAY_AGG`.
+///
+/// Filtered rows become null list entries, which `merge_batch` skips. With
+/// `ignore_nulls`, null values also become null list entries; without it they
+/// stay as `[NULL]` so `merge_batch` retains them.
+fn single_value_lists(
+ input: &ArrayRef,
+ opt_filter: Option<&BooleanArray>,
+ ignore_nulls: bool,
+ datatype: &DataType,
+) -> ArrayRef {
+ // Offsets are [0, 1, 2, ..., n].
+ let offsets = OffsetBuffer::from_repeated_length(1, input.len());
+ let filter_nulls = opt_filter.map(filter_to_nulls);
+ let nulls = if ignore_nulls {
+ let logical = input.logical_nulls();
+ NullBuffer::union(filter_nulls.as_ref(), logical.as_ref())
+ } else {
+ filter_nulls
+ };
+ let field = Arc::new(Field::new_list_field(datatype.clone(), true));
+ Arc::new(ListArray::new(field, offsets, Arc::clone(input), nulls))
+}
+
+/// [`GroupsAccumulator`] for `ARRAY_AGG(DISTINCT x)`, optionally with
+/// `ORDER BY x`.
+///
+/// Distinct values are encoded once with a [`RowConverter`] and stored in a
+/// single [`Rows`] buffer shared by all groups, so a value that appears in
many
+/// groups is held once. Each group records the values it has seen as
+/// `(group_index, value_index)` pairs, deduplicated in one hash set rather
than
+/// one set per group.
+#[derive(Debug)]
+struct DistinctArrayAggGroupsAccumulator {
+ datatype: DataType,
+ sort_options: Option<SortOptions>,
+ ignore_nulls: bool,
+ /// Encodes values to rows. Created on the first non-empty input, once the
+ /// runtime type is known.
+ converter: Option<RowConverter>,
+ /// Every distinct value seen so far, indexed by value index.
+ values: Option<Rows>,
+ /// `(hash, value_index)` for every row in `values`.
+ value_map: HashTable<(u64, u32)>,
+ /// Heap size of `value_map` in bytes.
+ value_map_size: usize,
+ /// Distinct `(group_index, value_index)` pairs, in first-seen order.
+ pairs: Vec<(u32, u32)>,
+ /// Deduplicates `pairs`.
+ pair_set: HashSet<(u32, u32), RandomState>,
+ /// Total number of groups tracked.
+ num_groups: usize,
+ random_state: RandomState,
+ /// Time spent deduplicating values, reported as the DISTINCT submetric.
+ distinct_metric: Option<Arc<dyn AggregateMetric>>,
+}
+
+impl DistinctArrayAggGroupsAccumulator {
+ fn new(
+ datatype: DataType,
+ sort_options: Option<SortOptions>,
+ ignore_nulls: bool,
+ ) -> Self {
+ Self {
+ datatype,
+ sort_options,
+ ignore_nulls,
+ converter: None,
+ values: None,
+ value_map: HashTable::new(),
+ value_map_size: 0,
+ pairs: Vec::new(),
+ pair_set: HashSet::default(),
+ num_groups: 0,
+ random_state: RandomState::default(),
+ distinct_metric: None,
+ }
+ }
+
+ fn converter(&mut self, data_type: &DataType) -> Result<&RowConverter> {
+ if self.converter.is_none() {
+ let sort_field = match self.sort_options {
+ Some(opts) => SortField::new_with_options(data_type.clone(),
opts),
+ None => SortField::new(data_type.clone()),
+ };
+ let converter = RowConverter::new(vec![sort_field])?;
+ self.values = Some(converter.empty_rows(0, 0));
+ self.converter = Some(converter);
+ }
+ self.converter
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("row converter must be
set"))
+ }
+
+ /// Like [`Self::insert`], but first gathers the selected rows when they
+ /// are only part of `input`, so rejected values are never encoded.
+ fn insert_selected(
+ &mut self,
+ input: &ArrayRef,
+ rows: Vec<(usize, usize)>,
+ ) -> Result<()> {
+ if rows.len() == input.len() {
+ return self.insert(input, rows);
+ }
+ let indices = UInt32Array::from_iter_values(
+ rows.iter().map(|&(_, row_index)| row_index as u32),
+ );
+ let selected = take_array(input.as_ref(), &indices, None)?;
+ self.insert(
+ &selected,
+ rows.into_iter()
+ .enumerate()
+ .map(|(position, (group_index, _))| (group_index, position)),
+ )
+ }
+
+ /// Records that `group_index` contains each value of `input` selected by
+ /// `rows`, given as `(group_index, row_in_input)` pairs.
+ fn insert(
+ &mut self,
+ input: &ArrayRef,
+ rows: impl IntoIterator<Item = (usize, usize)>,
+ ) -> Result<()> {
+ let encoded = self
+ .converter(input.data_type())?
+ .convert_columns(std::slice::from_ref(input))?;
+ let Self {
+ values,
+ value_map,
+ value_map_size,
+ pairs,
+ pair_set,
+ random_state,
+ ..
+ } = self;
+ let values = values
+ .as_mut()
+ .ok_or_else(|| internal_datafusion_err!("value rows must be
set"))?;
+
+ for (group_index, row_index) in rows {
+ let row = encoded.row(row_index);
+ let hash = random_state.hash_one(row.as_ref());
+ let value_index = match value_map.find(hash, |&(h, index)| {
+ h == hash && values.row(index as usize) == row
+ }) {
+ Some(&(_, index)) => index,
+ None => {
+ let index = values.num_rows() as u32;
+ values.push(row);
+ value_map.insert_accounted(
+ (hash, index),
+ |&(h, _)| h,
+ value_map_size,
+ );
+ index
+ }
+ };
+ let pair = (group_index as u32, value_index);
+ if pair_set.insert(pair) {
+ pairs.push(pair);
+ }
+ }
+ Ok(())
+ }
+
+ /// Drops the state of the first `emit_groups` groups and renumbers the
+ /// rest so they start from 0. Values only referenced by emitted groups are
+ /// dropped too.
+ fn retain_after(&mut self, emit_groups: usize) {
+ let emit_groups = emit_groups as u32;
+ let old_pairs = take(&mut self.pairs);
+ let old_values = self.values.take();
+ self.value_map = HashTable::new();
+ self.value_map_size = 0;
+ self.pair_set = HashSet::default();
+ self.num_groups -= emit_groups as usize;
+
+ let (Some(converter), Some(old_values)) = (&self.converter,
old_values) else {
+ return;
+ };
+ let mut values = converter.empty_rows(0, 0);
+ let mut remap: HashMap<u32, u32> = HashMap::new();
+ let mut pairs = Vec::new();
+ for (group_index, value_index) in old_pairs {
+ if group_index < emit_groups {
+ continue;
+ }
+ let new_index = *remap.entry(value_index).or_insert_with(|| {
+ let row = old_values.row(value_index as usize);
+ let index = values.num_rows() as u32;
+ values.push(row);
+ let hash = self.random_state.hash_one(row.as_ref());
+ self.value_map.insert_accounted(
+ (hash, index),
+ |&(h, _)| h,
+ &mut self.value_map_size,
+ );
+ index
+ });
+ let pair = (group_index - emit_groups, new_index);
+ self.pair_set.insert(pair);
+ pairs.push(pair);
+ }
+ self.pairs = pairs;
+ self.values = Some(values);
+ }
+}
+
+impl GroupsAccumulator for DistinctArrayAggGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> Result<()> {
+ // With `ORDER BY x` the ordering column is passed too. It is the
+ // argument itself, so only the first column is needed.
+ let Some(input) = values.first() else {
+ return internal_err!("array_agg(DISTINCT) expects an argument");
+ };
+ self.num_groups = self.num_groups.max(total_num_groups);
+
+ let nulls = if self.ignore_nulls {
+ input.logical_nulls()
+ } else {
+ None
+ };
+ let rows = group_indices
+ .iter()
+ .enumerate()
+ .filter(|&(row_index, _)| {
+ let filtered_out = opt_filter.is_some_and(|filter| {
+ filter.is_null(row_index) || !filter.value(row_index)
+ });
+ let skipped_null =
+ nulls.as_ref().is_some_and(|nulls|
nulls.is_null(row_index));
+ !filtered_out && !skipped_null
+ })
+ .map(|(row_index, &group_index)| (group_index, row_index))
+ .collect::<Vec<_>>();
+ if rows.is_empty() {
+ return Ok(());
+ }
+ let mut metric_recorder =
+ AggregateMetricRecorder::new(self.distinct_metric.clone());
+ let _timer = metric_recorder.timer();
+ self.insert_selected(input, rows)
+ }
+
+ /// Produces a `ListArray` ordered by group index. Groups without values
+ /// are null, matching [`DistinctArrayAggAccumulator`].
+ fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
+ let emit_groups = match emit_to {
+ EmitTo::All => self.num_groups,
+ EmitTo::First(n) => n,
+ };
+
+ // Counting sort of the pairs into group order.
+ let mut counts = vec![0u32; emit_groups];
+ for &(group_index, _) in &self.pairs {
+ if let Some(count) = counts.get_mut(group_index as usize) {
+ *count += 1;
+ }
+ }
+ let mut offsets = Vec::<i32>::with_capacity(emit_groups + 1);
+ offsets.push(0);
+ let mut nulls = NullBufferBuilder::new(emit_groups);
+ let mut write_positions = Vec::with_capacity(emit_groups);
+ let mut total = 0u32;
+ for &count in &counts {
+ if count == 0 {
+ nulls.append_null();
+ } else {
+ nulls.append_non_null();
+ }
+ write_positions.push(total);
+ total += count;
+ offsets.push(total as i32);
+ }
+ let mut ordered = vec![0u32; total as usize];
+ for &(group_index, value_index) in &self.pairs {
+ if let Some(position) = write_positions.get_mut(group_index as
usize) {
+ ordered[*position as usize] = value_index;
+ *position += 1;
+ }
+ }
+
+ let flat_values = match (&self.converter, &self.values) {
+ (Some(converter), Some(values)) if total > 0 => {
+ if self.sort_options.is_some() {
+ // The sort direction is encoded in the row bytes, so
+ // comparing rows gives the requested order.
+ for window in offsets.windows(2) {
+ ordered[window[0] as usize..window[1] as
usize].sort_unstable_by(
+ |&a, &b| values.row(a as usize).cmp(&values.row(b
as usize)),
+ );
+ }
+ }
+ let decoded = converter
+ .convert_rows(ordered.iter().map(|&i| values.row(i as
usize)))?
+ .swap_remove(0);
+ // `convert_rows` always returns the physical (non-dictionary)
+ // type, so cast back when the declared type has a dictionary.
+ if decoded.data_type() != &self.datatype
+ && datatype_contains_dictionary(&self.datatype)
+ {
+ cast(decoded.as_ref(), &self.datatype)?
+ } else {
+ decoded
+ }
+ }
+ _ => new_empty_array(&self.datatype),
+ };
+
+ match emit_to {
+ EmitTo::All => {
+ let converter = self.converter.take();
+ let distinct_metric = self.distinct_metric.take();
+ *self = Self::new(
+ self.datatype.clone(),
+ self.sort_options,
+ self.ignore_nulls,
+ );
+ self.distinct_metric = distinct_metric;
+ // Keep the converter: the runtime type does not change.
+ if let Some(converter) = converter {
+ self.values = Some(converter.empty_rows(0, 0));
+ self.converter = Some(converter);
+ }
+ }
+ EmitTo::First(_) => self.retain_after(emit_groups),
+ }
+
+ let field = Arc::new(Field::new_list_field(self.datatype.clone(),
true));
+ Ok(Arc::new(ListArray::new(
+ field,
+ OffsetBuffer::new(ScalarBuffer::from(offsets)),
+ flat_values,
+ nulls.finish(),
+ )))
+ }
+
+ fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+ Ok(vec![self.evaluate(emit_to)?])
+ }
+
+ fn merge_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ total_num_groups: usize,
+ ) -> Result<()> {
+ assert_eq_or_internal_err!(values.len(), 1, "one argument to
merge_batch");
+ let lists = values[0].as_list::<i32>();
+ self.num_groups = self.num_groups.max(total_num_groups);
+
+ // A sliced list keeps its whole child array, so encode only the span
+ // its offsets cover.
+ let offsets = lists.offsets();
+ let (span_start, span_len) = offset_span(offsets);
+ let child = lists.values().slice(span_start, span_len);
+ let rows = group_indices
+ .iter()
+ .enumerate()
+ .filter(|&(list_index, _)| lists.is_valid(list_index))
+ .flat_map(|(list_index, &group_index)| {
+ let start = offsets[list_index] as usize - span_start;
+ let end = offsets[list_index + 1] as usize - span_start;
+ (start..end).map(move |position| (group_index, position))
+ })
+ .collect::<Vec<_>>();
+ if rows.is_empty() {
+ return Ok(());
+ }
+ let mut metric_recorder =
+ AggregateMetricRecorder::new(self.distinct_metric.clone());
+ let _timer = metric_recorder.timer();
+ self.insert_selected(&child, rows)
+ }
+
+ fn convert_to_state(
+ &self,
+ values: &[ArrayRef],
+ opt_filter: Option<&BooleanArray>,
+ ) -> Result<Vec<ArrayRef>> {
+ let Some(input) = values.first() else {
+ return internal_err!("array_agg(DISTINCT) expects an argument");
+ };
+ Ok(vec![single_value_lists(
+ input,
+ opt_filter,
+ self.ignore_nulls,
+ &self.datatype,
+ )])
+ }
+
+ fn set_metrics(&mut self, metrics: Arc<dyn AggregateMetrics>) {
+ self.distinct_metric = Some(metrics.metric("distinct"));
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self)
+ + self.converter.as_ref().map_or(0, RowConverter::size)
+ + self.values.as_ref().map_or(0, Rows::size)
+ + self.value_map_size
+ + self.pairs.capacity() * size_of::<(u32, u32)>()
+ + self.pair_set.capacity() * (size_of::<(u32, u32)>() + 1)
Review Comment:
```suggestion
+ estimate_memory_size::<(u32, u32)>(self.pair_set.capacity(),
0).unwrap_or_default()
```
##########
datafusion/functions-aggregate/src/array_agg.rs:
##########
@@ -832,6 +834,418 @@ impl GroupsAccumulator for ArrayAggGroupsAccumulator {
}
}
+/// Converts each input row into a one-element list, the state shape of
+/// `ARRAY_AGG`.
+///
+/// Filtered rows become null list entries, which `merge_batch` skips. With
+/// `ignore_nulls`, null values also become null list entries; without it they
+/// stay as `[NULL]` so `merge_batch` retains them.
+fn single_value_lists(
+ input: &ArrayRef,
+ opt_filter: Option<&BooleanArray>,
+ ignore_nulls: bool,
+ datatype: &DataType,
+) -> ArrayRef {
+ // Offsets are [0, 1, 2, ..., n].
+ let offsets = OffsetBuffer::from_repeated_length(1, input.len());
+ let filter_nulls = opt_filter.map(filter_to_nulls);
+ let nulls = if ignore_nulls {
+ let logical = input.logical_nulls();
+ NullBuffer::union(filter_nulls.as_ref(), logical.as_ref())
+ } else {
+ filter_nulls
+ };
+ let field = Arc::new(Field::new_list_field(datatype.clone(), true));
+ Arc::new(ListArray::new(field, offsets, Arc::clone(input), nulls))
+}
+
+/// [`GroupsAccumulator`] for `ARRAY_AGG(DISTINCT x)`, optionally with
+/// `ORDER BY x`.
+///
+/// Distinct values are encoded once with a [`RowConverter`] and stored in a
+/// single [`Rows`] buffer shared by all groups, so a value that appears in
many
+/// groups is held once. Each group records the values it has seen as
+/// `(group_index, value_index)` pairs, deduplicated in one hash set rather
than
+/// one set per group.
+#[derive(Debug)]
+struct DistinctArrayAggGroupsAccumulator {
+ datatype: DataType,
+ sort_options: Option<SortOptions>,
+ ignore_nulls: bool,
+ /// Encodes values to rows. Created on the first non-empty input, once the
+ /// runtime type is known.
+ converter: Option<RowConverter>,
+ /// Every distinct value seen so far, indexed by value index.
+ values: Option<Rows>,
+ /// `(hash, value_index)` for every row in `values`.
+ value_map: HashTable<(u64, u32)>,
+ /// Heap size of `value_map` in bytes.
+ value_map_size: usize,
+ /// Distinct `(group_index, value_index)` pairs, in first-seen order.
+ pairs: Vec<(u32, u32)>,
+ /// Deduplicates `pairs`.
+ pair_set: HashSet<(u32, u32), RandomState>,
+ /// Total number of groups tracked.
+ num_groups: usize,
+ random_state: RandomState,
+ /// Time spent deduplicating values, reported as the DISTINCT submetric.
+ distinct_metric: Option<Arc<dyn AggregateMetric>>,
+}
+
+impl DistinctArrayAggGroupsAccumulator {
+ fn new(
+ datatype: DataType,
+ sort_options: Option<SortOptions>,
+ ignore_nulls: bool,
+ ) -> Self {
+ Self {
+ datatype,
+ sort_options,
+ ignore_nulls,
+ converter: None,
+ values: None,
+ value_map: HashTable::new(),
+ value_map_size: 0,
+ pairs: Vec::new(),
+ pair_set: HashSet::default(),
+ num_groups: 0,
+ random_state: RandomState::default(),
+ distinct_metric: None,
+ }
+ }
+
+ fn converter(&mut self, data_type: &DataType) -> Result<&RowConverter> {
+ if self.converter.is_none() {
+ let sort_field = match self.sort_options {
+ Some(opts) => SortField::new_with_options(data_type.clone(),
opts),
+ None => SortField::new(data_type.clone()),
+ };
+ let converter = RowConverter::new(vec![sort_field])?;
+ self.values = Some(converter.empty_rows(0, 0));
+ self.converter = Some(converter);
+ }
+ self.converter
+ .as_ref()
+ .ok_or_else(|| internal_datafusion_err!("row converter must be
set"))
+ }
+
+ /// Like [`Self::insert`], but first gathers the selected rows when they
+ /// are only part of `input`, so rejected values are never encoded.
+ fn insert_selected(
+ &mut self,
+ input: &ArrayRef,
+ rows: Vec<(usize, usize)>,
+ ) -> Result<()> {
+ if rows.len() == input.len() {
+ return self.insert(input, rows);
+ }
+ let indices = UInt32Array::from_iter_values(
+ rows.iter().map(|&(_, row_index)| row_index as u32),
+ );
+ let selected = take_array(input.as_ref(), &indices, None)?;
+ self.insert(
+ &selected,
+ rows.into_iter()
+ .enumerate()
+ .map(|(position, (group_index, _))| (group_index, position)),
+ )
+ }
+
+ /// Records that `group_index` contains each value of `input` selected by
+ /// `rows`, given as `(group_index, row_in_input)` pairs.
+ fn insert(
+ &mut self,
+ input: &ArrayRef,
+ rows: impl IntoIterator<Item = (usize, usize)>,
+ ) -> Result<()> {
+ let encoded = self
+ .converter(input.data_type())?
+ .convert_columns(std::slice::from_ref(input))?;
+ let Self {
+ values,
+ value_map,
+ value_map_size,
+ pairs,
+ pair_set,
+ random_state,
+ ..
+ } = self;
+ let values = values
+ .as_mut()
+ .ok_or_else(|| internal_datafusion_err!("value rows must be
set"))?;
+
+ for (group_index, row_index) in rows {
+ let row = encoded.row(row_index);
+ let hash = random_state.hash_one(row.as_ref());
+ let value_index = match value_map.find(hash, |&(h, index)| {
+ h == hash && values.row(index as usize) == row
+ }) {
+ Some(&(_, index)) => index,
+ None => {
+ let index = values.num_rows() as u32;
+ values.push(row);
+ value_map.insert_accounted(
+ (hash, index),
+ |&(h, _)| h,
+ value_map_size,
+ );
+ index
+ }
+ };
+ let pair = (group_index as u32, value_index);
+ if pair_set.insert(pair) {
+ pairs.push(pair);
+ }
+ }
+ Ok(())
+ }
+
+ /// Drops the state of the first `emit_groups` groups and renumbers the
+ /// rest so they start from 0. Values only referenced by emitted groups are
+ /// dropped too.
+ fn retain_after(&mut self, emit_groups: usize) {
+ let emit_groups = emit_groups as u32;
+ let old_pairs = take(&mut self.pairs);
+ let old_values = self.values.take();
+ self.value_map = HashTable::new();
+ self.value_map_size = 0;
+ self.pair_set = HashSet::default();
+ self.num_groups -= emit_groups as usize;
+
+ let (Some(converter), Some(old_values)) = (&self.converter,
old_values) else {
+ return;
+ };
+ let mut values = converter.empty_rows(0, 0);
+ let mut remap: HashMap<u32, u32> = HashMap::new();
+ let mut pairs = Vec::new();
+ for (group_index, value_index) in old_pairs {
+ if group_index < emit_groups {
+ continue;
+ }
+ let new_index = *remap.entry(value_index).or_insert_with(|| {
+ let row = old_values.row(value_index as usize);
+ let index = values.num_rows() as u32;
+ values.push(row);
+ let hash = self.random_state.hash_one(row.as_ref());
+ self.value_map.insert_accounted(
+ (hash, index),
+ |&(h, _)| h,
+ &mut self.value_map_size,
+ );
+ index
+ });
+ let pair = (group_index - emit_groups, new_index);
+ self.pair_set.insert(pair);
+ pairs.push(pair);
+ }
+ self.pairs = pairs;
+ self.values = Some(values);
+ }
+}
+
+impl GroupsAccumulator for DistinctArrayAggGroupsAccumulator {
+ fn update_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ opt_filter: Option<&BooleanArray>,
+ total_num_groups: usize,
+ ) -> Result<()> {
+ // With `ORDER BY x` the ordering column is passed too. It is the
+ // argument itself, so only the first column is needed.
+ let Some(input) = values.first() else {
+ return internal_err!("array_agg(DISTINCT) expects an argument");
+ };
+ self.num_groups = self.num_groups.max(total_num_groups);
+
+ let nulls = if self.ignore_nulls {
+ input.logical_nulls()
+ } else {
+ None
+ };
+ let rows = group_indices
+ .iter()
+ .enumerate()
+ .filter(|&(row_index, _)| {
+ let filtered_out = opt_filter.is_some_and(|filter| {
+ filter.is_null(row_index) || !filter.value(row_index)
+ });
+ let skipped_null =
+ nulls.as_ref().is_some_and(|nulls|
nulls.is_null(row_index));
+ !filtered_out && !skipped_null
+ })
+ .map(|(row_index, &group_index)| (group_index, row_index))
+ .collect::<Vec<_>>();
+ if rows.is_empty() {
+ return Ok(());
+ }
+ let mut metric_recorder =
+ AggregateMetricRecorder::new(self.distinct_metric.clone());
+ let _timer = metric_recorder.timer();
+ self.insert_selected(input, rows)
+ }
+
+ /// Produces a `ListArray` ordered by group index. Groups without values
+ /// are null, matching [`DistinctArrayAggAccumulator`].
+ fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
+ let emit_groups = match emit_to {
+ EmitTo::All => self.num_groups,
+ EmitTo::First(n) => n,
+ };
+
+ // Counting sort of the pairs into group order.
+ let mut counts = vec![0u32; emit_groups];
+ for &(group_index, _) in &self.pairs {
+ if let Some(count) = counts.get_mut(group_index as usize) {
+ *count += 1;
+ }
+ }
+ let mut offsets = Vec::<i32>::with_capacity(emit_groups + 1);
+ offsets.push(0);
+ let mut nulls = NullBufferBuilder::new(emit_groups);
+ let mut write_positions = Vec::with_capacity(emit_groups);
+ let mut total = 0u32;
+ for &count in &counts {
+ if count == 0 {
+ nulls.append_null();
+ } else {
+ nulls.append_non_null();
+ }
+ write_positions.push(total);
+ total += count;
+ offsets.push(total as i32);
+ }
+ let mut ordered = vec![0u32; total as usize];
+ for &(group_index, value_index) in &self.pairs {
+ if let Some(position) = write_positions.get_mut(group_index as
usize) {
+ ordered[*position as usize] = value_index;
+ *position += 1;
+ }
+ }
+
+ let flat_values = match (&self.converter, &self.values) {
+ (Some(converter), Some(values)) if total > 0 => {
+ if self.sort_options.is_some() {
+ // The sort direction is encoded in the row bytes, so
+ // comparing rows gives the requested order.
+ for window in offsets.windows(2) {
+ ordered[window[0] as usize..window[1] as
usize].sort_unstable_by(
+ |&a, &b| values.row(a as usize).cmp(&values.row(b
as usize)),
+ );
+ }
+ }
+ let decoded = converter
+ .convert_rows(ordered.iter().map(|&i| values.row(i as
usize)))?
+ .swap_remove(0);
+ // `convert_rows` always returns the physical (non-dictionary)
+ // type, so cast back when the declared type has a dictionary.
+ if decoded.data_type() != &self.datatype
+ && datatype_contains_dictionary(&self.datatype)
+ {
+ cast(decoded.as_ref(), &self.datatype)?
+ } else {
+ decoded
+ }
+ }
+ _ => new_empty_array(&self.datatype),
+ };
+
+ match emit_to {
+ EmitTo::All => {
+ let converter = self.converter.take();
+ let distinct_metric = self.distinct_metric.take();
+ *self = Self::new(
+ self.datatype.clone(),
+ self.sort_options,
+ self.ignore_nulls,
+ );
+ self.distinct_metric = distinct_metric;
+ // Keep the converter: the runtime type does not change.
+ if let Some(converter) = converter {
+ self.values = Some(converter.empty_rows(0, 0));
+ self.converter = Some(converter);
+ }
+ }
+ EmitTo::First(_) => self.retain_after(emit_groups),
+ }
+
+ let field = Arc::new(Field::new_list_field(self.datatype.clone(),
true));
+ Ok(Arc::new(ListArray::new(
+ field,
+ OffsetBuffer::new(ScalarBuffer::from(offsets)),
+ flat_values,
+ nulls.finish(),
+ )))
+ }
+
+ fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+ Ok(vec![self.evaluate(emit_to)?])
+ }
+
+ fn merge_batch(
+ &mut self,
+ values: &[ArrayRef],
+ group_indices: &[usize],
+ total_num_groups: usize,
+ ) -> Result<()> {
+ assert_eq_or_internal_err!(values.len(), 1, "one argument to
merge_batch");
+ let lists = values[0].as_list::<i32>();
+ self.num_groups = self.num_groups.max(total_num_groups);
+
+ // A sliced list keeps its whole child array, so encode only the span
+ // its offsets cover.
+ let offsets = lists.offsets();
+ let (span_start, span_len) = offset_span(offsets);
+ let child = lists.values().slice(span_start, span_len);
+ let rows = group_indices
+ .iter()
+ .enumerate()
+ .filter(|&(list_index, _)| lists.is_valid(list_index))
+ .flat_map(|(list_index, &group_index)| {
+ let start = offsets[list_index] as usize - span_start;
+ let end = offsets[list_index + 1] as usize - span_start;
+ (start..end).map(move |position| (group_index, position))
+ })
+ .collect::<Vec<_>>();
+ if rows.is_empty() {
+ return Ok(());
+ }
+ let mut metric_recorder =
+ AggregateMetricRecorder::new(self.distinct_metric.clone());
+ let _timer = metric_recorder.timer();
+ self.insert_selected(&child, rows)
+ }
+
+ fn convert_to_state(
+ &self,
+ values: &[ArrayRef],
+ opt_filter: Option<&BooleanArray>,
+ ) -> Result<Vec<ArrayRef>> {
+ let Some(input) = values.first() else {
+ return internal_err!("array_agg(DISTINCT) expects an argument");
+ };
+ Ok(vec![single_value_lists(
+ input,
+ opt_filter,
+ self.ignore_nulls,
+ &self.datatype,
+ )])
+ }
+
+ fn set_metrics(&mut self, metrics: Arc<dyn AggregateMetrics>) {
+ self.distinct_metric = Some(metrics.metric("distinct"));
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self)
+ + self.converter.as_ref().map_or(0, RowConverter::size)
+ + self.values.as_ref().map_or(0, Rows::size)
+ + self.value_map_size
+ + self.pairs.capacity() * size_of::<(u32, u32)>()
+ + self.pair_set.capacity() * (size_of::<(u32, u32)>() + 1)
Review Comment:
Should we use this one?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
