com.tencent.angel.split.SplitClassification Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of angel-ps-core Show documentation
There is a newer version: 3.2.0
/*
 * Tencent is pleased to support the open source community by making Angel available.
 *
 * Copyright (C) 2017-2018 THL A29 Limited, a Tencent company. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in 
 * compliance with the License. You may obtain a copy of the License at
 *
 * https://opensource.org/licenses/Apache-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License
 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 * or implied. See the License for the specific language governing permissions and limitations under
 * the License.
 *
 */


package com.tencent.angel.split;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.serializer.Deserializer;
import org.apache.hadoop.io.serializer.SerializationFactory;
import org.apache.hadoop.io.serializer.Serializer;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class SplitClassification {

  private String[] locations;
  private boolean useNewAPI;
  private List splitsOldAPI = null;
  private List splitsNewAPI = null;

  // private Configuration conf;

  public SplitClassification() {

  }

  public SplitClassification(List splitsOldAPI,
    List splitsNewAPI, String[] locations,
    boolean useNewAPI) {
    this.locations = locations;
    this.splitsOldAPI = splitsOldAPI;
    this.splitsNewAPI = splitsNewAPI;
    this.useNewAPI = useNewAPI;
  }

  public SplitClassification(List splitsOldAPI,
    List splitsNewAPI, boolean useNewAPI) {
    this.splitsOldAPI = splitsOldAPI;
    this.splitsNewAPI = splitsNewAPI;
    this.useNewAPI = useNewAPI;
  }

  public String[] getLocations() {
    return locations;
  }

  public void setLocations(String[] locations) {
    this.locations = locations;
  }

  public boolean isUseNewAPI() {
    return useNewAPI;
  }

  public void setUseNewAPI(boolean useNewAPI) {
    this.useNewAPI = useNewAPI;
  }

  public List getSplitsOldAPI() {
    return splitsOldAPI;
  }

  public void setSplitsOldAPI(List splitsOldAPI) {
    this.splitsOldAPI = splitsOldAPI;
  }

  public List getSplitsNewAPI() {
    return splitsNewAPI;
  }

  public void setSplitsNewAPI(List splitsNewAPI) {
    this.splitsNewAPI = splitsNewAPI;
  }

  public void addSplit(org.apache.hadoop.mapreduce.InputSplit split) {
    splitsNewAPI.add(split);
  }

  public void addSplit(org.apache.hadoop.mapred.InputSplit split) {
    splitsOldAPI.add(split);
  }

  public int getSplitNum() {
    if (useNewAPI) {
      return splitsNewAPI.size();
    } else {
      return splitsOldAPI.size();
    }
  }

  public org.apache.hadoop.mapreduce.InputSplit getSplitNewAPI(int splitIndex) {
    return splitsNewAPI.get(splitIndex);
  }

  public org.apache.hadoop.mapred.InputSplit getSplitOldAPI(int splitIndex) {
    return splitsOldAPI.get(splitIndex);
  }

  @Override public String toString() {
    return "SplitClassification [locations=" + Arrays.toString(locations) + ", useNewAPI="
      + useNewAPI + ", splitsOldAPI=" + splitsOldAPI + ", splitsNewAPI=" + splitsNewAPI + "]";
  }

  @SuppressWarnings({"rawtypes", "unchecked"}) public void serialize(DataOutputStream output)
    throws IOException {
    output.writeBoolean(useNewAPI);
    if (useNewAPI) {
      int size = splitsNewAPI.size();
      output.writeInt(size);
      output.writeInt(locations.length);
      for (int i = 0; i < locations.length; i++) {
        output.writeUTF(locations[i]);
      }
      if (size > 0) {
        output.writeUTF(splitsNewAPI.get(0).getClass().getName());
        SerializationFactory factory = new SerializationFactory(new Configuration());
        Serializer serializer = factory.getSerializer(splitsNewAPI.get(0).getClass());
        serializer.open(output);
        for (int i = 0; i < size; i++) {
          serializer.serialize(splitsNewAPI.get(i));
        }
      }
    } else {
      int size = splitsOldAPI.size();
      output.writeInt(size);
      output.writeInt(locations.length);
      for (int i = 0; i < size; i++) {
        output.writeUTF(locations[i]);
      }
      if (size > 0) {
        output.writeUTF(splitsOldAPI.get(0).getClass().getName());
        SerializationFactory factory = new SerializationFactory(new Configuration());
        Serializer serializer = factory.getSerializer(splitsOldAPI.get(0).getClass());
        serializer.open(output);
        for (int i = 0; i < size; i++) {
          serializer.serialize(splitsOldAPI.get(i));
        }
      }
    }
  }

  @SuppressWarnings("unchecked") public void deserialize(DataInputStream input)
    throws IOException, ClassNotFoundException {
    useNewAPI = input.readBoolean();
    int size = input.readInt();
    int locSize = input.readInt();
    locations = new String[locSize];
    for (int i = 0; i < locSize; i++) {
      locations[i] = input.readUTF();
    }

    if (useNewAPI) {
      if (size > 0) {
        String splitClass = input.readUTF();
        splitsNewAPI = new ArrayList(size);
        SerializationFactory factory = new SerializationFactory(new Configuration());
        Deserializer deSerializer = factory
          .getDeserializer(
            (Class) Class.forName(splitClass));

        deSerializer.open(input);
        for (int i = 0; i < size; i++) {
          splitsNewAPI.add(deSerializer.deserialize(null));
        }
      }

    } else {
      if (size > 0) {
        String splitClass = input.readUTF();
        splitsOldAPI = new ArrayList(size);
        SerializationFactory factory = new SerializationFactory(new Configuration());
        Deserializer deSerializer = factory
          .getDeserializer(
            (Class) Class.forName(splitClass));

        deSerializer.open(input);
        for (int i = 0; i < size; i++) {
          splitsOldAPI.add(deSerializer.deserialize(null));
        }
      }
    }
  }
}