<commit_before>package com.linkedin.uif.source.extractor.extract;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.Map.Entry;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import com.google.common.collect.Lists;
import com.linkedin.uif.configuration.ConfigurationKeys;
import com.linkedin.uif.configuration.SourceState;
import com.linkedin.uif.configuration.WorkUnitState;
import com.linkedin.uif.configuration.WorkUnitState.WorkingState;
import com.linkedin.uif.source.Source;
import com.linkedin.uif.source.extractor.extract.restapi.SalesforceSource;
import com.linkedin.uif.source.extractor.partition.Partitioner;
import com.linkedin.uif.source.extractor.utils.Utils;
import com.linkedin.uif.source.extractor.watermark.WatermarkPredicate;
import com.linkedin.uif.source.extractor.watermark.WatermarkType;
import com.linkedin.uif.source.workunit.Extract;
import com.linkedin.uif.source.workunit.WorkUnit;
import com.linkedin.uif.source.workunit.Extract.TableType;

/**
 * An implementation of Base source to get work units
 */
public abstract class BaseSource<S, D> implements Source<S, D> {
	private static final Log LOG = LogFactory.getLog(BaseSource.class);

    /**
     * get work units using Source state
     *
     * @param SourceState
     * @return list of work units
     */
	@Override
	public List<WorkUnit> getWorkunits(SourceState state) {
		LOG.info("Get work units");
		List<WorkUnit> workUnits = Lists.newArrayList();
		String nameSpaceName = state.getProp(ConfigurationKeys.EXTRACT_NAMESPACE_NAME_KEY);
		String entityName = state.getProp(ConfigurationKeys.SOURCE_ENTITY);
		TableType tableType = TableType.valueOf(state.getProp(ConfigurationKeys.EXTRACT_TABLE_TYPE_KEY).toUpperCase());
		long previousWatermark = this.getLatestWatermarkFromMetadata(state);

		Partitioner partitioner = new Partitioner(state);
		HashMap<Long, Long> partitions = partitioner.getPartitions(previousWatermark);
		Map<Long, Long> sortedPartitions = new TreeMap<Long, Long>(partitions);
		int workUnitCount = 0;
		for (Entry<Long, Long> entry : sortedPartitions.entrySet()) {
			SourceState partitionState = new SourceState();
			partitionState.addAll(state);
			partitionState.setProp(ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY, entry.getKey());
			partitionState.setProp(ConfigurationKeys.WORK_UNIT_HIGH_WATER_MARK_KEY, entry.getValue());
			Extract extract = partitionState.createExtract(tableType, nameSpaceName, entityName);
			workUnits.add(partitionState.createWorkUnit(extract));
			workUnitCount++;
		}
		LOG.info("Total number of work units for the current run: "+workUnitCount);
		
		List<WorkUnit> previousWorkUnits = this.getPreviousIncompleteWorkUnits(state);
		LOG.info("Total number of work units from the previous failed runs: "+previousWorkUnits.size());
		
		workUnits.addAll(previousWorkUnits);
		return workUnits;
	}

    /**
     * get all the previous work units which are in incomplete state
     *
     * @param SourceState
     * @return list of work units
     */
	private List<WorkUnit> getPreviousIncompleteWorkUnits(SourceState state) {
		LOG.debug("Getting previous unsuccessful work units");
		List<WorkUnit> previousWorkUnits = new ArrayList<WorkUnit>();
		List<WorkUnitState> previousWorkUnitStates = state.getPreviousStates();
		if(previousWorkUnitStates.size() == 0) {
			LOG.debug("Previous states are not found");
			return previousWorkUnits;
		}
		
		for(WorkUnitState workUnitState : previousWorkUnitStates) {
			if(workUnitState.getWorkingState() == WorkingState.FAILED || workUnitState.getWorkingState() == WorkingState.ABORTED) {
				previousWorkUnits.add(workUnitState.getWorkunit());
			}
		}
		
		return previousWorkUnits;
	}

    /**
     * get latest water mark from previous work unit states
     *
     * @param SourceState
     * @return latest water mark(high water mark of WorkUnitState)
     */
	private long getLatestWatermarkFromMetadata(SourceState state) {
		LOG.debug("Getting previous watermark");
		boolean hasDataInPreviousRun = false;
		long latestWaterMark = ConfigurationKeys.DEFAULT_WATERMARK_VALUE;
		
		List<WorkUnitState> previousWorkUnitStates = state.getPreviousStates();
		List<Long> previousWorkUnitStateHighWatermarks = new ArrayList<Long>(); 
		List<Long> previousWorkUnitLowWatermarks = new ArrayList<Long>();
		
		if(previousWorkUnitStates.size() == 0) {
			LOG.info("Previous states are not found; Latest watermark - Default watermark:"+latestWaterMark);
			return latestWaterMark;
		}
		
		for(WorkUnitState workUnitState : previousWorkUnitStates) {
			if(workUnitState.getWorkingState() == WorkingState.COMMITTED) {
				if(workUnitState.getHighWaterMark() != ConfigurationKeys.DEFAULT_WATERMARK_VALUE) {
					hasDataInPreviousRun = true;
				}
				previousWorkUnitStateHighWatermarks.add(workUnitState.getHighWaterMark());
				previousWorkUnitLowWatermarks.add(this.getLowWatermarkFromWorkUnit(workUnitState.getWorkunit(), workUnitState.getProp(ConfigurationKeys.SOURCE_WATERMARK_TYPE)));
			}
		}
		
		if(hasDataInPreviousRun) {
			latestWaterMark = Collections.max(previousWorkUnitStateHighWatermarks);
			LOG.info("Previous run has data; Latest watermark - Max watermark from WorkUnitStates:"+latestWaterMark);
		} else {
			latestWaterMark = Collections.min(previousWorkUnitLowWatermarks);
			LOG.info("Previous run has no data; Latest watermark - Min watermark from WorkUnits:"+latestWaterMark);
		}
		
		return latestWaterMark;
	}

    /**
     * get latest water mark from previous work units
     *
     * @param work unit
     * @param watermark type
     * @return latest water mark(low water mark of WorkUnit)
     */
	private long getLowWatermarkFromWorkUnit(WorkUnit workunit, String watermarkType) {
		if(workunit.getLowWaterMark() == ConfigurationKeys.DEFAULT_WATERMARK_VALUE) {
			return ConfigurationKeys.DEFAULT_WATERMARK_VALUE;
		}
		
		WatermarkType wmType = WatermarkType.valueOf(watermarkType.toUpperCase());
		int deltaNum = new WatermarkPredicate(null, wmType).getDeltaNumForNextWatermark();
		
		if(wmType == WatermarkType.SIMPLE) {
			return workunit.getLowWaterMark() - deltaNum;
		} else {
			return Long.parseLong(Utils.dateToString(Utils.addSecondsToDate(Utils.toDate(workunit.getLowWaterMark(), "yyyyMMddHHmmss"),deltaNum*-1), "yyyyMMddHHmmss"));
		}
	}

	@Override
	public void shutdown(SourceState state) {
	}
}
<commit_msg>Get default watermark from BaseSource when there are failed or aborted previous states without committed states<commit_after>package com.linkedin.uif.source.extractor.extract;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.Map.Entry;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import com.google.common.collect.Lists;
import com.linkedin.uif.configuration.ConfigurationKeys;
import com.linkedin.uif.configuration.SourceState;
import com.linkedin.uif.configuration.WorkUnitState;
import com.linkedin.uif.configuration.WorkUnitState.WorkingState;
import com.linkedin.uif.source.Source;
import com.linkedin.uif.source.extractor.extract.restapi.SalesforceSource;
import com.linkedin.uif.source.extractor.partition.Partitioner;
import com.linkedin.uif.source.extractor.utils.Utils;
import com.linkedin.uif.source.extractor.watermark.WatermarkPredicate;
import com.linkedin.uif.source.extractor.watermark.WatermarkType;
import com.linkedin.uif.source.workunit.Extract;
import com.linkedin.uif.source.workunit.WorkUnit;
import com.linkedin.uif.source.workunit.Extract.TableType;

/**
 * An implementation of Base source to get work units
 */
public abstract class BaseSource<S, D> implements Source<S, D> {
	private static final Log LOG = LogFactory.getLog(BaseSource.class);

    /**
     * get work units using Source state
     *
     * @param SourceState
     * @return list of work units
     */
	@Override
	public List<WorkUnit> getWorkunits(SourceState state) {
		LOG.info("Get work units");
		List<WorkUnit> workUnits = Lists.newArrayList();
		String nameSpaceName = state.getProp(ConfigurationKeys.EXTRACT_NAMESPACE_NAME_KEY);
		String entityName = state.getProp(ConfigurationKeys.SOURCE_ENTITY);
		TableType tableType = TableType.valueOf(state.getProp(ConfigurationKeys.EXTRACT_TABLE_TYPE_KEY).toUpperCase());
		long previousWatermark = this.getLatestWatermarkFromMetadata(state);

		Partitioner partitioner = new Partitioner(state);
		HashMap<Long, Long> partitions = partitioner.getPartitions(previousWatermark);
		Map<Long, Long> sortedPartitions = new TreeMap<Long, Long>(partitions);
		int workUnitCount = 0;
		for (Entry<Long, Long> entry : sortedPartitions.entrySet()) {
			SourceState partitionState = new SourceState();
			partitionState.addAll(state);
			partitionState.setProp(ConfigurationKeys.WORK_UNIT_LOW_WATER_MARK_KEY, entry.getKey());
			partitionState.setProp(ConfigurationKeys.WORK_UNIT_HIGH_WATER_MARK_KEY, entry.getValue());
			Extract extract = partitionState.createExtract(tableType, nameSpaceName, entityName);
			workUnits.add(partitionState.createWorkUnit(extract));
			workUnitCount++;
		}
		LOG.info("Total number of work units for the current run: "+workUnitCount);
		
		List<WorkUnit> previousWorkUnits = this.getPreviousIncompleteWorkUnits(state);
		LOG.info("Total number of work units from the previous failed runs: "+previousWorkUnits.size());
		
		workUnits.addAll(previousWorkUnits);
		return workUnits;
	}

    /**
     * get all the previous work units which are in incomplete state
     *
     * @param SourceState
     * @return list of work units
     */
	private List<WorkUnit> getPreviousIncompleteWorkUnits(SourceState state) {
		LOG.debug("Getting previous unsuccessful work units");
		List<WorkUnit> previousWorkUnits = new ArrayList<WorkUnit>();
		List<WorkUnitState> previousWorkUnitStates = state.getPreviousStates();
		if(previousWorkUnitStates.size() == 0) {
			LOG.debug("Previous states are not found");
			return previousWorkUnits;
		}
		
		for(WorkUnitState workUnitState : previousWorkUnitStates) {
			if(workUnitState.getWorkingState() == WorkingState.FAILED || workUnitState.getWorkingState() == WorkingState.ABORTED) {
				previousWorkUnits.add(workUnitState.getWorkunit());
			}
		}
		
		return previousWorkUnits;
	}

    /**
     * get latest water mark from previous work unit states
     *
     * @param SourceState
     * @return latest water mark(high water mark of WorkUnitState)
     */
	private long getLatestWatermarkFromMetadata(SourceState state) {
		LOG.debug("Getting previous watermark");
		boolean hasDataInPreviousRun = false;
		long latestWaterMark = ConfigurationKeys.DEFAULT_WATERMARK_VALUE;
		
		List<WorkUnitState> previousWorkUnitStates = state.getPreviousStates();
		List<Long> previousWorkUnitStateHighWatermarks = new ArrayList<Long>(); 
		List<Long> previousWorkUnitLowWatermarks = new ArrayList<Long>();
		
		if(previousWorkUnitStates.size() == 0) {
			LOG.info("Previous states are not found; Latest watermark - Default watermark:"+latestWaterMark);
			return latestWaterMark;
		}
		
		for(WorkUnitState workUnitState : previousWorkUnitStates) {
			if(workUnitState.getWorkingState() == WorkingState.COMMITTED) {
				if(workUnitState.getHighWaterMark() != ConfigurationKeys.DEFAULT_WATERMARK_VALUE) {
					hasDataInPreviousRun = true;
				}
				previousWorkUnitStateHighWatermarks.add(workUnitState.getHighWaterMark());
				previousWorkUnitLowWatermarks.add(this.getLowWatermarkFromWorkUnit(workUnitState.getWorkunit(), workUnitState.getProp(ConfigurationKeys.SOURCE_WATERMARK_TYPE)));
			}
		}
		
		if(previousWorkUnitStateHighWatermarks.isEmpty()) {
			latestWaterMark = ConfigurationKeys.DEFAULT_WATERMARK_VALUE;
			LOG.info("Previous Committed states are not found; Latest watermark - Default watermark:"+latestWaterMark);
		} else if(hasDataInPreviousRun) {
			latestWaterMark = Collections.max(previousWorkUnitStateHighWatermarks);
			LOG.info("Previous run has data; Latest watermark - Max watermark from WorkUnitStates:"+latestWaterMark);
		} else {
			latestWaterMark = Collections.min(previousWorkUnitLowWatermarks);
			LOG.info("Previous run has no data; Latest watermark - Min watermark from WorkUnits:"+latestWaterMark);
		}
		
		return latestWaterMark;
	}

    /**
     * get latest water mark from previous work units
     *
     * @param work unit
     * @param watermark type
     * @return latest water mark(low water mark of WorkUnit)
     */
	private long getLowWatermarkFromWorkUnit(WorkUnit workunit, String watermarkType) {
		if(workunit.getLowWaterMark() == ConfigurationKeys.DEFAULT_WATERMARK_VALUE) {
			return ConfigurationKeys.DEFAULT_WATERMARK_VALUE;
		}
		
		WatermarkType wmType = WatermarkType.valueOf(watermarkType.toUpperCase());
		int deltaNum = new WatermarkPredicate(null, wmType).getDeltaNumForNextWatermark();
		
		if(wmType == WatermarkType.SIMPLE) {
			return workunit.getLowWaterMark() - deltaNum;
		} else {
			return Long.parseLong(Utils.dateToString(Utils.addSecondsToDate(Utils.toDate(workunit.getLowWaterMark(), "yyyyMMddHHmmss"),deltaNum*-1), "yyyyMMddHHmmss"));
		}
	}

	@Override
	public void shutdown(SourceState state) {
	}
}

{'fraction_non_alphanumeric': 0.09892086330935251, 'fraction_numerical': 0.0005995203836930455, 'mean_word_length': 5.306710775047259, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3526112', 'n_tokens_mistral': 4211, 'n_tokens_neox': 3947, 'n_words': 783}