@@ -1040,26 +1040,29 @@ def get_stat(self, column_id: str, stat: agg_ops.UnaryAggregateOp):
1040
1040
self ._stats_cache [column_id ].update (stats_map )
1041
1041
return stats_map [stat .name ]
1042
1042
1043
- def get_corr_stat (self , column_id_left : str , column_id_right : str ):
1043
+ def get_binary_stat (
1044
+ self , column_id_left : str , column_id_right : str , stat : agg_ops .BinaryAggregateOp
1045
+ ):
1044
1046
# TODO(kemppeterson): Clean up the column names for DataFrames.corr support
1045
1047
# TODO(kemppeterson): Add a cache here.
1046
- corr_aggregations = [
1048
+ aggregations = [
1047
1049
(
1048
- column_id_left ,
1049
- column_id_right ,
1050
- "corr_" + column_id_left + column_id_right ,
1050
+ ex .BinaryAggregation (
1051
+ stat , ex .free_var (column_id_left ), ex .free_var (column_id_right )
1052
+ ),
1053
+ f"{ stat .name } _{ column_id_left } { column_id_right } " ,
1051
1054
)
1052
1055
]
1053
- expr = self .expr .corr_aggregate ( corr_aggregations )
1056
+ expr = self .expr .aggregate ( aggregations )
1054
1057
offset_index_id = guid .generate_guid ()
1055
1058
expr = expr .promote_offsets (offset_index_id )
1056
1059
block = Block (
1057
1060
expr ,
1058
1061
index_columns = [offset_index_id ],
1059
- column_labels = [a [2 ] for a in corr_aggregations ],
1062
+ column_labels = [a [1 ] for a in aggregations ],
1060
1063
)
1061
1064
df , _ = block .to_pandas ()
1062
- return df .loc [0 , "corr_" + column_id_left + column_id_right ]
1065
+ return df .loc [0 , f" { stat . name } _ { column_id_left } { column_id_right } " ]
1063
1066
1064
1067
def summarize (
1065
1068
self ,
0 commit comments