com.streamsets.pipeline.sdk.BatchMakerImpl Maven / Gradle / Ivy

Go to download
/*
 * Copyright (c) 2021 StreamSets Inc.
 */
package com.streamsets.pipeline.sdk;

import com.streamsets.datacollector.record.RecordImpl;
import com.streamsets.pipeline.api.BatchMaker;
import com.streamsets.pipeline.api.Record;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Implementation of BatchMaker interface for SDK
 */
class BatchMakerImpl implements BatchMaker {

  private final List outputLanes;
  private final Map> laneToRecordsMap;
  private final String singleLaneOutput;

  public BatchMakerImpl(List outputLanes) {
    this.outputLanes = outputLanes;
    if(outputLanes.size() == 1) {
      singleLaneOutput = outputLanes.iterator().next();
    } else {
      singleLaneOutput = null;
    }
    laneToRecordsMap = new HashMap<>();
    //The output map should always have a key for all the defined output lanes, if the stage did not produce any record
    // for a lane, the value in the map should be an empty record list.
    for(String lane : outputLanes) {
      laneToRecordsMap.put(lane, new ArrayList());
    }
  }

  @Override
  public List getLanes() {
    return outputLanes;
  }

  @Override
  public void addRecord(Record record, String... lanes) {
    if(lanes == null || lanes.length == 0) {
      List records = laneToRecordsMap.get(singleLaneOutput);
      if(records == null) {
        records = new ArrayList<>();
        laneToRecordsMap.put(singleLaneOutput, records);
      }
      records.add(((RecordImpl)record).clone());
      return;
    }
    for(String lane : lanes) {
      List records = laneToRecordsMap.get(lane);
      if(records == null) {
        records = new ArrayList<>();
        laneToRecordsMap.put(lane, records);
      }
      records.add(((RecordImpl)record).clone());
    }
  }

  public Map> getOutput() {
    return laneToRecordsMap;
  }
}