spark_connect/
group.rs

1//! GroupedData implementation for aggregations.
2//!
3//! Mirroring `pyspark.sql.GroupedData`.
4
5use spark_connect_core::error::Result;
6
7use crate::column::Column;
8use crate::dataframe::DataFrame;
9use crate::expression::Expression;
10use crate::plan::{AggregateGroupType, LogicalPlan};
11use crate::types::DataType;
12use crate::udf::CommonInlineUserDefinedFunctionExpression;
13
14/// Convert a [`crate::row::Value`] into a literal [`Expression`] (for pivot values).
15fn value_to_lit_expr(v: crate::row::Value) -> Expression {
16    use crate::expression::LiteralExpression as L;
17    use crate::row::Value;
18    let lit = match v {
19        Value::Bool(b) => L::boolean(b),
20        Value::Byte(x) => L::int(x as i32),
21        Value::Short(x) => L::int(x as i32),
22        Value::Integer(x) => L::int(x),
23        Value::Long(x) => L::long(x),
24        Value::Float(x) => L::double(x as f64),
25        Value::Double(x) => L::double(x),
26        Value::String(s) => L::string(s),
27        other => L::string(format!("{:?}", other)),
28    };
29    Expression::Literal(lit)
30}
31
32/// Grouped data for performing aggregations.
33///
34/// Mirrors `pyspark.sql.GroupedData`.
35#[derive(Clone)]
36pub struct GroupedData {
37    dataframe: DataFrame,
38    group_cols: Vec<Column>,
39    group_type: AggregateGroupType,
40    /// Pivot column, set by [`GroupedData::pivot`].
41    pivot_col: Option<Expression>,
42    /// Explicit pivot values (literals), set by [`GroupedData::pivot`].
43    pivot_values: Vec<Expression>,
44    /// Explicit grouping sets, set by [`GroupedData::new_grouping_sets`]. Each inner
45    /// Vec is one grouping set. Only meaningful when `group_type == GroupingSets`.
46    grouping_sets: Vec<Vec<Column>>,
47}
48
49impl GroupedData {
50    /// Create a new GroupedData.
51    pub(crate) fn new(
52        dataframe: DataFrame,
53        group_cols: Vec<Column>,
54        group_type: AggregateGroupType,
55    ) -> Self {
56        GroupedData {
57            dataframe,
58            group_cols,
59            group_type,
60            pivot_col: None,
61            pivot_values: vec![],
62            grouping_sets: vec![],
63        }
64    }
65
66    /// Create a GroupedData for `GROUPING SETS`. `grouping_sets` holds each explicit
67    /// set of grouping columns; the grouping expressions become the (de-duplicated)
68    /// union of all columns referenced across the sets.
69    pub(crate) fn new_grouping_sets(dataframe: DataFrame, grouping_sets: Vec<Vec<Column>>) -> Self {
70        // Union of all columns across the sets, de-duplicated by their proto encoding
71        // (Expression is not Eq/Hash), preserving first-seen order.
72        let mut seen: Vec<Vec<u8>> = Vec::new();
73        let mut group_cols: Vec<Column> = Vec::new();
74        for set in &grouping_sets {
75            for col in set {
76                let key = prost::Message::encode_to_vec(&col.expression().clone().to_proto());
77                if !seen.contains(&key) {
78                    seen.push(key);
79                    group_cols.push(col.clone());
80                }
81            }
82        }
83        GroupedData {
84            dataframe,
85            group_cols,
86            group_type: AggregateGroupType::GroupingSets,
87            pivot_col: None,
88            pivot_values: vec![],
89            grouping_sets,
90        }
91    }
92
93    /// Pivot on a column, optionally with explicit values (`GroupedData.pivot`).
94    ///
95    /// Mirrors `df.groupBy(...).pivot(col, values)`. When `values` is `None` the
96    /// server computes the distinct values; when supplied they are serialized as
97    /// pivot-value literals (previously they were dropped entirely).
98    pub fn pivot(&self, pivot_col: Column, values: Option<Vec<crate::row::Value>>) -> GroupedData {
99        let pivot_values = values
100            .unwrap_or_default()
101            .into_iter()
102            .map(value_to_lit_expr)
103            .collect();
104        GroupedData {
105            dataframe: self.dataframe.clone(),
106            group_cols: self.group_cols.clone(),
107            group_type: AggregateGroupType::Pivot,
108            pivot_col: Some(pivot_col.expression().clone()),
109            pivot_values,
110            grouping_sets: vec![],
111        }
112    }
113
114    /// Column names of the underlying (pre-grouping) DataFrame. These are passed as the
115    /// map/apply UDF's argument columns so the worker's pandas/Arrow input is named.
116    pub fn input_columns(&self) -> Result<Vec<String>> {
117        self.dataframe.columns()
118    }
119
120    /// Grouping-key expressions for this grouped data.
121    fn grouping_expressions(&self) -> Vec<Expression> {
122        self.group_cols
123            .iter()
124            .map(|col| col.expression().clone())
125            .collect()
126    }
127
128    /// Apply a pandas UDF to each group (`GroupedData.applyInPandas`).
129    ///
130    /// `func` is built on the Python side (cloudpickled, eval type
131    /// `SQL_GROUPED_MAP_PANDAS_UDF`).
132    pub fn apply_in_pandas(&self, func: CommonInlineUserDefinedFunctionExpression) -> DataFrame {
133        self.group_map(func)
134    }
135
136    /// Apply an Arrow UDF to each group (`GroupedData.applyInArrow`).
137    pub fn apply_in_arrow(&self, func: CommonInlineUserDefinedFunctionExpression) -> DataFrame {
138        self.group_map(func)
139    }
140
141    fn group_map(&self, func: CommonInlineUserDefinedFunctionExpression) -> DataFrame {
142        let plan = LogicalPlan::GroupMap {
143            input: Box::new(self.dataframe.plan.clone()),
144            grouping_expressions: self.grouping_expressions(),
145            func,
146            sorting_expressions: vec![],
147            initial_input: None,
148            initial_grouping_expressions: vec![],
149            is_map_groups_with_state: None,
150            output_mode: None,
151            timeout_conf: None,
152            state_schema: None,
153            transform_with_state_info: None,
154        };
155        DataFrame::new(self.dataframe.session.clone(), plan)
156    }
157
158    /// Apply a stateful pandas UDF to each group (`GroupedData.applyInPandasWithState`).
159    ///
160    /// `func` (built on the Python side with eval type
161    /// `SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE` and carrying the output schema as its
162    /// return type) is combined with the state schema, output mode, and timeout.
163    pub fn apply_in_pandas_with_state(
164        &self,
165        func: CommonInlineUserDefinedFunctionExpression,
166        state_schema: DataType,
167        output_mode: &str,
168        timeout_conf: &str,
169    ) -> DataFrame {
170        let plan = LogicalPlan::GroupMap {
171            input: Box::new(self.dataframe.plan.clone()),
172            grouping_expressions: self.grouping_expressions(),
173            func,
174            sorting_expressions: vec![],
175            initial_input: None,
176            initial_grouping_expressions: vec![],
177            is_map_groups_with_state: Some(false),
178            output_mode: Some(output_mode.to_string()),
179            timeout_conf: Some(timeout_conf.to_string()),
180            state_schema: Some(state_schema),
181            transform_with_state_info: None,
182        };
183        DataFrame::new(self.dataframe.session.clone(), plan)
184    }
185
186    /// `GroupedData.transformWithState` (row output). `func` is the cloudpickled
187    /// stateful processor (eval type in the 211-214 range).
188    pub fn transform_with_state(
189        &self,
190        func: CommonInlineUserDefinedFunctionExpression,
191        output_mode: &str,
192        time_mode: &str,
193        event_time_column_name: Option<&str>,
194        initial_state: Option<&GroupedData>,
195    ) -> DataFrame {
196        self.transform_with_state_impl(
197            func,
198            output_mode,
199            time_mode,
200            event_time_column_name,
201            initial_state,
202            None,
203        )
204    }
205
206    /// `GroupedData.transformWithStateInPandas` (carries the output schema in the
207    /// `TransformWithStateInfo`).
208    pub fn transform_with_state_in_pandas(
209        &self,
210        func: CommonInlineUserDefinedFunctionExpression,
211        output_schema: DataType,
212        output_mode: &str,
213        time_mode: &str,
214        event_time_column_name: Option<&str>,
215        initial_state: Option<&GroupedData>,
216    ) -> DataFrame {
217        self.transform_with_state_impl(
218            func,
219            output_mode,
220            time_mode,
221            event_time_column_name,
222            initial_state,
223            Some(output_schema),
224        )
225    }
226
227    #[allow(clippy::too_many_arguments)]
228    fn transform_with_state_impl(
229        &self,
230        func: CommonInlineUserDefinedFunctionExpression,
231        output_mode: &str,
232        time_mode: &str,
233        event_time_column_name: Option<&str>,
234        initial_state: Option<&GroupedData>,
235        output_schema: Option<DataType>,
236    ) -> DataFrame {
237        let (initial_input, initial_grouping_expressions) = match initial_state {
238            Some(gd) => (
239                Some(Box::new(gd.dataframe.plan.clone())),
240                gd.grouping_expressions(),
241            ),
242            None => (None, vec![]),
243        };
244        let plan = LogicalPlan::GroupMap {
245            input: Box::new(self.dataframe.plan.clone()),
246            grouping_expressions: self.grouping_expressions(),
247            func,
248            sorting_expressions: vec![],
249            initial_input,
250            initial_grouping_expressions,
251            is_map_groups_with_state: None,
252            output_mode: Some(output_mode.to_string()),
253            timeout_conf: None,
254            state_schema: None,
255            transform_with_state_info: Some(crate::plan::TransformWithStateInfo {
256                time_mode: time_mode.to_string(),
257                event_time_column_name: event_time_column_name.map(|s| s.to_string()),
258                output_schema,
259            }),
260        };
261        DataFrame::new(self.dataframe.session.clone(), plan)
262    }
263
264    /// Cogroup this grouped data with another (`GroupedData.cogroup`).
265    pub fn cogroup(&self, other: &GroupedData) -> CoGroupedData {
266        CoGroupedData {
267            left: self.clone(),
268            right: other.clone(),
269        }
270    }
271
272    /// Perform an aggregation.
273    pub fn agg(&self, expressions: Vec<Expression>) -> DataFrame {
274        // Convert group columns to expressions
275        let grouping_expressions = self
276            .group_cols
277            .iter()
278            .map(|col| col.expression().clone())
279            .collect();
280
281        // Explicit grouping sets → each set as its own list of grouping expressions.
282        let grouping_sets: Vec<Vec<Expression>> = self
283            .grouping_sets
284            .iter()
285            .map(|set| set.iter().map(|c| c.expression().clone()).collect())
286            .collect();
287
288        let plan = LogicalPlan::Aggregate {
289            input: Box::new(self.dataframe.plan.clone()),
290            group_type: self.group_type,
291            grouping_expressions,
292            aggregate_expressions: expressions,
293            pivot_col: self.pivot_col.clone(),
294            pivot_values: self.pivot_values.clone(),
295            grouping_sets,
296        };
297
298        DataFrame::new(self.dataframe.session.clone(), plan)
299    }
300
301    /// Count rows in each group.
302    pub fn count(&self) -> DataFrame {
303        use crate::functions;
304        let count_expr = functions::count(Column::new(Expression::Literal(
305            crate::expression::LiteralExpression::int(1),
306        )))
307        .expression()
308        .clone();
309
310        self.agg(vec![count_expr])
311    }
312
313    /// Sum values in each group.
314    pub fn sum(&self, columns: Vec<&str>) -> DataFrame {
315        use crate::functions;
316        let expressions: Vec<_> = columns
317            .iter()
318            .map(|col| functions::sum(crate::column::col(col)).expression().clone())
319            .collect();
320
321        self.agg(expressions)
322    }
323
324    /// Average values in each group.
325    pub fn avg(&self, columns: Vec<&str>) -> DataFrame {
326        use crate::functions;
327        let expressions: Vec<_> = columns
328            .iter()
329            .map(|col| functions::avg(crate::column::col(col)).expression().clone())
330            .collect();
331
332        self.agg(expressions)
333    }
334
335    /// Minimum values in each group.
336    pub fn min(&self, columns: Vec<&str>) -> DataFrame {
337        use crate::functions;
338        let expressions: Vec<_> = columns
339            .iter()
340            .map(|col| functions::min(crate::column::col(col)).expression().clone())
341            .collect();
342
343        self.agg(expressions)
344    }
345
346    /// Maximum values in each group.
347    pub fn max(&self, columns: Vec<&str>) -> DataFrame {
348        use crate::functions;
349        let expressions: Vec<_> = columns
350            .iter()
351            .map(|col| functions::max(crate::column::col(col)).expression().clone())
352            .collect();
353
354        self.agg(expressions)
355    }
356
357    /// Mean (alias for avg).
358    pub fn mean(&self, columns: Vec<&str>) -> DataFrame {
359        self.avg(columns)
360    }
361}
362
363/// Statistical functions for DataFrames.
364///
365/// Mirrors `pyspark.sql.DataFrameStatFunctions`.
366pub struct StatFunctions {
367    dataframe: DataFrame,
368}
369
370impl StatFunctions {
371    /// Create a new StatFunctions.
372    pub(crate) fn new(dataframe: DataFrame) -> Self {
373        StatFunctions { dataframe }
374    }
375
376    /// Compute a cross-tabulation of two columns.
377    pub fn crosstab(&self, col1: &str, col2: &str) -> DataFrame {
378        let plan = LogicalPlan::StatCrosstab {
379            input: Box::new(self.dataframe.plan.clone()),
380            col1: col1.to_string(),
381            col2: col2.to_string(),
382        };
383        DataFrame::new(self.dataframe.session.clone(), plan)
384    }
385
386    /// Find frequent items.
387    pub fn freq_items(&self, columns: Vec<&str>, support: f64) -> DataFrame {
388        let plan = LogicalPlan::StatFreqItems {
389            input: Box::new(self.dataframe.plan.clone()),
390            columns: columns.iter().map(|s| s.to_string()).collect(),
391            support,
392        };
393        DataFrame::new(self.dataframe.session.clone(), plan)
394    }
395
396    /// Compute approximate quantiles.
397    pub fn approx_quantile(
398        &self,
399        columns: Vec<&str>,
400        probabilities: Vec<f64>,
401        relative_error: f64,
402    ) -> DataFrame {
403        let plan = LogicalPlan::StatApproxQuantile {
404            input: Box::new(self.dataframe.plan.clone()),
405            columns: columns.iter().map(|s| s.to_string()).collect(),
406            probabilities,
407            relative_error,
408        };
409        DataFrame::new(self.dataframe.session.clone(), plan)
410    }
411
412    /// Compute correlation between two columns.
413    ///
414    /// Executes the stat aggregation and returns the resulting `f64`
415    /// (NaN when the correlation is undefined, e.g. an empty relation).
416    pub fn corr(&self, col1: &str, col2: &str) -> Result<f64> {
417        let plan = LogicalPlan::StatCorr {
418            input: Box::new(self.dataframe.plan.clone()),
419            col1: col1.to_string(),
420            col2: col2.to_string(),
421        };
422        let df = DataFrame::new(self.dataframe.session.clone(), plan);
423        Ok(df.scalar()?.and_then(|v| v.as_f64()).unwrap_or(f64::NAN))
424    }
425
426    /// Compute covariance between two columns.
427    ///
428    /// Executes the stat aggregation and returns the resulting `f64`.
429    pub fn cov(&self, col1: &str, col2: &str) -> Result<f64> {
430        let plan = LogicalPlan::StatCov {
431            input: Box::new(self.dataframe.plan.clone()),
432            col1: col1.to_string(),
433            col2: col2.to_string(),
434        };
435        let df = DataFrame::new(self.dataframe.session.clone(), plan);
436        Ok(df.scalar()?.and_then(|v| v.as_f64()).unwrap_or(f64::NAN))
437    }
438
439    /// Sample by values in a column.
440    pub fn sample_by(
441        &self,
442        col: &str,
443        fractions: Vec<(Expression, f64)>,
444        seed: Option<i64>,
445    ) -> DataFrame {
446        let plan = LogicalPlan::StatSampleBy {
447            input: Box::new(self.dataframe.plan.clone()),
448            col: col.to_string(),
449            fractions,
450            seed,
451        };
452        DataFrame::new(self.dataframe.session.clone(), plan)
453    }
454}
455
456/// Methods for handling missing data, accessed via [`DataFrame::na`].
457///
458/// Mirrors `pyspark.sql.DataFrameNaFunctions`.
459pub struct NaFunctions {
460    dataframe: DataFrame,
461}
462
463impl NaFunctions {
464    /// Create a new NaFunctions.
465    pub(crate) fn new(dataframe: DataFrame) -> Self {
466        NaFunctions { dataframe }
467    }
468
469    /// Drop rows containing null values. Mirrors `DataFrameNaFunctions.drop`.
470    pub fn drop(
471        &self,
472        how: Option<&str>,
473        thresh: Option<i32>,
474        subset: Option<Vec<&str>>,
475    ) -> DataFrame {
476        self.dataframe.dropna(how, thresh, subset)
477    }
478
479    /// Fill null values. Mirrors `DataFrameNaFunctions.fill`.
480    pub fn fill(&self, value: i64, subset: Option<Vec<&str>>) -> DataFrame {
481        self.dataframe.fillna(value, subset)
482    }
483
484    /// Replace values. Mirrors `DataFrameNaFunctions.replace`.
485    pub fn replace(
486        &self,
487        to_replace: Vec<(String, String)>,
488        subset: Option<Vec<&str>>,
489    ) -> DataFrame {
490        self.dataframe.replace(to_replace, subset)
491    }
492}
493
494/// A pair of cogrouped [`GroupedData`], mirroring `pyspark.sql.PandasCogroupedOps`.
495///
496/// Created via [`GroupedData::cogroup`].
497#[derive(Clone)]
498pub struct CoGroupedData {
499    left: GroupedData,
500    right: GroupedData,
501}
502
503impl CoGroupedData {
504    /// Column names of both sides (left then right), passed as the cogroup UDF's
505    /// argument columns so the worker's two pandas/Arrow inputs are named.
506    pub fn input_columns(&self) -> Result<Vec<String>> {
507        let mut cols = self.left.input_columns()?;
508        cols.extend(self.right.input_columns()?);
509        Ok(cols)
510    }
511
512    /// Apply a pandas UDF to each cogroup (`cogroup(...).applyInPandas`).
513    ///
514    /// `func` is built on the Python side (cloudpickled, eval type
515    /// `SQL_COGROUPED_MAP_PANDAS_UDF`).
516    pub fn apply_in_pandas(&self, func: CommonInlineUserDefinedFunctionExpression) -> DataFrame {
517        self.cogroup_map(func)
518    }
519
520    /// Apply an Arrow UDF to each cogroup (`cogroup(...).applyInArrow`).
521    pub fn apply_in_arrow(&self, func: CommonInlineUserDefinedFunctionExpression) -> DataFrame {
522        self.cogroup_map(func)
523    }
524
525    fn cogroup_map(&self, func: CommonInlineUserDefinedFunctionExpression) -> DataFrame {
526        let plan = LogicalPlan::CoGroupMap {
527            input: Box::new(self.left.dataframe.plan.clone()),
528            input_grouping_expressions: self.left.grouping_expressions(),
529            other: Box::new(self.right.dataframe.plan.clone()),
530            other_grouping_expressions: self.right.grouping_expressions(),
531            func,
532        };
533        DataFrame::new(self.left.dataframe.session.clone(), plan)
534    }
535}
536
537#[cfg(test)]
538mod tests {
539    use super::*;
540    use crate::session::SparkSession;
541
542    fn session() -> SparkSession {
543        SparkSession::builder()
544            .remote("sc://localhost:15002")
545            .get_or_create()
546            .expect("failed to build session")
547    }
548
549    #[test]
550    fn group_count_plan() {
551        let spark = session();
552        let df = spark.range(10).unwrap();
553        let grouped = df.group_by(vec![crate::column::col("id")]);
554        let result = grouped.count();
555        match &result.plan {
556            LogicalPlan::Aggregate {
557                group_type: AggregateGroupType::GroupBy,
558                ..
559            } => {
560                // Plan structure is correct
561            }
562            _ => panic!("expected Aggregate plan with GroupBy"),
563        }
564    }
565
566    #[test]
567    fn group_sum_plan() {
568        let spark = session();
569        let df = spark.range(10).unwrap();
570        let grouped = df.group_by(vec![crate::column::col("id")]);
571        let result = grouped.sum(vec!["id"]);
572        match &result.plan {
573            LogicalPlan::Aggregate {
574                group_type: AggregateGroupType::GroupBy,
575                aggregate_expressions,
576                ..
577            } => {
578                assert!(!aggregate_expressions.is_empty());
579            }
580            _ => panic!("expected Aggregate plan"),
581        }
582    }
583
584    #[test]
585    fn group_agg_plan() {
586        let spark = session();
587        let df = spark.range(10).unwrap();
588        let grouped = df.group_by(vec![crate::column::col("id")]);
589        let exprs = vec![crate::functions::sum(crate::column::col("id"))
590            .expression()
591            .clone()];
592        let result = grouped.agg(exprs);
593        match &result.plan {
594            LogicalPlan::Aggregate {
595                group_type: AggregateGroupType::GroupBy,
596                ..
597            } => {
598                // Plan structure is correct
599            }
600            _ => panic!("expected Aggregate plan"),
601        }
602    }
603
604    #[test]
605    fn pivot_plan() {
606        let spark = session();
607        let df = spark.range(10).unwrap();
608        let grouped = df.group_by(vec![crate::column::col("id")]);
609        let pivot_grouped = grouped.pivot(crate::column::col("category"), None);
610        assert_eq!(pivot_grouped.group_type, AggregateGroupType::Pivot);
611        assert!(pivot_grouped.pivot_col.is_some());
612    }
613
614    #[test]
615    fn stat_crosstab_plan() {
616        let spark = session();
617        let df = spark.range(10).unwrap();
618        let stats = df.stat();
619        let result = stats.crosstab("col1", "col2");
620        match &result.plan {
621            LogicalPlan::StatCrosstab { .. } => {
622                // Plan is correct
623            }
624            _ => panic!("expected StatCrosstab plan"),
625        }
626    }
627
628    #[test]
629    fn stat_freq_items_plan() {
630        let spark = session();
631        let df = spark.range(10).unwrap();
632        let stats = df.stat();
633        let result = stats.freq_items(vec!["col1", "col2"], 0.25);
634        match &result.plan {
635            LogicalPlan::StatFreqItems {
636                columns, support, ..
637            } => {
638                assert_eq!(columns.len(), 2);
639                assert_eq!(*support, 0.25);
640            }
641            _ => panic!("expected StatFreqItems plan"),
642        }
643    }
644
645    #[test]
646    fn stat_approx_quantile_plan() {
647        let spark = session();
648        let df = spark.range(10).unwrap();
649        let stats = df.stat();
650        let result = stats.approx_quantile(vec!["col1"], vec![0.25, 0.75], 0.05);
651        match &result.plan {
652            LogicalPlan::StatApproxQuantile {
653                columns,
654                probabilities,
655                relative_error,
656                ..
657            } => {
658                assert_eq!(columns.len(), 1);
659                assert_eq!(probabilities.len(), 2);
660                assert_eq!(*relative_error, 0.05);
661            }
662            _ => panic!("expected StatApproxQuantile plan"),
663        }
664    }
665
666    #[test]
667    fn stat_sample_by_plan() {
668        let spark = session();
669        let df = spark.range(10).unwrap();
670        let stats = df.stat();
671        let fractions = vec![
672            (
673                Expression::Literal(crate::expression::LiteralExpression::string("A")),
674                0.5,
675            ),
676            (
677                Expression::Literal(crate::expression::LiteralExpression::string("B")),
678                0.3,
679            ),
680        ];
681        let result = stats.sample_by("category", fractions, Some(42));
682        match &result.plan {
683            LogicalPlan::StatSampleBy {
684                col,
685                seed,
686                fractions,
687                ..
688            } => {
689                assert_eq!(col, "category");
690                assert_eq!(*seed, Some(42));
691                assert_eq!(fractions.len(), 2);
692            }
693            _ => panic!("expected StatSampleBy plan"),
694        }
695    }
696}
spark_connect/group.rs

spark_connect/
group.rs
