Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
Migrate.
Browse files Browse the repository at this point in the history
  • Loading branch information
@jorgecarleitao
jorgecarleitao committed Aug 31, 2021
1 parent 1343b1d commit c4c34fb
Show file tree
Hide file tree
Showing 14 changed files with 93 additions and 115 deletions.
2 changes: 1 addition & 1 deletion src/array/binary/mod.rs
View file
Expand Up @@ -56,7 +56,7 @@ impl<O: Offset> BinaryArray<O> {
assert_eq!(offsets.len() - 1, validity.len());
}

if data_type != Self::default_data_type() {
if data_type.to_physical_type() != Self::default_data_type().to_physical_type() {
panic!("BinaryArray can only be initialized with DataType::Binary or DataType::LargeBinary")
}

Expand Down
8 changes: 7 additions & 1 deletion src/array/boolean/mod.rs
View file
@@ -1,4 +1,7 @@
use crate::{bitmap::Bitmap, datatypes::DataType};
use crate::{
bitmap::Bitmap,
datatypes::{DataType, PhysicalType},
};

use super::{display_fmt, Array};

Expand Down Expand Up @@ -40,6 +43,9 @@ impl BooleanArray {
if let Some(ref validity) = validity {
assert_eq!(values.len(), validity.len());
}
if data_type.to_physical_type() != PhysicalType::Boolean {
panic!("BooleanArray can only be initialized with DataType::Boolean")
}
Self {
data_type,
values,
Expand Down
5 changes: 4 additions & 1 deletion src/array/boolean/mutable.rs
View file
Expand Up @@ -4,7 +4,7 @@ use std::sync::Arc;
use crate::{
array::{Array, MutableArray, TryExtend, TryPush},
bitmap::MutableBitmap,
datatypes::DataType,
datatypes::{DataType, PhysicalType},
error::Result,
trusted_len::TrustedLen,
};
Expand Down Expand Up @@ -74,6 +74,9 @@ impl MutableBooleanArray {
values: MutableBitmap,
validity: Option<MutableBitmap>,
) -> Self {
if data_type.to_physical_type() != PhysicalType::Boolean {
panic!("MutableBooleanArray can only be initialized with DataType::Boolean")
}
Self {
data_type,
values,
Expand Down
9 changes: 4 additions & 5 deletions src/array/fixed_size_binary/mod.rs
View file
Expand Up @@ -40,7 +40,7 @@ impl FixedSizeBinaryArray {

Self {
size,
data_type: DataType::FixedSizeBinary(size),
data_type,
values,
validity,
offset: 0,
Expand Down Expand Up @@ -97,10 +97,9 @@ impl FixedSizeBinaryArray {

impl FixedSizeBinaryArray {
pub(crate) fn get_size(data_type: &DataType) -> &i32 {
if let DataType::FixedSizeBinary(size) = data_type {
size
} else {
panic!("Wrong DataType")
match data_type {
DataType::FixedSizeBinary(size) => size,
_ => panic!("Wrong DataType"),
}
}
}
Expand Down
18 changes: 11 additions & 7 deletions src/array/fixed_size_list/mod.rs
View file
Expand Up @@ -26,13 +26,18 @@ pub struct FixedSizeListArray {
impl FixedSizeListArray {
/// Returns a new empty [`FixedSizeListArray`].
pub fn new_empty(data_type: DataType) -> Self {
let values = new_empty_array(Self::get_child_and_size(&data_type).0.clone()).into();
let values =
new_empty_array(Self::get_child_and_size(&data_type).0.data_type().clone()).into();
Self::from_data(data_type, values, None)
}

/// Returns a new null [`FixedSizeListArray`].
pub fn new_null(data_type: DataType, length: usize) -> Self {
let values = new_null_array(Self::get_child_and_size(&data_type).0.clone(), length).into();
let values = new_null_array(
Self::get_child_and_size(&data_type).0.data_type().clone(),
length,
)
.into();
Self::from_data(data_type, values, Some(Bitmap::new_zeroed(length)))
}

Expand Down Expand Up @@ -88,11 +93,10 @@ impl FixedSizeListArray {
}

impl FixedSizeListArray {
pub(crate) fn get_child_and_size(data_type: &DataType) -> (&DataType, &i32) {
if let DataType::FixedSizeList(field, size) = data_type {
(field.data_type(), size)
} else {
panic!("Wrong DataType")
pub(crate) fn get_child_and_size(data_type: &DataType) -> (&Field, &i32) {
match data_type {
DataType::FixedSizeList(child, size) => (child.as_ref(), size),
_ => panic!("Wrong DataType"),
}
}

Expand Down
70 changes: 28 additions & 42 deletions src/array/growable/mod.rs
View file
Expand Up @@ -96,9 +96,10 @@ pub fn make_growable<'a>(
let data_type = arrays[0].data_type();
assert!(arrays.iter().all(|&item| item.data_type() == data_type));

match data_type {
DataType::Null => Box::new(null::GrowableNull::new(data_type.clone())),
DataType::Boolean => {
use PhysicalType::*;
match data_type.to_physical_type() {
Null => Box::new(null::GrowableNull::new(data_type.clone())),
Boolean => {
let arrays = arrays
.iter()
.map(|array| array.as_any().downcast_ref().unwrap())
Expand All @@ -109,33 +110,19 @@ pub fn make_growable<'a>(
capacity,
))
}
DataType::Int8 => dyn_growable!(i8, arrays, use_validity, capacity),
DataType::Int16 => dyn_growable!(i16, arrays, use_validity, capacity),
DataType::Int32
| DataType::Date32
| DataType::Time32(_)
| DataType::Interval(IntervalUnit::YearMonth) => {
dyn_growable!(i32, arrays, use_validity, capacity)
}
DataType::Int64
| DataType::Date64
| DataType::Time64(_)
| DataType::Timestamp(_, _)
| DataType::Duration(_) => {
dyn_growable!(i64, arrays, use_validity, capacity)
}
DataType::Interval(IntervalUnit::DayTime) => {
dyn_growable!(days_ms, arrays, use_validity, capacity)
}
DataType::Decimal(_, _) => dyn_growable!(i128, arrays, use_validity, capacity),
DataType::UInt8 => dyn_growable!(u8, arrays, use_validity, capacity),
DataType::UInt16 => dyn_growable!(u16, arrays, use_validity, capacity),
DataType::UInt32 => dyn_growable!(u32, arrays, use_validity, capacity),
DataType::UInt64 => dyn_growable!(u64, arrays, use_validity, capacity),
DataType::Float16 => unreachable!(),
DataType::Float32 => dyn_growable!(f32, arrays, use_validity, capacity),
DataType::Float64 => dyn_growable!(f64, arrays, use_validity, capacity),
DataType::Utf8 => {
Int8 => dyn_growable!(i8, arrays, use_validity, capacity),
Int16 => dyn_growable!(i16, arrays, use_validity, capacity),
Int32 => dyn_growable!(i32, arrays, use_validity, capacity),
Int64 => dyn_growable!(i64, arrays, use_validity, capacity),
Int128 => dyn_growable!(i128, arrays, use_validity, capacity),
DaysMs => dyn_growable!(days_ms, arrays, use_validity, capacity),
UInt8 => dyn_growable!(u8, arrays, use_validity, capacity),
UInt16 => dyn_growable!(u16, arrays, use_validity, capacity),
UInt32 => dyn_growable!(u32, arrays, use_validity, capacity),
UInt64 => dyn_growable!(u64, arrays, use_validity, capacity),
Float32 => dyn_growable!(f32, arrays, use_validity, capacity),
Float64 => dyn_growable!(f64, arrays, use_validity, capacity),
Utf8 => {
let arrays = arrays
.iter()
.map(|array| array.as_any().downcast_ref().unwrap())
Expand All @@ -146,7 +133,7 @@ pub fn make_growable<'a>(
capacity,
))
}
DataType::LargeUtf8 => {
LargeUtf8 => {
let arrays = arrays
.iter()
.map(|array| array.as_any().downcast_ref().unwrap())
Expand All @@ -157,7 +144,7 @@ pub fn make_growable<'a>(
capacity,
))
}
DataType::Binary => {
Binary => {
let arrays = arrays
.iter()
.map(|array| array.as_any().downcast_ref().unwrap())
Expand All @@ -168,7 +155,7 @@ pub fn make_growable<'a>(
capacity,
))
}
DataType::LargeBinary => {
LargeBinary => {
let arrays = arrays
.iter()
.map(|array| array.as_any().downcast_ref().unwrap())
Expand All @@ -179,7 +166,7 @@ pub fn make_growable<'a>(
capacity,
))
}
DataType::FixedSizeBinary(_) => {
FixedSizeBinary => {
let arrays = arrays
.iter()
.map(|array| array.as_any().downcast_ref().unwrap())
Expand All @@ -190,8 +177,7 @@ pub fn make_growable<'a>(
capacity,
))
}

DataType::List(_) => {
List => {
let arrays = arrays
.iter()
.map(|array| array.as_any().downcast_ref().unwrap())
Expand All @@ -202,7 +188,7 @@ pub fn make_growable<'a>(
capacity,
))
}
DataType::LargeList(_) => {
LargeList => {
let arrays = arrays
.iter()
.map(|array| array.as_any().downcast_ref().unwrap())
Expand All @@ -213,7 +199,7 @@ pub fn make_growable<'a>(
capacity,
))
}
DataType::Struct(_) => {
Struct => {
let arrays = arrays
.iter()
.map(|array| array.as_any().downcast_ref().unwrap())
Expand All @@ -224,10 +210,10 @@ pub fn make_growable<'a>(
capacity,
))
}
DataType::FixedSizeList(_, _) => todo!(),
DataType::Union(_, _, _) => todo!(),
DataType::Dictionary(key_type, _) => {
with_match_dictionary_key_type!(key_type.as_ref(), |$T| {
FixedSizeList => todo!(),
Union => todo!(),
Dictionary(key_type) => {
with_match_physical_dictionary_key_type!(key_type, |$T| {
dyn_dict_growable!($T, arrays, use_validity, capacity)
})
}
Expand Down
14 changes: 7 additions & 7 deletions src/array/list/mod.rs
View file
Expand Up @@ -133,15 +133,15 @@ impl<O: Offset> ListArray<O> {
#[inline]
pub fn get_child_field(data_type: &DataType) -> &Field {
if O::is_large() {
if let DataType::LargeList(child) = data_type {
child.as_ref()
} else {
panic!("Wrong DataType")
match data_type {
DataType::LargeList(child) => child.as_ref(),
_ => panic!("Wrong DataType"),
}
} else if let DataType::List(child) = data_type {
child.as_ref()
} else {
panic!("Wrong DataType")
match data_type {
DataType::List(child) => child.as_ref(),
_ => panic!("Wrong DataType"),
}
}
}

Expand Down
2 changes: 1 addition & 1 deletion src/array/struct_.rs
View file
Expand Up @@ -24,7 +24,7 @@ use super::{ffi::ToFfi, new_empty_array, new_null_array, Array, FromFfi};
/// Field::new("c", DataType::Int32, false),
/// ];
///
/// let array = StructArray::from_data(fields, vec![boolean, int], None);
/// let array = StructArray::from_data(DataType::Struct(fields), vec![boolean, int], None);
/// ```
#[derive(Debug, Clone)]
pub struct StructArray {
Expand Down
6 changes: 5 additions & 1 deletion src/array/utf8/mod.rs
View file
Expand Up @@ -72,7 +72,7 @@ impl<O: Offset> Utf8Array<O> {
assert_eq!(offsets.len() - 1, validity.len());
}

if data_type != Self::default_data_type() {
if data_type.to_physical_type() != Self::default_data_type().to_physical_type() {
panic!("Utf8Array can only be initialized with DataType::Utf8 or DataType::LargeUtf8")
}

Expand Down Expand Up @@ -105,6 +105,10 @@ impl<O: Offset> Utf8Array<O> {
) -> Self {
check_offsets(&offsets, values.len());

if data_type.to_physical_type() != Self::default_data_type().to_physical_type() {
panic!("Utf8Array can only be initialized with DataType::Utf8 or DataType::LargeUtf8")
}

Self {
data_type,
offsets,
Expand Down
24 changes: 6 additions & 18 deletions src/array/utf8/mutable.rs
View file
Expand Up @@ -46,11 +46,7 @@ impl<O: Offset> MutableUtf8Array<O> {
let mut offsets = MutableBuffer::<O>::new();
offsets.push(O::default());
Self {
data_type: if O::is_large() {
DataType::LargeUtf8
} else {
DataType::Utf8
},
data_type: Self::default_data_type(),
offsets,
values: MutableBuffer::<u8>::new(),
validity: None,
Expand All @@ -73,10 +69,8 @@ impl<O: Offset> MutableUtf8Array<O> {
if let Some(ref validity) = validity {
assert_eq!(offsets.len() - 1, validity.len());
}
if O::is_large() {
assert_eq!(data_type, DataType::LargeUtf8)
} else {
assert_eq!(data_type, DataType::Utf8)
if data_type.to_physical_type() != Self::default_data_type().to_physical_type() {
panic!("MutableUtf8Array can only be initialized with DataType::Utf8 or DataType::LargeUtf8")
}
Self {
data_type,
Expand All @@ -103,10 +97,8 @@ impl<O: Offset> MutableUtf8Array<O> {
if let Some(ref validity) = validity {
assert_eq!(offsets.len() - 1, validity.len());
}
if O::is_large() {
assert_eq!(data_type, DataType::LargeUtf8)
} else {
assert_eq!(data_type, DataType::Utf8)
if data_type.to_physical_type() != Self::default_data_type().to_physical_type() {
panic!("MutableUtf8Array can only be initialized with DataType::Utf8 or DataType::LargeUtf8")
}
Self {
data_type,
Expand All @@ -117,11 +109,7 @@ impl<O: Offset> MutableUtf8Array<O> {
}

fn default_data_type() -> DataType {
if O::is_large() {
DataType::LargeUtf8
} else {
DataType::Utf8
}
Utf8Array::<O>::default_data_type()
}

/// Initializes a new [`MutableUtf8Array`] with a pre-allocated capacity of slots.
Expand Down
6 changes: 3 additions & 3 deletions src/datatypes/mod.rs
View file
Expand Up @@ -4,7 +4,7 @@ mod physical_type;
mod schema;

pub use field::Field;
pub(crate) use physical_type::*;
pub use physical_type::*;
pub use schema::Schema;

/// The set of datatypes that are supported by this implementation of Apache Arrow.
Expand Down Expand Up @@ -172,8 +172,8 @@ impl DataType {
}
}

/// Returns the physical type of the logical type
pub(crate) fn to_physical_type(&self) -> PhysicalType {
/// the [`PhysicalType`] of this [`DataType`].
pub fn to_physical_type(&self) -> PhysicalType {
use DataType::*;
match self {
Null => PhysicalType::Null,
Expand Down

0 comments on commit c4c34fb

Please sign in to comment.