Skip to content

Commit

Permalink
Columnar APIs: autogenerate helpers for unit-length batches (#8775)
Browse files Browse the repository at this point in the history
This makes the API far easier to use in the very common case where the
caller wants to split all of their data in unit-length batches.
I think it's worth the extra generated code.

Notes:
* The naming sucks.
* I disabled columnar APIs entirely for all blueprint stuff. I don't see
any reason to have all this code?

```diff
let times = TimeColumn::new_sequence("step", 0..STEPS);
let scalars = (0..STEPS).map(|step| (step as f64 / 10.0).sin());

rec.send_columns_v2(
    "scalars",
    [times],
    rerun::Scalar::update_fields()
        .with_many_scalar(scalars)
-        .columns(std::iter::repeat(1).take(STEPS as _))?,
+        .unary_columns()?,
)?;
```

* Follow-up to #8753
  • Loading branch information
teh-cmc authored Jan 23, 2025
1 parent 67df6b1 commit 12d26c8
Show file tree
Hide file tree
Showing 51 changed files with 591 additions and 758 deletions.
38 changes: 36 additions & 2 deletions crates/build/re_types_builder/src/codegen/rust/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1944,7 +1944,7 @@ fn quote_builder_from_obj(reporter: &Reporter, objects: &Objects, obj: &Object)
}
});

let columnar_methods = obj.is_eager_rust_archetype().then(|| {
let columnar_methods = (obj.is_eager_rust_archetype() && obj.scope().is_none()).then(|| {
let columns_doc = unindent::unindent("\
Partitions the component data into multiple sub-batches.
Expand All @@ -1959,15 +1959,32 @@ fn quote_builder_from_obj(reporter: &Reporter, objects: &Objects, obj: &Object)
");
let columns_doc = quote_doc_lines(&columns_doc.lines().map(|l| l.to_owned()).collect_vec());

let has_indicator = obj.fqname.as_str() != "rerun.archetypes.Scalar";
let columns_unary_doc = unindent::unindent("\
Helper to partition the component data into unit-length sub-batches.
This is semantically similar to calling [`Self::columns`] with `std::iter::take(1).repeat(n)`,
where `n` is automatically guessed.
");
let columns_unary_doc = quote_doc_lines(&columns_unary_doc.lines().map(|l| l.to_owned()).collect_vec());

let has_indicator = obj.fqname.as_str() != "rerun.archetypes.Scalar";
let num_fields = required.iter().chain(optional.iter()).count();

let fields = required.iter().chain(optional.iter()).map(|field| {
let field_name = format_ident!("{}", field.name);
let clone = if num_fields == 1 && !has_indicator { quote!(.into_iter()) } else { quote!(.clone()) };
quote!(self.#field_name.map(|#field_name| #field_name.partitioned(_lengths #clone)).transpose()?)
});

let field_lengths = required.iter().chain(optional.iter()).map(|field| {
format_ident!("len_{}", field.name)
}).collect_vec();
let unary_fields = required.iter().chain(optional.iter()).map(|field| {
let field_name = format_ident!("{}", field.name);
let len_field_name = format_ident!("len_{}", field.name);
quote!(let #len_field_name = self.#field_name.as_ref().map(|b| b.array.len()))
});

let indicator_column = if !has_indicator {
// NOTE(#8768): Scalar indicators are extremely wasteful, and not actually used for anything.
quote!(None)
Expand All @@ -1989,6 +2006,23 @@ fn quote_builder_from_obj(reporter: &Reporter, objects: &Objects, obj: &Object)
let indicator_column = #indicator_column;
Ok(columns.into_iter().chain([indicator_column]).flatten())
}

#columns_unary_doc
#[inline]
pub fn columns_of_unit_batches(
self,
) -> SerializationResult<impl Iterator<Item = ::re_types_core::SerializedComponentColumn>>
{
#(#unary_fields);*;

let len = None
#(.or(#field_lengths))*
.unwrap_or(0);

// NOTE: This will return an error if the different batches have different lengths,
// which is fine.
self.columns(std::iter::repeat(1).take(len))
}
}
});

Expand Down
29 changes: 29 additions & 0 deletions crates/store/re_types/src/archetypes/arrows2d.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

27 changes: 27 additions & 0 deletions crates/store/re_types/src/archetypes/arrows3d.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

14 changes: 14 additions & 0 deletions crates/store/re_types/src/archetypes/bar_chart.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

31 changes: 31 additions & 0 deletions crates/store/re_types/src/archetypes/capsules3d.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

27 changes: 27 additions & 0 deletions crates/store/re_types/src/archetypes/depth_image.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

33 changes: 33 additions & 0 deletions crates/store/re_types/src/archetypes/ellipsoids3d.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

21 changes: 21 additions & 0 deletions crates/store/re_types/src/archetypes/encoded_image.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

19 changes: 19 additions & 0 deletions crates/store/re_types/src/archetypes/geo_line_strings.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

21 changes: 21 additions & 0 deletions crates/store/re_types/src/archetypes/geo_points.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

14 changes: 14 additions & 0 deletions crates/store/re_types/src/archetypes/graph_edges.rs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading

0 comments on commit 12d26c8

Please sign in to comment.