import duckdb
import pyarrow as pa
import time
from matplotlib import pyplot as plt
import timeit
from functools import partial

def get_table(col_size):
    "Creates a one row arrow table with the given column size"
    row = {f"col_{idx}": f"my_value_{idx}" for idx in range(col_size)}
    return pa.Table.from_pydict({"data": pa.array([row])})

def run_unnest(col_size):    
    "Run unnest natively"
    with duckdb.connect() as con:
        con.register("my_view", get_table(col_size))
        start_time = time.time()
        con.sql("select unnest(data) from my_view").fetchall()
        end_time = time.time()
    return end_time-start_time

def run_python_unnest(col_size):
    "Read the data unnest and put it back, then read again"
    with duckdb.connect() as con:
        con.register("my_view", get_table(col_size))
        start_time = time.time()
        data = con.sql("select data from my_view").fetchall()
        con.register("my_new_view", pa.Table.from_pylist([k[0] for k in data]))
        con.execute("select * from my_new_view").fetchall()
        end_time = time.time()
    return end_time - start_time

def time_function(fn, col_size):
    "Take an average of a 5 run. Return milliseconds"
    return timeit.timeit(stmt=partial(fn, col_size=col_size), number=5)/5*1000

x = list(range(1,402,20))
y1 = [time_function(run_unnest, col_size=col_size) for col_size in x]
y2 = [time_function(run_python_unnest, col_size=col_size) for col_size in x]

plt.plot(x, y1, "r--", label="duckdb unnest")
plt.plot(x, y2, "b-", label="python unnest")
plt.legend(loc="upper left")
plt.xlabel("Number of columns")
plt.ylabel("Milliseconds to run")
plt.show()