Support pull/push value as array

Author: zhipeng93

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

@@ -137,14 +137,14 @@ public LogisticRegressionModel fit(Table... inputs) {
         IterationStageList<LogisticRegressionWithFtrlTrainingContext> iterationStages =
                 new IterationStageList<>(trainingContext);
         iterationStages
-                .addTrainingStage(new ComputeIndices())
-                .addTrainingStage(
+                .addStage(new ComputeIndices())
+                .addStage(
                         new PullStage(
                                 (SerializableSupplier<long[]>) () -> trainingContext.pullIndices,
                                 (SerializableConsumer<double[]>)
                                         x -> trainingContext.pulledValues = x))
-                .addTrainingStage(new ComputeGradients(BinaryLogisticLoss.INSTANCE))
-                .addTrainingStage(
+                .addStage(new ComputeGradients(BinaryLogisticLoss.INSTANCE))
+                .addStage(
                         new PushStage(
                                 (SerializableSupplier<long[]>) () -> trainingContext.pushIndices,
                                 (SerializableSupplier<double[]>) () -> trainingContext.pushValues))
@@ -160,13 +160,7 @@ public LogisticRegressionModel fit(Table... inputs) {
                         trainData.getParallelism());
 
         DataStream<Tuple3<Long, Long, double[]>> rawModelData =
-                TrainingUtils.train(
-                        modelDim,
-                        trainData,
-                        ftrl,
-                        iterationStages,
-                        getNumServers(),
-                        getNumServerCores());
+                TrainingUtils.train(modelDim, trainData, ftrl, iterationStages, getNumServers());
 
         final long modelVersion = 0L;
 
@@ -341,8 +335,8 @@ public void setWorldInfo(int workerId, int numWorkers) {
     }
 
     @Override
-    public void setTrainData(ResettableIterator<?> trainData) {
-        this.trainData = (ResettableIterator<LabeledLargePointWithWeight>) trainData;
+    public void setInputData(ResettableIterator<?> inputData) {
+        this.trainData = (ResettableIterator<LabeledLargePointWithWeight>) inputData;
     }
 
     @Override

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

@@ -51,13 +51,6 @@ public interface LogisticRegressionWithFtrlParams<T>
                     1,
                     ParamValidators.gtEq(1));
 
-    Param<Integer> NUM_SERVER_CORES =
-            new IntParam(
-                    "numServerCores",
-                    "number of cores that a server can use.",
-                    1,
-                    ParamValidators.gtEq(1));
-
     Param<Double> ALPHA =
             new DoubleParam(
                     "alpha",
@@ -81,14 +74,6 @@ default T setNumServers(Integer value) {
         return set(NUM_SERVERS, value);
     }
 
-    default int getNumServerCores() {
-        return get(NUM_SERVER_CORES);
-    }
-
-    default T setNumServerCores(int value) {
-        return set(NUM_SERVER_CORES, value);
-    }
-
     default double getAlpha() {
         return get(ALPHA);
     }

flink-ml-lib/src/main/java/org/apache/flink/ml/common/ps/RangePartitioner.java

@@ -67,6 +67,12 @@ private static class RequestsIterator implements Iterator<Tuple3<Integer, long[]
         private final int numServers;
         private final long[] indices;
         private final double[] values;
+        /**
+         * Number of values per key. If the model data is a vector, numValuesPerKey is one. If the
+         * model data is a matrix, numValuesPerKey is the number of columns.
+         */
+        private final int numValuesPerKey;
+
         private final long[] ranges;
 
         private int serverId = 0;
@@ -75,11 +81,18 @@ private static class RequestsIterator implements Iterator<Tuple3<Integer, long[]
 
         public RequestsIterator(
                 int numPss, long[] indices, @Nullable double[] values, long[] ranges) {
-            // Preconditions.checkArgument(values == null || values.length % indices.length == 0);
             this.numServers = numPss;
             this.indices = indices;
             this.values = values;
             this.ranges = ranges;
+            if (indices.length != 0 && values != null) {
+                numValuesPerKey = values.length / indices.length;
+                Preconditions.checkArgument(
+                        numValuesPerKey * indices.length == values.length,
+                        "The size of values cannot be divided by size of keys.");
+            } else {
+                numValuesPerKey = 1;
+            }
         }
 
         @Override
@@ -98,7 +111,11 @@ public Tuple3<Integer, long[], double[]> next() {
             double[] splitValues = values == null ? null : new double[0];
             if (s < e) {
                 splitIndices = Arrays.copyOfRange(indices, s, e);
-                splitValues = values == null ? null : Arrays.copyOfRange(values, s, e);
+                splitValues =
+                        values == null
+                                ? null
+                                : Arrays.copyOfRange(
+                                        values, s * numValuesPerKey, e * numValuesPerKey);
             }
             s = e;
             serverId++;