Skip to content

Commit 924155c

Browse files
committed
Move validation logic to ShreddingState::try_new
1 parent 6b66650 commit 924155c

File tree

1 file changed

+104
-100
lines changed

1 file changed

+104
-100
lines changed

parquet-variant-compute/src/variant_array.rs

Lines changed: 104 additions & 100 deletions
Original file line numberDiff line numberDiff line change
@@ -52,85 +52,6 @@ pub struct VariantArray {
5252
shredding_state: ShreddingState,
5353
}
5454

55-
/// Variant arrays can be shredded in one of three states, encoded here
56-
#[derive(Debug)]
57-
pub enum ShreddingState {
58-
/// This variant has no typed_value field
59-
Unshredded {
60-
metadata: BinaryViewArray,
61-
value: BinaryViewArray,
62-
},
63-
/// This variant has a typed_value field and no value field
64-
/// meaning it is fully shredded (aka the value is stored in typed_value)
65-
FullyShredded {
66-
metadata: BinaryViewArray,
67-
typed_value: ArrayRef,
68-
},
69-
/// This variant has both a value field and a typed_value field
70-
/// meaning it is partially shredded: first the typed_value is used, and
71-
/// if that is null, the value field is used.
72-
PartiallyShredded {
73-
metadata: BinaryViewArray,
74-
value: BinaryViewArray,
75-
typed_value: ArrayRef,
76-
},
77-
}
78-
79-
impl ShreddingState {
80-
/// Return a reference to the metadata field
81-
pub fn metadata_field(&self) -> &BinaryViewArray {
82-
match self {
83-
ShreddingState::Unshredded { metadata, .. } => metadata,
84-
ShreddingState::FullyShredded { metadata, .. } => metadata,
85-
ShreddingState::PartiallyShredded { metadata, .. } => metadata,
86-
}
87-
}
88-
89-
/// Return a reference to the value field, if present
90-
pub fn value_field(&self) -> Option<&BinaryViewArray> {
91-
match self {
92-
ShreddingState::Unshredded { value, .. } => Some(value),
93-
ShreddingState::FullyShredded { .. } => None,
94-
ShreddingState::PartiallyShredded { value, .. } => Some(value),
95-
}
96-
}
97-
98-
/// Return a reference to the typed_value field, if present
99-
pub fn typed_value_field(&self) -> Option<&ArrayRef> {
100-
match self {
101-
ShreddingState::Unshredded { .. } => None,
102-
ShreddingState::FullyShredded { typed_value, .. } => Some(typed_value),
103-
ShreddingState::PartiallyShredded { typed_value, .. } => Some(typed_value),
104-
}
105-
}
106-
107-
/// Slice all the underlying arrays
108-
pub fn slice(&self, offset: usize, length: usize) -> Self {
109-
match self {
110-
ShreddingState::Unshredded { metadata, value } => ShreddingState::Unshredded {
111-
metadata: metadata.slice(offset, length),
112-
value: value.slice(offset, length),
113-
},
114-
ShreddingState::FullyShredded {
115-
metadata,
116-
typed_value,
117-
} => ShreddingState::FullyShredded {
118-
metadata: metadata.slice(offset, length),
119-
typed_value: typed_value.slice(offset, length),
120-
},
121-
ShreddingState::PartiallyShredded {
122-
metadata,
123-
value,
124-
typed_value,
125-
} => ShreddingState::PartiallyShredded {
126-
metadata: metadata.slice(offset, length),
127-
value: value.slice(offset, length),
128-
typed_value: typed_value.slice(offset, length),
129-
},
130-
}
131-
}
132-
}
133-
13455
impl VariantArray {
13556
/// Creates a new `VariantArray` from a [`StructArray`].
13657
///
@@ -200,27 +121,8 @@ impl VariantArray {
200121

201122
// Note these clones are cheap, they just bump the ref count
202123
let inner = inner.clone();
203-
let metadata = metadata.clone();
204-
let value = value.cloned();
205-
let typed_value = typed_value.cloned();
206-
207-
let shredding_state = match (metadata, value, typed_value) {
208-
(metadata, Some(value), Some(typed_value)) => ShreddingState::PartiallyShredded {
209-
metadata,
210-
value,
211-
typed_value,
212-
},
213-
(metadata, Some(value), None) => ShreddingState::Unshredded { metadata, value },
214-
(metadata, None, Some(typed_value)) => ShreddingState::FullyShredded {
215-
metadata,
216-
typed_value,
217-
},
218-
(_metadata_field, None, None) => {
219-
return Err(ArrowError::InvalidArgumentError(String::from(
220-
"VariantArray has neither value nor typed_value field",
221-
)));
222-
}
223-
};
124+
let shredding_state =
125+
ShreddingState::try_new(metadata.clone(), value.cloned(), typed_value.cloned())?;
224126

225127
Ok(Self {
226128
inner,
@@ -308,6 +210,108 @@ impl VariantArray {
308210
}
309211
}
310212

213+
/// Variant arrays can be shredded in one of three states, encoded here
214+
#[derive(Debug)]
215+
pub enum ShreddingState {
216+
/// This variant has no typed_value field
217+
Unshredded {
218+
metadata: BinaryViewArray,
219+
value: BinaryViewArray,
220+
},
221+
/// This variant has a typed_value field and no value field
222+
/// meaning it is fully shredded (aka the value is stored in typed_value)
223+
FullyShredded {
224+
metadata: BinaryViewArray,
225+
typed_value: ArrayRef,
226+
},
227+
/// This variant has both a value field and a typed_value field
228+
/// meaning it is partially shredded: first the typed_value is used, and
229+
/// if that is null, the value field is used.
230+
PartiallyShredded {
231+
metadata: BinaryViewArray,
232+
value: BinaryViewArray,
233+
typed_value: ArrayRef,
234+
},
235+
}
236+
237+
impl ShreddingState {
238+
/// try to create a new `ShreddingState` from the given fields
239+
pub fn try_new(
240+
metadata: BinaryViewArray,
241+
value: Option<BinaryViewArray>,
242+
typed_value: Option<ArrayRef>,
243+
) -> Result<Self, ArrowError> {
244+
match (metadata, value, typed_value) {
245+
(metadata, Some(value), Some(typed_value)) => Ok(Self::PartiallyShredded {
246+
metadata,
247+
value,
248+
typed_value,
249+
}),
250+
(metadata, Some(value), None) => Ok(Self::Unshredded { metadata, value }),
251+
(metadata, None, Some(typed_value)) => Ok(Self::FullyShredded {
252+
metadata,
253+
typed_value,
254+
}),
255+
(_metadata_field, None, None) => Err(ArrowError::InvalidArgumentError(String::from(
256+
"VariantArray has neither value nor typed_value field",
257+
))),
258+
}
259+
}
260+
261+
/// Return a reference to the metadata field
262+
pub fn metadata_field(&self) -> &BinaryViewArray {
263+
match self {
264+
ShreddingState::Unshredded { metadata, .. } => metadata,
265+
ShreddingState::FullyShredded { metadata, .. } => metadata,
266+
ShreddingState::PartiallyShredded { metadata, .. } => metadata,
267+
}
268+
}
269+
270+
/// Return a reference to the value field, if present
271+
pub fn value_field(&self) -> Option<&BinaryViewArray> {
272+
match self {
273+
ShreddingState::Unshredded { value, .. } => Some(value),
274+
ShreddingState::FullyShredded { .. } => None,
275+
ShreddingState::PartiallyShredded { value, .. } => Some(value),
276+
}
277+
}
278+
279+
/// Return a reference to the typed_value field, if present
280+
pub fn typed_value_field(&self) -> Option<&ArrayRef> {
281+
match self {
282+
ShreddingState::Unshredded { .. } => None,
283+
ShreddingState::FullyShredded { typed_value, .. } => Some(typed_value),
284+
ShreddingState::PartiallyShredded { typed_value, .. } => Some(typed_value),
285+
}
286+
}
287+
288+
/// Slice all the underlying arrays
289+
pub fn slice(&self, offset: usize, length: usize) -> Self {
290+
match self {
291+
ShreddingState::Unshredded { metadata, value } => ShreddingState::Unshredded {
292+
metadata: metadata.slice(offset, length),
293+
value: value.slice(offset, length),
294+
},
295+
ShreddingState::FullyShredded {
296+
metadata,
297+
typed_value,
298+
} => ShreddingState::FullyShredded {
299+
metadata: metadata.slice(offset, length),
300+
typed_value: typed_value.slice(offset, length),
301+
},
302+
ShreddingState::PartiallyShredded {
303+
metadata,
304+
value,
305+
typed_value,
306+
} => ShreddingState::PartiallyShredded {
307+
metadata: metadata.slice(offset, length),
308+
value: value.slice(offset, length),
309+
typed_value: typed_value.slice(offset, length),
310+
},
311+
}
312+
}
313+
}
314+
311315
/// returns the non-null element at index as a Variant
312316
fn typed_value_to_variant(typed_value: &ArrayRef, index: usize) -> Variant {
313317
match typed_value.data_type() {

0 commit comments

Comments
 (0)