edit schema check(#4668)

Signed-off-by: Yu Wu <yolandawu131@gmail.com>
Signed-off-by: weiwee <wbwmat@gmail.com>

python/fate/ml/preprocessing/union.py

@@ -26,10 +26,6 @@ def __init__(self, axis=0):
         self.axis = axis
 
     def fit(self, ctx: Context, train_data_list):
-        label_name_list = [data.schema.label_name for data in train_data_list]
-        if sum([name != label_name_list[0] for name in label_name_list]):
-            raise ValueError(f"Data sets should all have the same label_name for union.")
-
         sample_id_name_list = [data.schema.sample_id_name for data in train_data_list]
         if sum([name != sample_id_name_list[0] for name in sample_id_name_list]):
             raise ValueError(f"Data sets should all have the same sample_id_name for union.")
@@ -39,13 +35,13 @@ def fit(self, ctx: Context, train_data_list):
             raise ValueError(f"Data sets should all have the same match_id_name for union.")
 
         if self.axis == 0:
-            data_cols = set()
-            for data in train_data_list:
-                if data_cols:
-                    if set(data.schema.columns) != data_cols:
-                        raise ValueError(f"Data sets should all have the same columns for union on 0 axis.")
-                else:
-                    data_cols = set(data.schema.columns)
+            label_name_list = [data.schema.label_name for data in train_data_list]
+            if sum([name != label_name_list[0] for name in label_name_list]):
+                raise ValueError(f"Data sets should all have the same label_name for union.")
+
+            column_name_list = [set(data.schema.columns) for data in train_data_list]
+            if sum([col_names != column_name_list[0] for col_names in column_name_list]):
+                raise ValueError(f"Data sets should all have the same columns for union on 0 axis.")
             result_data = DataFrame.vstack(train_data_list)
         elif self.axis == 1:
             if sum([data.label is not None for data in train_data_list]) > 1: