Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

wip: smol datatypes #1847

Closed
wants to merge 3 commits into from
Closed

wip: smol datatypes #1847

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5a0d009
arrow2 0.17, arrow2_convert 0.5 and update to new datatypes
teh-cmc Apr 14, 2023
cfad837
polars 0.28.0 patched with new datatypes
teh-cmc Apr 14, 2023
0e5138b
wip
teh-cmc Apr 17, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
323 changes: 71 additions & 252 deletions Cargo.lock
View file
Open in desktop

Large diffs are not rendered by default.

14 changes: 9 additions & 5 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,8 +52,8 @@ rerun = { path = "crates/rerun", version = "0.4.0" }

ahash = "0.8"
anyhow = "1.0"
arrow2 = "0.16"
arrow2_convert = "0.4.2"
arrow2 = "0.17"
arrow2_convert = "0.5.0"
clap = "4.0"
comfy-table = { version = "6.1", default-features = false }
ctrlc = { version = "3.0", features = ["termination"] }
Expand All @@ -76,9 +76,9 @@ macaw = "0.18"
mimalloc = "0.1.29"
ndarray = "0.15"
parking_lot = "0.12"
polars-core = "0.27.1"
polars-lazy = "0.27.1"
polars-ops = "0.27.1"
polars-core = "0.28.0"
polars-lazy = "0.28.0"
polars-ops = "0.28.0"
puffin = "0.14"
smallvec = { version = "1.0", features = ["const_generics", "union"] }
thiserror = "1.0"
Expand Down Expand Up @@ -112,3 +112,7 @@ debug = true
# If that is not possible, patch to a branch that has a PR open on the upstream repo.
# As a last resport, patch with a commit to our own repository.
# ALWAYS document what PR the commit hash is part of, or when it was merged into the upstream trunk.
arrow2 = { git = "https://github.com/rerun-io/arrow2", branch = "cmc/arc_datatype" }
arrow2_convert = { git = "https://github.com/rerun-io/arrow2-convert", branch = "cmc/arc_datatype" }
polars-core = { git = "https://github.com/rerun-io/polars", branch = "cmc/arc_datatype" }
polars-ops = { git = "https://github.com/rerun-io/polars", branch = "cmc/arc_datatype" }
2 changes: 1 addition & 1 deletion crates/re_arrow_store/benches/data_store.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -322,7 +322,7 @@ fn build_table(n: usize, packed: bool) -> DataTable {
// Do a serialization roundtrip to pack everything in contiguous memory.
if packed {
let (schema, columns) = table.serialize().unwrap();
table = DataTable::deserialize(TableId::ZERO, &schema, &columns).unwrap();
table = DataTable::deserialize(TableId::ZERO, &schema, columns, None).unwrap();
}

table
Expand Down
69 changes: 36 additions & 33 deletions crates/re_arrow_store/src/arrow_util.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
use std::sync::Arc;

use ahash::HashSet;
use arrow2::{
array::{
growable::make_growable, Array, FixedSizeListArray, ListArray, StructArray, UnionArray,
Expand Down Expand Up @@ -63,7 +66,7 @@ impl ArrayExt for dyn Array {
// Recursively clean the contents
let typed_arr = self.as_any().downcast_ref::<ListArray<i32>>().unwrap();
let clean_vals = typed_arr.values().as_ref().clean_for_polars();
let clean_data = DataType::List(Box::new(Field::new(
let clean_data = DataType::List(Arc::new(Field::new(
&field.name,
clean_vals.data_type().clone(),
field.is_nullable,
Expand All @@ -81,7 +84,7 @@ impl ArrayExt for dyn Array {
// Recursively clean the contents
let typed_arr = self.as_any().downcast_ref::<ListArray<i64>>().unwrap();
let clean_vals = typed_arr.values().as_ref().clean_for_polars();
let clean_data = DataType::LargeList(Box::new(Field::new(
let clean_data = DataType::LargeList(Arc::new(Field::new(
&field.name,
clean_vals.data_type().clone(),
field.is_nullable,
Expand All @@ -99,7 +102,7 @@ impl ArrayExt for dyn Array {
// Recursively clean the contents and convert `FixedSizeListArray` -> `ListArray`
let typed_arr = self.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
let clean_vals = typed_arr.values().as_ref().clean_for_polars();
let clean_data = DataType::List(Box::new(Field::new(
let clean_data = DataType::List(Arc::new(Field::new(
&field.name,
clean_vals.data_type().clone(),
field.is_nullable,
Expand All @@ -123,10 +126,10 @@ impl ArrayExt for dyn Array {
.iter()
.map(|v| v.as_ref().clean_for_polars())
.collect_vec();
let clean_fields = itertools::izip!(fields, &clean_vals)
let clean_fields = itertools::izip!(fields.iter(), &clean_vals)
.map(|(f, v)| Field::new(&f.name, v.data_type().clone(), f.is_nullable))
.collect_vec();
let clean_data = DataType::Struct(clean_fields);
let clean_data = DataType::Struct(Arc::new(clean_fields));
StructArray::try_new(clean_data, clean_vals, typed_arr.validity().cloned())
.unwrap()
.boxed()
Expand All @@ -144,12 +147,12 @@ impl ArrayExt for dyn Array {

let ids = ids
.clone()
.unwrap_or_else(|| (0i32..(clean_vals.len() as i32)).collect_vec());
.unwrap_or_else(|| Arc::new((0i32..(clean_vals.len() as i32)).collect_vec()));

// For Dense Unions, the value-arrays need to be padded to the
// correct length, which we do by growing using the existing type
// table.
let padded_vals = itertools::izip!(&clean_vals, &ids)
let padded_vals = itertools::izip!(clean_vals.iter(), ids.iter())
.map(|(dense, id)| {
let mut next = 0;
let mut grow = make_growable(&[dense.as_ref()], true, self.len());
Expand All @@ -165,12 +168,12 @@ impl ArrayExt for dyn Array {
})
.collect_vec();

let clean_field_types = itertools::izip!(fields, &clean_vals)
let clean_field_types = itertools::izip!(fields.iter(), &clean_vals)
.map(|(f, v)| Field::new(&f.name, v.data_type().clone(), f.is_nullable))
.collect_vec();

// The new type will be a struct
let clean_data = DataType::Struct(clean_field_types);
let clean_data = DataType::Struct(Arc::new(clean_field_types));

StructArray::try_new(clean_data, padded_vals, typed_arr.validity().cloned())
.unwrap()
Expand All @@ -189,11 +192,11 @@ impl ArrayExt for dyn Array {

let ids = ids
.clone()
.unwrap_or_else(|| (0i32..(clean_vals.len() as i32)).collect_vec());
.unwrap_or_else(|| Arc::new((0i32..(clean_vals.len() as i32)).collect_vec()));

// For Sparse Unions, the value-arrays is already the right
// correct length, but should have a validity derived from the types array.
let padded_vals = itertools::izip!(&clean_vals, &ids)
let padded_vals = itertools::izip!(&clean_vals, ids.iter())
.map(|(sparse, id)| {
let validity = Bitmap::from(
typed_arr
Expand All @@ -206,12 +209,12 @@ impl ArrayExt for dyn Array {
})
.collect_vec();

let clean_field_types = itertools::izip!(fields, &clean_vals)
let clean_field_types = itertools::izip!(fields.iter(), &clean_vals)
.map(|(f, v)| Field::new(&f.name, v.data_type().clone(), f.is_nullable))
.collect_vec();

// The new type will be a struct
let clean_data = DataType::Struct(clean_field_types);
let clean_data = DataType::Struct(Arc::new(clean_field_types));

StructArray::try_new(clean_data, padded_vals, typed_arr.validity().cloned())
.unwrap()
Expand Down Expand Up @@ -246,53 +249,53 @@ fn test_clean_for_polars_modify() {
assert_eq!(
*cell.datatype(),
DataType::Union(
vec![
Arc::new(vec![
Field::new("Unknown", DataType::Boolean, false),
Field::new(
"Rigid3",
DataType::Struct(vec![
DataType::Struct(Arc::new(vec![
Field::new(
"rotation",
DataType::FixedSizeList(
Box::new(Field::new("item", DataType::Float32, false)),
Arc::new(Field::new("item", DataType::Float32, false)),
4
),
false
),
Field::new(
"translation",
DataType::FixedSizeList(
Box::new(Field::new("item", DataType::Float32, false)),
Arc::new(Field::new("item", DataType::Float32, false)),
3
),
false
)
]),
])),
false
),
Field::new(
"Pinhole",
DataType::Struct(vec![
DataType::Struct(Arc::new(vec![
Field::new(
"image_from_cam",
DataType::FixedSizeList(
Box::new(Field::new("item", DataType::Float32, false)),
Arc::new(Field::new("item", DataType::Float32, false)),
9
),
false,
),
Field::new(
"resolution",
DataType::FixedSizeList(
Box::new(Field::new("item", DataType::Float32, false)),
Arc::new(Field::new("item", DataType::Float32, false)),
2
),
true,
),
]),
])),
false
)
],
]),
None,
UnionMode::Dense
),
Expand All @@ -302,40 +305,40 @@ fn test_clean_for_polars_modify() {

assert_eq!(
*cleaned.data_type(),
DataType::Struct(vec![
DataType::Struct(Arc::new(vec![
Field::new("Unknown", DataType::Boolean, false),
Field::new(
"Rigid3",
DataType::Struct(vec![
DataType::Struct(Arc::new(vec![
Field::new(
"rotation",
DataType::List(Box::new(Field::new("item", DataType::Float32, false)),),
DataType::List(Arc::new(Field::new("item", DataType::Float32, false)),),
false
),
Field::new(
"translation",
DataType::List(Box::new(Field::new("item", DataType::Float32, false)),),
DataType::List(Arc::new(Field::new("item", DataType::Float32, false)),),
false
)
]),
])),
false
),
Field::new(
"Pinhole",
DataType::Struct(vec![
DataType::Struct(Arc::new(vec![
Field::new(
"image_from_cam",
DataType::List(Box::new(Field::new("item", DataType::Float32, false))),
DataType::List(Arc::new(Field::new("item", DataType::Float32, false))),
false,
),
Field::new(
"resolution",
DataType::List(Box::new(Field::new("item", DataType::Float32, false))),
DataType::List(Arc::new(Field::new("item", DataType::Float32, false))),
true,
),
]),
])),
false
)
],),
],)),
);
}
Loading