()
| 23 | |
| 24 | @pytest.mark.gandiva |
| 25 | def test_tree_exp_builder(): |
| 26 | import pyarrow.gandiva as gandiva |
| 27 | |
| 28 | builder = gandiva.TreeExprBuilder() |
| 29 | |
| 30 | field_a = pa.field('a', pa.int32()) |
| 31 | field_b = pa.field('b', pa.int32()) |
| 32 | |
| 33 | schema = pa.schema([field_a, field_b]) |
| 34 | |
| 35 | field_result = pa.field('res', pa.int32()) |
| 36 | |
| 37 | node_a = builder.make_field(field_a) |
| 38 | node_b = builder.make_field(field_b) |
| 39 | |
| 40 | assert node_a.return_type() == field_a.type |
| 41 | |
| 42 | condition = builder.make_function("greater_than", [node_a, node_b], |
| 43 | pa.bool_()) |
| 44 | if_node = builder.make_if(condition, node_a, node_b, pa.int32()) |
| 45 | |
| 46 | expr = builder.make_expression(if_node, field_result) |
| 47 | |
| 48 | assert expr.result().type == pa.int32() |
| 49 | |
| 50 | config = gandiva.Configuration(dump_ir=True) |
| 51 | projector = gandiva.make_projector( |
| 52 | schema, [expr], pa.default_memory_pool(), "NONE", config) |
| 53 | |
| 54 | # Gandiva generates compute kernel function named `@expr_X` |
| 55 | assert projector.llvm_ir.find("@expr_") != -1 |
| 56 | |
| 57 | a = pa.array([10, 12, -20, 5], type=pa.int32()) |
| 58 | b = pa.array([5, 15, 15, 17], type=pa.int32()) |
| 59 | e = pa.array([10, 15, 15, 17], type=pa.int32()) |
| 60 | input_batch = pa.RecordBatch.from_arrays([a, b], names=['a', 'b']) |
| 61 | |
| 62 | r, = projector.evaluate(input_batch) |
| 63 | assert r.equals(e) |
| 64 | |
| 65 | |
| 66 | @pytest.mark.gandiva |
nothing calls this directly
no test coverage detected