Skip to content

Commit

Permalink
Codegen more of our Python arrow serialization (#6154)
Browse files Browse the repository at this point in the history
### What
This implements codegen of python arrow serialization in cases where we
can just delegate to our datatypes.

### Checklist
* [x] I have read and agree to [Contributor
Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and
the [Code of
Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md)
* [x] I've included a screenshot or gif (if applicable)
* [x] I have tested the web demo (if applicable):
* Using examples from latest `main` build:
[rerun.io/viewer](https://rerun.io/viewer/pr/6154?manifest_url=https://app.rerun.io/version/main/examples_manifest.json)
* Using full set of examples from `nightly` build:
[rerun.io/viewer](https://rerun.io/viewer/pr/6154?manifest_url=https://app.rerun.io/version/nightly/examples_manifest.json)
* [x] The PR title and labels are set such as to maximize their
usefulness for the next release's CHANGELOG
* [x] If applicable, add a new check to the [release
checklist](https://github.com/rerun-io/rerun/blob/main/tests/python/release_checklist)!

- [PR Build Summary](https://build.rerun.io/pr/6154)
- [Recent benchmark results](https://build.rerun.io/graphs/crates.html)
- [Wasm size tracking](https://build.rerun.io/graphs/sizes.html)

To run all checks from `main`, comment on the PR with `@rerun-bot
full-check`.
  • Loading branch information
emilk committed Apr 29, 2024
1 parent e4a86aa commit 8760aa9
Show file tree
Hide file tree
Showing 31 changed files with 142 additions and 88 deletions.
48 changes: 46 additions & 2 deletions crates/re_types_builder/src/codegen/python/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1841,7 +1841,7 @@ fn np_dtype_from_type(t: &Type) -> Option<&'static str> {
/// Only implemented for some cases.
fn quote_arrow_serialization(
reporter: &Reporter,
_objects: &Objects,
objects: &Objects,
obj: &Object,
) -> Result<String, String> {
let Object { name, .. } = obj;
Expand Down Expand Up @@ -1888,7 +1888,51 @@ fn quote_arrow_serialization(
}
}

Err("We lack codegen for arrow-serialization of structs".to_owned())
let mut code = String::new();

for field in &obj.fields {
let Type::Object(field_fqname) = &field.typ else {
return Err(
"We lack codegen for arrow-serialization of general structs".to_owned()
);
};
if let Some(last_dot) = field_fqname.rfind('.') {
let mod_path = &field_fqname[..last_dot];
let field_type_name = &field_fqname[last_dot + 1..];
code.push_indented(
0,
format!("from {mod_path} import {field_type_name}Batch"),
2,
);
}
}

code.push_indented(0, &format!("if isinstance(data, {name}):"), 1);
code.push_indented(1, "data = [data]", 2);

code.push_indented(0, "return pa.StructArray.from_arrays(", 1);
code.push_indented(1, "[", 1);
for field in &obj.fields {
let Type::Object(field_fqname) = &field.typ else {
return Err(
"We lack codegen for arrow-serialization of general structs".to_owned()
);
};
let field_obj = &objects[field_fqname];
let field_type_name = &field_obj.name;
let field_name = &field.name;
let field_batch_type = format!("{field_type_name}Batch");
// let field_batch_type = format!("datatypes.{field_type_name}Batch");
let field_array = format!("[x.{field_name} for x in data]");
let field_fwd =
format!("{field_batch_type}({field_array}).as_arrow_array().storage,");
code.push_indented(2, &field_fwd, 1);
}
code.push_indented(1, "],", 1);
code.push_indented(1, "fields=list(data_type),", 1);
code.push_indented(0, ")", 1);

Ok(code)
}

ObjectClass::Enum => {
Expand Down

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

15 changes: 12 additions & 3 deletions rerun_py/rerun_sdk/rerun/blueprint/datatypes/visible_time_range.py

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

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

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

16 changes: 13 additions & 3 deletions rerun_py/rerun_sdk/rerun/datatypes/aabb2d.py

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

27 changes: 0 additions & 27 deletions rerun_py/rerun_sdk/rerun/datatypes/aabb2d_ext.py

This file was deleted.

13 changes: 12 additions & 1 deletion rerun_py/rerun_sdk/rerun/datatypes/rotation_axis_angle.py

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

22 changes: 1 addition & 21 deletions rerun_py/rerun_sdk/rerun/datatypes/rotation_axis_angle_ext.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,8 @@

from typing import TYPE_CHECKING, Any

import pyarrow as pa

if TYPE_CHECKING:
from . import Angle, AngleLike, RotationAxisAngleArrayLike, Vec3DLike
from . import Angle, AngleLike, Vec3DLike


class RotationAxisAngleExt:
Expand Down Expand Up @@ -58,21 +56,3 @@ def angle__field_converter_override(x: AngleLike) -> Angle:
return x
else:
return Angle(rad=x)

@staticmethod
def native_to_pa_array_override(data: RotationAxisAngleArrayLike, data_type: pa.DataType) -> pa.Array:
from . import AngleBatch, RotationAxisAngle, Vec3DBatch

if isinstance(data, RotationAxisAngle):
data = [data]

axis_pa_array = Vec3DBatch._native_to_pa_array([rotation.axis for rotation in data], data_type["axis"].type)
angle_pa_arr = AngleBatch._native_to_pa_array([rotation.angle for rotation in data], data_type["angle"].type)

return pa.StructArray.from_arrays(
[
axis_pa_array,
angle_pa_arr,
],
fields=list(data_type),
)
2 changes: 1 addition & 1 deletion rerun_py/rerun_sdk/rerun/datatypes/tensor_dimension.py

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

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

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

2 changes: 1 addition & 1 deletion rerun_py/rerun_sdk/rerun/datatypes/uvec2d.py

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

2 changes: 1 addition & 1 deletion rerun_py/rerun_sdk/rerun/datatypes/uvec3d.py

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

2 changes: 1 addition & 1 deletion rerun_py/rerun_sdk/rerun/datatypes/uvec4d.py

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

2 changes: 1 addition & 1 deletion rerun_py/tests/test_types/components/affix_fuzzer11.py

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

2 changes: 1 addition & 1 deletion rerun_py/tests/test_types/components/affix_fuzzer12.py

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

2 changes: 1 addition & 1 deletion rerun_py/tests/test_types/components/affix_fuzzer13.py

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

2 changes: 1 addition & 1 deletion rerun_py/tests/test_types/components/affix_fuzzer16.py

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

2 changes: 1 addition & 1 deletion rerun_py/tests/test_types/components/affix_fuzzer17.py

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

2 changes: 1 addition & 1 deletion rerun_py/tests/test_types/components/affix_fuzzer18.py

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

2 changes: 1 addition & 1 deletion rerun_py/tests/test_types/components/affix_fuzzer7.py

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

2 changes: 1 addition & 1 deletion rerun_py/tests/test_types/components/affix_fuzzer8.py

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

2 changes: 1 addition & 1 deletion rerun_py/tests/test_types/datatypes/affix_fuzzer1.py

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

2 changes: 1 addition & 1 deletion rerun_py/tests/test_types/datatypes/affix_fuzzer2.py

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

15 changes: 12 additions & 3 deletions rerun_py/tests/test_types/datatypes/affix_fuzzer20.py

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

2 changes: 1 addition & 1 deletion rerun_py/tests/test_types/datatypes/affix_fuzzer21.py

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

0 comments on commit 8760aa9

Please sign in to comment.