mirror of
https://github.com/langchain-ai/datafusion.git
synced 2026-07-21 03:05:28 -04:00
Define the unittests using pytest (#493)
* Use pytest * Formatting * Update GHA conf * Remove TODO note * Format * Test requirements file * Update workflow file * Merge requirements file * Update workflow file
This commit is contained in:
Krisztián Szűcs
+151
-255
@@ -15,286 +15,182 @@
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
import unittest
|
||||
import tempfile
|
||||
import datetime
|
||||
import os.path
|
||||
import shutil
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from datafusion import ExecutionContext
|
||||
|
||||
import numpy
|
||||
import pyarrow
|
||||
import datafusion
|
||||
|
||||
# used to write parquet files
|
||||
import pyarrow.parquet
|
||||
|
||||
from tests.generic import *
|
||||
from . import generic as helpers
|
||||
|
||||
|
||||
class TestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
# Create a temporary directory
|
||||
self.test_dir = tempfile.mkdtemp()
|
||||
numpy.random.seed(1)
|
||||
@pytest.fixture
|
||||
def ctx():
|
||||
return ExecutionContext()
|
||||
|
||||
def tearDown(self):
|
||||
# Remove the directory after the test
|
||||
shutil.rmtree(self.test_dir)
|
||||
|
||||
def test_no_table(self):
|
||||
with self.assertRaises(Exception):
|
||||
datafusion.Context().sql("SELECT a FROM b").collect()
|
||||
def test_no_table(ctx):
|
||||
with pytest.raises(Exception, match="DataFusion error"):
|
||||
ctx.sql("SELECT a FROM b").collect()
|
||||
|
||||
def test_register(self):
|
||||
ctx = datafusion.ExecutionContext()
|
||||
|
||||
path = write_parquet(os.path.join(self.test_dir, "a.parquet"), data())
|
||||
def test_register(ctx, tmp_path):
|
||||
path = helpers.write_parquet(tmp_path / "a.parquet", helpers.data())
|
||||
ctx.register_parquet("t", path)
|
||||
|
||||
ctx.register_parquet("t", path)
|
||||
assert ctx.tables() == {"t"}
|
||||
|
||||
self.assertEqual(ctx.tables(), {"t"})
|
||||
|
||||
def test_execute(self):
|
||||
data = [1, 1, 2, 2, 3, 11, 12]
|
||||
def test_execute(ctx, tmp_path):
|
||||
data = [1, 1, 2, 2, 3, 11, 12]
|
||||
|
||||
ctx = datafusion.ExecutionContext()
|
||||
# single column, "a"
|
||||
path = helpers.write_parquet(tmp_path / "a.parquet", pa.array(data))
|
||||
ctx.register_parquet("t", path)
|
||||
|
||||
# single column, "a"
|
||||
path = write_parquet(
|
||||
os.path.join(self.test_dir, "a.parquet"), pyarrow.array(data)
|
||||
assert ctx.tables() == {"t"}
|
||||
|
||||
# count
|
||||
result = ctx.sql("SELECT COUNT(a) FROM t").collect()
|
||||
|
||||
expected = pa.array([7], pa.uint64())
|
||||
expected = [pa.RecordBatch.from_arrays([expected], ["COUNT(a)"])]
|
||||
assert result == expected
|
||||
|
||||
# where
|
||||
expected = pa.array([2], pa.uint64())
|
||||
expected = [pa.RecordBatch.from_arrays([expected], ["COUNT(a)"])]
|
||||
result = ctx.sql("SELECT COUNT(a) FROM t WHERE a > 10").collect()
|
||||
assert result == expected
|
||||
|
||||
# group by
|
||||
results = ctx.sql(
|
||||
"SELECT CAST(a as int), COUNT(a) FROM t GROUP BY CAST(a as int)"
|
||||
).collect()
|
||||
|
||||
# group by returns batches
|
||||
result_keys = []
|
||||
result_values = []
|
||||
for result in results:
|
||||
pydict = result.to_pydict()
|
||||
result_keys.extend(pydict["CAST(a AS Int32)"])
|
||||
result_values.extend(pydict["COUNT(a)"])
|
||||
|
||||
result_keys, result_values = (
|
||||
list(t) for t in zip(*sorted(zip(result_keys, result_values)))
|
||||
)
|
||||
|
||||
assert result_keys == [1, 2, 3, 11, 12]
|
||||
assert result_values == [2, 2, 1, 1, 1]
|
||||
|
||||
# order by
|
||||
result = ctx.sql(
|
||||
"SELECT a, CAST(a AS int) FROM t ORDER BY a DESC LIMIT 2"
|
||||
).collect()
|
||||
expected_a = pa.array([50.0219, 50.0152], pa.float64())
|
||||
expected_cast = pa.array([50, 50], pa.int32())
|
||||
expected = [
|
||||
pa.RecordBatch.from_arrays(
|
||||
[expected_a, expected_cast], ["a", "CAST(a AS Int32)"]
|
||||
)
|
||||
ctx.register_parquet("t", path)
|
||||
]
|
||||
np.testing.assert_equal(expected[0].column(1), expected[0].column(1))
|
||||
|
||||
self.assertEqual(ctx.tables(), {"t"})
|
||||
|
||||
# count
|
||||
result = ctx.sql("SELECT COUNT(a) FROM t").collect()
|
||||
def test_cast(ctx, tmp_path):
|
||||
"""
|
||||
Verify that we can cast
|
||||
"""
|
||||
path = helpers.write_parquet(tmp_path / "a.parquet", helpers.data())
|
||||
ctx.register_parquet("t", path)
|
||||
|
||||
expected = pyarrow.array([7], pyarrow.uint64())
|
||||
expected = [pyarrow.RecordBatch.from_arrays([expected], ["COUNT(a)"])]
|
||||
self.assertEqual(expected, result)
|
||||
valid_types = [
|
||||
"smallint",
|
||||
"int",
|
||||
"bigint",
|
||||
"float(32)",
|
||||
"float(64)",
|
||||
"float",
|
||||
]
|
||||
|
||||
# where
|
||||
expected = pyarrow.array([2], pyarrow.uint64())
|
||||
expected = [pyarrow.RecordBatch.from_arrays([expected], ["COUNT(a)"])]
|
||||
self.assertEqual(
|
||||
expected, ctx.sql("SELECT COUNT(a) FROM t WHERE a > 10").collect()
|
||||
)
|
||||
select = ", ".join(
|
||||
[f"CAST(9 AS {t}) AS A{i}" for i, t in enumerate(valid_types)]
|
||||
)
|
||||
|
||||
# group by
|
||||
results = ctx.sql(
|
||||
"SELECT CAST(a as int), COUNT(a) FROM t GROUP BY CAST(a as int)"
|
||||
).collect()
|
||||
# can execute, which implies that we can cast
|
||||
ctx.sql(f"SELECT {select} FROM t").collect()
|
||||
|
||||
# group by returns batches
|
||||
result_keys = []
|
||||
result_values = []
|
||||
for result in results:
|
||||
pydict = result.to_pydict()
|
||||
result_keys.extend(pydict["CAST(a AS Int32)"])
|
||||
result_values.extend(pydict["COUNT(a)"])
|
||||
|
||||
result_keys, result_values = (
|
||||
list(t) for t in zip(*sorted(zip(result_keys, result_values)))
|
||||
)
|
||||
|
||||
self.assertEqual(result_keys, [1, 2, 3, 11, 12])
|
||||
self.assertEqual(result_values, [2, 2, 1, 1, 1])
|
||||
|
||||
# order by
|
||||
result = ctx.sql(
|
||||
"SELECT a, CAST(a AS int) FROM t ORDER BY a DESC LIMIT 2"
|
||||
).collect()
|
||||
expected_a = pyarrow.array([50.0219, 50.0152], pyarrow.float64())
|
||||
expected_cast = pyarrow.array([50, 50], pyarrow.int32())
|
||||
expected = [
|
||||
pyarrow.RecordBatch.from_arrays(
|
||||
[expected_a, expected_cast], ["a", "CAST(a AS Int32)"]
|
||||
)
|
||||
]
|
||||
numpy.testing.assert_equal(expected[0].column(1), expected[0].column(1))
|
||||
|
||||
def test_cast(self):
|
||||
"""
|
||||
Verify that we can cast
|
||||
"""
|
||||
ctx = datafusion.ExecutionContext()
|
||||
|
||||
path = write_parquet(os.path.join(self.test_dir, "a.parquet"), data())
|
||||
ctx.register_parquet("t", path)
|
||||
|
||||
valid_types = [
|
||||
"smallint",
|
||||
"int",
|
||||
"bigint",
|
||||
"float(32)",
|
||||
"float(64)",
|
||||
"float",
|
||||
]
|
||||
|
||||
select = ", ".join(
|
||||
[f"CAST(9 AS {t}) AS A{i}" for i, t in enumerate(valid_types)]
|
||||
)
|
||||
|
||||
# can execute, which implies that we can cast
|
||||
ctx.sql(f"SELECT {select} FROM t").collect()
|
||||
|
||||
def _test_udf(self, udf, args, return_type, array, expected):
|
||||
ctx = datafusion.ExecutionContext()
|
||||
|
||||
# write to disk
|
||||
path = write_parquet(os.path.join(self.test_dir, "a.parquet"), array)
|
||||
ctx.register_parquet("t", path)
|
||||
|
||||
ctx.register_udf("udf", udf, args, return_type)
|
||||
|
||||
batches = ctx.sql("SELECT udf(a) AS tt FROM t").collect()
|
||||
|
||||
result = batches[0].column(0)
|
||||
|
||||
self.assertEqual(expected, result)
|
||||
|
||||
def test_udf_identity(self):
|
||||
self._test_udf(
|
||||
@pytest.mark.parametrize(
|
||||
("fn", "input_types", "output_type", "input_values", "expected_values"),
|
||||
[
|
||||
(
|
||||
lambda x: x,
|
||||
[pyarrow.float64()],
|
||||
pyarrow.float64(),
|
||||
pyarrow.array([-1.2, None, 1.2]),
|
||||
pyarrow.array([-1.2, None, 1.2]),
|
||||
)
|
||||
|
||||
def test_udf(self):
|
||||
self._test_udf(
|
||||
[pa.float64()],
|
||||
pa.float64(),
|
||||
[-1.2, None, 1.2],
|
||||
[-1.2, None, 1.2],
|
||||
),
|
||||
(
|
||||
lambda x: x.is_null(),
|
||||
[pyarrow.float64()],
|
||||
pyarrow.bool_(),
|
||||
pyarrow.array([-1.2, None, 1.2]),
|
||||
pyarrow.array([False, True, False]),
|
||||
)
|
||||
[pa.float64()],
|
||||
pa.bool_(),
|
||||
[-1.2, None, 1.2],
|
||||
[False, True, False],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_udf(
|
||||
ctx, tmp_path, fn, input_types, output_type, input_values, expected_values
|
||||
):
|
||||
# write to disk
|
||||
path = helpers.write_parquet(
|
||||
tmp_path / "a.parquet", pa.array(input_values)
|
||||
)
|
||||
ctx.register_parquet("t", path)
|
||||
ctx.register_udf("udf", fn, input_types, output_type)
|
||||
|
||||
batches = ctx.sql("SELECT udf(a) AS tt FROM t").collect()
|
||||
result = batches[0].column(0)
|
||||
|
||||
assert result == pa.array(expected_values)
|
||||
|
||||
|
||||
class TestIO(unittest.TestCase):
|
||||
def setUp(self):
|
||||
# Create a temporary directory
|
||||
self.test_dir = tempfile.mkdtemp()
|
||||
_null_mask = np.array([False, True, False])
|
||||
|
||||
def tearDown(self):
|
||||
# Remove the directory after the test
|
||||
shutil.rmtree(self.test_dir)
|
||||
|
||||
def _test_data(self, data):
|
||||
ctx = datafusion.ExecutionContext()
|
||||
@pytest.mark.parametrize(
|
||||
"arr",
|
||||
[
|
||||
pa.array(["a", "b", "c"], pa.utf8(), _null_mask),
|
||||
pa.array(["a", "b", "c"], pa.large_utf8(), _null_mask),
|
||||
pa.array([b"1", b"2", b"3"], pa.binary(), _null_mask),
|
||||
pa.array([b"1111", b"2222", b"3333"], pa.large_binary(), _null_mask),
|
||||
pa.array([False, True, True], None, _null_mask),
|
||||
pa.array([0, 1, 2], None),
|
||||
helpers.data_binary_other(),
|
||||
helpers.data_date32(),
|
||||
helpers.data_with_nans(),
|
||||
# C data interface missing
|
||||
pytest.param(
|
||||
pa.array([b"1111", b"2222", b"3333"], pa.binary(4), _null_mask),
|
||||
marks=pytest.mark.xfail,
|
||||
),
|
||||
pytest.param(helpers.data_datetime("s"), marks=pytest.mark.xfail),
|
||||
pytest.param(helpers.data_datetime("ms"), marks=pytest.mark.xfail),
|
||||
pytest.param(helpers.data_datetime("us"), marks=pytest.mark.xfail),
|
||||
pytest.param(helpers.data_datetime("ns"), marks=pytest.mark.xfail),
|
||||
# Not writtable to parquet
|
||||
pytest.param(helpers.data_timedelta("s"), marks=pytest.mark.xfail),
|
||||
pytest.param(helpers.data_timedelta("ms"), marks=pytest.mark.xfail),
|
||||
pytest.param(helpers.data_timedelta("us"), marks=pytest.mark.xfail),
|
||||
pytest.param(helpers.data_timedelta("ns"), marks=pytest.mark.xfail),
|
||||
],
|
||||
)
|
||||
def test_simple_select(ctx, tmp_path, arr):
|
||||
path = helpers.write_parquet(tmp_path / "a.parquet", arr)
|
||||
ctx.register_parquet("t", path)
|
||||
|
||||
# write to disk
|
||||
path = write_parquet(os.path.join(self.test_dir, "a.parquet"), data)
|
||||
ctx.register_parquet("t", path)
|
||||
batches = ctx.sql("SELECT a AS tt FROM t").collect()
|
||||
result = batches[0].column(0)
|
||||
|
||||
batches = ctx.sql("SELECT a AS tt FROM t").collect()
|
||||
|
||||
result = batches[0].column(0)
|
||||
|
||||
numpy.testing.assert_equal(data, result)
|
||||
|
||||
def test_nans(self):
|
||||
self._test_data(data_with_nans())
|
||||
|
||||
def test_utf8(self):
|
||||
array = pyarrow.array(
|
||||
["a", "b", "c"], pyarrow.utf8(), numpy.array([False, True, False])
|
||||
)
|
||||
self._test_data(array)
|
||||
|
||||
def test_large_utf8(self):
|
||||
array = pyarrow.array(
|
||||
["a", "b", "c"], pyarrow.large_utf8(), numpy.array([False, True, False])
|
||||
)
|
||||
self._test_data(array)
|
||||
|
||||
# Error from Arrow
|
||||
@unittest.expectedFailure
|
||||
def test_datetime_s(self):
|
||||
self._test_data(data_datetime("s"))
|
||||
|
||||
# C data interface missing
|
||||
@unittest.expectedFailure
|
||||
def test_datetime_ms(self):
|
||||
self._test_data(data_datetime("ms"))
|
||||
|
||||
# C data interface missing
|
||||
@unittest.expectedFailure
|
||||
def test_datetime_us(self):
|
||||
self._test_data(data_datetime("us"))
|
||||
|
||||
# Not writtable to parquet
|
||||
@unittest.expectedFailure
|
||||
def test_datetime_ns(self):
|
||||
self._test_data(data_datetime("ns"))
|
||||
|
||||
# Not writtable to parquet
|
||||
@unittest.expectedFailure
|
||||
def test_timedelta_s(self):
|
||||
self._test_data(data_timedelta("s"))
|
||||
|
||||
# Not writtable to parquet
|
||||
@unittest.expectedFailure
|
||||
def test_timedelta_ms(self):
|
||||
self._test_data(data_timedelta("ms"))
|
||||
|
||||
# Not writtable to parquet
|
||||
@unittest.expectedFailure
|
||||
def test_timedelta_us(self):
|
||||
self._test_data(data_timedelta("us"))
|
||||
|
||||
# Not writtable to parquet
|
||||
@unittest.expectedFailure
|
||||
def test_timedelta_ns(self):
|
||||
self._test_data(data_timedelta("ns"))
|
||||
|
||||
def test_date32(self):
|
||||
array = pyarrow.array(
|
||||
[
|
||||
datetime.date(2000, 1, 1),
|
||||
datetime.date(1980, 1, 1),
|
||||
datetime.date(2030, 1, 1),
|
||||
],
|
||||
pyarrow.date32(),
|
||||
numpy.array([False, True, False]),
|
||||
)
|
||||
self._test_data(array)
|
||||
|
||||
def test_binary_variable(self):
|
||||
array = pyarrow.array(
|
||||
[b"1", b"2", b"3"], pyarrow.binary(), numpy.array([False, True, False])
|
||||
)
|
||||
self._test_data(array)
|
||||
|
||||
# C data interface missing
|
||||
@unittest.expectedFailure
|
||||
def test_binary_fixed(self):
|
||||
array = pyarrow.array(
|
||||
[b"1111", b"2222", b"3333"],
|
||||
pyarrow.binary(4),
|
||||
numpy.array([False, True, False]),
|
||||
)
|
||||
self._test_data(array)
|
||||
|
||||
def test_large_binary(self):
|
||||
array = pyarrow.array(
|
||||
[b"1111", b"2222", b"3333"],
|
||||
pyarrow.large_binary(),
|
||||
numpy.array([False, True, False]),
|
||||
)
|
||||
self._test_data(array)
|
||||
|
||||
def test_binary_other(self):
|
||||
self._test_data(data_binary_other())
|
||||
|
||||
def test_bool(self):
|
||||
array = pyarrow.array(
|
||||
[False, True, True], None, numpy.array([False, True, False])
|
||||
)
|
||||
self._test_data(array)
|
||||
|
||||
def test_u32(self):
|
||||
array = pyarrow.array([0, 1, 2], None, numpy.array([False, True, False]))
|
||||
self._test_data(array)
|
||||
np.testing.assert_equal(result, arr)
|
||||
|
||||
Reference in New Issue
Block a user