add intersect count for intersection component

Signed-off-by: mgqa34 <mgq3374541@163.com>
FederatedAI · Jan 17, 2023 · 3f0ea5e · 3f0ea5e
1 parent 4f879a5
commit 3f0ea5e
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/python/fate/ml/intersection/raw_intersection.py b/python/fate/ml/intersection/raw_intersection.py
@@ -33,6 +33,10 @@ def fit(self, ctx: Context, train_data, validate_data=None):
         for intersect_index in intersect_indexes:
             intersect_data = intersect_data.loc(intersect_index)
 
+        intersect_count = intersect_data.count()
+        ctx.hosts.put("intersect_count", intersect_count)
+
+        logger.info(f"intersect count={intersect_count}")
         return intersect_data
 
 
@@ -46,4 +50,6 @@ def fit(self, ctx: Context, train_data, validate_data=None):
         # ctx.guest.put("intersect_index", intersect_data.index.tolist())
         ctx.guest.put("intersect_index", intersect_data.index.values)
 
+        intersect_count = ctx.guest.get("intersect_count")
+        logger.info(f"intersect count={intersect_count}")
         return intersect_data