mirror of
https://github.com/langchain-ai/datafusion.git
synced 2026-07-16 12:04:27 -04:00
Andy Grove
295 lines
8.5 KiB
Python
295 lines
8.5 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
import unittest
|
|
import tempfile
|
|
import datetime
|
|
import os.path
|
|
import shutil
|
|
|
|
import numpy
|
|
import pyarrow
|
|
import datafusion
|
|
|
|
# used to write parquet files
|
|
import pyarrow.parquet
|
|
|
|
from tests.generic import *
|
|
|
|
|
|
class TestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
# Create a temporary directory
|
|
self.test_dir = tempfile.mkdtemp()
|
|
numpy.random.seed(1)
|
|
|
|
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_register(self):
|
|
ctx = datafusion.ExecutionContext()
|
|
|
|
path = write_parquet(os.path.join(self.test_dir, "a.parquet"), data())
|
|
|
|
ctx.register_parquet("t", path)
|
|
|
|
self.assertEqual(ctx.tables(), {"t"})
|
|
|
|
def test_execute(self):
|
|
data = [1, 1, 2, 2, 3, 11, 12]
|
|
|
|
ctx = datafusion.ExecutionContext()
|
|
|
|
# single column, "a"
|
|
path = write_parquet(
|
|
os.path.join(self.test_dir, "a.parquet"), pyarrow.array(data)
|
|
)
|
|
ctx.register_parquet("t", path)
|
|
|
|
self.assertEqual(ctx.tables(), {"t"})
|
|
|
|
# count
|
|
result = ctx.sql("SELECT COUNT(a) FROM t").collect()
|
|
|
|
expected = pyarrow.array([7], pyarrow.uint64())
|
|
expected = [pyarrow.RecordBatch.from_arrays([expected], ["COUNT(a)"])]
|
|
self.assertEqual(expected, result)
|
|
|
|
# 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()
|
|
)
|
|
|
|
# group by
|
|
result = ctx.sql(
|
|
"SELECT CAST(a as int), COUNT(a) FROM t GROUP BY CAST(a as int)"
|
|
).collect()
|
|
|
|
result_keys = result[0].to_pydict()["CAST(a AS Int32)"]
|
|
result_values = result[0].to_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(
|
|
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(
|
|
lambda x: x.is_null(),
|
|
[pyarrow.float64()],
|
|
pyarrow.bool_(),
|
|
pyarrow.array([-1.2, None, 1.2]),
|
|
pyarrow.array([False, True, False]),
|
|
)
|
|
|
|
|
|
class TestIO(unittest.TestCase):
|
|
def setUp(self):
|
|
# Create a temporary directory
|
|
self.test_dir = tempfile.mkdtemp()
|
|
|
|
def tearDown(self):
|
|
# Remove the directory after the test
|
|
shutil.rmtree(self.test_dir)
|
|
|
|
def _test_data(self, data):
|
|
ctx = datafusion.ExecutionContext()
|
|
|
|
# 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)
|
|
|
|
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)
|