Skip to content

Commit a9c75d2

Browse files
authored
fix infering arrow lists with nullable elements (#2112)
1 parent 9f8e257 commit a9c75d2

File tree

2 files changed

+31
-3
lines changed

2 files changed

+31
-3
lines changed

vortex-array/src/arrow/dtype.rs

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,8 @@ pub fn infer_schema(dtype: &DType) -> VortexResult<Schema> {
116116
}
117117

118118
/// Try to convert a Vortex [`DType`] into an Arrow [`DataType`]
119+
/// Top level nulltability from the DType is dropped, Arrow represents
120+
/// nullability for a DataType in [`Field`]
119121
pub fn infer_data_type(dtype: &DType) -> VortexResult<DataType> {
120122
Ok(match dtype {
121123
DType::Null => DataType::Null,
@@ -150,9 +152,9 @@ pub fn infer_data_type(dtype: &DType) -> VortexResult<DataType> {
150152
// There are four kinds of lists: List (32-bit offsets), Large List (64-bit), List View
151153
// (32-bit), Large List View (64-bit). We cannot both guarantee zero-copy and commit to an
152154
// Arrow dtype because we do not how large our offsets are.
153-
DType::List(l, null) => DataType::List(FieldRef::new(Field::new_list_field(
155+
DType::List(l, _) => DataType::List(FieldRef::new(Field::new_list_field(
154156
infer_data_type(l.as_ref())?,
155-
(*null).into(),
157+
l.nullability().into(),
156158
))),
157159
DType::Extension(ext_dtype) => {
158160
// Try and match against the known extension DTypes.
@@ -214,6 +216,32 @@ mod test {
214216
);
215217
}
216218

219+
#[test]
220+
fn infer_nullable_list_element() {
221+
let list_non_nullable = DType::List(
222+
Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)),
223+
Nullability::Nullable,
224+
);
225+
226+
let arrow_list_non_nullable = infer_data_type(&list_non_nullable).unwrap();
227+
228+
let list_nullable = DType::List(
229+
Arc::new(DType::Primitive(PType::I64, Nullability::Nullable)),
230+
Nullability::Nullable,
231+
);
232+
let arrow_list_nullable = infer_data_type(&list_nullable).unwrap();
233+
234+
assert_ne!(arrow_list_non_nullable, arrow_list_nullable);
235+
assert_eq!(
236+
arrow_list_nullable,
237+
DataType::new_list(DataType::Int64, true)
238+
);
239+
assert_eq!(
240+
arrow_list_non_nullable,
241+
DataType::new_list(DataType::Int64, false)
242+
);
243+
}
244+
217245
#[test]
218246
#[should_panic]
219247
fn test_dtype_conversion_panics() {

vortex-array/src/canonical.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ pub enum Canonical {
6969
impl Canonical {
7070
// Create an empty canonical array of the given dtype.
7171
pub fn empty(dtype: &DType) -> Canonical {
72-
Self::try_empty(dtype).vortex_expect("Cannot build an empty array")
72+
Self::try_empty(dtype).vortex_expect("Cannot fail to build an empty array")
7373
}
7474

7575
pub fn try_empty(dtype: &DType) -> VortexResult<Canonical> {

0 commit comments

Comments
 (0)