<commit_before>package simpledb;

import java.util.*;

import static simpledb.ImputationType.DROP;
import static simpledb.ImputationType.MAXIMAL;
import static simpledb.ImputationType.MINIMAL;

/**
 * Adding imputation on top of a plan that already has imputed values along its
 * tree. We need this to add drop, impute min, impute max to plans that already
 * include joins.
 */
public class LogicalComposeImputation extends ImputedPlan {
	private static final double LOSS_FACTOR = 1.01;

	private final DbIterator physicalPlan;
	private final Set<QuantifiedName> dirtySet;
	private final double loss;
	private final double time;
	private TableStats tableStats;

	/**
	 * Constructor intended only for internal construction of updated plan
	 * @param tableStats
	 * @param physicalPlan
	 * @param dirtySet
	 * @param loss
	 * @param time
	 */
	private LogicalComposeImputation(TableStats tableStats, DbIterator physicalPlan, Set<QuantifiedName> dirtySet,
			double loss, double time) {
		super();
		this.physicalPlan = physicalPlan;
		this.dirtySet = dirtySet;
		this.loss = loss;
		this.time = time;
		this.tableStats = tableStats;
	}

	public TableStats getTableStats() {
		return this.tableStats;
	}

	public static ImputedPlan create(ImputedPlan subplan, ImputationType imp, Set<QuantifiedName> required,
			Map<String, Integer> tableMap) throws BadImputation {
		// don't change anything if there are no dirty columns
		if (subplan.getDirtySet().isEmpty()) {
			return subplan;
		}

		if (Collections.disjoint(subplan.getDirtySet(), required) && (imp == MINIMAL || imp == DROP)) {
			return subplan;
		}

		// which ones do we actually need to impute
		Set<QuantifiedName> impute = new HashSet<>();
		impute.addAll(subplan.getDirtySet());
		impute.retainAll(required);

		final Set<QuantifiedName> dirtySet = new HashSet<>();
		dirtySet.addAll(subplan.getDirtySet());

		TupleDesc schema = subplan.getPlan().getTupleDesc();
		double totalData = subplan.cardinality() * schema.numFields();

		DbIterator physicalPlan;
		double loss, time;
		// table stats for subplan
		TableStats subplanTableStats = subplan.getTableStats();
		// holds tablestats adjusted for imputation
		TableStats adjustedTableStats;

		switch (imp) {
		case DROP: {
			Collection<Integer> imputeIndices = schema.fieldNamesToIndices(impute);
			physicalPlan = new Drop(toNames(impute), subplan.getPlan());
			dirtySet.removeAll(impute);
			loss = estimateNumNulls(subplan, imputeIndices);
			time = subplan.cardinality();
			adjustedTableStats = subplanTableStats.adjustForImpute(DROP, imputeIndices);
			return new LogicalComposeImputation(adjustedTableStats, physicalPlan, dirtySet, loss, time);
		}
		case MINIMAL: {
			Collection<Integer> imputeIndices = schema.fieldNamesToIndices(impute);
			physicalPlan = new ImputeRegressionTree(toNames(impute), subplan.getPlan());
			dirtySet.removeAll(impute);
			loss = estimateNumNulls(subplan, imputeIndices) * Math.pow(LOSS_FACTOR, -totalData);
			time = subplan.cardinality() * imputeIndices.size();
			adjustedTableStats = subplanTableStats.adjustForImpute(MINIMAL, imputeIndices);
			return new LogicalComposeImputation(adjustedTableStats, physicalPlan, dirtySet, loss, time);
		}
		case MAXIMAL: {
			Set<QuantifiedName> i = subplan.getDirtySet();
			Collection<Integer> imputeIndices = schema.fieldNamesToIndices(i);
			physicalPlan = new ImputeRegressionTree(toNames(i), subplan.getPlan());
			dirtySet.clear();
			loss = estimateNumNulls(subplan, imputeIndices) * Math.pow(LOSS_FACTOR, -totalData);
			time = subplan.cardinality() * imputeIndices.size();
			adjustedTableStats = subplanTableStats.adjustForImpute(MAXIMAL, imputeIndices);
			return new LogicalComposeImputation(adjustedTableStats, physicalPlan, dirtySet, loss, time);
		}
		case NONE:
			if (impute.isEmpty()) {
				return subplan;
			}
			throw new BadImputation();
		default:
			throw new RuntimeException("Unexpected ImputationType.");
		}
	}

	public static double estimateNumNulls(ImputedPlan subplan, Collection<Integer> indices) {
		return subplan.getTableStats().estimateTotalNull(indices);
	}

	public DbIterator getPlan() {
		return physicalPlan;
	}

	public Set<QuantifiedName> getDirtySet() {
		return dirtySet;
	}

	public double cost(double lossWeight) {
		return lossWeight * loss + (1 - lossWeight) * time;
	}

	public double cardinality() {
		return tableStats.totalTuples();
	}

	private static Set<String> toNames(Set<QuantifiedName> attrs) {
		Set<String> names = new HashSet<>();
		for (QuantifiedName attr : attrs) {
			names.add(attr.toString());
		}
		return names;
	}

}
<commit_msg>fix: fixed compose imputation cost<commit_after>package simpledb;

import java.util.*;

import static simpledb.ImputationType.DROP;
import static simpledb.ImputationType.MAXIMAL;
import static simpledb.ImputationType.MINIMAL;

/**
 * Adding imputation on top of a plan that already has imputed values along its
 * tree. We need this to add drop, impute min, impute max to plans that already
 * include joins.
 */
public class LogicalComposeImputation extends ImputedPlan {
	private static final double LOSS_FACTOR = 1.01;

	private final DbIterator physicalPlan;
	private final ImputedPlan subplan;
	private final Set<QuantifiedName> dirtySet;
	private final double loss;
	private final double time;
	private TableStats tableStats;

	/**
	 * Constructor intended only for internal construction of updated plan
	 * @param tableStats
	 * @param physicalPlan
	 * @param dirtySet
	 * @param loss
	 * @param time
	 */
	private LogicalComposeImputation(TableStats tableStats, DbIterator physicalPlan, ImputedPlan subplan, Set<QuantifiedName> dirtySet,
			double loss, double time) {
		super();
		this.physicalPlan = physicalPlan;
		this.dirtySet = dirtySet;
		this.loss = loss;
		this.time = time;
		this.tableStats = tableStats;
		this.subplan = subplan;
	}

	public TableStats getTableStats() {
		return this.tableStats;
	}

	public static ImputedPlan create(ImputedPlan subplan, ImputationType imp, Set<QuantifiedName> required,
			Map<String, Integer> tableMap) throws BadImputation {
		// don't change anything if there are no dirty columns
		if (subplan.getDirtySet().isEmpty()) {
			return subplan;
		}

		if (Collections.disjoint(subplan.getDirtySet(), required) && (imp == MINIMAL || imp == DROP)) {
			return subplan;
		}

		// which ones do we actually need to impute
		Set<QuantifiedName> impute = new HashSet<>();
		impute.addAll(subplan.getDirtySet());
		impute.retainAll(required);

		final Set<QuantifiedName> dirtySet = new HashSet<>();
		dirtySet.addAll(subplan.getDirtySet());

		TupleDesc schema = subplan.getPlan().getTupleDesc();
		double totalData = subplan.cardinality() * schema.numFields();

		DbIterator physicalPlan;
		double loss, time;
		// table stats for subplan
		TableStats subplanTableStats = subplan.getTableStats();
		// holds tablestats adjusted for imputation
		TableStats adjustedTableStats;

		switch (imp) {
		case DROP: {
			Collection<Integer> imputeIndices = schema.fieldNamesToIndices(impute);
			physicalPlan = new Drop(toNames(impute), subplan.getPlan());
			dirtySet.removeAll(impute);
			loss = estimateNumNulls(subplan, imputeIndices);
			time = subplan.cardinality();
			adjustedTableStats = subplanTableStats.adjustForImpute(DROP, imputeIndices);
			return new LogicalComposeImputation(adjustedTableStats, physicalPlan, subplan, dirtySet, loss, time);
		}
		case MINIMAL: {
			Collection<Integer> imputeIndices = schema.fieldNamesToIndices(impute);
			Impute imputeOp = new ImputeRegressionTree(toNames(impute), subplan.getPlan());
			physicalPlan = imputeOp;
			dirtySet.removeAll(impute);
			loss = estimateNumNulls(subplan, imputeIndices) * Math.pow(LOSS_FACTOR, -totalData);
			int numComplete = schema.numFields() - dirtySet.size();
			time = imputeOp.getEstimatedCost(imputeIndices.size(), numComplete, (int) subplan.cardinality());
			adjustedTableStats = subplanTableStats.adjustForImpute(MINIMAL, imputeIndices);
			return new LogicalComposeImputation(adjustedTableStats, physicalPlan, subplan, dirtySet, loss, time);
		}
		case MAXIMAL: {
			Set<QuantifiedName> i = subplan.getDirtySet();
			Collection<Integer> imputeIndices = schema.fieldNamesToIndices(i);
			Impute imputeOp = new ImputeRegressionTree(toNames(i), subplan.getPlan());
			physicalPlan = imputeOp;
			dirtySet.clear();
			loss = estimateNumNulls(subplan, imputeIndices) * Math.pow(LOSS_FACTOR, -totalData);
			int numComplete = schema.numFields() - dirtySet.size();
			time = imputeOp.getEstimatedCost(imputeIndices.size(), numComplete, (int) subplan.cardinality());
			adjustedTableStats = subplanTableStats.adjustForImpute(MAXIMAL, imputeIndices);
			return new LogicalComposeImputation(adjustedTableStats, physicalPlan, subplan, dirtySet, loss, time);
		}
		case NONE:
			if (impute.isEmpty()) {
				return subplan;
			}
			throw new BadImputation();
		default:
			throw new RuntimeException("Unexpected ImputationType.");
		}
	}

	public static double estimateNumNulls(ImputedPlan subplan, Collection<Integer> indices) {
		return subplan.getTableStats().estimateTotalNull(indices);
	}

	public DbIterator getPlan() {
		return physicalPlan;
	}

	public Set<QuantifiedName> getDirtySet() {
		return dirtySet;
	}

	public double cost(double lossWeight) {
		return lossWeight * loss + (1 - lossWeight) * time + subplan.cost(lossWeight);
	}

	public double cardinality() {
		return tableStats.totalTuples();
	}

	private static Set<String> toNames(Set<QuantifiedName> attrs) {
		Set<String> names = new HashSet<>();
		for (QuantifiedName attr : attrs) {
			names.add(attr.toString());
		}
		return names;
	}

}

{'fraction_non_alphanumeric': 0.10238589211618257, 'fraction_numerical': 0.0008298755186721991, 'mean_word_length': 5.0031133250311335, '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': '4623969', 'n_tokens_mistral': 3277, 'n_tokens_neox': 2947, 'n_words': 738}