io.ray.streaming.api.stream.JoinStream Maven / Gradle / Ivy

package io.ray.streaming.api.stream;

import io.ray.streaming.api.function.impl.JoinFunction;
import io.ray.streaming.api.function.impl.KeyFunction;
import io.ray.streaming.operator.impl.JoinOperator;
import java.io.Serializable;

/**
 * Represents a DataStream of two joined DataStream.
 *
 * @param  Type of the data in the left stream.
 * @param  Type of the data in the right stream.
 * @param  Type of the data in the joined stream.
 */
public class JoinStream extends DataStream {

  private final DataStream rightStream;

  public JoinStream(DataStream leftStream, DataStream rightStream) {
    super(leftStream, new JoinOperator<>());
    this.rightStream = rightStream;
  }

  public DataStream getRightStream() {
    return rightStream;
  }

  /** Apply key-by to the left join stream. */
  public  Where where(KeyFunction keyFunction) {
    return new Where<>(this, keyFunction);
  }

  /**
   * Where clause of the join transformation.
   *
   * @param  Type of the join key.
   */
  class Where implements Serializable {

    private JoinStream joinStream;
    private KeyFunction leftKeyByFunction;

    Where(JoinStream joinStream, KeyFunction leftKeyByFunction) {
      this.joinStream = joinStream;
      this.leftKeyByFunction = leftKeyByFunction;
    }

    public Equal equalTo(KeyFunction rightKeyFunction) {
      return new Equal<>(joinStream, leftKeyByFunction, rightKeyFunction);
    }
  }

  /**
   * Equal clause of the join transformation.
   *
   * @param  Type of the join key.
   */
  class Equal implements Serializable {

    private JoinStream joinStream;
    private KeyFunction leftKeyByFunction;
    private KeyFunction rightKeyByFunction;

    Equal(
        JoinStream joinStream,
        KeyFunction leftKeyByFunction,
        KeyFunction rightKeyByFunction) {
      this.joinStream = joinStream;
      this.leftKeyByFunction = leftKeyByFunction;
      this.rightKeyByFunction = rightKeyByFunction;
    }

    @SuppressWarnings("unchecked")
    public DataStream with(JoinFunction joinFunction) {
      JoinOperator joinOperator = (JoinOperator) joinStream.getOperator();
      joinOperator.setFunction(joinFunction);
      return (DataStream) joinStream;
    }
  }
}