fix union

Signed-off-by: nemirorox <yolandawu131@gmail.com>
FederatedAI · Mar 19, 2020 · 06d8c9c · 06d8c9c
1 parent 65ae8f1
commit 06d8c9c
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/federatedml/statistic/union/union.py b/federatedml/statistic/union/union.py
@@ -97,9 +97,11 @@ def check_is_data_instance(self, table):
         self.is_data_instance = isinstance(entry[1], Instance)
 
     def fit(self, data):
-        if len(data) <= 0:
+        if data is None:
             LOGGER.warning("Union receives no data input.")
             return
+        if not isinstance(data, dict):
+            data = {"data": data}
         empty_count = 0
         combined_table = None
         combined_schema = None
@@ -156,7 +158,6 @@ def fit(self, data):
                                      metric_name=self.metric_name,
                                      metric_meta=MetricMeta(name=self.metric_name, metric_type=self.metric_type))
 
-        LOGGER.debug("after union schema: {}".format(combined_table.schema))
-
         LOGGER.info("Union operation finished. Total {} empty tables encountered.".format(empty_count))
         return combined_table
+