Skip to content

Commit 58d532d

Browse files
Alex-PLACETAlex-PLACET
committed
wip
1 parent 293bf3e commit 58d532d

File tree

1 file changed

+127
-97
lines changed

1 file changed

+127
-97
lines changed

‎test/test_polars_integration.py‎

Lines changed: 127 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -26,20 +26,20 @@ def setup_module_path():
2626
if module_dir.exists():
2727
sys.path.insert(0, str(module_dir))
2828
return
29-
29+
3030
# Try to find in build directory
3131
test_dir = Path(__file__).parent
3232
build_dirs = [
3333
test_dir.parent / "build" / "bin" / "Debug",
3434
test_dir.parent / "build" / "bin" / "Release",
3535
test_dir.parent / "build" / "bin",
3636
]
37-
37+
3838
for build_dir in build_dirs:
3939
if build_dir.exists():
4040
sys.path.insert(0, str(build_dir))
4141
return
42-
42+
4343
raise ImportError(
4444
"Could not find test_polars_helper module. "
4545
"Build the project first or set TEST_POLARS_HELPER_PATH."
@@ -53,103 +53,133 @@ def setup_module_path():
5353
import test_polars_helper # noqa: E402
5454

5555

56-
def test_create_array_in_cpp():
56+
# =============================================================================
57+
# Test 1: C++ -> Python (Create array in sparrow, import to Polars)
58+
# =============================================================================
59+
60+
61+
class TestCppToPython:
5762
"""Test creating an array in C++ (sparrow) and importing to Python/Polars."""
58-
print("\n" + "=" * 70)
59-
print("Test 1: C++ -> Python (Create array in sparrow, import to Polars)")
60-
print("=" * 70)
61-
62-
print("\n1. Creating test array in C++ using sparrow::pycapsule...")
63-
schema_capsule, array_capsule = test_polars_helper.create_test_array_capsules()
64-
65-
assert schema_capsule is not None, "Received null schema capsule from C++"
66-
assert array_capsule is not None, "Received null array capsule from C++"
67-
print(" PyCapsules created by sparrow::pycapsule::export_array_to_capsules()")
68-
69-
print("\n2. Importing to PyArrow...")
70-
arrow_array = pa.Array._import_from_c_capsule(schema_capsule, array_capsule)
71-
print(f" Arrow type: {arrow_array.type}")
72-
print(f" Arrow values: {arrow_array.to_pylist()}")
73-
74-
print("\n3. Converting to Polars...")
75-
polars_series = pl.from_arrow(arrow_array)
76-
print(f" Polars series: {polars_series.to_list()}")
77-
78-
expected = [10, 20, None, 40, 50]
79-
actual = polars_series.to_list()
80-
81-
assert expected == actual, f"Data mismatch! Expected: {expected}, Actual: {actual}"
82-
print(" Data matches expected values!")
83-
print("\n" + "=" * 70)
84-
print("Test 1 PASSED")
85-
print("=" * 70)
86-
87-
88-
def test_polars_to_cpp():
63+
64+
@pytest.fixture(autouse=True)
65+
def setup(self):
66+
"""Create test array capsules from C++."""
67+
self.schema_capsule, self.array_capsule = (
68+
test_polars_helper.create_test_array_capsules()
69+
)
70+
71+
def test_step1_create_capsules_in_cpp(self):
72+
"""Step 1: Create PyCapsules in C++ using sparrow::pycapsule."""
73+
assert self.schema_capsule is not None, "Received null schema capsule from C++"
74+
assert self.array_capsule is not None, "Received null array capsule from C++"
75+
76+
def test_step2_import_to_pyarrow(self):
77+
"""Step 2: Import PyCapsules to PyArrow."""
78+
arrow_array = pa.Array._import_from_c_capsule(
79+
self.schema_capsule, self.array_capsule
80+
)
81+
assert arrow_array.type == pa.int32(), f"Expected int32, got {arrow_array.type}"
82+
assert arrow_array.to_pylist() == [10, 20, None, 40, 50]
83+
84+
def test_step3_convert_to_polars(self):
85+
"""Step 3: Convert PyArrow array to Polars series."""
86+
arrow_array = pa.Array._import_from_c_capsule(
87+
self.schema_capsule, self.array_capsule
88+
)
89+
polars_series = pl.from_arrow(arrow_array)
90+
91+
expected = [10, 20, None, 40, 50]
92+
actual = polars_series.to_list()
93+
assert expected == actual, f"Data mismatch! Expected: {expected}, Actual: {actual}"
94+
95+
96+
# =============================================================================
97+
# Test 2: Python -> C++ (Export Polars to sparrow)
98+
# =============================================================================
99+
100+
101+
class TestPythonToCpp:
89102
"""Test exporting Polars data to C++ (sparrow)."""
90-
print("\n" + "=" * 70)
91-
print("Test 2: Python -> C++ (Export Polars to sparrow)")
92-
print("=" * 70)
93-
94-
print("\n1. Creating Polars series...")
95-
test_series = pl.Series([100, 200, None, 400, 500], dtype=pl.Int32)
96-
print(f" Polars series: {test_series.to_list()}")
97-
98-
print("\n2. Exporting to Arrow PyCapsules...")
99-
arrow_array = test_series.to_arrow()
100-
schema_capsule, array_capsule = arrow_array.__arrow_c_array__()
101-
print(" PyCapsules created by PyArrow")
102-
103-
print("\n3. Importing and verifying in sparrow using sparrow::pycapsule...")
104-
result = test_polars_helper.verify_array_size_from_capsules(schema_capsule, array_capsule, 5)
105-
106-
assert result is True, "C++ verification failed"
107-
print(" sparrow::pycapsule::import_array_from_capsules() succeeded!")
108-
print(" sparrow successfully imported and verified the array!")
109-
print("\n" + "=" * 70)
110-
print("Test 2 PASSED")
111-
print("=" * 70)
112-
113-
114-
def test_roundtrip():
103+
104+
@pytest.fixture(autouse=True)
105+
def setup(self):
106+
"""Create Polars series and export to capsules."""
107+
self.test_series = pl.Series([100, 200, None, 400, 500], dtype=pl.Int32)
108+
self.arrow_array = self.test_series.to_arrow()
109+
self.schema_capsule, self.array_capsule = self.arrow_array.__arrow_c_array__()
110+
111+
def test_step1_create_polars_series(self):
112+
"""Step 1: Create Polars series."""
113+
assert self.test_series.to_list() == [100, 200, None, 400, 500]
114+
assert self.test_series.dtype == pl.Int32
115+
116+
def test_step2_export_to_capsules(self):
117+
"""Step 2: Export Polars series to Arrow PyCapsules."""
118+
assert self.schema_capsule is not None, "Schema capsule is None"
119+
assert self.array_capsule is not None, "Array capsule is None"
120+
121+
def test_step3_import_in_sparrow(self):
122+
"""Step 3: Import and verify in sparrow using sparrow::pycapsule."""
123+
result = test_polars_helper.verify_array_size_from_capsules(
124+
self.schema_capsule, self.array_capsule, 5
125+
)
126+
assert result is True, "C++ verification failed"
127+
128+
129+
# =============================================================================
130+
# Test 3: Round-trip (Python -> sparrow -> Python)
131+
# =============================================================================
132+
133+
134+
class TestRoundtrip:
115135
"""Test round-trip: Python -> C++ (sparrow) -> Python."""
116-
print("\n" + "=" * 70)
117-
print("Test 3: Round-trip (Python -> sparrow -> Python)")
118-
print("=" * 70)
119-
120-
print("\n1. Creating Polars series...")
121-
original_series = pl.Series([1, 2, None, 4, 5], dtype=pl.Int32)
122-
print(f" Original: {original_series.to_list()}")
123-
124-
print("\n2. Exporting to Arrow PyCapsules...")
125-
arrow_array = original_series.to_arrow()
126-
schema_capsule_in, array_capsule_in = arrow_array.__arrow_c_array__()
127-
print(" PyCapsules created by PyArrow")
128-
129-
print("\n3. Round-tripping through sparrow using sparrow::pycapsule...")
130-
schema_capsule_out, array_capsule_out = test_polars_helper.roundtrip_array_capsules(
131-
schema_capsule_in,
132-
array_capsule_in
133-
)
134-
135-
assert schema_capsule_out is not None, "Received null schema capsule from C++"
136-
assert array_capsule_out is not None, "Received null array capsule from C++"
137-
print(" sparrow::pycapsule import/export succeeded!")
138-
139-
print("\n4. Importing back to Python...")
140-
arrow_array_out = pa.Array._import_from_c_capsule(schema_capsule_out, array_capsule_out)
141-
result_series = pl.from_arrow(arrow_array_out)
142-
print(f" Result: {result_series.to_list()}")
143-
144-
original_data = original_series.to_list()
145-
result_data = result_series.to_list()
146-
assert original_data == result_data, f"Data mismatch! Original: {original_data}, Result: {result_data}"
147-
148-
print(" Round-trip successful - data matches!")
149-
print("\n" + "=" * 70)
150-
print("Test 3 PASSED")
151-
print("=" * 70)
136+
137+
@pytest.fixture(autouse=True)
138+
def setup(self):
139+
"""Create original series and export to capsules."""
140+
self.original_series = pl.Series([1, 2, None, 4, 5], dtype=pl.Int32)
141+
self.arrow_array = self.original_series.to_arrow()
142+
self.schema_capsule_in, self.array_capsule_in = (
143+
self.arrow_array.__arrow_c_array__()
144+
)
145+
146+
def test_step1_create_original_series(self):
147+
"""Step 1: Create original Polars series."""
148+
assert self.original_series.to_list() == [1, 2, None, 4, 5]
149+
150+
def test_step2_export_to_capsules(self):
151+
"""Step 2: Export to Arrow PyCapsules."""
152+
assert self.schema_capsule_in is not None
153+
assert self.array_capsule_in is not None
154+
155+
def test_step3_roundtrip_through_sparrow(self):
156+
"""Step 3: Round-trip through sparrow using sparrow::pycapsule."""
157+
schema_capsule_out, array_capsule_out = (
158+
test_polars_helper.roundtrip_array_capsules(
159+
self.schema_capsule_in, self.array_capsule_in
160+
)
161+
)
162+
assert schema_capsule_out is not None, "Received null schema capsule from C++"
163+
assert array_capsule_out is not None, "Received null array capsule from C++"
164+
165+
def test_step4_import_back_to_python(self):
166+
"""Step 4: Import back to Python and verify data matches."""
167+
schema_capsule_out, array_capsule_out = (
168+
test_polars_helper.roundtrip_array_capsules(
169+
self.schema_capsule_in, self.array_capsule_in
170+
)
171+
)
172+
arrow_array_out = pa.Array._import_from_c_capsule(
173+
schema_capsule_out, array_capsule_out
174+
)
175+
result_series = pl.from_arrow(arrow_array_out)
176+
177+
original_data = self.original_series.to_list()
178+
result_data = result_series.to_list()
179+
assert (
180+
original_data == result_data
181+
), f"Data mismatch! Original: {original_data}, Result: {result_data}"
152182

153183

154184
if __name__ == "__main__":
155-
sys.exit(pytest.main([__file__, "-v", "-s"]))
185+
sys.exit(pytest.main([__file__, "-v"]))

0 commit comments

Comments
 (0)