@@ -904,27 +904,11 @@ impl LogicalPlanBuilder {
904
904
group_expr : impl IntoIterator < Item = impl Into < Expr > > ,
905
905
aggr_expr : impl IntoIterator < Item = impl Into < Expr > > ,
906
906
) -> Result < Self > {
907
- let mut group_expr = normalize_cols ( group_expr, & self . plan ) ?;
907
+ let group_expr = normalize_cols ( group_expr, & self . plan ) ?;
908
908
let aggr_expr = normalize_cols ( aggr_expr, & self . plan ) ?;
909
909
910
- // Rewrite groupby exprs according to functional dependencies
911
- let group_by_expr_names = group_expr
912
- . iter ( )
913
- . map ( |group_by_expr| group_by_expr. display_name ( ) )
914
- . collect :: < Result < Vec < _ > > > ( ) ?;
915
- let schema = self . plan . schema ( ) ;
916
- if let Some ( target_indices) =
917
- get_target_functional_dependencies ( schema, & group_by_expr_names)
918
- {
919
- for idx in target_indices {
920
- let field = schema. field ( idx) ;
921
- let expr =
922
- Expr :: Column ( Column :: new ( field. qualifier ( ) . cloned ( ) , field. name ( ) ) ) ;
923
- if !group_expr. contains ( & expr) {
924
- group_expr. push ( expr) ;
925
- }
926
- }
927
- }
910
+ let group_expr =
911
+ add_group_by_exprs_from_dependencies ( group_expr, self . plan . schema ( ) ) ?;
928
912
Aggregate :: try_new ( Arc :: new ( self . plan ) , group_expr, aggr_expr)
929
913
. map ( LogicalPlan :: Aggregate )
930
914
. map ( Self :: from)
@@ -1189,6 +1173,42 @@ pub fn build_join_schema(
1189
1173
schema. with_functional_dependencies ( func_dependencies)
1190
1174
}
1191
1175
1176
+ /// Add additional "synthetic" group by expressions based on functional
1177
+ /// dependencies.
1178
+ ///
1179
+ /// For example, if we are grouping on `[c1]`, and we know from
1180
+ /// functional dependencies that column `c1` determines `c2`, this function
1181
+ /// adds `c2` to the group by list.
1182
+ ///
1183
+ /// This allows MySQL style selects like
1184
+ /// `SELECT col FROM t WHERE pk = 5` if col is unique
1185
+ fn add_group_by_exprs_from_dependencies (
1186
+ mut group_expr : Vec < Expr > ,
1187
+ schema : & DFSchemaRef ,
1188
+ ) -> Result < Vec < Expr > > {
1189
+ // Names of the fields produced by the GROUP BY exprs for example, `GROUP BY
1190
+ // c1 + 1` produces an output field named `"c1 + 1"`
1191
+ let mut group_by_field_names = group_expr
1192
+ . iter ( )
1193
+ . map ( |e| e. display_name ( ) )
1194
+ . collect :: < Result < Vec < _ > > > ( ) ?;
1195
+
1196
+ if let Some ( target_indices) =
1197
+ get_target_functional_dependencies ( schema, & group_by_field_names)
1198
+ {
1199
+ for idx in target_indices {
1200
+ let field = schema. field ( idx) ;
1201
+ let expr =
1202
+ Expr :: Column ( Column :: new ( field. qualifier ( ) . cloned ( ) , field. name ( ) ) ) ;
1203
+ let expr_name = expr. display_name ( ) ?;
1204
+ if !group_by_field_names. contains ( & expr_name) {
1205
+ group_by_field_names. push ( expr_name) ;
1206
+ group_expr. push ( expr) ;
1207
+ }
1208
+ }
1209
+ }
1210
+ Ok ( group_expr)
1211
+ }
1192
1212
/// Errors if one or more expressions have equal names.
1193
1213
pub ( crate ) fn validate_unique_names < ' a > (
1194
1214
node_name : & str ,
0 commit comments