Skip to content

Commit

Permalink
feat: Lateral unnest + additional list array logic (#3363)
Browse files Browse the repository at this point in the history
```
GlareDB Shell
Preview (0.0.91) - There will be bugs!

>> SELECT * FROM (VALUES ([1,2,3]), ([8,9])) v(a), unnest(v.a);
┌─────────────┬────────┐
│ a           │ unnest │
│ List[Int32] │ Int32  │
├─────────────┼────────┤
│ [1, 2, 3]   │      1 │
│ [1, 2, 3]   │      2 │
│ [1, 2, 3]   │      3 │
│ [8, 9]      │      8 │
│ [8, 9]      │      9 │
└─────────────┴────────┘
```
  • Loading branch information
scsmithr authored Dec 16, 2024
1 parent 3653724 commit 3117aea
Show file tree
Hide file tree
Showing 12 changed files with 1,012 additions and 118 deletions.
96 changes: 88 additions & 8 deletions crates/rayexec_bullet/src/executor/scalar/fill.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@ use std::borrow::Borrow;

use rayexec_error::{RayexecError, Result};

use crate::array::Array;
use crate::array::{Array, ArrayData};
use crate::bitmap::Bitmap;
use crate::datatype::DataType;
use crate::executor::builder::{
ArrayBuilder,
ArrayDataBuffer,
Expand All @@ -23,6 +24,7 @@ use crate::executor::physical_type::{
PhysicalI64,
PhysicalI8,
PhysicalInterval,
PhysicalList,
PhysicalStorage,
PhysicalType,
PhysicalU128,
Expand All @@ -32,8 +34,15 @@ use crate::executor::physical_type::{
PhysicalU8,
PhysicalUtf8,
};
use crate::executor::scalar::UnaryExecutor;
use crate::selection;
use crate::storage::{AddressableStorage, UntypedNullStorage};
use crate::storage::{
AddressableStorage,
ListItemMetadata,
ListStorage,
PrimitiveStorage,
UntypedNullStorage,
};

/// Singular mapping of a `from` index to a `to` index.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
Expand Down Expand Up @@ -269,13 +278,63 @@ pub(crate) fn concat_with_exact_total_len(arrays: &[&Array], total_len: usize) -
});
concat_with_fill_state::<PhysicalBinary, _>(arrays, state)
}
PhysicalType::List => {
// TODO: Very doable
Err(RayexecError::new(
"concatenating list arrays not yet supported",
))
}
PhysicalType::List => concat_lists(datatype.clone(), arrays, total_len),
}
}

fn concat_lists(datatype: DataType, arrays: &[&Array], total_len: usize) -> Result<Array> {
let inner_arrays = arrays
.iter()
.map(|arr| match arr.array_data() {
ArrayData::List(list) => {
if list.array.has_selection() {
return Err(RayexecError::new("List child array has selection"));
}
Ok(&list.array)
}
other => Err(RayexecError::new(format!(
"Invalid inner array data for concatenating lists, got {:?}",
other.physical_type()
))),
})
.collect::<Result<Vec<_>>>()?;

let concatenated = concat(&inner_arrays)?;

// Update metadata objects.
let mut metadatas = Vec::with_capacity(total_len);
let mut validity = Bitmap::new_with_all_true(total_len);

let mut acc_rows = 0;

for (array, child_array) in arrays.iter().zip(inner_arrays) {
UnaryExecutor::for_each::<PhysicalList, _>(array, |_row_num, metadata| match metadata {
Some(metadata) => {
metadatas.push(ListItemMetadata {
offset: metadata.offset + acc_rows,
len: metadata.len,
});
}
None => {
metadatas.push(ListItemMetadata::default());
validity.set_unchecked(metadatas.len() - 1, false);
}
})?;

acc_rows += child_array.logical_len() as i32;
}

let data = ListStorage {
metadata: PrimitiveStorage::from(metadatas),
array: concatenated,
};

Ok(Array {
datatype,
selection: None,
validity: Some(validity.into()),
data: data.into(),
})
}

fn concat_with_fill_state<'a, S, B>(
Expand Down Expand Up @@ -641,4 +700,25 @@ mod tests {
assert_eq!(ScalarValue::from(7), got.logical_value(3).unwrap());
assert_eq!(ScalarValue::from(8), got.logical_value(4).unwrap());
}

#[test]
fn concat_lists() {
let arr1 = ScalarValue::List(vec![1.into(), 2.into()])
.as_array(1)
.unwrap();
let arr2 = ScalarValue::List(vec![3.into(), 4.into(), 5.into()])
.as_array(1)
.unwrap();

let got = concat(&[&arr1, &arr2]).unwrap();

assert_eq!(
ScalarValue::List(vec![1.into(), 2.into()]),
got.logical_value(0).unwrap()
);
assert_eq!(
ScalarValue::List(vec![3.into(), 4.into(), 5.into()]),
got.logical_value(1).unwrap()
);
}
}
Loading

0 comments on commit 3117aea

Please sign in to comment.