Skip to content

Commit 4068d0c

Browse files
committed
Tests for Arrow batch adapter scalar type conversion
Add tests for record_batches_to_struct and struct_to_record_batches covering all scalar types, nested structs, schema validation, and round-trip conversion. Add RecordBatch pull input adapter tests. Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com>
1 parent e3adf9c commit 4068d0c

2 files changed

Lines changed: 1080 additions & 1 deletion

File tree

csp/tests/adapters/test_arrow.py

Lines changed: 62 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import pytest
77

88
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
1010

1111
_STARTTIME = datetime(2020, 1, 1, 9, 0, 0)
1212

@@ -405,3 +405,64 @@ def __iter__(self):
405405
results = csp.run(G_lazy_schema, "TsCol", lazy_iter, False, starttime=_STARTTIME)
406406
assert len(results["data"]) == 2
407407
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

Comments
 (0)