langchain-ai/datafusion
1
0
Fork 0
mirror of https://github.com/langchain-ai/datafusion.git synced 2026-07-16 04:03:28 -04:00
Code Issues 3 Packages Projects Releases Wiki Activity

Revert "Add datafusion-python (#69)" (#257)

Browse Source 
This reverts commit 46bde0bd14.
This commit is contained in:
Andy Grove
2021-05-04 08:51:44 -06:00
committed by GitHub
parent 46bde0bd14
commit d0af907652
28 changed files with 1 additions and 2244 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tests/__init__.py
-16
View File
@@ -1,16 +0,0 @@
# 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.
python/tests/generic.py
-75
View File
@@ -1,75 +0,0 @@
# 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
def data():
data = numpy.concatenate(
[numpy.random.normal(0, 0.01, size=50), numpy.random.normal(50, 0.01, size=50)]
)
return pyarrow.array(data)
def data_with_nans():
data = numpy.random.normal(0, 0.01, size=50)
mask = numpy.random.randint(0, 2, size=50)
data[mask == 0] = numpy.NaN
return data
def data_datetime(f):
data = [
datetime.datetime.now(),
datetime.datetime.now() - datetime.timedelta(days=1),
datetime.datetime.now() + datetime.timedelta(days=1),
]
return pyarrow.array(
data, type=pyarrow.timestamp(f), mask=numpy.array([False, True, False])
)
def data_timedelta(f):
data = [
datetime.timedelta(days=100),
datetime.timedelta(days=1),
datetime.timedelta(seconds=1),
]
return pyarrow.array(
data, type=pyarrow.duration(f), mask=numpy.array([False, True, False])
)
def data_binary_other():
return numpy.array([1, 0, 0], dtype="u4")
def write_parquet(path, data):
table = pyarrow.Table.from_arrays([data], names=["a"])
pyarrow.parquet.write_table(table, path)
return path
python/tests/test_df.py
-115
View File
@@ -1,115 +0,0 @@
# 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 pyarrow
import datafusion
f = datafusion.functions
class TestCase(unittest.TestCase):
def _prepare(self):
ctx = datafusion.ExecutionContext()
# create a RecordBatch and a new DataFrame from it
batch = pyarrow.RecordBatch.from_arrays(
[pyarrow.array([1, 2, 3]), pyarrow.array([4, 5, 6])],
names=["a", "b"],
)
return ctx.create_dataframe([[batch]])
def test_select(self):
df = self._prepare()
df = df.select(
f.col("a") + f.col("b"),
f.col("a") - f.col("b"),
)
# execute and collect the first (and only) batch
result = df.collect()[0]
self.assertEqual(result.column(0), pyarrow.array([5, 7, 9]))
self.assertEqual(result.column(1), pyarrow.array([-3, -3, -3]))
def test_filter(self):
df = self._prepare()
df = df \
.select(
f.col("a") + f.col("b"),
f.col("a") - f.col("b"),
) \
.filter(f.col("a") > f.lit(2))
# execute and collect the first (and only) batch
result = df.collect()[0]
self.assertEqual(result.column(0), pyarrow.array([9]))
self.assertEqual(result.column(1), pyarrow.array([-3]))
def test_limit(self):
df = self._prepare()
df = df.limit(1)
# execute and collect the first (and only) batch
result = df.collect()[0]
self.assertEqual(len(result.column(0)), 1)
self.assertEqual(len(result.column(1)), 1)
def test_udf(self):
df = self._prepare()
# is_null is a pyarrow function over arrays
udf = f.udf(lambda x: x.is_null(), [pyarrow.int64()], pyarrow.bool_())
df = df.select(udf(f.col("a")))
self.assertEqual(df.collect()[0].column(0), pyarrow.array([False, False, False]))
def test_join(self):
ctx = datafusion.ExecutionContext()
batch = pyarrow.RecordBatch.from_arrays(
[pyarrow.array([1, 2, 3]), pyarrow.array([4, 5, 6])],
names=["a", "b"],
)
df = ctx.create_dataframe([[batch]])
batch = pyarrow.RecordBatch.from_arrays(
[pyarrow.array([1, 2]), pyarrow.array([8, 10])],
names=["a", "c"],
)
df1 = ctx.create_dataframe([[batch]])
df = df.join(df1, on="a", how="inner")
# execute and collect the first (and only) batch
batch = df.collect()[0]
if batch.column(0) == pyarrow.array([1, 2]):
self.assertEqual(batch.column(0), pyarrow.array([1, 2]))
self.assertEqual(batch.column(1), pyarrow.array([8, 10]))
self.assertEqual(batch.column(2), pyarrow.array([4, 5]))
else:
self.assertEqual(batch.column(0), pyarrow.array([2, 1]))
self.assertEqual(batch.column(1), pyarrow.array([10, 8]))
self.assertEqual(batch.column(2), pyarrow.array([5, 4]))
python/tests/test_sql.py
-294
View File
@@ -1,294 +0,0 @@
# 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)
python/tests/test_udaf.py
-91
View File
@@ -1,91 +0,0 @@
# 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 pyarrow
import pyarrow.compute
import datafusion
f = datafusion.functions
class Accumulator:
"""
Interface of a user-defined accumulation.
"""
def __init__(self):
self._sum = pyarrow.scalar(0.0)
def to_scalars(self) -> [pyarrow.Scalar]:
return [self._sum]
def update(self, values: pyarrow.Array) -> None:
# not nice since pyarrow scalars can't be summed yet. This breaks on `None`
self._sum = pyarrow.scalar(
self._sum.as_py() + pyarrow.compute.sum(values).as_py()
)
def merge(self, states: pyarrow.Array) -> None:
# not nice since pyarrow scalars can't be summed yet. This breaks on `None`
self._sum = pyarrow.scalar(
self._sum.as_py() + pyarrow.compute.sum(states).as_py()
)
def evaluate(self) -> pyarrow.Scalar:
return self._sum
class TestCase(unittest.TestCase):
def _prepare(self):
ctx = datafusion.ExecutionContext()
# create a RecordBatch and a new DataFrame from it
batch = pyarrow.RecordBatch.from_arrays(
[pyarrow.array([1, 2, 3]), pyarrow.array([4, 4, 6])],
names=["a", "b"],
)
return ctx.create_dataframe([[batch]])
def test_aggregate(self):
df = self._prepare()
udaf = f.udaf(
Accumulator, pyarrow.float64(), pyarrow.float64(), [pyarrow.float64()]
)
df = df.aggregate([], [udaf(f.col("a"))])
# execute and collect the first (and only) batch
result = df.collect()[0]
self.assertEqual(result.column(0), pyarrow.array([1.0 + 2.0 + 3.0]))
def test_group_by(self):
df = self._prepare()
udaf = f.udaf(
Accumulator, pyarrow.float64(), pyarrow.float64(), [pyarrow.float64()]
)
df = df.aggregate([f.col("b")], [udaf(f.col("a"))])
# execute and collect the first (and only) batch
result = df.collect()[0]
self.assertEqual(result.column(1), pyarrow.array([1.0 + 2.0, 3.0]))