From 682669058e77756eeed045d54c6a662a76d0e417 Mon Sep 17 00:00:00 2001 From: Weston Pace Date: Fri, 22 Mar 2024 07:20:07 -0700 Subject: [PATCH] Added dictionary and ree to the one_way test --- python/pyarrow/tests/test_substrait.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/python/pyarrow/tests/test_substrait.py b/python/pyarrow/tests/test_substrait.py index 53e96f4540363..8579eb1ab2fc3 100644 --- a/python/pyarrow/tests/test_substrait.py +++ b/python/pyarrow/tests/test_substrait.py @@ -969,20 +969,28 @@ def test_arrow_one_way_types(): [ pa.field("binary_view", pa.binary_view()), pa.field("string_view", pa.string_view()), + pa.field("dictionary", pa.dictionary(pa.int32(), pa.string())), + pa.field("ree", pa.run_end_encoded(pa.int32(), pa.string())), ] ) alt_schema = pa.schema( - [pa.field("binary_view", pa.binary()), pa.field("string_view", pa.string())] + [ + pa.field("binary_view", pa.binary()), + pa.field("string_view", pa.string()), + pa.field("dictionary", pa.string()), + pa.field("ree", pa.string()) + ] ) - def check_one_way(expr): + def check_one_way(field): + expr = pc.is_null(pc.field(field.name)) buf = pa.substrait.serialize_expressions([expr], ["test_expr"], schema) returned = pa.substrait.deserialize_expressions(buf) assert alt_schema == returned.schema - check_one_way(pc.is_null(pc.field("binary_view"))) - check_one_way(pc.is_null(pc.field("string_view"))) - + for field in schema: + check_one_way(field) + def test_invalid_expression_ser_des(): schema = pa.schema([