Add safe zero-copy converion from bytes::Bytes (#4254)

Author: tustvold

arrow-buffer/Cargo.toml

@@ -33,7 +33,11 @@ name = "arrow_buffer"
 path = "src/lib.rs"
 bench = false
 
+[package.metadata.docs.rs]
+features = ["bytes"]
+
 [dependencies]
+bytes = { version = "1.4", optional = true }
 num = { version = "0.4", default-features = false, features = ["std"] }
 half = { version = "2.1", default-features = false }
 

arrow-buffer/src/bytes.rs

@@ -148,3 +148,32 @@ impl Debug for Bytes {
         write!(f, " }}")
     }
 }
+
+#[cfg(feature = "bytes")]
+impl From<bytes::Bytes> for Bytes {
+    fn from(value: bytes::Bytes) -> Self {
+        Self {
+            len: value.len(),
+            ptr: NonNull::new(value.as_ptr() as _).unwrap(),
+            deallocation: Deallocation::Custom(std::sync::Arc::new(value)),
+        }
+    }
+}
+
+#[cfg(all(test, feature = "bytes"))]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_from_bytes() {
+        let bytes = bytes::Bytes::from(vec![1, 2, 3, 4]);
+        let arrow_bytes: Bytes = bytes.clone().into();
+
+        assert_eq!(bytes.as_ptr(), arrow_bytes.as_ptr());
+
+        drop(bytes);
+        drop(arrow_bytes);
+
+        let _ = Bytes::from(bytes::Bytes::new());
+    }
+}

arrow-flight/Cargo.toml

@@ -28,7 +28,7 @@ license = { workspace = true }
 
 [dependencies]
 arrow-array = { workspace = true }
-arrow-buffer = { workspace = true }
+arrow-buffer = { workspace = true, features = ["bytes"] }
 # Cast is needed to work around https://github.com/apache/arrow-rs/issues/3389
 arrow-cast = { workspace = true }
 arrow-ipc = { workspace = true }

arrow-flight/src/decode.rs

@@ -17,6 +17,7 @@
 
 use crate::{utils::flight_data_to_arrow_batch, FlightData};
 use arrow_array::{ArrayRef, RecordBatch};
+use arrow_buffer::Buffer;
 use arrow_schema::{Schema, SchemaRef};
 use bytes::Bytes;
 use futures::{ready, stream::BoxStream, Stream, StreamExt};
@@ -258,7 +259,7 @@ impl FlightDataDecoder {
                     ));
                 };
 
-                let buffer: arrow_buffer::Buffer = data.data_body.into();
+                let buffer = Buffer::from_bytes(data.data_body.into());
                 let dictionary_batch =
                     message.header_as_dictionary_batch().ok_or_else(|| {
                         FlightError::protocol(

arrow-flight/src/sql/client.rs

@@ -520,7 +520,7 @@ pub fn arrow_data_from_flight_data(
 
             let dictionaries_by_field = HashMap::new();
             let record_batch = read_record_batch(
-                &Buffer::from(&flight_data.data_body),
+                &Buffer::from_bytes(flight_data.data_body.into()),
                 ipc_record_batch,
                 arrow_schema_ref.clone(),
                 &dictionaries_by_field,