graphrecords_query/operations/aggregation/
median.rs

1use crate::{
2    Bare, BareValueDomain, EvaluateOperand, Explain, Failure, IndexDomain, Indexed, Labeled,
3    Multiple, Operand, OrderState, QueryResult, Single,
4    capabilities::ValueMedian,
5    error::comparison::{IncomparableValues, IncomparableValuesAt},
6    execution::EvaluationCache,
7    operands::OperandHandle,
8    operations::{
9        Apply, BareStream, KeyedStream, LaneKernel, Operation, OperationContext, Prepare,
10    },
11    optimizer::{Estimate, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs, Stats},
12    registry::operation_manifest,
13    traits::Median,
14};
15use graphrecords_core::GraphRecord;
16use std::{
17    cmp::Ordering,
18    fmt::{Debug, Display},
19};
20
21#[derive(Clone, Explain, Operation, OperationInputs, OptimizerHints, PlanIdentity, PlanInputs)]
22#[operation(scope = Lane)]
23#[explain(label = "Median")]
24#[plan(optimizer_hints(empty = if_any))]
25pub struct MedianOperation;
26
27impl Prepare for MedianOperation {
28    type Prepared<'a> = ();
29
30    fn prepare<'a>(
31        &'a self,
32        _graphrecord: &'a GraphRecord,
33        _cache: &'a EvaluationCache<'a>,
34    ) -> QueryResult<Self::Prepared<'a>> {
35        Ok(())
36    }
37}
38
39fn middle_values<T, F>(mut values: Vec<T>, compare: F) -> Option<(T, Option<T>)>
40where
41    T: Clone,
42    F: Fn(&T, &T) -> Ordering,
43{
44    let length = values.len();
45
46    if length == 0 {
47        return None;
48    }
49
50    let middle = length / 2;
51    let (lower_values, middle_value, _) =
52        values.select_nth_unstable_by(middle, |left, right| compare(left, right));
53
54    if length.is_multiple_of(2) {
55        let lower = lower_values
56            .iter()
57            .max_by(|left, right| compare(left, right))
58            .expect("an even-length lane has a lower middle value");
59
60        Some((lower.clone(), Some(middle_value.clone())))
61    } else {
62        Some((middle_value.clone(), None))
63    }
64}
65
66impl<I, V, O> LaneKernel<Indexed<I, V>, Multiple<O>> for MedianOperation
67where
68    I: IndexDomain,
69    V: ValueMedian + BareValueDomain,
70    O: OrderState,
71    V::Owned: Debug + Display + Send + Sync,
72{
73    type Output = OperandHandle<Bare<V>, Single>;
74
75    fn execute<'a>(
76        _graphrecord: &'a GraphRecord,
77        values: KeyedStream<'a, I, V, Multiple<O>>,
78        _prepared: Self::Prepared<'a>,
79    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
80        let collected = match values
81            .map(|(index, value)| {
82                let value = value?;
83                V::validate_median(Self::LABEL, &value)
84                    .map_err(|failure| failure.at::<I>(&index))?;
85
86                Ok((index, value))
87            })
88            .collect::<QueryResult<Vec<_>>>()
89        {
90            Ok(collected) => collected,
91            Err(failure) => return Ok(Some(Err(failure))),
92        };
93
94        if let Some((first_position, second_position)) =
95            V::find_incomparable_median_values(collected.iter().map(|(_, value)| value))
96        {
97            let (first_index, first) = &collected[first_position];
98            let (second_index, second) = &collected[second_position];
99            let failure = Failure::new_at::<I, _>(
100                Self::LABEL,
101                IncomparableValuesAt::new(
102                    V::into_owned(first.clone()),
103                    V::into_owned(second.clone()),
104                    I::to_owned(first_index),
105                    I::to_owned(second_index),
106                ),
107                second_index,
108            );
109
110            return Ok(Some(Err(failure)));
111        }
112
113        let Some(((lower_index, lower), upper)) = middle_values(collected, |left, right| {
114            V::ordering(&left.1, &right.1).expect("median values were checked for comparability")
115        }) else {
116            return Ok(None);
117        };
118        let (upper_index, upper) = match upper {
119            Some((index, value)) => (Some(index), Some(value)),
120            None => (None, None),
121        };
122        let failure_index = upper_index.as_ref().unwrap_or(&lower_index);
123
124        Ok(Some(
125            V::median(Self::LABEL, lower, upper).map_err(|failure| failure.at::<I>(failure_index)),
126        ))
127    }
128
129    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
130        input.zero_or_one()
131    }
132}
133
134impl<V, O> LaneKernel<Bare<V>, Multiple<O>> for MedianOperation
135where
136    V: ValueMedian + BareValueDomain,
137    O: OrderState,
138    V::Owned: Debug + Display + Send + Sync,
139{
140    type Output = OperandHandle<Bare<V>, Single>;
141
142    fn execute<'a>(
143        _graphrecord: &'a GraphRecord,
144        values: BareStream<'a, V, Multiple<O>>,
145        _prepared: Self::Prepared<'a>,
146    ) -> QueryResult<<Self::Output as EvaluateOperand>::ReturnValue<'a>> {
147        let collected = match values
148            .map(|value| {
149                let value = value?;
150                V::validate_median(Self::LABEL, &value)?;
151
152                Ok(value)
153            })
154            .collect::<QueryResult<Vec<_>>>()
155        {
156            Ok(collected) => collected,
157            Err(failure) => return Ok(Some(Err(failure))),
158        };
159
160        if let Some((first_position, second_position)) =
161            V::find_incomparable_median_values(collected.iter())
162        {
163            let failure = Failure::new(
164                Self::LABEL,
165                IncomparableValues::new(
166                    V::into_owned(collected[first_position].clone()),
167                    V::into_owned(collected[second_position].clone()),
168                ),
169            );
170
171            return Ok(Some(Err(failure)));
172        }
173
174        let Some((lower, upper)) = middle_values(collected, |left, right| {
175            V::ordering(left, right).expect("median values were checked for comparability")
176        }) else {
177            return Ok(None);
178        };
179
180        Ok(Some(V::median(Self::LABEL, lower, upper)))
181    }
182
183    fn estimate(&self, input: Estimate, _stats: &Stats) -> Estimate {
184        input.zero_or_one()
185    }
186}
187
188impl<O: Apply<MedianOperation>> Median for O {
189    type ReturnOperand = O::Output;
190
191    fn median(&self) -> Self::ReturnOperand {
192        Self::ReturnOperand::new(OperationContext::new(self.clone(), MedianOperation))
193    }
194}
195
196operation_manifest! {
197    MedianOperation {
198        method: Median::median;
199        scope: lane;
200
201        kernel {
202            parameters: <
203                I: IndexDomain,
204                V: ValueMedian + BareValueDomain,
205                O: OrderState,
206            >;
207            input: (Indexed<I, V>, Multiple<O>);
208            output: OperandHandle<Bare<V>, Single>;
209            where V::Owned: Debug + Display + Send + Sync;
210        }
211
212        kernel {
213            parameters: <
214                V: ValueMedian + BareValueDomain,
215                O: OrderState,
216            >;
217            input: (Bare<V>, Multiple<O>);
218            output: OperandHandle<Bare<V>, Single>;
219            where V::Owned: Debug + Display + Send + Sync;
220        }
221    }
222}
graphrecords_query/operations/aggregation/median.rs

graphrecords_query/operations/aggregation/
median.rs
