[FLINK-27826] Support training very high dimensional logisticRegression

Author: zhipeng93

flink-ml-lib/pom.xml

@@ -138,6 +138,12 @@ under the License.
       <scope>test</scope>
       <type>test-jar</type>
     </dependency>
+    <dependency>
+      <groupId>fastutil</groupId>
+      <artifactId>fastutil</artifactId>
+      <version>5.0.9</version>
+    </dependency>
+
   </dependencies>
 
   <build>

flink-ml-lib/src/main/java/org/apache/flink/ml/classification/logisticregression/LogisticRegressionModel.java

@@ -43,6 +43,7 @@
 import java.io.IOException;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 /** A Model which classifies data using the model data computed by {@link LogisticRegression}. */
@@ -147,10 +148,16 @@ public PredictLabelFunction(String broadcastModelKey, Map<Param<?>, Object> para
         @Override
         public Row map(Row dataPoint) {
             if (servable == null) {
-                LogisticRegressionModelData modelData =
-                        (LogisticRegressionModelData)
-                                getRuntimeContext().getBroadcastVariable(broadcastModelKey).get(0);
-                servable = new LogisticRegressionModelServable(modelData);
+                List<LogisticRegressionModelData> modelData =
+                        getRuntimeContext().getBroadcastVariable(broadcastModelKey);
+
+                if (modelData.size() == 1) {
+                    servable = new LogisticRegressionModelServable(modelData.get(0));
+                } else {
+                    LogisticRegressionModelData mergedModel =
+                            LogisticRegressionModelServable.mergePieces(modelData);
+                    servable = new LogisticRegressionModelServable(mergedModel);
+                }
                 ParamUtils.updateExistingParams(servable, params);
             }
             Vector features = (Vector) dataPoint.getField(servable.getFeaturesCol());

flink-ml-lib/src/main/java/org/apache/flink/ml/classification/logisticregression/LogisticRegressionModelDataUtil.java

@@ -89,7 +89,13 @@ public static DataStream<LogisticRegressionModelData> getModelDataStream(Table m
         StreamTableEnvironment tEnv =
                 (StreamTableEnvironment) ((TableImpl) modelData).getTableEnvironment();
         return tEnv.toDataStream(modelData)
-                .map(x -> new LogisticRegressionModelData(x.getFieldAs(0), x.getFieldAs(1)));
+                .map(
+                        x ->
+                                new LogisticRegressionModelData(
+                                        x.getFieldAs(0),
+                                        x.getFieldAs(1),
+                                        x.getFieldAs(2),
+                                        x.getFieldAs(3)));
     }
 
     /**
@@ -107,7 +113,10 @@ public static DataStream<byte[]> getModelDataByteStream(Table modelDataTable) {
                         x -> {
                             LogisticRegressionModelData modelData =
                                     new LogisticRegressionModelData(
-                                            x.getFieldAs(0), x.getFieldAs(1));
+                                            x.getFieldAs(0),
+                                            x.getFieldAs(1),
+                                            x.getFieldAs(2),
+                                            x.getFieldAs(3));
 
                             ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
                             modelData.encode(outputStream);