<commit_before>package com.facebook.presto;

import com.facebook.presto.aggregation.CountAggregation;
import com.facebook.presto.aggregation.SumAggregation;
import com.facebook.presto.block.BlockStream;
import com.facebook.presto.block.Cursor;
import com.facebook.presto.ingest.RowSourceBuilder;
import com.facebook.presto.operator.AggregationOperator;
import com.facebook.presto.operator.GroupByBlockStream;
import com.facebook.presto.operator.HashAggregationBlockStream;
import com.facebook.presto.slice.Slices;
import com.google.common.base.Charsets;
import com.google.common.base.Function;
import com.google.common.base.Joiner;
import com.google.common.base.Splitter;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.io.CharStreams;
import com.google.common.io.InputSupplier;
import com.google.common.io.Resources;
import org.skife.jdbi.v2.DBI;
import org.skife.jdbi.v2.Handle;
import org.skife.jdbi.v2.PreparedBatch;
import org.skife.jdbi.v2.PreparedBatchPart;
import org.skife.jdbi.v2.StatementContext;
import org.skife.jdbi.v2.tweak.ResultSetMapper;
import org.testng.annotations.AfterSuite;
import org.testng.annotations.BeforeSuite;
import org.testng.annotations.Test;

import java.io.IOException;
import java.io.InputStreamReader;
import java.io.Reader;
import java.net.URL;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.zip.GZIPInputStream;

import static com.facebook.presto.TupleInfo.Type.FIXED_INT_64;
import static com.facebook.presto.TupleInfo.Type.VARIABLE_BINARY;
import static com.facebook.presto.ingest.RowSourceBuilder.RowGenerator;
import static com.google.common.base.Preconditions.checkArgument;
import static io.airlift.testing.Assertions.assertEqualsIgnoreOrder;
import static java.lang.String.format;
import static java.util.Collections.nCopies;

public class TestQueries
{
    private Handle handle;
    private List<List<String>> ordersData;
    private List<List<String>> lineitemData;

    @BeforeSuite
    public void setupDatabase()
            throws IOException
    {
        handle = DBI.open("jdbc:h2:mem:test" + System.nanoTime());

        ordersData = readTestData("tpch/orders.dat.gz", 15000);
        handle.execute("CREATE TABLE orders (\n" +
                "  orderkey BIGINT NOT NULL,\n" +
                "  custkey BIGINT NOT NULL,\n" +
                "  orderstatus CHAR(1) NOT NULL,\n" +
                "  totalprice DOUBLE NOT NULL,\n" +
                "  orderdate CHAR(10) NOT NULL,\n" +
                "  orderpriority CHAR(15) NOT NULL,\n" +
                "  shippriority BIGINT NOT NULL\n" +
                ")");
        insertRows("orders", handle, ordersData);

        lineitemData = readTestData("tpch/lineitem.dat.gz", 60175);
        handle.execute("CREATE TABLE lineitem (\n" +
                "  orderkey BIGINT NOT NULL,\n" +
                "  partkey BIGINT NOT NULL,\n" +
                "  suppkey BIGINT NOT NULL,\n" +
                "  linenumber BIGINT NOT NULL,\n" +
                "  quantity BIGINT NOT NULL,\n" +
                "  extendedprice DOUBLE NOT NULL,\n" +
                "  discount DOUBLE NOT NULL,\n" +
                "  tax DOUBLE NOT NULL,\n" +
                "  returnflag CHAR(1) NOT NULL,\n" +
                "  linestatus CHAR(1) NOT NULL,\n" +
                "  shipdate CHAR(10) NOT NULL,\n" +
                "  commitdate CHAR(10) NOT NULL,\n" +
                "  receiptdate CHAR(10) NOT NULL,\n" +
                "  shipinstruct VARCHAR(25) NOT NULL,\n" +
                "  shipmode VARCHAR(10) NOT NULL\n" +
                ")");
        insertRows("lineitem", handle, lineitemData);
    }

    @AfterSuite
    public void cleanupDatabase()
    {
        handle.close();
    }

    @Test
    public void testCountAll()
    {
        List<Tuple> expected = computeExpected("SELECT COUNT(*) FROM orders", FIXED_INT_64);

        BlockStream orders = createBlockStream(ordersData, 1, FIXED_INT_64);
        AggregationOperator aggregation = new AggregationOperator(orders, CountAggregation.PROVIDER);

        assertEqualsIgnoreOrder(tuples(aggregation), expected);
    }

    @Test
    public void testGroupBySum()
    {
        List<Tuple> expected = computeExpected(
                "SELECT orderstatus, SUM(custkey) FROM orders GROUP BY orderstatus",
                VARIABLE_BINARY, FIXED_INT_64);

        BlockStream groupBySource = createBlockStream(ordersData, 2, VARIABLE_BINARY);
        BlockStream aggregateSource = createBlockStream(ordersData, 1, FIXED_INT_64);

        GroupByBlockStream groupBy = new GroupByBlockStream(groupBySource);
        HashAggregationBlockStream aggregation = new HashAggregationBlockStream(groupBy, aggregateSource, SumAggregation.PROVIDER);

        assertEqualsIgnoreOrder(tuples(aggregation), expected);
    }

    private List<Tuple> computeExpected(final String sql, TupleInfo.Type... types)
    {
        TupleInfo tupleInfo = new TupleInfo(types);
        return handle.createQuery(sql)
                .map(tupleMapper(tupleInfo))
                .list();
    }

    private static List<Tuple> tuples(BlockStream blockStream)
    {
        Cursor cursor = blockStream.cursor();
        List<Tuple> list = new ArrayList<>();
        while (cursor.advanceNextPosition()) {
            list.add(cursor.getTuple());
        }
        return list;
    }

    @SuppressWarnings("UnusedDeclaration")
    private static Iterable<List<Object>> tupleValues(Iterable<Tuple> tuples)
    {
        return Iterables.transform(tuples, new Function<Tuple, List<Object>>()
        {
            @Override
            public List<Object> apply(Tuple input)
            {
                return input.toValues();
            }
        });
    }

    private static BlockStream createBlockStream(List<List<String>> data, final int index, final TupleInfo.Type type)
    {
        final Iterator<List<String>> iterator = data.iterator();
        return new RowSourceBuilder(new TupleInfo(type), new RowGenerator()
        {
            @Override
            public boolean generate(RowSourceBuilder.RowBuilder rowBuilder)
            {
                if (!iterator.hasNext()) {
                    return false;
                }
                String value = iterator.next().get(index);
                switch (type) {
                    case FIXED_INT_64:
                        rowBuilder.append(Long.parseLong(value));
                        break;
                    case VARIABLE_BINARY:
                        rowBuilder.append(value.getBytes(Charsets.UTF_8));
                        break;
                    default:
                        throw new AssertionError("unhandled type: " + type);
                }
                return true;
            }

            @Override
            public void close() {}
        });
    }

    private static void insertRows(String table, Handle handle, List<List<String>> data)
    {
        checkArgument(!data.isEmpty(), "no data to insert");
        int columns = data.get(0).size();<commit_msg>Improve readability by using an enum for column names<commit_after>package com.facebook.presto;

import com.facebook.presto.aggregation.CountAggregation;
import com.facebook.presto.aggregation.SumAggregation;
import com.facebook.presto.block.BlockStream;
import com.facebook.presto.block.Cursor;
import com.facebook.presto.ingest.RowSourceBuilder;
import com.facebook.presto.operator.AggregationOperator;
import com.facebook.presto.operator.GroupByBlockStream;
import com.facebook.presto.operator.HashAggregationBlockStream;
import com.facebook.presto.slice.Slices;
import com.google.common.base.Charsets;
import com.google.common.base.Function;
import com.google.common.base.Joiner;
import com.google.common.base.Splitter;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.io.CharStreams;
import com.google.common.io.InputSupplier;
import com.google.common.io.Resources;
import org.skife.jdbi.v2.DBI;
import org.skife.jdbi.v2.Handle;
import org.skife.jdbi.v2.PreparedBatch;
import org.skife.jdbi.v2.PreparedBatchPart;
import org.skife.jdbi.v2.StatementContext;
import org.skife.jdbi.v2.tweak.ResultSetMapper;
import org.testng.annotations.AfterSuite;
import org.testng.annotations.BeforeSuite;
import org.testng.annotations.Test;

import java.io.IOException;
import java.io.InputStreamReader;
import java.io.Reader;
import java.net.URL;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.zip.GZIPInputStream;

import static com.facebook.presto.TupleInfo.Type.FIXED_INT_64;
import static com.facebook.presto.TupleInfo.Type.VARIABLE_BINARY;
import static com.facebook.presto.TupleInfo.Type.DOUBLE;
import static com.facebook.presto.ingest.RowSourceBuilder.RowGenerator;
import static com.google.common.base.Preconditions.checkArgument;
import static io.airlift.testing.Assertions.assertEqualsIgnoreOrder;
import static java.lang.String.format;
import static java.util.Collections.nCopies;

public class TestQueries
{
    private Handle handle;
    private List<List<String>> ordersData;
    private List<List<String>> lineitemData;

    private enum Column
    {
        ORDER_ORDERKEY(0),
        ORDER_CUSTKEY(1),
        ORDER_ORDERSTATUS(2),
        ORDER_TOTALPRICE(3),
        ORDER_ORDERDATE(4),
        ORDER_ORDERPRIORITY(5),
        ORDER_SHIPPRIORITY(6),

        LINEITEM_ORDERKEY(0),
        LINEITEM_DISCOUNT(6),
        LINEITEM_TAX(7);

        private final int index;

        Column(int index)
        {
            this.index = index;
        }

        public int getIndex()
        {
            return index;
        }
    }

    @BeforeSuite
    public void setupDatabase()
            throws IOException
    {
        handle = DBI.open("jdbc:h2:mem:test" + System.nanoTime());

        ordersData = readTestData("tpch/orders.dat.gz", 15000);
        handle.execute("CREATE TABLE orders (\n" +
                "  orderkey BIGINT NOT NULL,\n" +
                "  custkey BIGINT NOT NULL,\n" +
                "  orderstatus CHAR(1) NOT NULL,\n" +
                "  totalprice DOUBLE NOT NULL,\n" +
                "  orderdate CHAR(10) NOT NULL,\n" +
                "  orderpriority CHAR(15) NOT NULL,\n" +
                "  shippriority BIGINT NOT NULL\n" +
                ")");
        insertRows("orders", handle, ordersData);

        lineitemData = readTestData("tpch/lineitem.dat.gz", 60175);
        handle.execute("CREATE TABLE lineitem (\n" +
                "  orderkey BIGINT NOT NULL,\n" +
                "  partkey BIGINT NOT NULL,\n" +
                "  suppkey BIGINT NOT NULL,\n" +
                "  linenumber BIGINT NOT NULL,\n" +
                "  quantity BIGINT NOT NULL,\n" +
                "  extendedprice DOUBLE NOT NULL,\n" +
                "  discount DOUBLE NOT NULL,\n" +
                "  tax DOUBLE NOT NULL,\n" +
                "  returnflag CHAR(1) NOT NULL,\n" +
                "  linestatus CHAR(1) NOT NULL,\n" +
                "  shipdate CHAR(10) NOT NULL,\n" +
                "  commitdate CHAR(10) NOT NULL,\n" +
                "  receiptdate CHAR(10) NOT NULL,\n" +
                "  shipinstruct VARCHAR(25) NOT NULL,\n" +
                "  shipmode VARCHAR(10) NOT NULL\n" +
                ")");
        insertRows("lineitem", handle, lineitemData);
    }

    @AfterSuite
    public void cleanupDatabase()
    {
        handle.close();
    }

    @Test
    public void testCountAll()
    {
        List<Tuple> expected = computeExpected("SELECT COUNT(*) FROM orders", FIXED_INT_64);

        BlockStream orders = createBlockStream(ordersData, Column.ORDER_ORDERKEY, FIXED_INT_64);
        AggregationOperator aggregation = new AggregationOperator(orders, CountAggregation.PROVIDER);

        assertEqualsIgnoreOrder(tuples(aggregation), expected);
    }

    @Test
    public void testGroupBySum()
    {
        List<Tuple> expected = computeExpected(
                "SELECT orderstatus, SUM(custkey) FROM orders GROUP BY orderstatus",
                VARIABLE_BINARY, FIXED_INT_64);

        BlockStream groupBySource = createBlockStream(ordersData, Column.ORDER_ORDERSTATUS, VARIABLE_BINARY);
        BlockStream aggregateSource = createBlockStream(ordersData, Column.ORDER_TOTALPRICE, DOUBLE);

        GroupByBlockStream groupBy = new GroupByBlockStream(groupBySource);
        HashAggregationBlockStream aggregation = new HashAggregationBlockStream(groupBy, aggregateSource, SumAggregation.PROVIDER);

        assertEqualsIgnoreOrder(tuples(aggregation), expected);
    }

    private List<Tuple> computeExpected(final String sql, TupleInfo.Type... types)
    {
        TupleInfo tupleInfo = new TupleInfo(types);
        return handle.createQuery(sql)
                .map(tupleMapper(tupleInfo))
                .list();
    }

    private static List<Tuple> tuples(BlockStream blockStream)
    {
        Cursor cursor = blockStream.cursor();
        List<Tuple> list = new ArrayList<>();
        while (cursor.advanceNextPosition()) {
            list.add(cursor.getTuple());
        }
        return list;
    }

    @SuppressWarnings("UnusedDeclaration")
    private static Iterable<List<Object>> tupleValues(Iterable<Tuple> tuples)
    {
        return Iterables.transform(tuples, new Function<Tuple, List<Object>>()
        {
            @Override
            public List<Object> apply(Tuple input)
            {
                return input.toValues();
            }
        });
    }

    private static BlockStream createBlockStream(List<List<String>> data, final Column column, final TupleInfo.Type type)
    {
        final Iterator<List<String>> iterator = data.iterator();
        return new RowSourceBuilder(new TupleInfo(type), new RowGenerator()
        {
            @Override
            public boolean generate(RowSourceBuilder.RowBuilder rowBuilder)
            {
                if (!iterator.hasNext()) {
                    return false;
                }
                String value = iterator.next().get(column.getIndex());
                switch (type) {
                    case FIXED_INT_64:
                        rowBuilder.append(Long.parseLong(value));
                        break;
                    case VARIABLE_BINARY:
                        rowBuilder.append(value.getBytes(Charsets.UTF_8));
                        break;
                    default:
                        throw new AssertionError("unhandled type: " + type);
                }
                return true;
            }

            @Override
            public void close() {}
        });
    }

    private static void insertRows(String table, Handle handle, List<List<String>> data)
    {
        checkArgument(!data.isEmpty(), "no data to insert");
        int columns = data.get(0).size();
{'fraction_non_alphanumeric': 0.10133527234267584, 'fraction_numerical': 0.00707297726070862, 'mean_word_length': 2.4151241534988714, 'pattern_counts': {'":': 0, '<': 47, '<?xml version=': 0, '>': 47, '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': '5622833', 'n_tokens_mistral': 4434, 'n_tokens_neox': 4215, 'n_words': 870}