<commit_before>package edu.umich.verdict.dbms;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Set;

import com.google.common.base.Joiner;
import com.google.common.base.Optional;

import edu.umich.verdict.VerdictContext;
import edu.umich.verdict.datatypes.SampleParam;
import edu.umich.verdict.datatypes.SampleSizeInfo;
import edu.umich.verdict.datatypes.TableUniqueName;
import edu.umich.verdict.exceptions.VerdictException;
import edu.umich.verdict.relation.ExactRelation;
import edu.umich.verdict.relation.Relation;
import edu.umich.verdict.relation.SingleRelation;
import edu.umich.verdict.relation.condition.CompCond;
import edu.umich.verdict.relation.expr.BinaryOpExpr;
import edu.umich.verdict.relation.expr.ConstantExpr;
import edu.umich.verdict.relation.expr.Expr;
import edu.umich.verdict.relation.expr.SubqueryExpr;
import edu.umich.verdict.util.VerdictLogger;

public class DbmsImpala extends DbmsJDBC {

	public DbmsImpala(VerdictContext vc, String dbName, String host, String port, String schema, String user,
			String password, String jdbcClassName) throws VerdictException {
		super(vc, dbName, host, port, schema, user, password, jdbcClassName);
		// TODO Auto-generated constructor stub
	}

	@Override
	public ResultSet getDatabaseNamesInResultSet() throws VerdictException {
		return executeJdbcQuery("show databases");
	}
	
	@Override
	public ResultSet getTablesInResultSet(String schema) throws VerdictException {
		return executeJdbcQuery("show tables in " + schema);
	}

	@Override
	public List<String> getTables(String schema) throws VerdictException {
		List<String> tables = new ArrayList<String>();
		try {
			ResultSet rs = executeJdbcQuery("show tables in " + schema);
			while (rs.next()) {
				String table = rs.getString(1);
				tables.add(table);
			}
		} catch (SQLException e) {
			VerdictLogger.error(this, "Failed to access the database: " + schema);
			throw new VerdictException(e);
		}
		return tables;
	}
	
	public List<String> getColumns(TableUniqueName table) throws VerdictException {
		List<String> columns = new ArrayList<String>();
		try {
			ResultSet rs = executeJdbcQuery("describe " + table);
			while (rs.next()) {
				String column = rs.getString(1);
				columns.add(column);
			}
		} catch (SQLException e) {
			throw new VerdictException(e);
		}
		return columns;
	}

	@Override
	public ResultSet describeTableInResultSet(TableUniqueName tableUniqueName)  throws VerdictException {
		return executeJdbcQuery(String.format("describe %s", tableUniqueName));
	}

	/**
	 * Impala does not support the standard JDBC protocol {@link java.sql.Connection#setCatalog(String) setCatalog}
	 * function for changing the current database. This is a workaround.
	 */
	@Override
	public void changeDatabase(String schemaName) throws VerdictException {
		execute(String.format("use %s", schemaName));
		currentSchema = Optional.fromNullable(schemaName);
		VerdictLogger.info("Database changed to: " + schemaName);
	}

	@Override
	public void createMetaTablesInDMBS(
			TableUniqueName originalTableName,
			TableUniqueName sizeTableName,
			TableUniqueName nameTableName) throws VerdictException {
		VerdictLogger.debug(this, "Creates meta tables if not exist.");
		String sql = String.format("CREATE TABLE IF NOT EXISTS %s", sizeTableName)
				+ " (schemaname STRING, "
				+ " tablename STRING, "
				+ " samplesize BIGINT, "
				+ " originaltablesize BIGINT)";
		executeUpdate(sql);
	
		sql = String.format("CREATE TABLE IF NOT EXISTS %s", nameTableName)
				+ " (originalschemaname STRING, "
				+ " originaltablename STRING, "
				+ " sampleschemaaname STRING, "
				+ " sampletablename STRING, "
				+ " sampletype STRING, "
				+ " samplingratio DOUBLE, "
				+ " columnnames STRING)";
		executeUpdate(sql);
		
		VerdictLogger.debug(this, "Meta tables created.");
	}

	@Override
	protected void justCreateUniformRandomSampleTableOf(SampleParam param) throws VerdictException {
		TableUniqueName temp = createTableWithRandNumbers(param);
		createUniformRandomSampleFromRandTable(param, temp);
		dropTable(temp);
//		String samplingProbCol = vc.getDbms().samplingProbabilityColumnName();
//		Set<String> colNames = vc.getMeta().getColumns(param.originalTable);
		
//		// This query exploits the fact that if subquery is combined with a regular column in a comparison condition,
//		// Impala properly generates a random number for every row and makes a comparison.
//		Expr threshold = ConstantExpr.from(String.format("one * (select %f from %s limit 1)", param.samplingRatio, param.originalTable));
//		ExactRelation sampled = SingleRelation.from(vc, param.originalTable)
//				                .select("*, 1 AS one, count(*) OVER () AS " + quote("__total_size"))
//						        .where(CompCond.from(Expr.from("rand(unix_timestamp())"), "<", threshold));
//		sampled = sampled.select(
//					Joiner.on(", ").join(colNames) +
//					", count(*) over () / " + quote("__total_size") + " AS " + samplingProbCol + ", " +  // attach sampling prob
//					randomPartitionColumn());										 // attach partition number
//		String sql = String.format("create table %s AS ", param.sampleTableName()) + sampled.toSql();
//		VerdictLogger.debug(this, "The query used for creating a uniform random sample:");
//		VerdictLogger.debugPretty(this, Relation.prettyfySql(sql), "  ");
//		
//		executeUpdate(sql);
	}
	
	private TableUniqueName createTableWithRandNumbers(SampleParam param) throws VerdictException {
		ExactRelation withRand = SingleRelation.from(vc, param.originalTable)
								 .select("*, rand(unix_timestamp()) AS __rand");
		TableUniqueName temp = Relation.getTempTableName(vc, param.sampleTableName().getSchemaName());
		String sql = String.format("create table %s AS %s", temp, withRand.toSql());
		executeUpdate(sql);
		return temp;
	}
	
	private void createUniformRandomSampleFromRandTable(SampleParam param, TableUniqueName temp) throws VerdictException {
		String samplingProbCol = vc.getDbms().samplingProbabilityColumnName();
		Set<String> colNames = vc.getMeta().getColumns(param.originalTable);
		
		ExactRelation withRand = SingleRelation.from(vc, temp);
		long total_size = withRand.countValue();
		
		ExactRelation sampled = withRand.select(
							      Joiner.on(", ").join(colNames) +
							      ", count(*) over () / " + total_size + " AS " + samplingProbCol + ", " + // attach sampling prob
							      randomPartitionColumn());
		String sql = String.format("create table %s AS %s", param.sampleTableName() + sampled.toSql());
		VerdictLogger.debug(this, "The query used for creating a uniform random sample:");
		VerdictLogger.debugPretty(this, Relation.prettyfySql(sql), "  ");
		executeUpdate(sql);
	}
	
	/**
	 * 
	 */
	@Override
	protected void justCreateStratifiedSampleTableof(SampleParam param) throws VerdictException {
		SampleSizeInfo info = vc.getMeta().getSampleSizeOf(new SampleParam(vc, param.originalTable, "uniform", null, new ArrayList<String>()));
		if (info == null) {
			String msg = "A uniform random must first be created before creating a stratified sample.";
			VerdictLogger.error(this, msg);
			throw new VerdictException(msg);
		}
		
		long originalTableSize = info.originalTableSize;
		double samplingProbability = param.samplingRatio;
		String groupName = Joiner.on(", ").join(param.columnNames);
		String samplingProbColName = vc.getDbms().samplingProbabilityColumnName();
		TableUniqueName sampleTable = param.sampleTableName();<commit_msg>minor bug fix<commit_after>package edu.umich.verdict.dbms;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Set;

import com.google.common.base.Joiner;
import com.google.common.base.Optional;

import edu.umich.verdict.VerdictContext;
import edu.umich.verdict.datatypes.SampleParam;
import edu.umich.verdict.datatypes.SampleSizeInfo;
import edu.umich.verdict.datatypes.TableUniqueName;
import edu.umich.verdict.exceptions.VerdictException;
import edu.umich.verdict.relation.ExactRelation;
import edu.umich.verdict.relation.Relation;
import edu.umich.verdict.relation.SingleRelation;
import edu.umich.verdict.relation.condition.CompCond;
import edu.umich.verdict.relation.expr.BinaryOpExpr;
import edu.umich.verdict.relation.expr.ConstantExpr;
import edu.umich.verdict.relation.expr.Expr;
import edu.umich.verdict.relation.expr.SubqueryExpr;
import edu.umich.verdict.util.VerdictLogger;

public class DbmsImpala extends DbmsJDBC {

	public DbmsImpala(VerdictContext vc, String dbName, String host, String port, String schema, String user,
			String password, String jdbcClassName) throws VerdictException {
		super(vc, dbName, host, port, schema, user, password, jdbcClassName);
		// TODO Auto-generated constructor stub
	}

	@Override
	public ResultSet getDatabaseNamesInResultSet() throws VerdictException {
		return executeJdbcQuery("show databases");
	}
	
	@Override
	public ResultSet getTablesInResultSet(String schema) throws VerdictException {
		return executeJdbcQuery("show tables in " + schema);
	}

	@Override
	public List<String> getTables(String schema) throws VerdictException {
		List<String> tables = new ArrayList<String>();
		try {
			ResultSet rs = executeJdbcQuery("show tables in " + schema);
			while (rs.next()) {
				String table = rs.getString(1);
				tables.add(table);
			}
		} catch (SQLException e) {
			VerdictLogger.error(this, "Failed to access the database: " + schema);
			throw new VerdictException(e);
		}
		return tables;
	}
	
	public List<String> getColumns(TableUniqueName table) throws VerdictException {
		List<String> columns = new ArrayList<String>();
		try {
			ResultSet rs = executeJdbcQuery("describe " + table);
			while (rs.next()) {
				String column = rs.getString(1);
				columns.add(column);
			}
		} catch (SQLException e) {
			throw new VerdictException(e);
		}
		return columns;
	}

	@Override
	public ResultSet describeTableInResultSet(TableUniqueName tableUniqueName)  throws VerdictException {
		return executeJdbcQuery(String.format("describe %s", tableUniqueName));
	}

	/**
	 * Impala does not support the standard JDBC protocol {@link java.sql.Connection#setCatalog(String) setCatalog}
	 * function for changing the current database. This is a workaround.
	 */
	@Override
	public void changeDatabase(String schemaName) throws VerdictException {
		execute(String.format("use %s", schemaName));
		currentSchema = Optional.fromNullable(schemaName);
		VerdictLogger.info("Database changed to: " + schemaName);
	}

	@Override
	public void createMetaTablesInDMBS(
			TableUniqueName originalTableName,
			TableUniqueName sizeTableName,
			TableUniqueName nameTableName) throws VerdictException {
		VerdictLogger.debug(this, "Creates meta tables if not exist.");
		String sql = String.format("CREATE TABLE IF NOT EXISTS %s", sizeTableName)
				+ " (schemaname STRING, "
				+ " tablename STRING, "
				+ " samplesize BIGINT, "
				+ " originaltablesize BIGINT)";
		executeUpdate(sql);
	
		sql = String.format("CREATE TABLE IF NOT EXISTS %s", nameTableName)
				+ " (originalschemaname STRING, "
				+ " originaltablename STRING, "
				+ " sampleschemaaname STRING, "
				+ " sampletablename STRING, "
				+ " sampletype STRING, "
				+ " samplingratio DOUBLE, "
				+ " columnnames STRING)";
		executeUpdate(sql);
		
		VerdictLogger.debug(this, "Meta tables created.");
	}

	@Override
	protected void justCreateUniformRandomSampleTableOf(SampleParam param) throws VerdictException {
		TableUniqueName temp = createTableWithRandNumbers(param);
		createUniformRandomSampleFromRandTable(param, temp);
		dropTable(temp);
//		String samplingProbCol = vc.getDbms().samplingProbabilityColumnName();
//		Set<String> colNames = vc.getMeta().getColumns(param.originalTable);
		
//		// This query exploits the fact that if subquery is combined with a regular column in a comparison condition,
//		// Impala properly generates a random number for every row and makes a comparison.
//		Expr threshold = ConstantExpr.from(String.format("one * (select %f from %s limit 1)", param.samplingRatio, param.originalTable));
//		ExactRelation sampled = SingleRelation.from(vc, param.originalTable)
//				                .select("*, 1 AS one, count(*) OVER () AS " + quote("__total_size"))
//						        .where(CompCond.from(Expr.from("rand(unix_timestamp())"), "<", threshold));
//		sampled = sampled.select(
//					Joiner.on(", ").join(colNames) +
//					", count(*) over () / " + quote("__total_size") + " AS " + samplingProbCol + ", " +  // attach sampling prob
//					randomPartitionColumn());										 // attach partition number
//		String sql = String.format("create table %s AS ", param.sampleTableName()) + sampled.toSql();
//		VerdictLogger.debug(this, "The query used for creating a uniform random sample:");
//		VerdictLogger.debugPretty(this, Relation.prettyfySql(sql), "  ");
//		
//		executeUpdate(sql);
	}
	
	private TableUniqueName createTableWithRandNumbers(SampleParam param) throws VerdictException {
		ExactRelation withRand = SingleRelation.from(vc, param.originalTable)
								 .select("*, rand(unix_timestamp()) AS __rand");
		TableUniqueName temp = Relation.getTempTableName(vc, param.sampleTableName().getSchemaName());
		String sql = String.format("create table %s AS %s", temp, withRand.toSql());
		VerdictLogger.debug(this, "The query used for creating a temporary table with a column containing random numbers:");
		VerdictLogger.debugPretty(this, Relation.prettyfySql(sql), "  ");
		executeUpdate(sql);
		return temp;
	}
	
	private void createUniformRandomSampleFromRandTable(SampleParam param, TableUniqueName temp) throws VerdictException {
		String samplingProbCol = vc.getDbms().samplingProbabilityColumnName();
		Set<String> colNames = vc.getMeta().getColumns(param.originalTable);
		
		ExactRelation withRand = SingleRelation.from(vc, temp);
		long total_size = withRand.countValue();
		
		ExactRelation sampled = withRand
							   .where("__rand < " + param.samplingRatio)
							   .select(
							      Joiner.on(", ").join(colNames) +
							      ", count(*) over () / " + total_size + " AS " + samplingProbCol + ", " + // attach sampling prob
							      randomPartitionColumn());
		String sql = String.format("create table %s AS %s", param.sampleTableName(), sampled.toSql());
		VerdictLogger.debug(this, "The query used for creating a uniform random sample from a temporary table:");
		VerdictLogger.debugPretty(this, Relation.prettyfySql(sql), "  ");
		executeUpdate(sql);
	}
	
	/**
	 * 
	 */
	@Override
	protected void justCreateStratifiedSampleTableof(SampleParam param) throws VerdictException {
		SampleSizeInfo info = vc.getMeta().getSampleSizeOf(new SampleParam(vc, param.originalTable, "uniform", null, new ArrayList<String>()));
		if (info == null) {
			String msg = "A uniform random must first be created before creating a stratified sample.";
			VerdictLogger.error(this, msg);
			throw new VerdictException(msg);
		}
		
		long originalTableSize = info.originalTableSize;
		double samplingProbability = param.samplingRatio;
		String groupName = Joiner.on(", ").join(param.columnNames);
		String samplingProbColName = vc.getDbms().samplingProbabilityColumnName();
		TableUniqueName sampleTable = param.sampleTableName();
{'fraction_non_alphanumeric': 0.11609584581325694, 'fraction_numerical': 0.0005209011590050788, 'mean_word_length': 4.591190389515836, 'pattern_counts': {'":': 0, '<': 24, '<?xml version=': 0, '>': 21, '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': '6644245', 'n_tokens_mistral': 5093, 'n_tokens_neox': 4546, 'n_words': 1097}