wip codegen for to_arrow

rerun-io · Jul 20, 2023 · f0b409f · f0b409f
1 parent 73b1054
commit f0b409f
Show file tree

Hide file tree

Showing 87 changed files with 518 additions and 56 deletions.
diff --git 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.
-4945c6a5462b379e74846b23eb272fadc9b00bf029ae55e902345ee24bc887e5
+ac036433ca594e3fb9458dc9c2d12d4eb913fd336f20b9ca0f3c4b5b2c981976
diff --git a/crates/re_types_builder/src/codegen/cpp/mod.rs b/crates/re_types_builder/src/codegen/cpp/mod.rs
@@ -276,16 +276,27 @@ impl QuotedObject {
             }
         };
 
+        let datatype = arrow_registry.get(&obj.fqname);
+
         match obj.kind {
             ObjectKind::Datatype | ObjectKind::Component => {
                 methods.push(arrow_data_type_method(
-                    &arrow_registry.get(&obj.fqname),
+                    &datatype,
                     &mut hpp_includes,
                     &mut cpp_includes,
                     &mut hpp_declarations,
                 ));
+
+                methods.push(to_arrow_method(
+                    &datatype,
+                    &mut hpp_includes,
+                    &mut cpp_includes,
+                    &mut hpp_declarations,
+                ));
+            }
+            ObjectKind::Archetype => {
+                // TODO(andreas): Should also be convertable to arrow?
             }
-            ObjectKind::Archetype => {}
         };
 
         let hpp_method_section = if methods.is_empty() {
@@ -595,8 +606,8 @@ fn arrow_data_type_method(
     hpp_declarations: &mut ForwardDecls,
 ) -> Method {
     hpp_declarations.insert("arrow", ForwardDecl::Class("DataType".to_owned()));
-    cpp_includes.system.insert("arrow/api.h".to_owned());
     hpp_includes.system.insert("memory".to_owned()); // std::shared_ptr
+    cpp_includes.system.insert("arrow/api.h".to_owned());
 
     let quoted_datatype = quote_arrow_data_type(datatype, cpp_includes, true);
 
@@ -612,6 +623,29 @@ fn arrow_data_type_method(
     }
 }
 
+fn to_arrow_method(
+    datatype: &DataType,
+    hpp_includes: &mut Includes,
+    cpp_includes: &mut Includes,
+    hpp_declarations: &mut ForwardDecls,
+) -> Method {
+    hpp_declarations.insert("arrow", ForwardDecl::Class("Array".to_owned()));
+    hpp_includes.system.insert("arrow/result.h".to_owned());
+    hpp_includes.system.insert("memory".to_owned()); // std::shared_ptr
+    cpp_includes.system.insert("arrow/api.h".to_owned());
+
+    Method {
+        docs: "Converts an array of this type into an arrow array.".into(),
+        declaration: MethodDeclaration {
+            is_static: true,
+            return_type: quote! { arrow::Result<std::shared_ptr<arrow::Array>> },
+            name_and_parameters: quote! { to_arrow() },
+        },
+        definition_body: quote! { /* TODO */ return arrow::Result<std::shared_ptr<arrow::Array>>(); },
+        inline: false,
+    }
+}
+
 /// e.g. `static Angle radians(float radians);` -> `auto angle = Angle::radians(radians);`
 fn static_constructor_for_enum_type(
     objects: &Objects,

diff --git a/examples/cpp/minimal/main.cpp b/examples/cpp/minimal/main.cpp
@@ -1,6 +1,37 @@
 #include <loguru.hpp>
 #include <rerun.hpp>
 
+#include <components/point2d.hpp>
+
+arrow::Result<std::shared_ptr<arrow::Table>> points2(size_t num_points, const float* xy) {
+    arrow::MemoryPool* pool = arrow::default_memory_pool();
+
+    auto x_builder = std::make_shared<arrow::FloatBuilder>(pool);
+    auto y_builder = std::make_shared<arrow::FloatBuilder>(pool);
+
+    auto nullable = false;
+
+    auto data_type = rr::components::Point2D::to_arrow_datatype();
+    auto struct_builder =
+        arrow::StructBuilder(data_type, pool, {x_builder, y_builder});
+
+    for (size_t i = 0; i < num_points; ++i) {
+        ARROW_RETURN_NOT_OK(struct_builder.Append());
+        ARROW_RETURN_NOT_OK(x_builder->Append(xy[2 * i + 0]));
+        ARROW_RETURN_NOT_OK(y_builder->Append(xy[2 * i + 1]));
+    }
+
+    std::shared_ptr<arrow::Array> array;
+    ARROW_RETURN_NOT_OK(struct_builder.Finish(&array));
+
+    auto name = "points"; // Unused, but should be the name of the field in the archetype
+    auto schema = arrow::schema({arrow::field(name, data_type, nullable)});
+
+    return arrow::Table::Make(schema, {array});
+}
+
+
+
 int main(int argc, char** argv) {
     loguru::g_preamble_uptime = false;
     loguru::g_preamble_thread = false;
@@ -10,18 +41,37 @@ int main(int argc, char** argv) {
 
     auto rr_stream = rr::RecordingStream{"c-example-app", "127.0.0.1:9876"};
 
-    float xyz[9] = {0.0, 0.0, 0.0, 1.0, 2.0, 3.0, 5.0, 5.0, 5.0};
-    auto points = rr::points3(3, xyz).ValueOrDie();
-    auto buffer = rr::ipc_from_table(*points).ValueOrDie();
+    // Points3D, fully custom.
+    {
+        float xyz[9] = {0.0, 0.0, 0.0, 1.0, 2.0, 3.0, 5.0, 5.0, 5.0};
+        auto points = rr::points3(3, xyz).ValueOrDie();
+        auto buffer = rr::ipc_from_table(*points).ValueOrDie();
+
+        const rr::DataCell data_cells[1] = {rr::DataCell{
+            .component_name = "rerun.point3d",
+            .num_bytes = static_cast<size_t>(buffer->size()),
+            .bytes = buffer->data(),
+        }};
+
+        uint32_t num_instances = 3;
+        rr_stream.log_data_row("3d/points", num_instances, 1, data_cells);
+    }
+
+    // Points2D, via generated code.
+    {
+        float xy[6] = {0.0, 0.0, 1.0, 3.0, 5.0, 5.0};
+        auto points = points2(3, xy).ValueOrDie();
+        auto buffer = rr::ipc_from_table(*points).ValueOrDie();
 
-    const rr::DataCell data_cells[1] = {rr::DataCell{
-        .component_name = "rerun.point3d",
-        .num_bytes = static_cast<size_t>(buffer->size()),
-        .bytes = buffer->data(),
-    }};
+        const rr::DataCell data_cells[1] = {rr::DataCell{
+            .component_name = "rerun.point2d",
+            .num_bytes = static_cast<size_t>(buffer->size()),
+            .bytes = buffer->data(),
+        }};
 
-    uint32_t num_instances = 3;
-    rr_stream.log_data_row("points", num_instances, 1, data_cells);
+        uint32_t num_instances = 3;
+        rr_stream.log_data_row("2d/points", num_instances, 1, data_cells);
+    }
 
     // Test some type instantiation
     auto tls = rr::datatypes::TranslationRotationScale3D{};

diff --git a/rerun_cpp/src/components/affix_fuzzer1.cpp b/rerun_cpp/src/components/affix_fuzzer1.cpp
@@ -11,5 +11,9 @@ namespace rr {
         std::shared_ptr<arrow::DataType> AffixFuzzer1::to_arrow_datatype() {
             return rr::datatypes::AffixFuzzer1::to_arrow_datatype();
         }
+
+        arrow::Result<std::shared_ptr<arrow::Array>> AffixFuzzer1::to_arrow() {
+            return arrow::Result<std::shared_ptr<arrow::Array>>();
+        }
     } // namespace components
 } // namespace rr
diff --git a/rerun_cpp/src/components/affix_fuzzer1.hpp b/rerun_cpp/src/components/affix_fuzzer1.hpp
@@ -3,15 +3,18 @@
 
 #pragma once
 
+#include <arrow/result.h>
+
 #include <cstdint>
 #include <memory>
 #include <utility>
 
 #include "../datatypes/affix_fuzzer1.hpp"
 
 namespace arrow {
+    class Array;
     class DataType;
-}
+} // namespace arrow
 
 namespace rr {
     namespace components {
@@ -24,6 +27,8 @@ namespace rr {
 
             /// Returns the arrow data type this type corresponds to.
             static std::shared_ptr<arrow::DataType> to_arrow_datatype();
+            /// Converts an array of this type into an arrow array.
+            static arrow::Result<std::shared_ptr<arrow::Array>> to_arrow();
         };
     } // namespace components
 } // namespace rr
diff --git a/rerun_cpp/src/components/affix_fuzzer10.cpp b/rerun_cpp/src/components/affix_fuzzer10.cpp
@@ -10,5 +10,9 @@ namespace rr {
         std::shared_ptr<arrow::DataType> AffixFuzzer10::to_arrow_datatype() {
             return arrow::utf8();
         }
+
+        arrow::Result<std::shared_ptr<arrow::Array>> AffixFuzzer10::to_arrow() {
+            return arrow::Result<std::shared_ptr<arrow::Array>>();
+        }
     } // namespace components
 } // namespace rr
diff --git a/rerun_cpp/src/components/affix_fuzzer10.hpp b/rerun_cpp/src/components/affix_fuzzer10.hpp
@@ -3,15 +3,18 @@
 
 #pragma once
 
+#include <arrow/result.h>
+
 #include <cstdint>
 #include <memory>
 #include <optional>
 #include <string>
 #include <utility>
 
 namespace arrow {
+    class Array;
     class DataType;
-}
+} // namespace arrow
 
 namespace rr {
     namespace components {
@@ -24,6 +27,8 @@ namespace rr {
 
             /// Returns the arrow data type this type corresponds to.
             static std::shared_ptr<arrow::DataType> to_arrow_datatype();
+            /// Converts an array of this type into an arrow array.
+            static arrow::Result<std::shared_ptr<arrow::Array>> to_arrow();
         };
     } // namespace components
 } // namespace rr
diff --git a/rerun_cpp/src/components/affix_fuzzer11.cpp b/rerun_cpp/src/components/affix_fuzzer11.cpp
@@ -10,5 +10,9 @@ namespace rr {
         std::shared_ptr<arrow::DataType> AffixFuzzer11::to_arrow_datatype() {
             return arrow::list(arrow::field("item", arrow::float32(), true, nullptr));
         }
+
+        arrow::Result<std::shared_ptr<arrow::Array>> AffixFuzzer11::to_arrow() {
+            return arrow::Result<std::shared_ptr<arrow::Array>>();
+        }
     } // namespace components
 } // namespace rr
diff --git a/rerun_cpp/src/components/affix_fuzzer11.hpp b/rerun_cpp/src/components/affix_fuzzer11.hpp
@@ -3,15 +3,18 @@
 
 #pragma once
 
+#include <arrow/result.h>
+
 #include <cstdint>
 #include <memory>
 #include <optional>
 #include <utility>
 #include <vector>
 
 namespace arrow {
+    class Array;
     class DataType;
-}
+} // namespace arrow
 
 namespace rr {
     namespace components {
@@ -24,6 +27,8 @@ namespace rr {
 
             /// Returns the arrow data type this type corresponds to.
             static std::shared_ptr<arrow::DataType> to_arrow_datatype();
+            /// Converts an array of this type into an arrow array.
+            static arrow::Result<std::shared_ptr<arrow::Array>> to_arrow();
         };
     } // namespace components
 } // namespace rr
diff --git a/rerun_cpp/src/components/affix_fuzzer12.cpp b/rerun_cpp/src/components/affix_fuzzer12.cpp
@@ -10,5 +10,9 @@ namespace rr {
         std::shared_ptr<arrow::DataType> AffixFuzzer12::to_arrow_datatype() {
             return arrow::list(arrow::field("item", arrow::utf8(), false, nullptr));
         }
+
+        arrow::Result<std::shared_ptr<arrow::Array>> AffixFuzzer12::to_arrow() {
+            return arrow::Result<std::shared_ptr<arrow::Array>>();
+        }
     } // namespace components
 } // namespace rr
diff --git a/rerun_cpp/src/components/affix_fuzzer12.hpp b/rerun_cpp/src/components/affix_fuzzer12.hpp
@@ -3,15 +3,18 @@
 
 #pragma once
 
+#include <arrow/result.h>
+
 #include <cstdint>
 #include <memory>
 #include <string>
 #include <utility>
 #include <vector>
 
 namespace arrow {
+    class Array;
     class DataType;
-}
+} // namespace arrow
 
 namespace rr {
     namespace components {
@@ -24,6 +27,8 @@ namespace rr {
 
             /// Returns the arrow data type this type corresponds to.
             static std::shared_ptr<arrow::DataType> to_arrow_datatype();
+            /// Converts an array of this type into an arrow array.
+            static arrow::Result<std::shared_ptr<arrow::Array>> to_arrow();
         };
     } // namespace components
 } // namespace rr
diff --git a/rerun_cpp/src/components/affix_fuzzer13.cpp b/rerun_cpp/src/components/affix_fuzzer13.cpp
@@ -10,5 +10,9 @@ namespace rr {
         std::shared_ptr<arrow::DataType> AffixFuzzer13::to_arrow_datatype() {
             return arrow::list(arrow::field("item", arrow::utf8(), true, nullptr));
         }
+
+        arrow::Result<std::shared_ptr<arrow::Array>> AffixFuzzer13::to_arrow() {
+            return arrow::Result<std::shared_ptr<arrow::Array>>();
+        }
     } // namespace components
 } // namespace rr
diff --git a/rerun_cpp/src/components/affix_fuzzer13.hpp b/rerun_cpp/src/components/affix_fuzzer13.hpp
@@ -3,6 +3,8 @@
 
 #pragma once
 
+#include <arrow/result.h>
+
 #include <cstdint>
 #include <memory>
 #include <optional>
@@ -11,8 +13,9 @@
 #include <vector>
 
 namespace arrow {
+    class Array;
     class DataType;
-}
+} // namespace arrow
 
 namespace rr {
     namespace components {
@@ -25,6 +28,8 @@ namespace rr {
 
             /// Returns the arrow data type this type corresponds to.
             static std::shared_ptr<arrow::DataType> to_arrow_datatype();
+            /// Converts an array of this type into an arrow array.
+            static arrow::Result<std::shared_ptr<arrow::Array>> to_arrow();
         };
     } // namespace components
 } // namespace rr
diff --git a/rerun_cpp/src/components/affix_fuzzer14.cpp b/rerun_cpp/src/components/affix_fuzzer14.cpp
@@ -11,5 +11,9 @@ namespace rr {
         std::shared_ptr<arrow::DataType> AffixFuzzer14::to_arrow_datatype() {
             return rr::datatypes::AffixFuzzer3::to_arrow_datatype();
         }
+
+        arrow::Result<std::shared_ptr<arrow::Array>> AffixFuzzer14::to_arrow() {
+            return arrow::Result<std::shared_ptr<arrow::Array>>();
+        }
     } // namespace components
 } // namespace rr
diff --git a/rerun_cpp/src/components/affix_fuzzer14.hpp b/rerun_cpp/src/components/affix_fuzzer14.hpp
@@ -3,15 +3,18 @@
 
 #pragma once
 
+#include <arrow/result.h>
+
 #include <cstdint>
 #include <memory>
 #include <utility>
 
 #include "../datatypes/affix_fuzzer3.hpp"
 
 namespace arrow {
+    class Array;
     class DataType;
-}
+} // namespace arrow
 
 namespace rr {
     namespace components {
@@ -24,6 +27,8 @@ namespace rr {
 
             /// Returns the arrow data type this type corresponds to.
             static std::shared_ptr<arrow::DataType> to_arrow_datatype();
+            /// Converts an array of this type into an arrow array.
+            static arrow::Result<std::shared_ptr<arrow::Array>> to_arrow();
         };
     } // namespace components
 } // namespace rr