Copilot commented on code in PR #26357:
URL: https://github.com/apache/datafusion/pull/26357#discussion_r4279668509
##########
datafusion/functions-aggregate/src/array_agg.rs:
##########
@@ -832,6 +834,377 @@ 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,
+}
+
+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(),
+ }
+ }
+
+ 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"))
+ }
+
+ /// 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(
Review Comment:
The native groups accumulator inherits the no-op `set_metrics` and never
records a DISTINCT duration. Switching from the adapter therefore removes
`agg_expr_0_internal_distinct_time`, which
`grouped_hash_stream_reports_distinct_array_agg_submetric` requires. The two
`legacy_grouped_distinct_*` tests also still expect adapter-specific metric
counts. Implement the native metric hook and update/merge timing, preserve the
metric handle when `evaluate(All)` resets the accumulator, and make the
adapter-specific tests construct the adapter explicitly.
##########
datafusion/functions-aggregate/src/array_agg.rs:
##########
@@ -832,6 +834,377 @@ 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,
+}
+
+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(),
+ }
+ }
+
+ 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"))
+ }
+
+ /// 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(());
+ }
+ self.insert(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();
+ *self = Self::new(
+ self.datatype.clone(),
+ self.sort_options,
+ self.ignore_nulls,
+ );
+ // 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);
+
+ let offsets = lists.offsets();
+ 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;
+ let end = offsets[list_index + 1] as usize;
+ (start..end).map(move |position| (group_index, position))
+ })
+ .collect::<Vec<_>>();
+ if rows.is_empty() {
+ return Ok(());
+ }
+ self.insert(lists.values(), rows)
Review Comment:
A sliced `ListArray` retains its full child array, and `insert` encodes that
entire array before selecting rows. Partial aggregation emits batch-sized
slices of one materialized result (`aggregate_hash_table/common.rs:654–663`),
so merging successive slices repeatedly encodes the complete partial result.
For N one-element groups emitted in batches of B, this causes O(N²/B) encoding
work and a full-state temporary allocation per batch. Slice the children to the
current offset span and rebase the selected positions before insertion.
##########
datafusion/functions-aggregate/src/array_agg.rs:
##########
@@ -832,6 +834,377 @@ 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,
+}
+
+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(),
+ }
+ }
+
+ 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"))
+ }
+
+ /// 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(());
+ }
+ self.insert(input, rows)
Review Comment:
When callers pass a sparse `opt_filter`, `rows` excludes rejected values but
`insert` still encodes the entire input. The legacy grouped stream passes
filters this way (`grouped_hash_stream.rs:946–957`), so retaining one small
string also copies all discarded large strings into temporary row storage. The
newer hash-table path already compacts FILTER inputs upstream. Gather selected
rows before insertion on this path, while retaining the no-copy fast path when
every row is selected.
--
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]
