Skip to content

Commit b1d18a8

Browse files
committed
test: add validation for schema mismatch in Arrow C stream
1 parent f7a2407 commit b1d18a8

File tree

2 files changed

+40
-1
lines changed

2 files changed

+40
-1
lines changed

python/tests/test_dataframe.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1675,6 +1675,29 @@ def test_arrow_c_stream_schema_selection(fail_collect):
16751675
assert batches[0].equals(expected_batch)
16761676

16771677

1678+
def test_arrow_c_stream_schema_mismatch(fail_collect):
1679+
ctx = SessionContext()
1680+
1681+
batch = pa.RecordBatch.from_arrays(
1682+
[pa.array([1, 2]), pa.array([3, 4])], names=["a", "b"]
1683+
)
1684+
df = ctx.create_dataframe([[batch]])
1685+
1686+
bad_schema = pa.schema([("a", pa.string())])
1687+
1688+
c_schema = pa_cffi.ffi.new("struct ArrowSchema*")
1689+
address = int(pa_cffi.ffi.cast("uintptr_t", c_schema))
1690+
bad_schema._export_to_c(address)
1691+
1692+
capsule_new = ctypes.pythonapi.PyCapsule_New
1693+
capsule_new.restype = ctypes.py_object
1694+
capsule_new.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p]
1695+
bad_capsule = capsule_new(ctypes.c_void_p(address), b"arrow_schema", None)
1696+
1697+
with pytest.raises(Exception, match="Fail to merge schema"):
1698+
df.__arrow_c_stream__(bad_capsule)
1699+
1700+
16781701
def test_to_pylist(df):
16791702
# Convert datafusion dataframe to Python list
16801703
pylist = df.to_pylist()

python/tests/utils.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,26 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
118
"""Testing-only helpers for datafusion-python.
219
320
This module contains utilities used by the test-suite that should not be
421
exposed as part of the public API. Keep the implementation minimal and
522
documented so reviewers can easily see it's test-only.
623
"""
7-
824
from __future__ import annotations
925

1026
from typing import TYPE_CHECKING

0 commit comments

Comments
 (0)