import duckdb
import os

try:
    import pyarrow
    import pyarrow.parquet
    import numpy as np

    can_run = True
except:
    can_run = False


class TestArrowParallel(object):
    def test_parallel_run(self, duckdb_cursor):
        if not can_run:
            return
        duckdb_conn = duckdb.connect()
        duckdb_conn.execute("PRAGMA threads=4")
        duckdb_conn.execute("PRAGMA verify_parallelism")
        data = pyarrow.array(np.random.randint(800, size=1000000), type=pyarrow.int32())
        tbl = pyarrow.Table.from_batches(pyarrow.Table.from_arrays([data], ['a']).to_batches(10000))
        rel = duckdb_conn.from_arrow(tbl)
        # Also test multiple reads
        assert rel.aggregate("(count(a))::INT").execute().fetchone()[0] == 1000000
        assert rel.aggregate("(count(a))::INT").execute().fetchone()[0] == 1000000

    def test_parallel_types_and_different_batches(self, duckdb_cursor):
        if not can_run:
            return
        duckdb_conn = duckdb.connect()
        duckdb_conn.execute("PRAGMA threads=4")
        duckdb_conn.execute("PRAGMA verify_parallelism")

        parquet_filename = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'data', 'userdata1.parquet')
        cols = 'id, first_name, last_name, email, gender, ip_address, cc, country, birthdate, salary, title, comments'

        userdata_parquet_table = pyarrow.parquet.read_table(parquet_filename)
        for i in [7, 51, 99, 100, 101, 500, 1000, 2000]:
            data = pyarrow.array(np.arange(3, 7), type=pyarrow.int32())
            tbl = pyarrow.Table.from_arrays([data], ['a'])
            rel_id = duckdb_conn.from_arrow(tbl)
            userdata_parquet_table2 = pyarrow.Table.from_batches(userdata_parquet_table.to_batches(i))
            rel = duckdb_conn.from_arrow(userdata_parquet_table2)
            result = rel.filter("first_name=\'Jose\' and salary > 134708.82").aggregate('count(*)')
            assert result.execute().fetchone()[0] == 4

    def test_parallel_fewer_batches_than_threads(self, duckdb_cursor):
        if not can_run:
            return
        duckdb_conn = duckdb.connect()
        duckdb_conn.execute("PRAGMA threads=4")
        duckdb_conn.execute("PRAGMA verify_parallelism")

        data = pyarrow.array(np.random.randint(800, size=1000), type=pyarrow.int32())
        tbl = pyarrow.Table.from_batches(pyarrow.Table.from_arrays([data], ['a']).to_batches(2))
        rel = duckdb_conn.from_arrow(tbl)
        # Also test multiple reads
        assert rel.aggregate("(count(a))::INT").execute().fetchone()[0] == 1000
