Skip to content

Commit 153b5f1

Browse files
committed
Add assertions for UUID extension type in tests
Ensure collected UUID results are extension arrays or chunked arrays with the UUID extension type before comparison to expected values, preserving end-to-end metadata validation.
1 parent bae5d54 commit 153b5f1

File tree

1 file changed

+5
-2
lines changed

1 file changed

+5
-2
lines changed

python/tests/test_udf.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,10 @@ def ensure_extension(values: pa.Array) -> pa.Array:
184184

185185
expected = uuid_type.wrap_array(storage)
186186

187-
if isinstance(result, pa.Array) and result.type.equals(uuid_type.storage_type):
188-
result = uuid_type.wrap_array(result)
187+
if isinstance(result, pa.ChunkedArray):
188+
assert result.type.equals(uuid_type)
189+
else:
190+
assert isinstance(result, pa.ExtensionArray)
191+
assert result.type.equals(uuid_type)
189192

190193
assert result.equals(expected)

0 commit comments

Comments
 (0)