Add dictionary support for C data interface #1407

sunchao · 2022-03-08T04:02:15Z

Which issue does this PR close?

Closes #1397.

Rationale for this change

Currently the Rust implementation of C data interface doesn't support dictionary type yet, which is necessary if we want to pass data between Rust and other languages such as Java/C++.

What changes are included in this PR?

Added dictionary support in Rust FFI.
Added roundtrip tests within Rust (i.e., Rust -> Rust -> Rust)
Added roundtrip tests between Python and Rust

I kept the DictionaryArray untouched so the dictionary (i.e., values) is still stored as the only child of the ArrayData it is associated to.

Are there any user-facing changes?

Yes, now we can use C data interface with dictionary type. There is one API change on ArrowSchema.try_new:

    pub fn try_new(
        format: &str,
        children: Vec<FFI_ArrowSchema>,
        dictionary: Option<FFI_ArrowSchema>,
    ) -> Result<Self>

This PR added one parameter dictionary.

codecov-commenter · 2022-03-08T05:06:04Z

Codecov Report

Merging #1407 (7ed4d37) into master (1fbe618) will increase coverage by 0.02%.
The diff coverage is 82.60%.

@@            Coverage Diff             @@
##           master    #1407      +/-   ##
==========================================
+ Coverage   83.13%   83.16%   +0.02%     
==========================================
  Files         182      182              
  Lines       53321    53394      +73     
==========================================
+ Hits        44330    44404      +74     
+ Misses       8991     8990       -1

Impacted Files	Coverage Δ
arrow/src/datatypes/ffi.rs	`73.65% <68.75%> (+2.80%)`	⬆️
arrow/src/ffi.rs	`85.78% <90.19%> (+1.16%)`	⬆️
arrow/src/array/ffi.rs	`100.00% <100.00%> (ø)`
arrow/src/datatypes/datatype.rs	`66.40% <0.00%> (-0.40%)`	⬇️
arrow/src/datatypes/field.rs	`53.79% <0.00%> (-0.31%)`	⬇️
parquet_derive/src/parquet_field.rs	`65.98% <0.00%> (-0.23%)`	⬇️
parquet/src/encodings/encoding.rs	`93.52% <0.00%> (-0.20%)`	⬇️
arrow/src/array/array_binary.rs	`93.28% <0.00%> (+0.05%)`	⬆️
arrow/src/array/transform/mod.rs	`86.43% <0.00%> (+0.11%)`	⬆️
arrow/src/array/equal/mod.rs	`93.24% <0.00%> (+0.33%)`	⬆️
... and 1 more

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 1fbe618...7ed4d37. Read the comment docs.

sunchao · 2022-03-08T19:50:06Z

cc @alamb @jorgecarleitao @kszucs could you help to review this? thanks!

sunchao · 2022-03-08T19:56:28Z

We should also be able to pass is_ordered flag via the C data interface, but it appears that property is not supported in the Rust implementation yet (perhaps it should be part of DataType::Dictionary).

alamb

Looks good to me -- thank you @sunchao

I wonder if @viirya may have some time to review this PR as well as he has been working in the FFI interface recently as well

alamb · 2022-03-08T20:02:00Z

arrow-pyarrow-integration-testing/tests/test_sql.py

@@ -122,14 +123,6 @@ def test_type_roundtrip_raises(pyarrow_type):
    with pytest.raises(pa.ArrowException):
        rust.round_trip_type(pyarrow_type)

-
-def test_dictionary_type_roundtrip():


alamb · 2022-03-08T20:04:31Z

arrow/src/ffi.rs

+        // verify
+        let new_values = vec!["a", "aaa", "aaa", "a", "aaa", "aaa"];
+        let expected: DictionaryArray<Int8Type> = new_values.into_iter().collect();
+        assert_eq!(actual, &expected);


viirya · 2022-03-08T21:21:08Z

Yea, I will review this PR. Thanks @alamb @sunchao

jorgecarleitao

Thanks a lot @sunchao . Overall looks great! Left 3 comments below.

jorgecarleitao · 2022-03-08T20:01:38Z

arrow/src/ffi.rs

-        let data_type = &self.data_type()?;
+        // Special handling for dictionary type as we only care about the key type in the case.
+        let data_type = match &self.data_type()? {
+            DataType::Dictionary(key_data_type, _) => key_data_type.as_ref().clone(),


maybe it would be possible to not clone here by returning a reference?

Yea we can do this. I need to move &self.data_type()? into a separate assignment though since otherwise I'm getting "temporary value dropped while borrowed" error.

jorgecarleitao · 2022-03-08T21:19:03Z

arrow/src/array/ffi.rs

@@ -127,4 +128,14 @@ mod tests {
        let data = array.data();
        test_round_trip(data)
    }
+
+    #[test]
+    fn test_dictionary() -> Result<()> {


would it be worth to have an example with validity (in both the keys and values), so that we cover the most complex use-case?

Good idea. Will add.

jorgecarleitao · 2022-03-08T21:19:49Z

arrow-pyarrow-integration-testing/tests/test_sql.py

+    """
+    Python -> Rust -> Python
+    """
+    a = pa.array(["a", "b", "a"], type=pa.dictionary(pa.int8(), pa.string()))


same here - with validities?

viirya · 2022-03-08T21:41:02Z

arrow/src/ffi.rs

            .map(|i| {
                let child = self.child(i);
                child.to_data()
            })
            .map(|d| d.unwrap())
            .collect();

+        if let Some(d) = self.dictionary() {
+            // For dictionary type there should only be a single child, so we don't need to worry if


Should we add assert for that?

viirya

Looks good.

viirya · 2022-03-08T22:10:57Z

Yes, now we can use C data interface with dictionary type. No API change.

FFI_ArrowSchema.try_new is a public function, so I wonder this should be an API change?

sunchao · 2022-03-08T23:00:13Z

FFI_ArrowSchema.try_new is a public function, so I wonder this should be an API change?

Hmm you are right. Looks like try_new doesn't have to be public though. I think people should just use ArrowArray directly.

sunchao · 2022-03-09T01:00:39Z

Thanks @jorgecarleitao @alamb @viirya ! I've addressed your comments. Let me know if you have more feedback. Otherwise I'll merge this tonight.

sunchao · 2022-03-09T07:27:52Z

Merged, thanks!

initial commit

5ee2bf7

github-actions bot added the arrow Changes to the arrow crate label Mar 8, 2022

add integration tests for python

b86a377

alamb approved these changes Mar 8, 2022

View reviewed changes

jorgecarleitao reviewed Mar 8, 2022

View reviewed changes

viirya reviewed Mar 8, 2022

View reviewed changes

viirya approved these changes Mar 8, 2022

View reviewed changes

address comments

7ed4d37

jorgecarleitao approved these changes Mar 9, 2022

View reviewed changes

sunchao merged commit f19d1ed into apache:master Mar 9, 2022

sunchao deleted the ffi-dict branch March 9, 2022 07:27

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add dictionary support for C data interface #1407

Add dictionary support for C data interface #1407

sunchao commented Mar 8, 2022 •

edited

Loading

codecov-commenter commented Mar 8, 2022 •

edited

Loading

sunchao commented Mar 8, 2022

sunchao commented Mar 8, 2022

alamb left a comment

alamb Mar 8, 2022

alamb Mar 8, 2022

viirya commented Mar 8, 2022

jorgecarleitao left a comment

jorgecarleitao Mar 8, 2022

sunchao Mar 8, 2022

jorgecarleitao Mar 8, 2022

sunchao Mar 8, 2022

jorgecarleitao Mar 8, 2022

viirya Mar 8, 2022

sunchao Mar 8, 2022

viirya left a comment

viirya commented Mar 8, 2022

sunchao commented Mar 8, 2022

sunchao commented Mar 9, 2022

sunchao commented Mar 9, 2022

Add dictionary support for C data interface #1407

Add dictionary support for C data interface #1407

Conversation

sunchao commented Mar 8, 2022 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

codecov-commenter commented Mar 8, 2022 • edited Loading

Codecov Report

sunchao commented Mar 8, 2022

sunchao commented Mar 8, 2022

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

viirya commented Mar 8, 2022

jorgecarleitao left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

viirya left a comment

Choose a reason for hiding this comment

viirya commented Mar 8, 2022

sunchao commented Mar 8, 2022

sunchao commented Mar 9, 2022

sunchao commented Mar 9, 2022

sunchao commented Mar 8, 2022 •

edited

Loading

codecov-commenter commented Mar 8, 2022 •

edited

Loading