Skip to content

Commit f6261b0

Browse files
add a test to capture the bug (#843)
1 parent a18ece8 commit f6261b0

File tree

1 file changed

+19
-0
lines changed

1 file changed

+19
-0
lines changed

python/datafusion/tests/test_context.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -372,6 +372,25 @@ def test_dataset_filter(ctx, capfd):
372372
assert result[0].column(1) == pa.array([-3])
373373

374374

375+
def test_dataset_count(ctx):
376+
# `datafusion-python` issue: https://github.com/apache/datafusion-python/issues/800
377+
batch = pa.RecordBatch.from_arrays(
378+
[pa.array([1, 2, 3]), pa.array([4, 5, 6])],
379+
names=["a", "b"],
380+
)
381+
dataset = ds.dataset([batch])
382+
ctx.register_dataset("t", dataset)
383+
384+
# Testing the dataframe API
385+
df = ctx.table("t")
386+
assert df.count() == 3
387+
388+
# Testing the SQL API
389+
count = ctx.sql("SELECT COUNT(*) FROM t")
390+
count = count.collect()
391+
assert count[0].column(0) == pa.array([3])
392+
393+
375394
def test_pyarrow_predicate_pushdown_is_null(ctx, capfd):
376395
"""Ensure that pyarrow filter gets pushed down for `IsNull`"""
377396
# create a RecordBatch and register it as a pyarrow.dataset.Dataset

0 commit comments

Comments
 (0)