|
6 | 6 | import pytest |
7 | 7 |
|
8 | 8 | import csp |
9 | | -from csp.adapters.arrow import RecordBatchPullInputAdapter, write_record_batches |
| 9 | +from csp.adapters.arrow import RecordBatchPullInputAdapter, struct_to_record_batches, write_record_batches |
10 | 10 |
|
11 | 11 | _STARTTIME = datetime(2020, 1, 1, 9, 0, 0) |
12 | 12 |
|
@@ -405,3 +405,64 @@ def __iter__(self): |
405 | 405 | results = csp.run(G_lazy_schema, "TsCol", lazy_iter, False, starttime=_STARTTIME) |
406 | 406 | assert len(results["data"]) == 2 |
407 | 407 | assert [len(r[1][0]) for r in results["data"]] == [5, 3] |
| 408 | + |
| 409 | + |
| 410 | +class TestStructConversion: |
| 411 | + """Tests for struct_to_record_batches conversion.""" |
| 412 | + |
| 413 | + def test_nested_struct_with_nulls(self): |
| 414 | + """Regression test: NestedStructWriter::writeNull must not double-append to child builders. |
| 415 | +
|
| 416 | + When a parent struct has a null nested struct field, the children of the |
| 417 | + struct column must still have exactly one entry per row. A prior bug caused |
| 418 | + AppendNull() to also append empty values to children (on top of the explicit |
| 419 | + writeNull per child), shifting all subsequent valid values to wrong indices. |
| 420 | + """ |
| 421 | + |
| 422 | + class Inner(csp.Struct): |
| 423 | + x: int |
| 424 | + y: float |
| 425 | + |
| 426 | + class Outer(csp.Struct): |
| 427 | + name: str |
| 428 | + inner: Inner |
| 429 | + |
| 430 | + @csp.node |
| 431 | + def validate_batches(batches: csp.ts[object]) -> csp.ts[object]: |
| 432 | + if csp.ticked(batches): |
| 433 | + return batches |
| 434 | + |
| 435 | + @csp.graph |
| 436 | + def g(): |
| 437 | + s1 = Outer(name="a", inner=Inner(x=1, y=1.5)) |
| 438 | + s2 = Outer(name="b") # inner is unset -> null |
| 439 | + s3 = Outer(name="c", inner=Inner(x=3, y=3.5)) |
| 440 | + s4 = Outer(name="d") # inner is unset -> null |
| 441 | + s5 = Outer(name="e", inner=Inner(x=5, y=5.5)) |
| 442 | + |
| 443 | + structs = csp.const([s1, s2, s3, s4, s5]) |
| 444 | + batches = struct_to_record_batches(structs, Outer, max_batch_size=100) |
| 445 | + out = validate_batches(batches) |
| 446 | + csp.add_graph_output("batches", out) |
| 447 | + |
| 448 | + results = csp.run(g, starttime=_STARTTIME, endtime=_STARTTIME + timedelta(seconds=1)) |
| 449 | + batch_list = results["batches"][0][1] |
| 450 | + batch = batch_list[0] |
| 451 | + |
| 452 | + assert batch.num_rows == 5 |
| 453 | + |
| 454 | + struct_col = batch.column("inner") |
| 455 | + x_child = struct_col.field("x") |
| 456 | + y_child = struct_col.field("y") |
| 457 | + |
| 458 | + # Child arrays must have exactly the same length as the struct column |
| 459 | + assert len(x_child) == 5, f"child x length {len(x_child)} != struct length 5" |
| 460 | + assert len(y_child) == 5, f"child y length {len(y_child)} != struct length 5" |
| 461 | + |
| 462 | + # Verify actual values are correct (not shifted by double-append) |
| 463 | + values = struct_col.to_pylist() |
| 464 | + assert values[0] == {"x": 1, "y": 1.5} |
| 465 | + assert values[1] is None |
| 466 | + assert values[2] == {"x": 3, "y": 3.5} |
| 467 | + assert values[3] is None |
| 468 | + assert values[4] == {"x": 5, "y": 5.5} |
0 commit comments