diff --git a/crates/re_log_types/src/data_row.rs b/crates/re_log_types/src/data_row.rs index b1091e4868cd..c2b48554e9fd 100644 --- a/crates/re_log_types/src/data_row.rs +++ b/crates/re_log_types/src/data_row.rs @@ -46,7 +46,7 @@ pub type DataRowResult = ::std::result::Result; // --- -type DataCellVec = SmallVec<[DataCell; 4]>; +pub type DataCellVec = SmallVec<[DataCell; 4]>; /// A row's worth of [`DataCell`]s: a collection of independent [`DataCell`]s with different /// underlying datatypes and pointing to different parts of the heap. diff --git a/crates/re_log_types/src/lib.rs b/crates/re_log_types/src/lib.rs index 0bf8eb04ce1a..f43efb1d7dad 100644 --- a/crates/re_log_types/src/lib.rs +++ b/crates/re_log_types/src/lib.rs @@ -44,7 +44,7 @@ use std::sync::Arc; pub use self::arrow_msg::ArrowMsg; pub use self::component::{Component, DeserializableComponent, SerializableComponent}; pub use self::data_cell::{DataCell, DataCellError, DataCellInner, DataCellResult}; -pub use self::data_row::{DataRow, DataRowError, DataRowResult, RowId}; +pub use self::data_row::{DataCellVec, DataRow, DataRowError, DataRowResult, RowId}; pub use self::data_table::{ DataCellColumn, DataCellOptVec, DataTable, DataTableError, DataTableResult, EntityPathVec, ErasedTimeVec, NumInstancesVec, RowIdVec, TableId, TimePointVec, COLUMN_ENTITY_PATH, diff --git a/crates/re_types/definitions/rerun/archetypes/points2d.fbs b/crates/re_types/definitions/rerun/archetypes/points2d.fbs index 6903b3d26562..a483faa3aae5 100644 --- a/crates/re_types/definitions/rerun/archetypes/points2d.fbs +++ b/crates/re_types/definitions/rerun/archetypes/points2d.fbs @@ -11,6 +11,7 @@ namespace rerun.archetypes; // TODO(#2372): archetype IDL definitions must refer to objects of kind component // TODO(#2373): `attr.rerun.component_required` implies `required` // TODO(#2427): distinguish optional vs. recommended in language backends +// TODO(#2521): always derive debug & clone for rust backend /// A 2D point cloud with positions and optional colors, radii, labels, etc. table Points2D ( diff --git a/crates/re_types/source_hash.txt b/crates/re_types/source_hash.txt index 3f01ac270f24..672ccc20c301 100644 --- a/crates/re_types/source_hash.txt +++ b/crates/re_types/source_hash.txt @@ -1,4 +1,4 @@ # This is a sha256 hash for all direct and indirect dependencies of this crate's build script. # It can be safely removed at anytime to force the build script to run again. # Check out build.rs to see how it's computed. -0960d9b4f6df9136f7857a7b7280a4803f3eba7a085c98aa1ce7c95dcd88539e \ No newline at end of file +c10dc39333002ce5c62d9e88a7feb4fca76098528fe643012a53665a9934581e \ No newline at end of file diff --git a/crates/re_types/src/components/class_id.rs b/crates/re_types/src/components/class_id.rs index e5c03e11b5d7..24520b7e5040 100644 --- a/crates/re_types/src/components/class_id.rs +++ b/crates/re_types/src/components/class_id.rs @@ -14,6 +14,10 @@ impl crate::Component for ClassId { #[allow(clippy::wildcard_imports)] fn to_arrow_datatype() -> arrow2::datatypes::DataType { use ::arrow2::datatypes::*; - DataType::UInt16 + DataType::Extension( + "rerun.components.ClassId".to_owned(), + Box::new(DataType::UInt16), + None, + ) } } diff --git a/crates/re_types/src/components/color.rs b/crates/re_types/src/components/color.rs index 1de749cdb19a..5c7e6744b545 100644 --- a/crates/re_types/src/components/color.rs +++ b/crates/re_types/src/components/color.rs @@ -24,6 +24,10 @@ impl crate::Component for Color { #[allow(clippy::wildcard_imports)] fn to_arrow_datatype() -> arrow2::datatypes::DataType { use ::arrow2::datatypes::*; - DataType::UInt32 + DataType::Extension( + "rerun.components.Color".to_owned(), + Box::new(DataType::UInt32), + None, + ) } } diff --git a/crates/re_types/src/components/draw_order.rs b/crates/re_types/src/components/draw_order.rs index 1ed1c97acedf..b69822b74457 100644 --- a/crates/re_types/src/components/draw_order.rs +++ b/crates/re_types/src/components/draw_order.rs @@ -19,6 +19,10 @@ impl crate::Component for DrawOrder { #[allow(clippy::wildcard_imports)] fn to_arrow_datatype() -> arrow2::datatypes::DataType { use ::arrow2::datatypes::*; - DataType::Float32 + DataType::Extension( + "rerun.components.DrawOrder".to_owned(), + Box::new(DataType::Float32), + None, + ) } } diff --git a/crates/re_types/src/components/instance_key.rs b/crates/re_types/src/components/instance_key.rs index 229858d46cca..a1df2a09f429 100644 --- a/crates/re_types/src/components/instance_key.rs +++ b/crates/re_types/src/components/instance_key.rs @@ -12,6 +12,10 @@ impl crate::Component for InstanceKey { #[allow(clippy::wildcard_imports)] fn to_arrow_datatype() -> arrow2::datatypes::DataType { use ::arrow2::datatypes::*; - DataType::UInt64 + DataType::Extension( + "rerun.components.InstanceKey".to_owned(), + Box::new(DataType::UInt64), + None, + ) } } diff --git a/crates/re_types/src/components/keypoint_id.rs b/crates/re_types/src/components/keypoint_id.rs index dc052c879909..b5646db5f0aa 100644 --- a/crates/re_types/src/components/keypoint_id.rs +++ b/crates/re_types/src/components/keypoint_id.rs @@ -16,6 +16,10 @@ impl crate::Component for KeypointId { #[allow(clippy::wildcard_imports)] fn to_arrow_datatype() -> arrow2::datatypes::DataType { use ::arrow2::datatypes::*; - DataType::UInt16 + DataType::Extension( + "rerun.components.KeypointId".to_owned(), + Box::new(DataType::UInt16), + None, + ) } } diff --git a/crates/re_types/src/components/label.rs b/crates/re_types/src/components/label.rs index 782e45f32196..bfc8421e231f 100644 --- a/crates/re_types/src/components/label.rs +++ b/crates/re_types/src/components/label.rs @@ -13,6 +13,10 @@ impl crate::Component for Label { #[allow(clippy::wildcard_imports)] fn to_arrow_datatype() -> arrow2::datatypes::DataType { use ::arrow2::datatypes::*; - DataType::Utf8 + DataType::Extension( + "rerun.components.Label".to_owned(), + Box::new(DataType::Utf8), + None, + ) } } diff --git a/crates/re_types/src/components/point2d.rs b/crates/re_types/src/components/point2d.rs index 98f057872886..19df9df7dc73 100644 --- a/crates/re_types/src/components/point2d.rs +++ b/crates/re_types/src/components/point2d.rs @@ -21,13 +21,13 @@ impl crate::Component for Point2D { Field { name: "x".to_owned(), data_type: DataType::Float32, - is_nullable: true, + is_nullable: false, metadata: [].into(), }, Field { name: "y".to_owned(), data_type: DataType::Float32, - is_nullable: true, + is_nullable: false, metadata: [].into(), }, ])), diff --git a/crates/re_types/src/components/radius.rs b/crates/re_types/src/components/radius.rs index 02fffa83b587..2dec841b3dca 100644 --- a/crates/re_types/src/components/radius.rs +++ b/crates/re_types/src/components/radius.rs @@ -12,6 +12,10 @@ impl crate::Component for Radius { #[allow(clippy::wildcard_imports)] fn to_arrow_datatype() -> arrow2::datatypes::DataType { use ::arrow2::datatypes::*; - DataType::Float32 + DataType::Extension( + "rerun.components.Radius".to_owned(), + Box::new(DataType::Float32), + None, + ) } } diff --git a/crates/re_types/src/datatypes/vec2d.rs b/crates/re_types/src/datatypes/vec2d.rs index 0a1338fd216b..efdeb64097cd 100644 --- a/crates/re_types/src/datatypes/vec2d.rs +++ b/crates/re_types/src/datatypes/vec2d.rs @@ -12,14 +12,18 @@ impl crate::Datatype for Vec2D { #[allow(clippy::wildcard_imports)] fn to_arrow_datatype() -> arrow2::datatypes::DataType { use ::arrow2::datatypes::*; - DataType::FixedSizeList( - Box::new(Field { - name: "item".to_owned(), - data_type: DataType::Float32, - is_nullable: false, - metadata: [].into(), - }), - 2usize, + DataType::Extension( + "rerun.datatypes.Vec2D".to_owned(), + Box::new(DataType::FixedSizeList( + Box::new(Field { + name: "item".to_owned(), + data_type: DataType::Float32, + is_nullable: false, + metadata: [].into(), + }), + 2usize, + )), + None, ) } } diff --git a/crates/re_types_builder/src/arrow_registry.rs b/crates/re_types_builder/src/arrow_registry.rs index 4a59b18bac0c..041eda4a7c67 100644 --- a/crates/re_types_builder/src/arrow_registry.rs +++ b/crates/re_types_builder/src/arrow_registry.rs @@ -4,7 +4,9 @@ use anyhow::Context as _; use arrow2::datatypes::{DataType, Field, UnionMode}; use std::collections::{BTreeMap, HashMap}; -use crate::{ElementType, Object, Type, ATTR_ARROW_SPARSE_UNION, ATTR_ARROW_TRANSPARENT}; +use crate::{ + ElementType, Object, ObjectField, Type, ATTR_ARROW_SPARSE_UNION, ATTR_ARROW_TRANSPARENT, +}; // --- Registry --- @@ -18,7 +20,7 @@ pub struct ArrowRegistry { impl ArrowRegistry { /// Computes the Arrow datatype for the specified object and stores it in the registry, to be /// resolved later on. - pub fn register(&mut self, obj: &Object) { + pub fn register(&mut self, obj: &mut Object) { let (fqname, datatype) = (obj.fqname.clone(), self.arrow_datatype_from_object(obj)); self.registry.insert(fqname, datatype); } @@ -48,7 +50,7 @@ impl ArrowRegistry { // --- - fn arrow_datatype_from_object(&self, obj: &Object) -> LazyDatatype { + fn arrow_datatype_from_object(&mut self, obj: &mut Object) -> LazyDatatype { let is_struct = obj.is_struct(); let is_transparent = obj.try_get_attr::(ATTR_ARROW_TRANSPARENT).is_some(); let num_fields = obj.fields.len(); @@ -59,18 +61,25 @@ impl ArrowRegistry { obj.fqname, ); - if is_transparent { - self.arrow_datatype_from_type(&obj.fields[0].typ) + let datatype = if is_transparent { + LazyDatatype::Extension( + obj.fqname.clone(), + Box::new( + self.arrow_datatype_from_type(obj.fields[0].typ.clone(), &mut obj.fields[0]), + ), + None, + ) } else if is_struct { LazyDatatype::Extension( obj.fqname.clone(), Box::new(LazyDatatype::Struct( obj.fields - .iter() - .map(|field| LazyField { - name: field.name.clone(), - datatype: self.arrow_datatype_from_type(&field.typ), - is_nullable: field.required, + .iter_mut() + .map(|obj_field| LazyField { + name: obj_field.name.clone(), + datatype: self + .arrow_datatype_from_type(obj_field.typ.clone(), obj_field), + is_nullable: obj_field.is_nullable, metadata: Default::default(), }) .collect(), @@ -85,10 +94,10 @@ impl ArrowRegistry { obj.fqname.clone(), Box::new(LazyDatatype::Union( obj.fields - .iter() + .iter_mut() .map(|field| LazyField { name: field.name.clone(), - datatype: self.arrow_datatype_from_type(&field.typ), + datatype: self.arrow_datatype_from_type(field.typ.clone(), field), is_nullable: false, metadata: Default::default(), }) @@ -102,11 +111,18 @@ impl ArrowRegistry { )), None, ) + }; + + // NOTE: Arrow-transparent objects by definition don't have a datatype of their own. + if !is_transparent { + obj.datatype = datatype.clone().into(); } + + datatype } - fn arrow_datatype_from_type(&self, typ: &Type) -> LazyDatatype { - match typ { + fn arrow_datatype_from_type(&mut self, typ: Type, field: &mut ObjectField) -> LazyDatatype { + let datatype = match typ { Type::UInt8 => LazyDatatype::UInt8, Type::UInt16 => LazyDatatype::UInt16, Type::UInt32 => LazyDatatype::UInt32, @@ -124,22 +140,27 @@ impl ArrowRegistry { Box::new(LazyField { name: "item".into(), datatype: self.arrow_datatype_from_element_type(elem_type), - is_nullable: false, + is_nullable: field.is_nullable, metadata: Default::default(), }), - *length, + length, ), Type::Vector { elem_type } => LazyDatatype::List(Box::new(LazyField { name: "item".into(), datatype: self.arrow_datatype_from_element_type(elem_type), - is_nullable: false, + is_nullable: field.is_nullable, metadata: Default::default(), })), - Type::Object(fqname) => LazyDatatype::Unresolved(fqname.clone()), - } + Type::Object(fqname) => LazyDatatype::Unresolved(fqname), + }; + + field.datatype = datatype.clone().into(); + self.registry.insert(field.fqname.clone(), datatype.clone()); + + datatype } - fn arrow_datatype_from_element_type(&self, typ: &ElementType) -> LazyDatatype { + fn arrow_datatype_from_element_type(&self, typ: ElementType) -> LazyDatatype { _ = self; match typ { ElementType::UInt8 => LazyDatatype::UInt8, @@ -155,7 +176,7 @@ impl ArrowRegistry { ElementType::Float32 => LazyDatatype::Float32, ElementType::Float64 => LazyDatatype::Float64, ElementType::String => LazyDatatype::Utf8, - ElementType::Object(fqname) => LazyDatatype::Unresolved(fqname.clone()), + ElementType::Object(fqname) => LazyDatatype::Unresolved(fqname), } } } diff --git a/crates/re_types_builder/src/codegen/python.rs b/crates/re_types_builder/src/codegen/python.rs index f459b3ce2441..6c0a1404893d 100644 --- a/crates/re_types_builder/src/codegen/python.rs +++ b/crates/re_types_builder/src/codegen/python.rs @@ -2,7 +2,7 @@ use anyhow::Context as _; use std::{ - collections::{BTreeMap, HashMap}, + collections::{BTreeMap, HashMap, HashSet}, io::Write, path::{Path, PathBuf}, }; @@ -41,7 +41,7 @@ impl CodeGenerator for PythonCodeGenerator { datatypes_path, arrow_registry, objs, - &objs.ordered_datatypes(), + &objs.ordered_objects(ObjectKind::Datatype.into()), ) .0, ); @@ -55,7 +55,7 @@ impl CodeGenerator for PythonCodeGenerator { components_path, arrow_registry, objs, - &objs.ordered_components(), + &objs.ordered_objects(ObjectKind::Component.into()), ) .0, ); @@ -68,7 +68,7 @@ impl CodeGenerator for PythonCodeGenerator { archetypes_path, arrow_registry, objs, - &objs.ordered_archetypes(), + &objs.ordered_objects(ObjectKind::Archetype.into()), ); filepaths.extend(paths); @@ -145,9 +145,9 @@ fn quote_objects( for (filepath, objs) in files { let names = objs .iter() - .flat_map(|obj| match obj.kind { + .flat_map(|obj| match obj.object.kind { ObjectKind::Datatype | ObjectKind::Component => { - let name = &obj.name; + let name = &obj.object.name; vec![ format!("{name}"), @@ -157,7 +157,7 @@ fn quote_objects( format!("{name}Type"), ] } - ObjectKind::Archetype => vec![obj.name.clone()], + ObjectKind::Archetype => vec![obj.object.name.clone()], }) .collect::>(); @@ -181,6 +181,13 @@ fn quote_objects( " from __future__ import annotations + import numpy as np + import numpy.typing as npt + import pyarrow as pa + + from dataclasses import dataclass + from typing import Any, Dict, Iterable, Optional, Sequence, Set, Tuple, Union + __all__ = [{manifest}] ", @@ -188,6 +195,15 @@ fn quote_objects( 0, ); + let import_clauses: HashSet<_> = objs + .iter() + .flat_map(|obj| obj.object.fields.iter()) + .filter_map(quote_import_clauses_from_field) + .collect(); + for clause in import_clauses { + code.push_text(&clause, 1, 0); + } + for obj in objs { code.push_text(&obj.code, 1, 0); } @@ -235,9 +251,8 @@ fn quote_objects( #[derive(Debug, Clone)] struct QuotedObject { + object: Object, filepath: PathBuf, - name: String, - kind: ObjectKind, code: String, } @@ -255,24 +270,17 @@ impl QuotedObject { attrs: _, fields, specifics: _, + datatype: _, } = obj; let mut code = String::new(); - code.push_text("e_module_prelude(), 0, 0); - - for clause in obj - .fields - .iter() - .filter_map(quote_import_clauses_from_field) - { - code.push_text(&clause, 1, 0); - } - code.push_unindented_text( format!( r#" + ## --- {name} --- ## + @dataclass class {name}: "# @@ -282,7 +290,13 @@ impl QuotedObject { code.push_text(quote_doc_from_docs(docs), 0, 4); - for field in fields { + // NOTE: We need to add required fields first, and then optional ones, otherwise mypy + // complains. + let fields_in_order = fields + .iter() + .filter(|field| !field.is_nullable) + .chain(fields.iter().filter(|field| field.is_nullable)); + for field in fields_in_order { let ObjectField { filepath: _, fqname: _, @@ -291,8 +305,9 @@ impl QuotedObject { docs, typ: _, attrs: _, - required: _, - deprecated: _, + is_nullable, + is_deprecated: _, + datatype: _, } = field; let (typ, _) = quote_field_type_from_field(objects, field, false); @@ -302,10 +317,10 @@ impl QuotedObject { } else { typ }; - let typ = if field.required { - typ - } else { + let typ = if *is_nullable { format!("{typ} | None = None") + } else { + typ }; code.push_text(format!("{name}: {typ}"), 1, 4); @@ -329,9 +344,8 @@ impl QuotedObject { filepath.set_extension("py"); Self { + object: obj.clone(), filepath, - name: obj.name.clone(), - kind: obj.kind, code, } } @@ -349,20 +363,11 @@ impl QuotedObject { attrs: _, fields, specifics: _, + datatype: _, } = obj; let mut code = String::new(); - code.push_text("e_module_prelude(), 0, 0); - - for clause in obj - .fields - .iter() - .filter_map(quote_import_clauses_from_field) - { - code.push_text(&clause, 1, 0); - } - code.push_unindented_text( format!( r#" @@ -385,8 +390,9 @@ impl QuotedObject { docs, typ: _, attrs: _, - required: _, - deprecated: _, + is_nullable: _, + is_deprecated: _, + datatype: _, } = field; let (typ, _) = quote_field_type_from_field(objects, field, false); @@ -407,9 +413,8 @@ impl QuotedObject { filepath.set_extension("py"); Self { + object: obj.clone(), filepath, - name: obj.name.clone(), - kind: obj.kind, code, } } @@ -427,21 +432,6 @@ fn quote_manifest(names: impl IntoIterator>) -> String { quoted_names.join(", ") } -fn quote_module_prelude() -> String { - // NOTE: All the extraneous stuff will be cleaned up courtesy of `ruff`. - unindent::unindent( - r#" - import numpy as np - import numpy.typing as npt - import pyarrow as pa - - from dataclasses import dataclass - from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple, Union - - "#, - ) -} - fn quote_doc_from_docs(docs: &Docs) -> String { let lines = crate::codegen::get_documentation(docs, &["py", "python"]); @@ -671,7 +661,7 @@ fn quote_field_type_from_field( unwrapped = true; typ } else { - format!("List[{typ}]") + format!("list[{typ}]") } } } @@ -823,12 +813,12 @@ fn quote_builder_from_obj(objects: &Objects, obj: &Object) -> String { let required = obj .fields .iter() - .filter(|field| field.required) + .filter(|field| !field.is_nullable) .collect::>(); let optional = obj .fields .iter() - .filter(|field| !field.required) + .filter(|field| field.is_nullable) .collect::>(); let mut code = String::new(); @@ -918,10 +908,17 @@ fn quote_arrow_datatype(datatype: &DataType) -> String { DataType::LargeBinary => "pa.large_binary()".to_owned(), DataType::Utf8 => "pa.utf8()".to_owned(), DataType::LargeUtf8 => "pa.large_utf8()".to_owned(), + + DataType::List(field) => { + let field = quote_arrow_field(field); + format!("pa.list_({field})") + } + DataType::FixedSizeList(field, length) => { let field = quote_arrow_field(field); format!("pa.list_({field}, {length})") } + DataType::Union(fields, _, mode) => { let fields = fields .iter() @@ -933,6 +930,7 @@ fn quote_arrow_datatype(datatype: &DataType) -> String { UnionMode::Sparse => format!(r#"pa.sparse_union([{fields}])"#), } } + DataType::Struct(fields) => { let fields = fields .iter() @@ -941,7 +939,9 @@ fn quote_arrow_datatype(datatype: &DataType) -> String { .join(", "); format!("pa.struct([{fields}])") } + DataType::Extension(_, datatype, _) => quote_arrow_datatype(datatype), + _ => unimplemented!("{datatype:#?}"), // NOLINT } } diff --git a/crates/re_types_builder/src/codegen/rust.rs b/crates/re_types_builder/src/codegen/rust.rs index 0b4002828d60..484f624d30c1 100644 --- a/crates/re_types_builder/src/codegen/rust.rs +++ b/crates/re_types_builder/src/codegen/rust.rs @@ -42,7 +42,7 @@ impl CodeGenerator for RustCodeGenerator { filepaths.extend(create_files( datatypes_path, arrow_registry, - &objects.ordered_datatypes(), + &objects.ordered_objects(ObjectKind::Datatype.into()), )); let components_path = self.crate_path.join("src/components"); @@ -52,7 +52,7 @@ impl CodeGenerator for RustCodeGenerator { filepaths.extend(create_files( components_path, arrow_registry, - &objects.ordered_components(), + &objects.ordered_objects(ObjectKind::Component.into()), )); let archetypes_path = self.crate_path.join("src/archetypes"); @@ -62,7 +62,7 @@ impl CodeGenerator for RustCodeGenerator { filepaths.extend(create_files( archetypes_path, arrow_registry, - &objects.ordered_archetypes(), + &objects.ordered_objects(ObjectKind::Archetype.into()), )); filepaths @@ -189,6 +189,7 @@ impl QuotedObject { attrs: _, fields, specifics: _, + datatype: _, } = obj; let name = format_ident!("{name}"); @@ -247,6 +248,7 @@ impl QuotedObject { attrs: _, fields, specifics: _, + datatype: _, } = obj; let name = format_ident!("{name}"); @@ -264,8 +266,9 @@ impl QuotedObject { docs, typ: _, attrs: _, - required: _, - deprecated: _, + is_nullable: _, + is_deprecated: _, + datatype: _, } = obj_field; let name = format_ident!("{name}"); @@ -320,19 +323,20 @@ impl quote::ToTokens for ObjectFieldTokenizer<'_> { docs, typ: _, attrs: _, - required, + is_nullable, // TODO(#2366): support for deprecation notices - deprecated: _, + is_deprecated: _, + datatype: _, } = obj_field; let quoted_docs = quote_doc_from_docs(docs); let name = format_ident!("{name}"); let (quoted_type, _) = quote_field_type_from_field(obj_field, false); - let quoted_type = if *required { - quoted_type - } else { + let quoted_type = if *is_nullable { quote!(Option<#quoted_type>) + } else { + quoted_type }; if is_tuple_struct_from_obj(obj) { @@ -460,6 +464,7 @@ fn quote_trait_impls_from_obj(arrow_registry: &ArrowRegistry, obj: &Object) -> T attrs: _, fields: _, specifics: _, + datatype: _, } = obj; let name = format_ident!("{name}"); @@ -564,6 +569,7 @@ fn quote_builder_from_obj(obj: &Object) -> TokenStream { attrs: _, fields, specifics: _, + datatype: _, } = obj; let name = format_ident!("{name}"); @@ -571,11 +577,11 @@ fn quote_builder_from_obj(obj: &Object) -> TokenStream { // NOTE: Collecting because we need to iterate them more than once. let required = fields .iter() - .filter(|field| field.required) + .filter(|field| !field.is_nullable) .collect::>(); let optional = fields .iter() - .filter(|field| !field.required) + .filter(|field| field.is_nullable) .collect::>(); // --- impl new() --- @@ -761,9 +767,17 @@ fn quote_fqname_as_type_path(fqname: impl AsRef) -> TokenStream { // --- Helpers --- fn is_tuple_struct_from_obj(obj: &Object) -> bool { - obj.is_struct() - && obj.fields.len() == 1 - && obj.try_get_attr::(ATTR_RUST_TUPLE_STRUCT).is_some() + let is_tuple_struct = + obj.is_struct() && obj.try_get_attr::(ATTR_RUST_TUPLE_STRUCT).is_some(); + + assert!( + !is_tuple_struct || obj.fields.len() == 1, + "`{ATTR_RUST_TUPLE_STRUCT}` is only supported for objects with a single field, but {} has {}", + obj.fqname, + obj.fields.len(), + ); + + is_tuple_struct } fn iter_archetype_components<'a>( diff --git a/crates/re_types_builder/src/lib.rs b/crates/re_types_builder/src/lib.rs index be53dce368b4..9da75e8df863 100644 --- a/crates/re_types_builder/src/lib.rs +++ b/crates/re_types_builder/src/lib.rs @@ -93,7 +93,7 @@ mod codegen; #[allow(clippy::unimplemented)] mod objects; -pub use self::arrow_registry::ArrowRegistry; +pub use self::arrow_registry::{ArrowRegistry, LazyDatatype, LazyField}; pub use self::codegen::{CodeGenerator, PythonCodeGenerator, RustCodeGenerator}; pub use self::objects::{ Attributes, Docs, ElementType, Object, ObjectField, ObjectKind, Objects, Type, @@ -184,7 +184,7 @@ fn generate_lang_agnostic( binary_entrypoint_path.set_extension("bfbs"); // semantic pass: high level objects from low-level reflection data - let objects = Objects::from_buf( + let mut objects = Objects::from_buf( sh.read_binary_file(tmp.path().join(binary_entrypoint_path)) .unwrap() .as_slice(), @@ -192,7 +192,7 @@ fn generate_lang_agnostic( // create and fill out arrow registry let mut arrow_registry = ArrowRegistry::default(); - for obj in objects.ordered_objects(None) { + for obj in objects.ordered_objects_mut(None) { arrow_registry.register(obj); } diff --git a/crates/re_types_builder/src/objects.rs b/crates/re_types_builder/src/objects.rs index 677779a7c7b3..069a27330e59 100644 --- a/crates/re_types_builder/src/objects.rs +++ b/crates/re_types_builder/src/objects.rs @@ -81,22 +81,18 @@ impl Objects { .unwrap() } - /// Returns all available datatypes, pre-sorted in ascending order based on their `order` + /// Returns all available objects, pre-sorted in ascending order based on their `order` /// attribute. - pub fn ordered_datatypes(&self) -> Vec<&Object> { - self.ordered_objects(ObjectKind::Datatype.into()) - } + pub fn ordered_objects_mut(&mut self, kind: Option) -> Vec<&mut Object> { + let objs = self + .objects + .values_mut() + .filter(|obj| kind.map_or(true, |kind| obj.kind == kind)); - /// Returns all available components, pre-sorted in ascending order based on their `order` - /// attribute. - pub fn ordered_components(&self) -> Vec<&Object> { - self.ordered_objects(ObjectKind::Component.into()) - } + let mut objs = objs.collect::>(); + objs.sort_by_key(|anyobj| anyobj.order()); - /// Returns all available archetypes, pre-sorted in ascending order based on their `order` - /// attribute. - pub fn ordered_archetypes(&self) -> Vec<&Object> { - self.ordered_objects(ObjectKind::Archetype.into()) + objs } /// Returns all available objects, pre-sorted in ascending order based on their `order` @@ -248,6 +244,12 @@ pub struct Object { /// Properties that only apply to either structs or unions. pub specifics: ObjectSpecifics, + + /// The Arrow datatype of this `Object`, or `None` if the object is Arrow-transparent. + /// + /// This is lazily computed when the parent object gets registered into the Arrow registry and + /// will be `None` until then. + pub datatype: Option, } impl Object { @@ -297,6 +299,7 @@ impl Object { attrs, fields, specifics: ObjectSpecifics::Struct {}, + datatype: None, } } @@ -356,6 +359,7 @@ impl Object { attrs, fields, specifics: ObjectSpecifics::Union { utype }, + datatype: None, } } @@ -444,13 +448,19 @@ pub struct ObjectField { /// Whether the field is required. /// /// Always true for IDL definitions using flatbuffers' `struct` type (as opposed to `table`). - pub required: bool, + pub is_nullable: bool, /// Whether the field is deprecated. // // TODO(#2366): do something with this // TODO(#2367): implement custom attr to specify deprecation reason - pub deprecated: bool, + pub is_deprecated: bool, + + /// The Arrow datatype of this `ObjectField`. + /// + /// This is lazily computed when the parent object gets registered into the Arrow registry and + /// will be `None` until then. + pub datatype: Option, } impl ObjectField { @@ -478,8 +488,8 @@ impl ObjectField { let typ = Type::from_raw_type(enums, objs, field.type_()); let attrs = Attributes::from_raw_attrs(field.attributes()); - let required = field.required() || obj.is_struct(); - let deprecated = field.deprecated(); + let is_nullable = !obj.is_struct() && !field.required(); + let is_deprecated = field.deprecated(); Self { filepath, @@ -489,8 +499,9 @@ impl ObjectField { docs, typ, attrs, - required, - deprecated, + is_nullable, + is_deprecated, + datatype: None, } } @@ -525,8 +536,8 @@ impl ObjectField { let attrs = Attributes::from_raw_attrs(val.attributes()); // TODO(cmc): not sure about this, but fbs unions are a bit weird that way - let required = true; - let deprecated = false; + let is_nullable = false; + let is_deprecated = false; Self { filepath, @@ -536,8 +547,9 @@ impl ObjectField { docs, typ, attrs, - required, - deprecated, + is_nullable, + is_deprecated, + datatype: None, } } @@ -666,6 +678,44 @@ impl Type { _ => unreachable!("{typ:#?}"), } } + + /// True if this is some kind of array/vector. + pub fn is_plural(&self) -> bool { + match self { + Type::Array { + elem_type: _, + length: _, + } + | Type::Vector { elem_type: _ } => true, + Type::UInt8 + | Type::UInt16 + | Type::UInt32 + | Type::UInt64 + | Type::Int8 + | Type::Int16 + | Type::Int32 + | Type::Int64 + | Type::Bool + | Type::Float16 + | Type::Float32 + | Type::Float64 + | Type::String + | Type::Object(_) => false, + } + } + + /// `Some(fqname)` if this is an `Object` or an `Array`/`Vector` of `Object`s. + pub fn fqname(&self) -> Option<&str> { + match self { + Type::Object(fqname) => Some(fqname.as_str()), + Type::Array { + elem_type, + length: _, + } + | Type::Vector { elem_type } => elem_type.fqname(), + _ => None, + } + } } /// The underlying element type for arrays/vectors/maps. @@ -725,6 +775,14 @@ impl ElementType { _ => unreachable!("{inner_type:#?}"), } } + + /// `Some(fqname)` if this is an `Object`. + pub fn fqname(&self) -> Option<&str> { + match self { + ElementType::Object(fqname) => Some(fqname.as_str()), + _ => None, + } + } } // --- Common --- diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/points2d.py b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/points2d.py index 89b567557c51..6fd3d30c8795 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/points2d.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/archetypes/points2d.py @@ -2,13 +2,14 @@ from __future__ import annotations -__all__ = ["Points2D"] - - from dataclasses import dataclass +__all__ = ["Points2D"] + from .. import components +## --- Points2D --- ## + @dataclass class Points2D: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/class_id.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/class_id.py index c61b49ca56b8..0cd3d86043c4 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/class_id.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/class_id.py @@ -2,8 +2,6 @@ from __future__ import annotations -__all__ = ["ClassId", "ClassIdArray", "ClassIdArrayLike", "ClassIdLike", "ClassIdType"] - from dataclasses import dataclass from typing import Any, Sequence, Union @@ -11,6 +9,11 @@ import numpy.typing as npt import pyarrow as pa +__all__ = ["ClassId", "ClassIdArray", "ClassIdArrayLike", "ClassIdLike", "ClassIdType"] + + +## --- ClassId --- ## + @dataclass class ClassId: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/color.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/color.py index dca8ab267845..327bf8bee058 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/color.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/color.py @@ -2,8 +2,6 @@ from __future__ import annotations -__all__ = ["Color", "ColorArray", "ColorArrayLike", "ColorLike", "ColorType"] - from dataclasses import dataclass from typing import Any, Sequence, Union @@ -11,6 +9,11 @@ import numpy.typing as npt import pyarrow as pa +__all__ = ["Color", "ColorArray", "ColorArrayLike", "ColorLike", "ColorType"] + + +## --- Color --- ## + @dataclass class Color: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/draw_order.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/draw_order.py index dd4c9314b7fc..b015c66feb4f 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/draw_order.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/draw_order.py @@ -2,8 +2,6 @@ from __future__ import annotations -__all__ = ["DrawOrder", "DrawOrderArray", "DrawOrderArrayLike", "DrawOrderLike", "DrawOrderType"] - from dataclasses import dataclass from typing import Any, Sequence, Union @@ -11,6 +9,11 @@ import numpy.typing as npt import pyarrow as pa +__all__ = ["DrawOrder", "DrawOrderArray", "DrawOrderArrayLike", "DrawOrderLike", "DrawOrderType"] + + +## --- DrawOrder --- ## + @dataclass class DrawOrder: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/instance_key.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/instance_key.py index f4733520cf7a..01b825e7beb5 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/instance_key.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/instance_key.py @@ -2,8 +2,6 @@ from __future__ import annotations -__all__ = ["InstanceKey", "InstanceKeyArray", "InstanceKeyArrayLike", "InstanceKeyLike", "InstanceKeyType"] - from dataclasses import dataclass from typing import Any, Sequence, Union @@ -11,6 +9,11 @@ import numpy.typing as npt import pyarrow as pa +__all__ = ["InstanceKey", "InstanceKeyArray", "InstanceKeyArrayLike", "InstanceKeyLike", "InstanceKeyType"] + + +## --- InstanceKey --- ## + @dataclass class InstanceKey: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/keypoint_id.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/keypoint_id.py index 9e406e06abfa..38f9b3ddbac1 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/keypoint_id.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/keypoint_id.py @@ -2,8 +2,6 @@ from __future__ import annotations -__all__ = ["KeypointId", "KeypointIdArray", "KeypointIdArrayLike", "KeypointIdLike", "KeypointIdType"] - from dataclasses import dataclass from typing import Any, Sequence, Union @@ -11,6 +9,11 @@ import numpy.typing as npt import pyarrow as pa +__all__ = ["KeypointId", "KeypointIdArray", "KeypointIdArrayLike", "KeypointIdLike", "KeypointIdType"] + + +## --- KeypointId --- ## + @dataclass class KeypointId: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/label.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/label.py index a48e2a92f971..5af6ea0c614b 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/label.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/label.py @@ -2,13 +2,16 @@ from __future__ import annotations -__all__ = ["Label", "LabelArray", "LabelArrayLike", "LabelLike", "LabelType"] - from dataclasses import dataclass from typing import Any, Sequence, Union import pyarrow as pa +__all__ = ["Label", "LabelArray", "LabelArrayLike", "LabelLike", "LabelType"] + + +## --- Label --- ## + @dataclass class Label: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/point2d.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/point2d.py index 3096ec2d1c5e..a7d5aaac503e 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/point2d.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/point2d.py @@ -2,8 +2,6 @@ from __future__ import annotations -__all__ = ["Point2D", "Point2DArray", "Point2DArrayLike", "Point2DLike", "Point2DType"] - from dataclasses import dataclass from typing import Any, Sequence, Tuple, Union @@ -11,6 +9,11 @@ import numpy.typing as npt import pyarrow as pa +__all__ = ["Point2D", "Point2DArray", "Point2DArrayLike", "Point2DLike", "Point2DType"] + + +## --- Point2D --- ## + @dataclass class Point2D: @@ -34,7 +37,7 @@ class Point2DType(pa.ExtensionType): # type: ignore[misc] def __init__(self: type[pa.ExtensionType]) -> None: pa.ExtensionType.__init__( self, - pa.struct([pa.field("x", pa.float32(), True, {}), pa.field("y", pa.float32(), True, {})]), + pa.struct([pa.field("x", pa.float32(), False, {}), pa.field("y", pa.float32(), False, {})]), "rerun.point2d", ) diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/components/radius.py b/rerun_py/rerun_sdk/rerun/_rerun2/components/radius.py index 7508475d1ad9..ec81b9eeffe9 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/components/radius.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/components/radius.py @@ -2,8 +2,6 @@ from __future__ import annotations -__all__ = ["Radius", "RadiusArray", "RadiusArrayLike", "RadiusLike", "RadiusType"] - from dataclasses import dataclass from typing import Any, Sequence, Union @@ -11,6 +9,11 @@ import numpy.typing as npt import pyarrow as pa +__all__ = ["Radius", "RadiusArray", "RadiusArrayLike", "RadiusLike", "RadiusType"] + + +## --- Radius --- ## + @dataclass class Radius: diff --git a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/vec2d.py b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/vec2d.py index bc6320a233bb..2efcf56fbcf2 100644 --- a/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/vec2d.py +++ b/rerun_py/rerun_sdk/rerun/_rerun2/datatypes/vec2d.py @@ -2,8 +2,6 @@ from __future__ import annotations -__all__ = ["Vec2D", "Vec2DArray", "Vec2DArrayLike", "Vec2DLike", "Vec2DType"] - from dataclasses import dataclass from typing import Any, Sequence, Union @@ -11,6 +9,11 @@ import numpy.typing as npt import pyarrow as pa +__all__ = ["Vec2D", "Vec2DArray", "Vec2DArrayLike", "Vec2DLike", "Vec2DType"] + + +## --- Vec2D --- ## + @dataclass class Vec2D: diff --git a/scripts/lint.py b/scripts/lint.py index 5eb57c32b7f5..dd0f2e7c4d31 100755 --- a/scripts/lint.py +++ b/scripts/lint.py @@ -438,7 +438,7 @@ def main() -> None: root_dirpath = os.path.abspath(f"{script_dirpath}/..") os.chdir(root_dirpath) - extensions = ["html", "js", "md", "py", "rs", "sh", "toml", "wgsl", "yml"] + extensions = ["fbs", "html", "js", "md", "py", "rs", "sh", "toml", "wgsl", "yml"] exclude_dirs = {"env", "renv", "venv", "target", "target_ra", "target_wasm"}