Skip to content

Commit

Permalink
Convert example components to arrow-rs
Browse files Browse the repository at this point in the history
  • Loading branch information
@emilk
emilk committed Dec 18, 2024
1 parent 3e792d2 commit 10d58b6
Show file tree
Hide file tree
Showing 3 changed files with 102 additions and 80 deletions.
1 change: 1 addition & 0 deletions Cargo.lock
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6014,6 +6014,7 @@ version = "0.21.0-alpha.1+dev"
dependencies = [
"ahash",
"anyhow",
"arrow",
"backtrace",
"bytemuck",
"clean-path",
Expand Down
1 change: 1 addition & 0 deletions crates/store/re_log_types/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ re_types_core.workspace = true
# External
ahash.workspace = true
anyhow.workspace = true
arrow.workspace = true
arrow2 = { workspace = true, features = [
"io_ipc",
"io_print",
Expand Down
180 changes: 100 additions & 80 deletions crates/store/re_log_types/src/example_components.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,68 +75,78 @@ impl SizeBytes for MyPoint {
}

impl Loggable for MyPoint {
fn arrow2_datatype() -> arrow2::datatypes::DataType {
use arrow2::datatypes::DataType::Float32;
arrow2::datatypes::DataType::Struct(Arc::new(vec![
arrow2::datatypes::Field::new("x", Float32, false),
arrow2::datatypes::Field::new("y", Float32, false),
fn arrow_datatype() -> arrow::datatypes::DataType {
use arrow::datatypes::DataType::Float32;
arrow::datatypes::DataType::Struct(arrow::datatypes::Fields::from(vec![
arrow::datatypes::Field::new("x", Float32, false),
arrow::datatypes::Field::new("y", Float32, false),
]))
}

fn to_arrow2_opt<'a>(
fn to_arrow_opt<'a>(
data: impl IntoIterator<Item = Option<impl Into<std::borrow::Cow<'a, Self>>>>,
) -> re_types_core::SerializationResult<Box<dyn arrow2::array::Array>>
) -> re_types_core::SerializationResult<arrow::array::ArrayRef>
where
Self: 'a,
{
use arrow::datatypes::DataType::Float32;

let (xs, ys): (Vec<_>, Vec<_>) = data
.into_iter()
.map(Option::unwrap)
.map(Into::into)
.map(|p| (p.x, p.y))
.unzip();

let x_array = arrow2::array::Float32Array::from_vec(xs).boxed();
let y_array = arrow2::array::Float32Array::from_vec(ys).boxed();
let x_array = Arc::new(arrow::array::Float32Array::from(xs));
let y_array = Arc::new(arrow::array::Float32Array::from(ys));

Ok(
arrow2::array::StructArray::new(Self::arrow2_datatype(), vec![x_array, y_array], None)
.boxed(),
)
Ok(Arc::new(arrow::array::StructArray::new(
arrow::datatypes::Fields::from(vec![
arrow::datatypes::Field::new("x", Float32, false),
arrow::datatypes::Field::new("y", Float32, false),
]),
vec![x_array, y_array],
None,
)))
}

fn from_arrow2_opt(
data: &dyn arrow2::array::Array,
fn from_arrow_opt(
data: &dyn arrow::array::Array,
) -> re_types_core::DeserializationResult<Vec<Option<Self>>> {
let array = data
.as_any()
.downcast_ref::<arrow2::array::StructArray>()
.downcast_ref::<arrow::array::StructArray>()
.ok_or(DeserializationError::downcast_error::<
arrow2::array::StructArray,
arrow::array::StructArray,
>())?;

let x_array = array.values()[0].as_ref();
let y_array = array.values()[1].as_ref();
let x_array = array.columns()[0].as_ref();
let y_array = array.columns()[1].as_ref();

let xs = x_array
.as_any()
.downcast_ref::<arrow2::array::Float32Array>()
.downcast_ref::<arrow::array::Float32Array>()
.ok_or(DeserializationError::downcast_error::<
arrow2::array::Float32Array,
arrow::array::Float32Array,
>())?;
let ys = y_array
.as_any()
.downcast_ref::<arrow2::array::Float32Array>()
.downcast_ref::<arrow::array::Float32Array>()
.ok_or(DeserializationError::downcast_error::<
arrow2::array::Float32Array,
arrow::array::Float32Array,
>())?;

Ok(xs
.values_iter()
.copied()
.zip(ys.values_iter().copied())
.map(|(x, y)| Self { x, y })
.map(Some)
.iter()
.zip(ys.iter())
.map(|(x, y)| {
if let (Some(x), Some(y)) = (x, y) {
Some(Self { x, y })
} else {
None
}
})
.collect())
}
}
Expand Down Expand Up @@ -184,68 +194,78 @@ impl SizeBytes for MyPoint64 {
}

impl Loggable for MyPoint64 {
fn arrow2_datatype() -> arrow2::datatypes::DataType {
use arrow2::datatypes::DataType::Float64;
arrow2::datatypes::DataType::Struct(Arc::new(vec![
arrow2::datatypes::Field::new("x", Float64, false),
arrow2::datatypes::Field::new("y", Float64, false),
fn arrow_datatype() -> arrow::datatypes::DataType {
use arrow::datatypes::DataType::Float64;
arrow::datatypes::DataType::Struct(arrow::datatypes::Fields::from(vec![
arrow::datatypes::Field::new("x", Float64, false),
arrow::datatypes::Field::new("y", Float64, false),
]))
}

fn to_arrow2_opt<'a>(
fn to_arrow_opt<'a>(
data: impl IntoIterator<Item = Option<impl Into<std::borrow::Cow<'a, Self>>>>,
) -> re_types_core::SerializationResult<Box<dyn arrow2::array::Array>>
) -> re_types_core::SerializationResult<arrow::array::ArrayRef>
where
Self: 'a,
{
use arrow::datatypes::DataType::Float64;

let (xs, ys): (Vec<_>, Vec<_>) = data
.into_iter()
.map(Option::unwrap)
.map(Into::into)
.map(|p| (p.x, p.y))
.unzip();

let x_array = arrow2::array::Float64Array::from_vec(xs).boxed();
let y_array = arrow2::array::Float64Array::from_vec(ys).boxed();
let x_array = Arc::new(arrow::array::Float64Array::from(xs));
let y_array = Arc::new(arrow::array::Float64Array::from(ys));

Ok(
arrow2::array::StructArray::new(Self::arrow2_datatype(), vec![x_array, y_array], None)
.boxed(),
)
Ok(Arc::new(arrow::array::StructArray::new(
arrow::datatypes::Fields::from(vec![
arrow::datatypes::Field::new("x", Float64, false),
arrow::datatypes::Field::new("y", Float64, false),
]),
vec![x_array, y_array],
None,
)))
}

fn from_arrow2_opt(
data: &dyn arrow2::array::Array,
fn from_arrow_opt(
data: &dyn arrow::array::Array,
) -> re_types_core::DeserializationResult<Vec<Option<Self>>> {
let array = data
.as_any()
.downcast_ref::<arrow2::array::StructArray>()
.downcast_ref::<arrow::array::StructArray>()
.ok_or(DeserializationError::downcast_error::<
arrow2::array::StructArray,
arrow::array::StructArray,
>())?;

let x_array = array.values()[0].as_ref();
let y_array = array.values()[1].as_ref();
let x_array = array.columns()[0].as_ref();
let y_array = array.columns()[1].as_ref();

let xs = x_array
.as_any()
.downcast_ref::<arrow2::array::Float64Array>()
.downcast_ref::<arrow::array::Float64Array>()
.ok_or(DeserializationError::downcast_error::<
arrow2::array::Float64Array,
arrow::array::Float64Array,
>())?;
let ys = y_array
.as_any()
.downcast_ref::<arrow2::array::Float64Array>()
.downcast_ref::<arrow::array::Float64Array>()
.ok_or(DeserializationError::downcast_error::<
arrow2::array::Float64Array,
arrow::array::Float64Array,
>())?;

Ok(xs
.values_iter()
.copied()
.zip(ys.values_iter().copied())
.map(|(x, y)| Self { x, y })
.map(Some)
.iter()
.zip(ys.iter())
.map(|(x, y)| {
if let (Some(x), Some(y)) = (x, y) {
Some(Self { x, y })
} else {
None
}
})
.collect())
}
}
Expand Down Expand Up @@ -296,28 +316,28 @@ impl SizeBytes for MyColor {
}

impl Loggable for MyColor {
fn arrow2_datatype() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::UInt32
fn arrow_datatype() -> arrow::datatypes::DataType {
arrow::datatypes::DataType::UInt32
}

fn to_arrow2_opt<'a>(
fn to_arrow_opt<'a>(
data: impl IntoIterator<Item = Option<impl Into<std::borrow::Cow<'a, Self>>>>,
) -> re_types_core::SerializationResult<Box<dyn arrow2::array::Array>>
) -> re_types_core::SerializationResult<arrow::array::ArrayRef>
where
Self: 'a,
{
use re_types_core::datatypes::UInt32;
UInt32::to_arrow2_opt(
UInt32::to_arrow_opt(
data.into_iter()
.map(|opt| opt.map(Into::into).map(|c| UInt32(c.0))),
)
}

fn from_arrow2_opt(
data: &dyn arrow2::array::Array,
fn from_arrow_opt(
data: &dyn arrow::array::Array,
) -> re_types_core::DeserializationResult<Vec<Option<Self>>> {
use re_types_core::datatypes::UInt32;
Ok(UInt32::from_arrow2_opt(data)?
Ok(UInt32::from_arrow_opt(data)?
.into_iter()
.map(|opt| opt.map(|v| Self(v.0)))
.collect())
Expand Down Expand Up @@ -347,28 +367,28 @@ impl SizeBytes for MyLabel {
}

impl Loggable for MyLabel {
fn arrow2_datatype() -> arrow2::datatypes::DataType {
re_types_core::datatypes::Utf8::arrow2_datatype()
fn arrow_datatype() -> arrow::datatypes::DataType {
re_types_core::datatypes::Utf8::arrow_datatype()
}

fn to_arrow2_opt<'a>(
fn to_arrow_opt<'a>(
data: impl IntoIterator<Item = Option<impl Into<std::borrow::Cow<'a, Self>>>>,
) -> re_types_core::SerializationResult<Box<dyn arrow2::array::Array>>
) -> re_types_core::SerializationResult<arrow::array::ArrayRef>
where
Self: 'a,
{
use re_types_core::datatypes::Utf8;
Utf8::to_arrow2_opt(
Utf8::to_arrow_opt(
data.into_iter()
.map(|opt| opt.map(Into::into).map(|l| Utf8(l.0.clone().into()))),
)
}

fn from_arrow2_opt(
data: &dyn arrow2::array::Array,
fn from_arrow_opt(
data: &dyn arrow::array::Array,
) -> re_types_core::DeserializationResult<Vec<Option<Self>>> {
use re_types_core::datatypes::Utf8;
Ok(Utf8::from_arrow2_opt(data)?
Ok(Utf8::from_arrow_opt(data)?
.into_iter()
.map(|opt| opt.map(|v| Self(v.0.to_string())))
.collect())
Expand Down Expand Up @@ -407,28 +427,28 @@ impl SizeBytes for MyIndex {
}

impl Loggable for MyIndex {
fn arrow2_datatype() -> arrow2::datatypes::DataType {
arrow2::datatypes::DataType::UInt64
fn arrow_datatype() -> arrow::datatypes::DataType {
arrow::datatypes::DataType::UInt64
}

fn to_arrow2_opt<'a>(
fn to_arrow_opt<'a>(
data: impl IntoIterator<Item = Option<impl Into<std::borrow::Cow<'a, Self>>>>,
) -> re_types_core::SerializationResult<Box<dyn arrow2::array::Array>>
) -> re_types_core::SerializationResult<arrow::array::ArrayRef>
where
Self: 'a,
{
use re_types_core::datatypes::UInt64;
UInt64::to_arrow2_opt(
UInt64::to_arrow_opt(
data.into_iter()
.map(|opt| opt.map(Into::into).map(|c| UInt64(c.0))),
)
}

fn from_arrow2_opt(
data: &dyn arrow2::array::Array,
fn from_arrow_opt(
data: &dyn arrow::array::Array,
) -> re_types_core::DeserializationResult<Vec<Option<Self>>> {
use re_types_core::datatypes::UInt64;
Ok(UInt64::from_arrow2_opt(data)?
Ok(UInt64::from_arrow_opt(data)?
.into_iter()
.map(|opt| opt.map(|v| Self(v.0)))
.collect())
Expand Down

0 comments on commit 10d58b6

Please sign in to comment.