()
| 313 | |
| 314 | |
| 315 | def test_hash_join(): |
| 316 | left = pa.table({'key': [1, 2, 3], 'a': [4, 5, 6]}) |
| 317 | left_source = Declaration("table_source", options=TableSourceNodeOptions(left)) |
| 318 | right = pa.table({'key': [2, 3, 4], 'b': [4, 5, 6]}) |
| 319 | right_source = Declaration("table_source", options=TableSourceNodeOptions(right)) |
| 320 | |
| 321 | # inner join |
| 322 | join_opts = HashJoinNodeOptions("inner", left_keys="key", right_keys="key") |
| 323 | joined = Declaration( |
| 324 | "hashjoin", options=join_opts, inputs=[left_source, right_source]) |
| 325 | result = joined.to_table() |
| 326 | expected = pa.table( |
| 327 | [[2, 3], [5, 6], [2, 3], [4, 5]], |
| 328 | names=["key", "a", "key", "b"]) |
| 329 | assert result.equals(expected) |
| 330 | |
| 331 | for keys in [field("key"), ["key"], [field("key")]]: |
| 332 | join_opts = HashJoinNodeOptions("inner", left_keys=keys, right_keys=keys) |
| 333 | joined = Declaration( |
| 334 | "hashjoin", options=join_opts, inputs=[left_source, right_source]) |
| 335 | result = joined.to_table() |
| 336 | assert result.equals(expected) |
| 337 | |
| 338 | # left join |
| 339 | join_opts = HashJoinNodeOptions( |
| 340 | "left outer", left_keys="key", right_keys="key") |
| 341 | joined = Declaration( |
| 342 | "hashjoin", options=join_opts, inputs=[left_source, right_source]) |
| 343 | result = joined.to_table() |
| 344 | expected = pa.table( |
| 345 | [[1, 2, 3], [4, 5, 6], [None, 2, 3], [None, 4, 5]], |
| 346 | names=["key", "a", "key", "b"] |
| 347 | ) |
| 348 | assert result.sort_by("a").equals(expected) |
| 349 | |
| 350 | # suffixes |
| 351 | join_opts = HashJoinNodeOptions( |
| 352 | "left outer", left_keys="key", right_keys="key", |
| 353 | output_suffix_for_left="_left", output_suffix_for_right="_right") |
| 354 | joined = Declaration( |
| 355 | "hashjoin", options=join_opts, inputs=[left_source, right_source]) |
| 356 | result = joined.to_table() |
| 357 | expected = pa.table( |
| 358 | [[1, 2, 3], [4, 5, 6], [None, 2, 3], [None, 4, 5]], |
| 359 | names=["key_left", "a", "key_right", "b"] |
| 360 | ) |
| 361 | assert result.sort_by("a").equals(expected) |
| 362 | |
| 363 | # manually specifying output columns |
| 364 | join_opts = HashJoinNodeOptions( |
| 365 | "left outer", left_keys="key", right_keys="key", |
| 366 | left_output=["key", "a"], right_output=[field("b")]) |
| 367 | joined = Declaration( |
| 368 | "hashjoin", options=join_opts, inputs=[left_source, right_source]) |
| 369 | result = joined.to_table() |
| 370 | expected = pa.table( |
| 371 | [[1, 2, 3], [4, 5, 6], [None, 4, 5]], |
| 372 | names=["key", "a", "b"] |
nothing calls this directly
no test coverage detected