-
Notifications
You must be signed in to change notification settings - Fork 1.2k
Add support for using ListView arrays and types through FFI #8822
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 4 commits
d0980b1
ff3dc75
7375a94
c0e5df9
cedca3b
6c0e7a4
5b00813
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -620,6 +620,16 @@ impl ArrayData { | |
| vec![ArrayData::new_null(f.data_type(), *list_len as usize * len)], | ||
| true, | ||
| ), | ||
| DataType::ListView(f) => ( | ||
| vec![zeroed(len * 4), zeroed(len * 4)], | ||
| vec![ArrayData::new_empty(f.data_type())], | ||
| true, | ||
| ), | ||
| DataType::LargeListView(f) => ( | ||
| vec![zeroed(len * 8), zeroed(len * 8)], | ||
| vec![ArrayData::new_empty(f.data_type())], | ||
| true, | ||
| ), | ||
| DataType::Struct(fields) => ( | ||
| vec![], | ||
| fields | ||
|
|
@@ -1783,7 +1793,7 @@ impl DataTypeLayout { | |
| }, | ||
| ], | ||
| can_contain_null_mask: true, | ||
| variadic: true, | ||
| variadic: false, | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This is a bug, see the spec, list views aren't represented by a variable number of buffers. |
||
| } | ||
| } | ||
| } | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -27,7 +27,9 @@ | |
|
|
||
| import arrow_pyarrow_integration_testing as rust | ||
|
|
||
| PYARROW_PRE_14 = int(pa.__version__.split('.')[0]) < 14 | ||
| PYARROW_MAJOR_VER = int(pa.__version__.split(".")[0]) | ||
| PYARROW_PRE_14 = PYARROW_MAJOR_VER < 14 | ||
| PYARROW_PRE_16 = PYARROW_MAJOR_VER < 16 | ||
|
|
||
|
|
||
| @contextlib.contextmanager | ||
|
|
@@ -112,8 +114,16 @@ def assert_pyarrow_leak(): | |
| ), | ||
| ] | ||
|
|
||
| _unsupported_pyarrow_types = [ | ||
| ] | ||
| if PYARROW_MAJOR_VER >= 16: | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Did this to get everything green, but I think that it might be a good opportunity to only test with pyarrow 16 and above, which supports all relevant features here -
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can you file a follow on ticket / PR to discuss dropping support for versions less than 16?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. will do
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||
| _supported_pyarrow_types.extend( | ||
| [ | ||
| pa.list_view(pa.uint64()), | ||
| pa.large_list_view(pa.uint64()), | ||
| pa.list_view(pa.string()), | ||
| pa.large_list_view(pa.string()), | ||
| ] | ||
| ) | ||
|
|
||
|
|
||
| # As of pyarrow 14, pyarrow implements the Arrow PyCapsule interface | ||
| # (https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html). | ||
|
|
@@ -158,12 +168,6 @@ def test_type_roundtrip_pycapsule(pyarrow_type): | |
| assert restored == pyarrow_type | ||
| assert restored is not pyarrow_type | ||
|
|
||
|
|
||
| @pytest.mark.parametrize("pyarrow_type", _unsupported_pyarrow_types, ids=str) | ||
| def test_type_roundtrip_raises(pyarrow_type): | ||
| with pytest.raises(pa.ArrowException): | ||
| rust.round_trip_type(pyarrow_type) | ||
|
|
||
| @pytest.mark.parametrize('pyarrow_type', _supported_pyarrow_types, ids=str) | ||
| def test_field_roundtrip(pyarrow_type): | ||
| pyarrow_field = pa.field("test", pyarrow_type, nullable=True) | ||
|
|
@@ -337,6 +341,21 @@ def test_list_array(): | |
| del a | ||
| del b | ||
|
|
||
|
|
||
| @pytest.mark.skipif(PYARROW_PRE_16, reason="requires pyarrow 16") | ||
| def test_list_view_array(): | ||
| """ | ||
| Python -> Rust -> Python | ||
| """ | ||
| a = pa.array([[], None, [1, 2], [4, 5, 6]], pa.list_view(pa.int64())) | ||
| b = rust.round_trip_array(a) | ||
| b.validate(full=True) | ||
| assert a.to_pylist() == b.to_pylist() | ||
| assert a.type == b.type | ||
| del a | ||
| del b | ||
|
|
||
|
|
||
| def test_map_array(): | ||
| """ | ||
| Python -> Rust -> Python | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -456,6 +456,14 @@ impl TryFrom<&FFI_ArrowSchema> for DataType { | |
| let c_child = c_schema.child(0); | ||
| DataType::LargeList(Arc::new(Field::try_from(c_child)?)) | ||
| } | ||
| "+vl" => { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Double checked that this matches https://arrow.apache.org/docs/format/CDataInterface.html#data-type-description-format-strings ✅ |
||
| let c_child = c_schema.child(0); | ||
| DataType::ListView(Arc::new(Field::try_from(c_child)?)) | ||
| } | ||
| "+vL" => { | ||
| let c_child = c_schema.child(0); | ||
| DataType::LargeListView(Arc::new(Field::try_from(c_child)?)) | ||
| } | ||
| "+s" => { | ||
| let fields = c_schema.children().map(Field::try_from); | ||
| DataType::Struct(fields.collect::<Result<_, ArrowError>>()?) | ||
|
|
@@ -657,6 +665,8 @@ impl TryFrom<&DataType> for FFI_ArrowSchema { | |
| let children = match dtype { | ||
| DataType::List(child) | ||
| | DataType::LargeList(child) | ||
| | DataType::ListView(child) | ||
| | DataType::LargeListView(child) | ||
| | DataType::FixedSizeList(child, _) | ||
| | DataType::Map(child, _) => { | ||
| vec![FFI_ArrowSchema::try_from(child.as_ref())?] | ||
|
|
@@ -746,6 +756,8 @@ fn get_format_string(dtype: &DataType) -> Result<Cow<'static, str>, ArrowError> | |
| DataType::Interval(IntervalUnit::MonthDayNano) => Ok("tin".into()), | ||
| DataType::List(_) => Ok("+l".into()), | ||
| DataType::LargeList(_) => Ok("+L".into()), | ||
| DataType::ListView(_) => Ok("+vl".into()), | ||
| DataType::LargeListView(_) => Ok("+vL".into()), | ||
| DataType::Struct(_) => Ok("+s".into()), | ||
| DataType::Map(_, _) => Ok("+m".into()), | ||
| DataType::RunEndEncoded(_, _) => Ok("+r".into()), | ||
|
|
@@ -874,6 +886,16 @@ mod tests { | |
| DataType::Int16, | ||
| false, | ||
| )))); | ||
| round_trip_type(DataType::ListView(Arc::new(Field::new( | ||
| "a", | ||
| DataType::Int16, | ||
| false, | ||
| )))); | ||
| round_trip_type(DataType::LargeListView(Arc::new(Field::new( | ||
| "a", | ||
| DataType::Int16, | ||
| false, | ||
| )))); | ||
| round_trip_type(DataType::Struct(Fields::from(vec![Field::new( | ||
| "a", | ||
| DataType::Utf8, | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just nicer when working with IDEs.