From 85478f7e1664a3bc302a6bdbf57a3d142d339e52 Mon Sep 17 00:00:00 2001 From: Jay Chia Date: Sat, 14 Dec 2024 21:40:32 +0800 Subject: [PATCH] Add arrow return test --- tests/memory/test_udf_project.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/memory/test_udf_project.py b/tests/memory/test_udf_project.py index 9ce2afff36..7ac1bfee98 100644 --- a/tests/memory/test_udf_project.py +++ b/tests/memory/test_udf_project.py @@ -1,5 +1,6 @@ import uuid +import pyarrow as pa import pytest from memray._memray import compute_statistics @@ -32,6 +33,12 @@ def to_pylist_identity_batched(s): return data +@daft.udf(return_dtype=str, batch_size=128) +def to_pylist_identity_batched_arrow_return(s): + data = s.to_pylist() + return pa.array(data) + + @pytest.mark.parametrize( "udf", [ @@ -39,6 +46,7 @@ def to_pylist_identity_batched(s): to_pylist_identity, to_arrow_identity_batched, to_pylist_identity_batched, + to_pylist_identity_batched_arrow_return, ], ) def test_string_identity_projection(udf):