sort duplicate cluster by size

Signed-off-by: Yang Yu <yayu@yayu-mlt.client.nvidia.com>

tutorials/pretraining-data-curation/helper.py

@@ -52,3 +52,15 @@ def convert_str_id_to_int(df, id_column="id"):
     df["doc_id"] = dx[1].astype("int64").values
     df["dataset_id"] = dx[0].hash_values()
     return df
+
+
+def get_dataframe_complement(original_df, filtered_df):
+    def partition_complement(part_original_df, partition_info=None):
+        if not partition_info:
+            return part_original_df
+        part_filtered_df = filtered_df.get_partition(partition_info["number"])
+        complement_mask = ~part_original_df.index.isin(part_filtered_df.index.persist())
+        complement_df = part_original_df[complement_mask]
+        return complement_df
+
+    return original_df.map_partitions(partition_complement)