| 116 | |
| 117 | @pytest.mark.gandiva |
| 118 | def test_in_expr(): |
| 119 | import pyarrow.gandiva as gandiva |
| 120 | |
| 121 | arr = pa.array(["ga", "an", "nd", "di", "iv", "va"]) |
| 122 | table = pa.Table.from_arrays([arr], ["a"]) |
| 123 | |
| 124 | # string |
| 125 | builder = gandiva.TreeExprBuilder() |
| 126 | node_a = builder.make_field(table.schema.field("a")) |
| 127 | cond = builder.make_in_expression(node_a, ["an", "nd"], pa.string()) |
| 128 | condition = builder.make_condition(cond) |
| 129 | filter = gandiva.make_filter(table.schema, condition) |
| 130 | result = filter.evaluate(table.to_batches()[0], pa.default_memory_pool()) |
| 131 | assert result.to_array().equals(pa.array([1, 2], type=pa.uint32())) |
| 132 | |
| 133 | # int32 |
| 134 | arr = pa.array([3, 1, 4, 1, 5, 9, 2, 6, 5, 4]) |
| 135 | table = pa.Table.from_arrays([arr.cast(pa.int32())], ["a"]) |
| 136 | node_a = builder.make_field(table.schema.field("a")) |
| 137 | cond = builder.make_in_expression(node_a, [1, 5], pa.int32()) |
| 138 | condition = builder.make_condition(cond) |
| 139 | filter = gandiva.make_filter(table.schema, condition) |
| 140 | result = filter.evaluate(table.to_batches()[0], pa.default_memory_pool()) |
| 141 | assert result.to_array().equals(pa.array([1, 3, 4, 8], type=pa.uint32())) |
| 142 | |
| 143 | # int64 |
| 144 | arr = pa.array([3, 1, 4, 1, 5, 9, 2, 6, 5, 4]) |
| 145 | table = pa.Table.from_arrays([arr], ["a"]) |
| 146 | node_a = builder.make_field(table.schema.field("a")) |
| 147 | cond = builder.make_in_expression(node_a, [1, 5], pa.int64()) |
| 148 | condition = builder.make_condition(cond) |
| 149 | filter = gandiva.make_filter(table.schema, condition) |
| 150 | result = filter.evaluate(table.to_batches()[0], pa.default_memory_pool()) |
| 151 | assert result.to_array().equals(pa.array([1, 3, 4, 8], type=pa.uint32())) |
| 152 | |
| 153 | |
| 154 | @pytest.mark.skip(reason="Gandiva C++ did not have *real* binary, " |