From f761d417addd2d0f0cb0278fc076ed097aa5f485 Mon Sep 17 00:00:00 2001 From: Ben Sully Date: Sun, 26 Sep 2021 07:39:45 +0100 Subject: [PATCH] Add `MutableArray::as_box` (#450) --- src/array/binary/mutable.rs | 9 +++++++++ src/array/boolean/mutable.rs | 8 ++++++++ src/array/dictionary/mutable.rs | 7 +++++++ src/array/fixed_size_binary/mutable.rs | 8 ++++++++ src/array/fixed_size_list/mutable.rs | 8 ++++++++ src/array/list/mutable.rs | 9 +++++++++ src/array/mod.rs | 12 ++++++++++-- src/array/primitive/mutable.rs | 8 ++++++++ src/array/utf8/mutable.rs | 9 +++++++++ 9 files changed, 76 insertions(+), 2 deletions(-) diff --git a/src/array/binary/mutable.rs b/src/array/binary/mutable.rs index 89c645de00a..1eee5e10406 100644 --- a/src/array/binary/mutable.rs +++ b/src/array/binary/mutable.rs @@ -146,6 +146,15 @@ impl MutableArray for MutableBinaryArray { self.validity.as_ref() } + fn as_box(&mut self) -> Box { + Box::new(BinaryArray::from_data( + self.data_type.clone(), + std::mem::take(&mut self.offsets).into(), + std::mem::take(&mut self.values).into(), + std::mem::take(&mut self.validity).map(|x| x.into()), + )) + } + fn as_arc(&mut self) -> Arc { Arc::new(BinaryArray::from_data( self.data_type.clone(), diff --git a/src/array/boolean/mutable.rs b/src/array/boolean/mutable.rs index 857c434a75d..a776b7a12bc 100644 --- a/src/array/boolean/mutable.rs +++ b/src/array/boolean/mutable.rs @@ -338,6 +338,14 @@ impl MutableArray for MutableBooleanArray { self.validity.as_ref() } + fn as_box(&mut self) -> Box { + Box::new(BooleanArray::from_data( + self.data_type.clone(), + std::mem::take(&mut self.values).into(), + std::mem::take(&mut self.validity).map(|x| x.into()), + )) + } + fn as_arc(&mut self) -> Arc { Arc::new(BooleanArray::from_data( self.data_type.clone(), diff --git a/src/array/dictionary/mutable.rs b/src/array/dictionary/mutable.rs index fc786bcf51b..17f6f4caa4d 100644 --- a/src/array/dictionary/mutable.rs +++ b/src/array/dictionary/mutable.rs @@ -115,6 +115,13 @@ impl MutableArray for MutableDictio self.keys.validity() } + fn as_box(&mut self) -> Box { + Box::new(DictionaryArray::::from_data( + std::mem::take(&mut self.keys).into(), + self.values.as_arc(), + )) + } + fn as_arc(&mut self) -> Arc { Arc::new(DictionaryArray::::from_data( std::mem::take(&mut self.keys).into(), diff --git a/src/array/fixed_size_binary/mutable.rs b/src/array/fixed_size_binary/mutable.rs index e60619c4e4c..9f3872c704b 100644 --- a/src/array/fixed_size_binary/mutable.rs +++ b/src/array/fixed_size_binary/mutable.rs @@ -173,6 +173,14 @@ impl MutableArray for MutableFixedSizeBinaryArray { self.validity.as_ref() } + fn as_box(&mut self) -> Box { + Box::new(FixedSizeBinaryArray::from_data( + DataType::FixedSizeBinary(self.size as i32), + std::mem::take(&mut self.values).into(), + std::mem::take(&mut self.validity).map(|x| x.into()), + )) + } + fn as_arc(&mut self) -> Arc { Arc::new(FixedSizeBinaryArray::from_data( DataType::FixedSizeBinary(self.size as i32), diff --git a/src/array/fixed_size_list/mutable.rs b/src/array/fixed_size_list/mutable.rs index 0cf7105f6c7..7409aaca66e 100644 --- a/src/array/fixed_size_list/mutable.rs +++ b/src/array/fixed_size_list/mutable.rs @@ -69,6 +69,14 @@ impl MutableArray for MutableFixedSizeListArray { self.validity.as_ref() } + fn as_box(&mut self) -> Box { + Box::new(FixedSizeListArray::from_data( + self.data_type.clone(), + self.values.as_arc(), + std::mem::take(&mut self.validity).map(|x| x.into()), + )) + } + fn as_arc(&mut self) -> Arc { Arc::new(FixedSizeListArray::from_data( self.data_type.clone(), diff --git a/src/array/list/mutable.rs b/src/array/list/mutable.rs index d4f873b7afd..bd268912a1d 100644 --- a/src/array/list/mutable.rs +++ b/src/array/list/mutable.rs @@ -177,6 +177,15 @@ impl MutableArray for MutableListArray Box { + Box::new(ListArray::from_data( + self.data_type.clone(), + std::mem::take(&mut self.offsets).into(), + self.values.as_arc(), + std::mem::take(&mut self.validity).map(|x| x.into()), + )) + } + fn as_arc(&mut self) -> Arc { Arc::new(ListArray::from_data( self.data_type.clone(), diff --git a/src/array/mod.rs b/src/array/mod.rs index 944752c683d..a9f0d2d82d2 100644 --- a/src/array/mod.rs +++ b/src/array/mod.rs @@ -121,8 +121,16 @@ pub trait MutableArray: std::fmt::Debug { /// The optional validity of the array. fn validity(&self) -> Option<&MutableBitmap>; - /// Convert itself to an (immutable) [`Array`]. - fn as_arc(&mut self) -> Arc; + /// Convert itself to an (immutable) ['Array']. + fn as_box(&mut self) -> Box; + + /// Convert itself to an (immutable) atomically reference counted [`Array`]. + // This provided implementation has an extra allocation as it first + // boxes `self`, then converts the box into an `Arc`. Implementors may wish + // to avoid an allocation by skipping the box completely. + fn as_arc(&mut self) -> Arc { + self.as_box().into() + } /// Convert to `Any`, to enable dynamic casting. fn as_any(&self) -> &dyn Any; diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs index 6e0a5fbdf80..bb5b611d638 100644 --- a/src/array/primitive/mutable.rs +++ b/src/array/primitive/mutable.rs @@ -326,6 +326,14 @@ impl MutableArray for MutablePrimitiveArray { self.validity.as_ref() } + fn as_box(&mut self) -> Box { + Box::new(PrimitiveArray::from_data( + self.data_type.clone(), + std::mem::take(&mut self.values).into(), + std::mem::take(&mut self.validity).map(|x| x.into()), + )) + } + fn as_arc(&mut self) -> Arc { Arc::new(PrimitiveArray::from_data( self.data_type.clone(), diff --git a/src/array/utf8/mutable.rs b/src/array/utf8/mutable.rs index 8ad0fd88192..78b3a7ba7d0 100644 --- a/src/array/utf8/mutable.rs +++ b/src/array/utf8/mutable.rs @@ -187,6 +187,15 @@ impl MutableArray for MutableUtf8Array { self.validity.as_ref() } + fn as_box(&mut self) -> Box { + Box::new(Utf8Array::from_data( + Self::default_data_type(), + std::mem::take(&mut self.offsets).into(), + std::mem::take(&mut self.values).into(), + std::mem::take(&mut self.validity).map(|x| x.into()), + )) + } + fn as_arc(&mut self) -> Arc { Arc::new(Utf8Array::from_data( Self::default_data_type(),