<commit_before>package edu.washington.escience.myria;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.Reader;
import java.io.Serializable;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.Arrays;
import java.util.BitSet;
import java.util.List;
import java.util.Objects;

import net.jcip.annotations.ThreadSafe;

import org.joda.time.DateTime;
import org.postgresql.PGConnection;
import org.postgresql.copy.CopyManager;

import com.almworks.sqlite4java.SQLiteException;
import com.almworks.sqlite4java.SQLiteStatement;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.hash.HashFunction;
import com.google.common.hash.Hasher;
import com.google.common.hash.Hashing;

import edu.washington.escience.myria.column.Column;
import edu.washington.escience.myria.parallel.PartitionFunction;
import edu.washington.escience.myria.proto.TransportProto.TransportMessage;
import edu.washington.escience.myria.util.IPCUtils;

/**
 * Container class for a batch of tuples. The goal is to amortize memory management overhead.
 * 
 */
@ThreadSafe
public class TupleBatch implements ReadableTable, Serializable {
  /** Required for Java serialization. */
  private static final long serialVersionUID = 1L;
  /** The hard-coded number of tuples in a batch. */
  public static final int BATCH_SIZE = 10 * 1000;
  /** Class-specific magic number used to generate the hash code. */
  private static final int MAGIC_HASHCODE = 243;
  /** The hash function for this class. */
  private static final HashFunction HASH_FUNCTION = Hashing.murmur3_32(MAGIC_HASHCODE);
  /** Schema of tuples in this batch. */
  private final Schema schema;
  /** Tuple data stored as columns in this batch. */
  private final ImmutableList<Column<?>> columns;
  /** Number of tuples in this TB. */
  private final int numTuples;
  /** Whether this TB is an EOI TB. */
  private final boolean isEOI;

  /**
   * EOI TB constructor.
   * 
   * @param schema schema of the tuples in this batch.
   * @param isEoi whether this TupleBatch is an EOI TupleBatch.
   * */
  private TupleBatch(final Schema schema, final boolean isEoi) {
    this.schema = schema;
    numTuples = 0;
    ImmutableList.Builder<Column<?>> b = ImmutableList.builder();
    for (Type type : schema.getColumnTypes()) {
      b.add(Column.emptyColumn(type));
    }
    columns = b.build();
    isEOI = isEoi;
  }

  /**
   * @param columnNames the new column names.
   * @return a shallow copy of the specified TupleBatch with the new column names.
   */
  public TupleBatch rename(final List<String> columnNames) {
    Schema newSchema = Schema.of(schema.getColumnTypes(), columnNames);
    return new TupleBatch(newSchema, columns, numTuples, isEOI);
  }

  /**
   * Standard immutable TupleBatch constructor. All fields must be populated before creation and cannot be changed.
   * 
   * @param schema schema of the tuples in this batch. Must match columns.
   * @param columns contains the column-stored data. Must match schema.
   * @param numTuples the number of tuples in this TupleBatch.
   */
  public TupleBatch(final Schema schema, final List<Column<?>> columns, final int numTuples) {
    this(schema, columns, numTuples, false);
  }

  /**
   * Constructor that gets the number of tuples from the columns.
   * 
   * @param schema schema of the tuples in this batch. Must match columns.
   * @param columns contains the column-stored data. Must match schema.
   */
  public TupleBatch(final Schema schema, final List<Column<?>> columns) {
    this(schema, columns, columns.get(0).size());
  }

  /**
   * Construct a TupleBatch from the specified components.
   * 
   * @param schema schema of the tuples in this batch. Must match columns.
   * @param columns schema of the tuples in this batch. Must match columns.
   * @param numTuples the number of tuples in this batch. Must match columns.
   * @param isEOI whether this is an EOI TupleBatch.
   */
  public TupleBatch(final Schema schema, final List<Column<?>> columns, final int numTuples, final boolean isEOI) {
    this.schema = Objects.requireNonNull(schema);
    Objects.requireNonNull(columns);
    Preconditions.checkArgument(columns.size() == schema.numColumns(),
        "Number of columns in data must equal the number of fields in schema");
    this.columns = ImmutableList.copyOf(columns);
    this.numTuples = numTuples;
    this.isEOI = isEOI;
  }

  /**
   * put the tuple batch into TBB by smashing it into cells and putting them one by one.
   * 
   * @param tbb the TBB buffer.
   * */
  public final void compactInto(final TupleBatchBuffer tbb) {
    if (isEOI()) {
      /* an EOI TB has no data */
      tbb.appendTB(this);
      return;
    }
    final int numColumns = columns.size();
    for (int i = 0; i < numTuples; i++) {
      for (int column = 0; column < numColumns; column++) {
        tbb.put(column, columns.get(column), i);
      }
    }
  }

  /**
   * Return a new TupleBatch that contains only the filtered rows of the current dataset. Note that if some of the
   * tuples in this batch are invalid, we will have to map the indices in the specified filter to the "real" indices in
   * the tuple.
   * 
   * @param filter the rows to be retained.
   * @return a TupleBatch that contains only the filtered rows of the current dataset.
   */
  public final TupleBatch filter(final BitSet filter) {
    Preconditions.checkArgument(filter.length() <= numTuples(),
        "Error: trying to filter a TupleBatch of length %s with a filter of length %s", numTuples(), filter.length());
    int newNumTuples = filter.cardinality();

    /* Shortcut: the filter is full, so all current tuples are retained. Just return this. */
    if (newNumTuples == numTuples) {
      return this;
    }

    ImmutableList.Builder<Column<?>> newColumns = ImmutableList.builder();
    for (Column<?> column : columns) {
      newColumns.add(column.filter(filter));
    }
    return new TupleBatch(schema, newColumns.build(), newNumTuples, isEOI);
  }

  @Override
  public final boolean getBoolean(final int column, final int row) {
    return columns.get(column).getBoolean(row);
  }

  @Override
  public final double getDouble(final int column, final int row) {
    return columns.get(column).getDouble(row);
  }

  @Override
  public final float getFloat(final int column, final int row) {
    return columns.get(column).getFloat(row);
  }

  @Override
  public final int getInt(final int column, final int row) {
    return columns.get(column).getInt(row);
  }

  /**
   * store this TB into JDBC.
   * 
   * @param statement JDBC statement.
   * @throws SQLException any exception caused by JDBC.
   * */
  public final void getIntoJdbc(final PreparedStatement statement) throws SQLException {
    for (int i = 0; i < numTuples; i++) {
      int column = 0;
      for (final Column<?> c : columns) {
        c.getIntoJdbc(i, statement, ++column);
      }
      statement.addBatch();
    }
  }

  /**
   * store this TB into SQLite.
   * 
   * @param statement SQLite statement.
   * @throws SQLiteException any exception caused by SQLite.
   * */
  public final void getIntoSQLite(final SQLiteStatement statement) throws SQLiteException {
    for (int i = 0; i < numTuples; i++) {
      int column = 0;
      for (final Column<?> c : columns) {
        c.getIntoSQLite(i, statement, ++column);
      }
      statement.step();
      statement.reset();
    }
  }

  /**
   * Copy data directly into postgres by creating a CSV.
   * 
   * @param connection the pg connection
   * @param table the table to import into
   * @throws SQLException exception raised when inserting into database
   * @throws IOException exception raised from buffer
   */
  public final void getIntoPostgres(final PGConnection connection, final RelationKey table) throws SQLException,
      IOException {
    CopyManager cpManager = connection.getCopyAPI();

    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    TupleWriter tw = new CsvTupleWriter('t', baos);
    tw.writeTuples(this);
    tw.done();

    // TODO: this seems complicated
    Reader reader = new InputStreamReader(new ByteArrayInputStream(baos.toByteArray()));
    cpManager.copyIn("COPY " + table.toString(MyriaConstants.STORAGE_SYSTEM_POSTGRESQL) + " FROM STDIN WITH CSV",
        reader);
  }

  @Override
  public final long getLong(final int column, final int row) {
    return columns.get(column).getLong(row);
  }

  @Override
  public final Object getObject(final int column, final int row) {
    return columns.get(column).getObject(row);
  }

  @Override
  public final Schema getSchema() {
    return schema;
  }

  @Override<commit_msg>Fix separator in csv generation<commit_after>package edu.washington.escience.myria;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.Reader;
import java.io.Serializable;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.Arrays;
import java.util.BitSet;
import java.util.List;
import java.util.Objects;

import net.jcip.annotations.ThreadSafe;

import org.joda.time.DateTime;
import org.postgresql.PGConnection;
import org.postgresql.copy.CopyManager;

import com.almworks.sqlite4java.SQLiteException;
import com.almworks.sqlite4java.SQLiteStatement;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.hash.HashFunction;
import com.google.common.hash.Hasher;
import com.google.common.hash.Hashing;

import edu.washington.escience.myria.column.Column;
import edu.washington.escience.myria.parallel.PartitionFunction;
import edu.washington.escience.myria.proto.TransportProto.TransportMessage;
import edu.washington.escience.myria.util.IPCUtils;

/**
 * Container class for a batch of tuples. The goal is to amortize memory management overhead.
 * 
 */
@ThreadSafe
public class TupleBatch implements ReadableTable, Serializable {
  /** Required for Java serialization. */
  private static final long serialVersionUID = 1L;
  /** The hard-coded number of tuples in a batch. */
  public static final int BATCH_SIZE = 10 * 1000;
  /** Class-specific magic number used to generate the hash code. */
  private static final int MAGIC_HASHCODE = 243;
  /** The hash function for this class. */
  private static final HashFunction HASH_FUNCTION = Hashing.murmur3_32(MAGIC_HASHCODE);
  /** Schema of tuples in this batch. */
  private final Schema schema;
  /** Tuple data stored as columns in this batch. */
  private final ImmutableList<Column<?>> columns;
  /** Number of tuples in this TB. */
  private final int numTuples;
  /** Whether this TB is an EOI TB. */
  private final boolean isEOI;

  /**
   * EOI TB constructor.
   * 
   * @param schema schema of the tuples in this batch.
   * @param isEoi whether this TupleBatch is an EOI TupleBatch.
   * */
  private TupleBatch(final Schema schema, final boolean isEoi) {
    this.schema = schema;
    numTuples = 0;
    ImmutableList.Builder<Column<?>> b = ImmutableList.builder();
    for (Type type : schema.getColumnTypes()) {
      b.add(Column.emptyColumn(type));
    }
    columns = b.build();
    isEOI = isEoi;
  }

  /**
   * @param columnNames the new column names.
   * @return a shallow copy of the specified TupleBatch with the new column names.
   */
  public TupleBatch rename(final List<String> columnNames) {
    Schema newSchema = Schema.of(schema.getColumnTypes(), columnNames);
    return new TupleBatch(newSchema, columns, numTuples, isEOI);
  }

  /**
   * Standard immutable TupleBatch constructor. All fields must be populated before creation and cannot be changed.
   * 
   * @param schema schema of the tuples in this batch. Must match columns.
   * @param columns contains the column-stored data. Must match schema.
   * @param numTuples the number of tuples in this TupleBatch.
   */
  public TupleBatch(final Schema schema, final List<Column<?>> columns, final int numTuples) {
    this(schema, columns, numTuples, false);
  }

  /**
   * Constructor that gets the number of tuples from the columns.
   * 
   * @param schema schema of the tuples in this batch. Must match columns.
   * @param columns contains the column-stored data. Must match schema.
   */
  public TupleBatch(final Schema schema, final List<Column<?>> columns) {
    this(schema, columns, columns.get(0).size());
  }

  /**
   * Construct a TupleBatch from the specified components.
   * 
   * @param schema schema of the tuples in this batch. Must match columns.
   * @param columns schema of the tuples in this batch. Must match columns.
   * @param numTuples the number of tuples in this batch. Must match columns.
   * @param isEOI whether this is an EOI TupleBatch.
   */
  public TupleBatch(final Schema schema, final List<Column<?>> columns, final int numTuples, final boolean isEOI) {
    this.schema = Objects.requireNonNull(schema);
    Objects.requireNonNull(columns);
    Preconditions.checkArgument(columns.size() == schema.numColumns(),
        "Number of columns in data must equal the number of fields in schema");
    this.columns = ImmutableList.copyOf(columns);
    this.numTuples = numTuples;
    this.isEOI = isEOI;
  }

  /**
   * put the tuple batch into TBB by smashing it into cells and putting them one by one.
   * 
   * @param tbb the TBB buffer.
   * */
  public final void compactInto(final TupleBatchBuffer tbb) {
    if (isEOI()) {
      /* an EOI TB has no data */
      tbb.appendTB(this);
      return;
    }
    final int numColumns = columns.size();
    for (int i = 0; i < numTuples; i++) {
      for (int column = 0; column < numColumns; column++) {
        tbb.put(column, columns.get(column), i);
      }
    }
  }

  /**
   * Return a new TupleBatch that contains only the filtered rows of the current dataset. Note that if some of the
   * tuples in this batch are invalid, we will have to map the indices in the specified filter to the "real" indices in
   * the tuple.
   * 
   * @param filter the rows to be retained.
   * @return a TupleBatch that contains only the filtered rows of the current dataset.
   */
  public final TupleBatch filter(final BitSet filter) {
    Preconditions.checkArgument(filter.length() <= numTuples(),
        "Error: trying to filter a TupleBatch of length %s with a filter of length %s", numTuples(), filter.length());
    int newNumTuples = filter.cardinality();

    /* Shortcut: the filter is full, so all current tuples are retained. Just return this. */
    if (newNumTuples == numTuples) {
      return this;
    }

    ImmutableList.Builder<Column<?>> newColumns = ImmutableList.builder();
    for (Column<?> column : columns) {
      newColumns.add(column.filter(filter));
    }
    return new TupleBatch(schema, newColumns.build(), newNumTuples, isEOI);
  }

  @Override
  public final boolean getBoolean(final int column, final int row) {
    return columns.get(column).getBoolean(row);
  }

  @Override
  public final double getDouble(final int column, final int row) {
    return columns.get(column).getDouble(row);
  }

  @Override
  public final float getFloat(final int column, final int row) {
    return columns.get(column).getFloat(row);
  }

  @Override
  public final int getInt(final int column, final int row) {
    return columns.get(column).getInt(row);
  }

  /**
   * store this TB into JDBC.
   * 
   * @param statement JDBC statement.
   * @throws SQLException any exception caused by JDBC.
   * */
  public final void getIntoJdbc(final PreparedStatement statement) throws SQLException {
    for (int i = 0; i < numTuples; i++) {
      int column = 0;
      for (final Column<?> c : columns) {
        c.getIntoJdbc(i, statement, ++column);
      }
      statement.addBatch();
    }
  }

  /**
   * store this TB into SQLite.
   * 
   * @param statement SQLite statement.
   * @throws SQLiteException any exception caused by SQLite.
   * */
  public final void getIntoSQLite(final SQLiteStatement statement) throws SQLiteException {
    for (int i = 0; i < numTuples; i++) {
      int column = 0;
      for (final Column<?> c : columns) {
        c.getIntoSQLite(i, statement, ++column);
      }
      statement.step();
      statement.reset();
    }
  }

  /**
   * Copy data directly into postgres by creating a CSV.
   * 
   * @param connection the pg connection
   * @param table the table to import into
   * @throws SQLException exception raised when inserting into database
   * @throws IOException exception raised from buffer
   */
  public final void getIntoPostgres(final PGConnection connection, final RelationKey table) throws SQLException,
      IOException {
    CopyManager cpManager = connection.getCopyAPI();

    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    TupleWriter tw = new CsvTupleWriter(',', baos);
    tw.writeTuples(this);
    tw.done();

    // TODO: this seems complicated
    Reader reader = new InputStreamReader(new ByteArrayInputStream(baos.toByteArray()));
    cpManager.copyIn("COPY " + table.toString(MyriaConstants.STORAGE_SYSTEM_POSTGRESQL) + " FROM STDIN WITH CSV",
        reader);
  }

  @Override
  public final long getLong(final int column, final int row) {
    return columns.get(column).getLong(row);
  }

  @Override
  public final Object getObject(final int column, final int row) {
    return columns.get(column).getObject(row);
  }

  @Override
  public final Schema getSchema() {
    return schema;
  }

  @Override
{'fraction_non_alphanumeric': 0.09850458819530984, 'fraction_numerical': 0.0026056417808995127, 'mean_word_length': 3.4214876033057853, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 35, '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': '128868', 'n_tokens_mistral': 5339, 'n_tokens_neox': 5166, 'n_words': 1739}