xudong963 commented on code in PR #15296:
URL: https://github.com/apache/datafusion/pull/15296#discussion_r2002959262
##########
datafusion/expr-common/src/statistics.rs:
##########
@@ -203,6 +203,121 @@ impl Distribution {
};
Ok(dt)
}
+
+ /// Merges two distributions into a single distribution that represents
their combined statistics.
+ /// This creates a more general distribution that approximates the mixture
of the input distributions.
+ pub fn merge(&self, other: &Self) -> Result<Self> {
+ let range_a = self.range()?;
+ let range_b = other.range()?;
+
+ // Determine data type and create combined range
+ let combined_range = range_a.union(&range_b)?;
+
+ // Calculate weights for the mixture distribution
+ let (weight_a, weight_b) = match (range_a.cardinality(),
range_b.cardinality()) {
+ (Some(ca), Some(cb)) => {
+ let total = (ca + cb) as f64;
+ ((ca as f64) / total, (cb as f64) / total)
+ }
+ _ => (0.5, 0.5), // Equal weights if cardinalities not available
+ };
+
+ // Get the original statistics
+ let mean_a = self.mean()?;
+ let mean_b = other.mean()?;
+ let median_a = self.median()?;
+ let median_b = other.median()?;
+ let var_a = self.variance()?;
+ let var_b = other.variance()?;
+
+ // Always use Float64 for intermediate calculations to avoid truncation
+ // I assume that the target type is always numeric
+ // Todo: maybe we can keep all `ScalarValue` as `Float64` in
`Distribution`?
+ let calc_type = DataType::Float64;
+
+ // Create weight scalars using Float64 to avoid truncation
+ let weight_a_scalar = ScalarValue::from(weight_a);
+ let weight_b_scalar = ScalarValue::from(weight_b);
+
+ // Calculate combined mean
+ let combined_mean = if mean_a.is_null() || mean_b.is_null() {
+ if mean_a.is_null() {
+ mean_b.clone()
+ } else {
+ mean_a.clone()
+ }
+ } else {
+ // Cast to Float64 for calculation
+ let mean_a_f64 = mean_a.cast_to(&calc_type)?;
+ let mean_b_f64 = mean_b.cast_to(&calc_type)?;
+
+ // Calculate weighted mean
+ mean_a_f64
+ .mul_checked(&weight_a_scalar)?
+ .add_checked(&mean_b_f64.mul_checked(&weight_b_scalar)?)?
+ };
+
+ // Calculate combined median
+ let combined_median = if median_a.is_null() || median_b.is_null() {
+ if median_a.is_null() {
+ median_b
+ } else {
+ median_a
+ }
+ } else {
+ // Cast to Float64 for calculation
+ let median_a_f64 = median_a.cast_to(&calc_type)?;
+ let median_b_f64 = median_b.cast_to(&calc_type)?;
+
+ // Calculate weighted median
+ median_a_f64
+ .mul_checked(&weight_a_scalar)?
+ .add_checked(&median_b_f64.mul_checked(&weight_b_scalar)?)?
Review Comment:
I add some notes for the method:
https://github.com/apache/datafusion/pull/15296/commits/5b98f4c5efe5404a6ae4c8609a3574cccc91e9ab
##########
datafusion/expr-common/src/statistics.rs:
##########
@@ -857,6 +857,143 @@ pub fn compute_variance(
ScalarValue::try_from(target_type)
}
+/// Merges two distributions into a single distribution that represents their
combined statistics.
+/// This creates a more general distribution that approximates the mixture of
the input distributions.
Review Comment:
I added some notes for the method:
https://github.com/apache/datafusion/pull/15296/commits/5b98f4c5efe5404a6ae4c8609a3574cccc91e9ab
--
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
