Merge pull request #2 from peisuke/fix/aggregated_column_name

Fix column name aggregation
glassonion1 · Aug 4, 2022 · f43ef85 · f43ef85
2 parents b0c764d + 7e9bb5b
commit f43ef85
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/anonypy/anonypy.py b/anonypy/anonypy.py
@@ -76,7 +76,7 @@ def anonymize(df, partitions, feature_columns, sensitive_column, max_partitions=
         sensitive_counts = (
             df.loc[partition].groupby(sensitive_column).agg({sensitive_column: "count"})
         )
-        values = grouped_columns.iloc[0].to_dict()
+        values = grouped_columns.apply(lambda x: x[0]).to_dict()
         for sensitive_value, count in sensitive_counts[sensitive_column].items():
             if count == 0:
                 continue
@@ -105,6 +105,6 @@ def count_anonymity(
         if max_partitions is not None and i > max_partitions:
             break
         grouped_columns = df.loc[partition].agg(aggregations, squeeze=False)
-        values = grouped_columns.iloc[0].to_dict()
+        values = grouped_columns.apply(lambda x: x[0]).to_dict()
         rows.append(values.copy())
     return rows