Expand LogisticRegressionModelData as many pieces

Author: zhipeng93

docs/content/docs/operators/classification/logisticregression.md

@@ -323,7 +323,7 @@ public class OnlineLogisticRegressionExample {
 
         // Creates an online LogisticRegression object and initializes its parameters and initial
         // model data.
-        Row initModelData = Row.of(Vectors.dense(0.41233679404769874, -0.18088118293232122), 0L);
+		Row initModelData = Row.of(Vectors.dense(0.41233679404769874, -0.18088118293232122), 0L, 2L, 0L);
         Table initModelDataTable = tEnv.fromDataStream(env.fromElements(initModelData));
         OnlineLogisticRegression olr =
                 new OnlineLogisticRegression()

flink-ml-examples/src/main/java/org/apache/flink/ml/examples/classification/OnlineLogisticRegressionExample.java

@@ -96,7 +96,8 @@ public static void main(String[] args) {
 
         // Creates an online LogisticRegression object and initializes its parameters and initial
         // model data.
-        Row initModelData = Row.of(Vectors.dense(0.41233679404769874, -0.18088118293232122), 0L);
+        Row initModelData =
+                Row.of(Vectors.dense(0.41233679404769874, -0.18088118293232122), 0L, 2L, 0L);
         Table initModelDataTable = tEnv.fromDataStream(env.fromElements(initModelData));
         OnlineLogisticRegression olr =
                 new OnlineLogisticRegression()

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);

flink-ml-lib/src/test/java/org/apache/flink/ml/classification/LogisticRegressionTest.java

@@ -316,7 +316,7 @@ public void testSaveLoadAndPredict() throws Exception {
                         tempFolder.newFolder().getAbsolutePath(),
                         LogisticRegressionModel::load);
         assertEquals(
-                Arrays.asList("coefficient", "modelVersion"),
+                Arrays.asList("coefficient", "startIndex", "endIndex", "modelVersion"),
                 model.getModelData()[0].getResolvedSchema().getColumnNames());
         Table output = model.transform(binomialDataTable)[0];
         verifyPredictionResult(

flink-ml-lib/src/test/java/org/apache/flink/ml/classification/OnlineLogisticRegressionTest.java

@@ -256,6 +256,8 @@ public void before() throws Exception {
                                                 new double[] {
                                                     0.41233679404769874, -0.18088118293232122
                                                 }),
+                                        0L,
+                                        2L,
                                         0L)));
         initSparseModel =
                 tEnv.fromDataStream(
@@ -266,6 +268,8 @@ public void before() throws Exception {
                                                     0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01,
                                                     0.01, 0.01
                                                 }),
+                                        0L,
+                                        10L,
                                         0L)));
     }
 

flink-ml-python/pyflink/ml/classification/tests/test_logisticregression.py

@@ -120,7 +120,7 @@ def test_save_load_and_predict(self):
         model = regression.fit(self.binomial_data_table)
         self.assertEqual(
             model.get_model_data()[0].get_schema().get_field_names(),
-            ['coefficient', 'modelVersion'])
+            ['coefficient', "startIndex", "endIndex", 'modelVersion'])
         output = model.transform(self.binomial_data_table)[0]
         field_names = output.get_schema().get_field_names()
         self.verify_predict_result(

flink-ml-servable-core/src/main/java/org/apache/flink/ml/util/ServableReadWriteUtils.java

@@ -25,13 +25,14 @@
 import org.apache.flink.ml.servable.api.TransformerServable;
 import org.apache.flink.ml.servable.builder.PipelineModelServable;
 import org.apache.flink.util.InstantiationUtil;
-import org.apache.flink.util.Preconditions;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.SequenceInputStream;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -143,10 +144,10 @@ public static InputStream loadModelData(String path) throws IOException {
         FileSystem fileSystem = modelDataPath.getFileSystem();
 
         FileStatus[] files = fileSystem.listStatus(modelDataPath);
-        Preconditions.checkState(
-                files.length == 1,
-                "Only one model data file is expected in the directory %s.",
-                path);
-        return fileSystem.open(files[0].getPath());
+        List<InputStream> inputStreams = new ArrayList<>();
+        for (FileStatus file : files) {
+            inputStreams.add(fileSystem.open(file.getPath()));
+        }
+        return new SequenceInputStream(Collections.enumeration(inputStreams));
     }
 }

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

@@ -33,12 +33,23 @@ public class LogisticRegressionModelData {
 
     public DenseVector coefficient;
 
+    public long startIndex;
+
+    public long endIndex;
+
     public long modelVersion;
 
     public LogisticRegressionModelData() {}
 
     public LogisticRegressionModelData(DenseVector coefficient, long modelVersion) {
+        this(coefficient, 0L, coefficient.size(), modelVersion);
+    }
+
+    public LogisticRegressionModelData(
+            DenseVector coefficient, long startIndex, long endIndex, long modelVersion) {
         this.coefficient = coefficient;
+        this.startIndex = startIndex;
+        this.endIndex = endIndex;
         this.modelVersion = modelVersion;
     }
 
@@ -54,6 +65,8 @@ public void encode(OutputStream outputStream) throws IOException {
 
         DenseVectorSerializer serializer = new DenseVectorSerializer();
         serializer.serialize(coefficient, dataOutputViewStreamWrapper);
+        dataOutputViewStreamWrapper.writeLong(startIndex);
+        dataOutputViewStreamWrapper.writeLong(endIndex);
         dataOutputViewStreamWrapper.writeLong(modelVersion);
     }
 
@@ -69,8 +82,10 @@ static LogisticRegressionModelData decode(InputStream inputStream) throws IOExce
 
         DenseVectorSerializer serializer = new DenseVectorSerializer();
         DenseVector coefficient = serializer.deserialize(dataInputViewStreamWrapper);
+        long startIndex = dataInputViewStreamWrapper.readLong();
+        long endIndex = dataInputViewStreamWrapper.readLong();
         long modelVersion = dataInputViewStreamWrapper.readLong();
 
-        return new LogisticRegressionModelData(coefficient, modelVersion);
+        return new LogisticRegressionModelData(coefficient, startIndex, endIndex, modelVersion);
     }
 }

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

@@ -18,6 +18,7 @@
 
 package org.apache.flink.ml.classification.logisticregression;
 
+import org.apache.flink.annotation.VisibleForTesting;
 import org.apache.flink.api.java.tuple.Tuple2;
 import org.apache.flink.ml.linalg.BLAS;
 import org.apache.flink.ml.linalg.DenseVector;
@@ -81,11 +82,49 @@ public DataFrame transform(DataFrame input) {
     public LogisticRegressionModelServable setModelData(InputStream... modelDataInputs)
             throws IOException {
         Preconditions.checkArgument(modelDataInputs.length == 1);
+        List<LogisticRegressionModelData> modelPieces = new ArrayList<>();
+        while (true) {
+            try {
+                LogisticRegressionModelData piece =
+                        LogisticRegressionModelData.decode(modelDataInputs[0]);
+                modelPieces.add(piece);
+            } catch (IOException e) {
+                // Reached the end of model stream.
+                break;
+            }
+        }
 
-        modelData = LogisticRegressionModelData.decode(modelDataInputs[0]);
+        modelData = mergePieces(modelPieces);
         return this;
     }
 
+    @VisibleForTesting
+    public static LogisticRegressionModelData mergePieces(
+            List<LogisticRegressionModelData> pieces) {
+        long dim = 0;
+        for (LogisticRegressionModelData piece : pieces) {
+            dim = Math.max(dim, piece.endIndex);
+        }
+        // TODO: Add distributed inference for very large models.
+        Preconditions.checkState(
+                dim < Integer.MAX_VALUE,
+                "The dimension of logistic regression model is larger than INT.MAX. Please consider using distributed inference.");
+        int intDim = (int) dim;
+        DenseVector mergedCoefficient = new DenseVector(intDim);
+        for (LogisticRegressionModelData piece : pieces) {
+            int startIndex = (int) piece.startIndex;
+            int endIndex = (int) piece.endIndex;
+            System.arraycopy(
+                    piece.coefficient.values,
+                    0,
+                    mergedCoefficient.values,
+                    startIndex,
+                    endIndex - startIndex);
+        }
+        return new LogisticRegressionModelData(
+                mergedCoefficient, 0, mergedCoefficient.size(), pieces.get(0).modelVersion);
+    }
+
     public static LogisticRegressionModelServable load(String path) throws IOException {
         LogisticRegressionModelServable servable =
                 ServableReadWriteUtils.loadServableParam(