# /// script
# requires-python = ">=3.12,<3.13"
# dependencies = ["datafusion==54.0.0", "duckdb==1.4.3", "pyarrow==25.0.1"]
# ///
import hashlib
import json
from pathlib import Path

import duckdb
import pyarrow as pa
from datafusion import SessionContext, SessionConfig

folder = Path(__file__).resolve().parent
fixture = json.loads((folder / 'plans.json').read_text())
file = folder / 'weather.parquet'
assert hashlib.sha256(file.read_bytes()).hexdigest() == fixture['fileSha256']
context = SessionContext(SessionConfig().with_target_partitions(2))
context.register_parquet('data', str(file))
context.register_record_batches('stations', [[pa.record_batch({'station':['North','South'], 'description':['Northern site','Southern site']})]])
con = duckdb.connect()
con.read_parquet(str(file)).create_view('data')
con.execute("CREATE VIEW stations AS SELECT * FROM (VALUES ('North','Northern site'),('South','Southern site')) AS lookup(station,description)")
for query in fixture['queries']:
    frame = context.sql(query['sql'])
    actual = frame.to_pylist()
    cursor = con.execute(query['sql'])
    names = [column[0] for column in cursor.description]
    assert actual == [dict(zip(names, row)) for row in cursor.fetchall()]
    assert [[str(row[name]) if row[name] is not None else None for name in names] for row in actual] == query['rows']
    print(query['label'], 'verified')
    print(frame.optimized_logical_plan().display_indent())
    print(frame.execution_plan().display_indent())