Skip to content

Commit

Permalink
Fix duckdb tests
Browse files Browse the repository at this point in the history
  • Loading branch information
goodwanghan authored Jan 21, 2024
1 parent 47d1960 commit 6e2bdcc
Showing 1 changed file with 40 additions and 26 deletions.
66 changes: 40 additions & 26 deletions tests/fugue_duckdb/test_execution_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,13 @@
import duckdb
import pandas as pd
import pyarrow as pa
import pytest
from pytest import raises

import fugue.api as fa
import fugue.test as ft
from fugue import ArrowDataFrame, DataFrame, FugueWorkflow, fsql
from fugue.api import engine_context
from fugue.plugins import infer_execution_engine
from fugue_duckdb import DuckExecutionEngine
from fugue_duckdb.dataframe import DuckDataFrame
from fugue_test.builtin_suite import BuiltInTests
from fugue_test.execution_suite import ExecutionEngineTests
Expand Down Expand Up @@ -109,39 +107,55 @@ def test_builtin_connection():


def test_configs():
dag = FugueWorkflow()
df = dag.df([[None], [1]], "a:double")
df = dag.select("SELECT * FROM ", df, "ORDER BY a LIMIT 1")
df.assert_eq(dag.df([[None]], "a:double"))

dag.run(
"duckdb",
{
"fugue.duckdb.pragma.threads": 2,
"fugue.duckdb.pragma.default_null_order": "NULLS FIRST",
},
df = fa.as_pandas(
fa.fugue_sql(
"""
SELECT name, value FROM duckdb_settings()
WHERE name IN ('threads')
""",
engine="duckdb",
engine_conf={"fugue.duckdb.pragma.threads": 1},
)
)
assert df.value.iloc[0] == "1"

dag = FugueWorkflow()
df = dag.df([[None], [1]], "a:double")
df = dag.select("SELECT * FROM ", df, "ORDER BY a LIMIT 1")
df.assert_eq(dag.df([[1]], "a:double"))

dag.run(
"duckdb",
{
"fugue.duckdb.pragma.threads": 2,
"fugue.duckdb.pragma.default_null_order": "NULLS LAST",
},
df = fa.as_pandas(
fa.fugue_sql(
"""
SELECT name, value FROM duckdb_settings()
WHERE name IN ('threads')
""",
engine="duckdb",
engine_conf={"fugue.duckdb.pragma.threads": 3},
)
)
assert df.value.iloc[0] == "3"

with raises(ValueError):
# invalid config format
dag.run("duckdb", {"fugue.duckdb.pragma.threads;xx": 2})
df = fa.as_pandas(
fa.fugue_sql(
"""
SELECT name, value FROM duckdb_settings()
WHERE name IN ('threads')
""",
engine="duckdb",
engine_conf={"fugue.duckdb.pragma.threads;xx": 3},
)
)

with raises(Exception):
# non-existent config
dag.run("duckdb", {"fugue.duckdb.pragma.threads_xx": 2})
df = fa.as_pandas(
fa.fugue_sql(
"""
SELECT name, value FROM duckdb_settings()
WHERE name IN ('threads')
""",
engine="duckdb",
engine_conf={"fugue.duckdb.pragma.threads_xx": 3},
)
)


def test_annotations():
Expand Down

0 comments on commit 6e2bdcc

Please sign in to comment.