| 281 | // Test that timestamp types work as inputs to in expressions |
| 282 | template <typename Type> |
| 283 | void TestDateOrTimeInExpression( |
| 284 | arrow::MemoryPool* pool, const DataTypePtr& type, |
| 285 | std::function<NodePtr(NodePtr, std::unordered_set<typename Type::c_type>)> |
| 286 | expression_factory) { |
| 287 | using c_type = typename Type::c_type; |
| 288 | |
| 289 | // schema for input fields |
| 290 | auto field0 = field("f0", type); |
| 291 | auto schema = arrow::schema({field0}); |
| 292 | |
| 293 | // Build f0 in (717629471, 717630471) |
| 294 | auto node_f0 = TreeExprBuilder::MakeField(field0); |
| 295 | std::unordered_set<c_type> in_constants({717629471, 717630471}); |
| 296 | auto in_expr = expression_factory(node_f0, in_constants); |
| 297 | auto condition = TreeExprBuilder::MakeCondition(in_expr); |
| 298 | |
| 299 | std::shared_ptr<Filter> filter; |
| 300 | auto status = Filter::Make(schema, condition, TestConfiguration(), &filter); |
| 301 | ASSERT_TRUE(status.ok()) << status.message(); |
| 302 | |
| 303 | // Create a row-batch with some sample data |
| 304 | int num_records = 5; |
| 305 | auto array0 = MakeArrowArray<Type, c_type>( |
| 306 | type, {717629471, 717630471, 717729471, 717629471, 717629571}, |
| 307 | {true, true, true, false, true}); |
| 308 | // expected output (indices for which condition matches) |
| 309 | auto exp = MakeArrowArrayUint16({0, 1}); |
| 310 | |
| 311 | // prepare input record batch |
| 312 | auto in_batch = arrow::RecordBatch::Make(schema, num_records, {array0}); |
| 313 | |
| 314 | std::shared_ptr<SelectionVector> selection_vector; |
| 315 | status = SelectionVector::MakeInt16(num_records, pool, &selection_vector); |
| 316 | ASSERT_TRUE(status.ok()) << status.message(); |
| 317 | |
| 318 | // Evaluate expression |
| 319 | status = filter->Evaluate(*in_batch, selection_vector); |
| 320 | ASSERT_TRUE(status.ok()) << status.message(); |
| 321 | |
| 322 | // Validate results |
| 323 | EXPECT_ARROW_ARRAY_EQUALS(exp, selection_vector->ToArray()); |
| 324 | } |
| 325 | |
| 326 | TEST_F(TestIn, TestInDate32) { |
| 327 | TestDateOrTimeInExpression<arrow::Date32Type>( |