support duration in ffi

apache · May 11, 2022 · 9c36b4e · 9c36b4e
1 parent b9a41f3
commit 9c36b4e
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 5 deletions.
diff --git a/arrow-pyarrow-integration-testing/tests/test_sql.py b/arrow-pyarrow-integration-testing/tests/test_sql.py
@@ -55,6 +55,10 @@ def assert_pyarrow_leak():
     pa.timestamp("us"),
     pa.timestamp("us", tz="UTC"),
     pa.timestamp("us", tz="Europe/Paris"),
+    pa.duration("s"),
+    pa.duration("ms"),
+    pa.duration("us"),
+    pa.duration("ns"),
     pa.float16(),
     pa.float32(),
     pa.float64(),
@@ -86,7 +90,6 @@ def assert_pyarrow_leak():
 
 _unsupported_pyarrow_types = [
     pa.decimal256(76, 38),
-    pa.duration("s"),
     pa.map_(pa.string(), pa.int32()),
     pa.union(
         [pa.field("a", pa.binary(10)), pa.field("b", pa.string())],

diff --git a/arrow/src/datatypes/ffi.rs b/arrow/src/datatypes/ffi.rs
@@ -52,6 +52,10 @@ impl TryFrom<&FFI_ArrowSchema> for DataType {
             "ttm" => DataType::Time32(TimeUnit::Millisecond),
             "ttu" => DataType::Time64(TimeUnit::Microsecond),
             "ttn" => DataType::Time64(TimeUnit::Nanosecond),
+            "tDs" => DataType::Duration(TimeUnit::Second),
+            "tDm" => DataType::Duration(TimeUnit::Millisecond),
+            "tDu" => DataType::Duration(TimeUnit::Microsecond),
+            "tDn" => DataType::Duration(TimeUnit::Nanosecond),
             "+l" => {
                 let c_child = c_schema.child(0);
                 DataType::List(Box::new(Field::try_from(c_child)?))
@@ -251,6 +255,10 @@ fn get_format_string(dtype: &DataType) -> Result<String> {
         DataType::Timestamp(TimeUnit::Millisecond, Some(tz)) => Ok(format!("tsm:{}", tz)),
         DataType::Timestamp(TimeUnit::Microsecond, Some(tz)) => Ok(format!("tsu:{}", tz)),
         DataType::Timestamp(TimeUnit::Nanosecond, Some(tz)) => Ok(format!("tsn:{}", tz)),
+        DataType::Duration(TimeUnit::Second) => Ok("tDs".to_string()),
+        DataType::Duration(TimeUnit::Millisecond) => Ok("tDm".to_string()),
+        DataType::Duration(TimeUnit::Microsecond) => Ok("tDu".to_string()),
+        DataType::Duration(TimeUnit::Nanosecond) => Ok("tDn".to_string()),
         DataType::List(_) => Ok("+l".to_string()),
         DataType::LargeList(_) => Ok("+L".to_string()),
         DataType::Struct(_) => Ok("+s".to_string()),

diff --git a/arrow/src/ffi.rs b/arrow/src/ffi.rs
@@ -314,6 +314,7 @@ fn bit_width(data_type: &DataType, i: usize) -> Result<usize> {
         (DataType::Float64, 1) => size_of::<f64>() * 8,
         (DataType::Decimal(..), 1) => size_of::<i128>() * 8,
         (DataType::Timestamp(..), 1) => size_of::<i64>() * 8,
+        (DataType::Duration(..), 1) => size_of::<i64>() * 8,
         // primitive types have a single buffer
         (DataType::Boolean, _) |
         (DataType::UInt8, _) |
@@ -327,7 +328,8 @@ fn bit_width(data_type: &DataType, i: usize) -> Result<usize> {
         (DataType::Float32, _) |
         (DataType::Float64, _) |
         (DataType::Decimal(..), _) |
-        (DataType::Timestamp(..), _) => {
+        (DataType::Timestamp(..), _) |
+        (DataType::Duration(..), _) => {
             return Err(ArrowError::CDataInterface(format!(
                 "The datatype \"{:?}\" expects 2 buffers, but requested {}. Please verify that the C data interface is correctly implemented.",
                 data_type, i
@@ -873,9 +875,9 @@ mod tests {
     use super::*;
     use crate::array::{
         export_array_into_raw, make_array, Array, ArrayData, BooleanArray, DecimalArray,
-        DictionaryArray, FixedSizeBinaryArray, FixedSizeListArray, GenericBinaryArray,
-        GenericListArray, GenericStringArray, Int32Array, OffsetSizeTrait,
-        Time32MillisecondArray, TimestampMillisecondArray,
+        DictionaryArray, DurationSecondArray, FixedSizeBinaryArray, FixedSizeListArray,
+        GenericBinaryArray, GenericListArray, GenericStringArray, Int32Array,
+        OffsetSizeTrait, Time32MillisecondArray, TimestampMillisecondArray,
     };
     use crate::compute::kernels;
     use crate::datatypes::{Field, Int8Type};
@@ -1358,4 +1360,40 @@ mod tests {
         }
         Ok(())
     }
+
+    #[test]
+    fn test_duration() -> Result<()> {
+        // create an array natively
+        let array = DurationSecondArray::from(vec![None, Some(1), Some(2)]);
+
+        // export it
+        let array = ArrowArray::try_from(array.data().clone())?;
+
+        // (simulate consumer) import it
+        let data = ArrayData::try_from(array)?;
+        let array = make_array(data);
+
+        // perform some operation
+        let array = kernels::concat::concat(&[array.as_ref(), array.as_ref()]).unwrap();
+        let array = array
+            .as_any()
+            .downcast_ref::<DurationSecondArray>()
+            .unwrap();
+
+        // verify
+        assert_eq!(
+            array,
+            &DurationSecondArray::from(vec![
+                None,
+                Some(1),
+                Some(2),
+                None,
+                Some(1),
+                Some(2)
+            ])
+        );
+
+        // (drop/release)
+        Ok(())
+    }
 }