private <U> SparkCollection<U> wrap(JavaDStream<U> stream) { return new DStreamCollection<>(sec, stream); } }
private <T, U> PairDStreamCollection<T, U> wrap(JavaPairDStream<T, U> pairStream) { return new PairDStreamCollection<>(sec, pairStream); } }
@Override public SparkCollection<RecordInfo<Object>> multiOutputTransform(StageSpec stageSpec, StageStatisticsCollector collector) { return wrap(stream.transform(new DynamicTransform<T>(new DynamicDriverContext(stageSpec, sec, collector), true))); }
@Override protected SparkPairCollection<Object, Object> addJoinKey(StageSpec stageSpec, String inputStageName, SparkCollection<Object> inputCollection, StageStatisticsCollector collector) throws Exception { DynamicDriverContext dynamicDriverContext = new DynamicDriverContext(stageSpec, sec, collector); JavaDStream<Object> dStream = inputCollection.getUnderlying(); JavaPairDStream<Object, Object> result = dStream.transformToPair(new DynamicJoinOn<>(dynamicDriverContext, inputStageName)); return new PairDStreamCollection<>(sec, result); }
@Override protected SparkCollection<Object> mergeJoinResults( StageSpec stageSpec, SparkPairCollection<Object, List<JoinElement<Object>>> joinedInputs, StageStatisticsCollector collector) throws Exception { DynamicDriverContext dynamicDriverContext = new DynamicDriverContext(stageSpec, sec, collector); JavaPairDStream<Object, List<JoinElement<Object>>> pairDStream = joinedInputs.getUnderlying(); JavaDStream<Object> result = pairDStream.transform(new DynamicJoinMerge<>(dynamicDriverContext)); return new DStreamCollection<>(sec, result); } }
@Override public SparkCollection<RecordInfo<Object>> aggregate(StageSpec stageSpec, @Nullable Integer partitions, StageStatisticsCollector collector) { DynamicDriverContext dynamicDriverContext = new DynamicDriverContext(stageSpec, sec, collector); JavaPairDStream<Object, T> keyedCollection = stream.transformToPair(new DynamicAggregatorGroupBy<Object, T>(dynamicDriverContext)); JavaPairDStream<Object, Iterable<T>> groupedCollection = partitions == null ? keyedCollection.groupByKey() : keyedCollection.groupByKey(partitions); return wrap(groupedCollection.transform(new DynamicAggregatorAggregate<Object, T, Object>(dynamicDriverContext))); }
@Override public <T> SparkPairCollection<K, T> mapValues(Function<V, T> function) { return wrap(pairStream.mapValues(function)); }
@Override public <U> SparkCollection<U> flatMap(StageSpec stageSpec, FlatMapFunction<T, U> function) { return wrap(stream.flatMap(function)); }
@Override public JavaRDD<RecordInfo<Object>> call(JavaPairRDD<GROUP_KEY, Iterable<GROUP_VAL>> input, Time batchTime) throws Exception { if (function == null) { function = Compat.convert( new AggregatorAggregateFunction<GROUP_KEY, GROUP_VAL, OUT>(dynamicDriverContext.getPluginFunctionContext())); } return input.flatMap(function); } }
@Override public SparkCollection<RecordInfo<Object>> transform(StageSpec stageSpec, StageStatisticsCollector collector) { return wrap(stream.transform(new DynamicTransform<T>(new DynamicDriverContext(stageSpec, sec, collector), false))); }
@Override public <T> SparkCollection<T> flatMap(FlatMapFunction<Tuple2<K, V>, T> function) { return new DStreamCollection<>(sec, pairStream.flatMap(function)); }
@Override public <K, V> SparkPairCollection<K, V> flatMapToPair(PairFlatMapFunction<T, K, V> function) { return new PairDStreamCollection<>(sec, stream.flatMapToPair(function)); }
@SuppressWarnings("unchecked") @Override public <T> SparkPairCollection<K, Tuple2<Optional<V>, Optional<T>>> fullOuterJoin(SparkPairCollection<K, T> other) { return wrap(StreamingCompat.fullOuterJoin(pairStream, (JavaPairDStream<K, T>) other.getUnderlying())); }
@Override public SparkCollection<T> cache() { return wrap(stream.cache()); }
@SuppressWarnings("unchecked") @Override public <T> SparkPairCollection<K, Tuple2<V, T>> join(SparkPairCollection<K, T> other) { return wrap(pairStream.join((JavaPairDStream<K, T>) other.getUnderlying())); }
@SuppressWarnings("unchecked") @Override public SparkCollection<T> union(SparkCollection<T> other) { return wrap(stream.union((JavaDStream<T>) other.getUnderlying())); }
@SuppressWarnings("unchecked") @Override public <T> SparkPairCollection<K, Tuple2<V, Optional<T>>> leftOuterJoin(SparkPairCollection<K, T> other, int numPartitions) { return wrap( StreamingCompat.leftOuterJoin(pairStream, (JavaPairDStream<K, T>) other.getUnderlying(), numPartitions)); }
@SuppressWarnings("unchecked") @Override public <T> SparkPairCollection<K, Tuple2<Optional<V>, Optional<T>>> fullOuterJoin(SparkPairCollection<K, T> other, int numPartitions) { return wrap( StreamingCompat.fullOuterJoin(pairStream, (JavaPairDStream<K, T>) other.getUnderlying(), numPartitions)); }
@SuppressWarnings("unchecked") @Override public <T> SparkPairCollection<K, Tuple2<V, T>> join(SparkPairCollection<K, T> other, int numPartitions) { return wrap(pairStream.join((JavaPairDStream<K, T>) other.getUnderlying(), numPartitions)); }
@SuppressWarnings("unchecked") @Override public <T> SparkPairCollection<K, Tuple2<V, Optional<T>>> leftOuterJoin(SparkPairCollection<K, T> other) { return wrap(StreamingCompat.leftOuterJoin(pairStream, (JavaPairDStream<K, T>) other.getUnderlying())); }