alamb commented on code in PR #10226:
URL: https://github.com/apache/datafusion/pull/10226#discussion_r1580917363
##########
datafusion/physical-expr/src/aggregate/median.rs:
##########
@@ -196,6 +184,192 @@ impl<T: ArrowNumericType> Accumulator for
MedianAccumulator<T> {
}
}
+/// MEDIAN(DISTINCT) aggregate expression. Similar to MEDIAN but computes
after taking
+/// all unique values. This may use a lot of memory if the cardinality is high.
+#[derive(Debug)]
+pub struct DistinctMedian {
+ name: String,
+ expr: Arc<dyn PhysicalExpr>,
+ data_type: DataType,
+}
+
+impl DistinctMedian {
+ /// Create a new MEDIAN(DISTINCT) aggregate function
+ pub fn new(
+ expr: Arc<dyn PhysicalExpr>,
+ name: impl Into<String>,
+ data_type: DataType,
+ ) -> Self {
+ Self {
+ name: name.into(),
+ expr,
+ data_type,
+ }
+ }
+}
+
+impl AggregateExpr for DistinctMedian {
+ /// Return a reference to Any that can be used for downcasting
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn field(&self) -> Result<Field> {
+ Ok(Field::new(&self.name, self.data_type.clone(), true))
+ }
+
+ fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
+ use arrow_array::types::*;
+ macro_rules! helper {
+ ($t:ty, $dt:expr) => {
+ Ok(Box::new(DistinctMedianAccumulator::<$t> {
+ data_type: $dt.clone(),
+ distinct_values: Default::default(),
+ }))
+ };
+ }
+ let dt = &self.data_type;
+ downcast_integer! {
+ dt => (helper, dt),
+ DataType::Float16 => helper!(Float16Type, dt),
+ DataType::Float32 => helper!(Float32Type, dt),
+ DataType::Float64 => helper!(Float64Type, dt),
+ DataType::Decimal128(_, _) => helper!(Decimal128Type, dt),
+ DataType::Decimal256(_, _) => helper!(Decimal256Type, dt),
+ _ => Err(DataFusionError::NotImplemented(format!(
+ "DistinctMedianAccumulator not supported for {} with {}",
+ self.name(),
+ self.data_type
+ ))),
+ }
+ }
+
+ fn state_fields(&self) -> Result<Vec<Field>> {
+ // Intermediate state is a list of the unique elements we have
+ // collected so far
+ let field = Field::new("item", self.data_type.clone(), true);
+ let data_type = DataType::List(Arc::new(field));
+
+ Ok(vec![Field::new(
+ format_state_name(&self.name, "distinct_median"),
+ data_type,
+ true,
+ )])
+ }
+
+ fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
+ vec![self.expr.clone()]
+ }
+
+ fn name(&self) -> &str {
+ &self.name
+ }
+}
+
+impl PartialEq<dyn Any> for DistinctMedian {
+ fn eq(&self, other: &dyn Any) -> bool {
+ down_cast_any_ref(other)
+ .downcast_ref::<Self>()
+ .map(|x| {
+ self.name == x.name
+ && self.data_type == x.data_type
+ && self.expr.eq(&x.expr)
+ })
+ .unwrap_or(false)
+ }
+}
+
+/// The distinct median accumulator accumulates the raw input values
+/// as `ScalarValue`s
+///
+/// The intermediate state is represented as a List of scalar values updated by
+/// `merge_batch` and a `Vec` of `ArrayRef` that are converted to scalar values
+/// in the final evaluation step so that we avoid expensive conversions and
+/// allocations during `update_batch`.
+struct DistinctMedianAccumulator<T: ArrowNumericType> {
+ data_type: DataType,
+ distinct_values: HashSet<Hashable<T::Native>>,
+}
+
+impl<T: ArrowNumericType> std::fmt::Debug for DistinctMedianAccumulator<T> {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "DistinctMedianAccumulator({})", self.data_type)
+ }
+}
+
+impl<T: ArrowNumericType> Accumulator for DistinctMedianAccumulator<T> {
+ fn state(&mut self) -> Result<Vec<ScalarValue>> {
+ let all_values = self
+ .distinct_values
+ .iter()
+ .map(|x| ScalarValue::new_primitive::<T>(Some(x.0),
&self.data_type))
+ .collect::<Result<Vec<_>>>()?;
+
+ let arr = ScalarValue::new_list(&all_values, &self.data_type);
+ Ok(vec![ScalarValue::List(arr)])
+ }
+
+ fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
+ if values.is_empty() {
+ return Ok(());
+ }
+
+ let array = values[0].as_primitive::<T>();
+ match array.nulls().filter(|x| x.null_count() > 0) {
Review Comment:
I think it is a common optimization (in the arrow-rs kernels and datafusion)
to special case the 'no nulls' case -- if you know there are no nulls in the
input you can avoid a branch (to check for null) in the inner loop, which gives
the compiler a better chance for auto-vectorization
--
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: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
