add tree-shaped data with primitives to deep ops benchmarks

PiperOrigin-RevId: 856257494
Change-Id: I984b2435abf5157a4dcdafdfa9d3e7d77a9ebcef

Author: Koladata Team

koladata/internal/op_utils/BUILD

@@ -395,7 +395,6 @@ cc_test(
     srcs = ["deep_clone_test.cc"],
     deps = [
         ":deep_clone",
-        ":presence_and",
         "//koladata:test_utils",
         "//koladata/internal:data_bag",
         "//koladata/internal:data_item",
@@ -406,10 +405,8 @@ cc_test(
         "//koladata/internal:uuid_object",
         "//koladata/internal/testing:deep_op_utils",
         "//koladata/internal/testing:matchers",
-        "@com_google_absl//absl/random",
         "@com_google_absl//absl/types:span",
         "@com_google_arolla//arolla/dense_array",
-        "@com_google_arolla//arolla/memory",
         "@com_google_arolla//arolla/qtype",
         "@com_google_arolla//arolla/util",
         "@com_google_googletest//:gtest_main",
@@ -755,11 +752,13 @@ cc_test(
     deps = [
         ":deep_op_benchmarks_util",
         ":extract",
+        "//koladata/internal:casting",
         "//koladata/internal:data_bag",
         "//koladata/internal:data_item",
         "//koladata/internal:data_slice",
         "//koladata/internal:dtype",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_arolla//arolla/qtype",
         "@com_google_arolla//arolla/util/testing",
         "@com_google_benchmark//:benchmark_main",
@@ -1274,6 +1273,7 @@ cc_library(
         "@com_google_absl//absl/random",
         "@com_google_absl//absl/random:distributions",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
         "@com_google_arolla//arolla/dense_array",
         "@com_google_arolla//arolla/memory",
         "@com_google_arolla//arolla/qtype",

koladata/internal/op_utils/deep_op_benchmarks_util.h

@@ -15,8 +15,10 @@
 #ifndef KOLADATA_INTERNAL_OP_UTILS_DEEP_OP_BENCHMARKS_UTIL_H_
 #define KOLADATA_INTERNAL_OP_UTILS_DEEP_OP_BENCHMARKS_UTIL_H_
 
+#include <cmath>
 #include <cstdint>
 #include <functional>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -27,6 +29,7 @@
 #include "absl/random/distributions.h"
 #include "absl/random/random.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
 #include "arolla/dense_array/dense_array.h"
 #include "arolla/memory/optional_value.h"
 #include "arolla/qtype/base_types.h"
@@ -71,10 +74,32 @@ constexpr auto kLayersBenchmarkFn = [](auto* b) {
       ->Args({20, 100, 10, 10});
 };
 
+constexpr auto kTreesBenchmarkFn = [](auto* b) {
+  // Number of layers, number of non-leaf childs per object, number of leaves
+  // per object, size of the slice.
+  b->Args({2, 10, 10, 1})
+      ->Args({2, 10, 10, 10})
+      ->Args({2, 10, 10, 100})
+      ->Args({2, 1000, 1000, 1})
+      ->Args({2, 1000, 1000, 10})
+      ->Args({5, 10, 10, 1})
+      ->Args({5, 10, 10, 10})
+      ->Args({5, 10, 10, 100})
+      ->Args({18, 2, 2, 1})
+      ->Args({18, 2, 2, 10})
+      ->Args({18, 2, 2, 100});
+};
+
 using RunBenchmarksFn =
     std::function<void(benchmark::State&, DataSliceImpl&, DataItem&,
                        DataBagImplPtr&, DataBagImpl::FallbackSpan)>;
 
+using RunCastingBenchmarksFn = std::function<void(
+    benchmark::State&, DataSliceImpl& ds_impl, DataItem& ds_schema,
+    DataBagImplPtr& ds_databag, const DataBagImpl::FallbackSpan ds_fallbacks,
+    DataItem& new_schema, DataBagImplPtr& new_schema_databag,
+    const DataBagImpl::FallbackSpan new_schema_fallbacks)>;
+
 inline DataSliceImpl ShuffleObjectsSlice(const DataSliceImpl& ds,
                                          absl::BitGen& gen) {
   // Generate random permutation.
@@ -108,7 +133,8 @@ inline DataSliceImpl ApplyRandomMask(const DataSliceImpl& ds,
   return result;
 }
 
-inline void BM_DisjointChains(benchmark::State& state, RunBenchmarksFn run_fn) {
+inline void BM_DisjointChains(benchmark::State& state,
+                              const RunBenchmarksFn& run_fn) {
   int64_t schema_depth = state.range(0);
   int64_t ds_size = state.range(1);
 
@@ -131,7 +157,7 @@ inline void BM_DisjointChains(benchmark::State& state, RunBenchmarksFn run_fn) {
 }
 
 inline void BM_DisjointChainsObjects(benchmark::State& state,
-                                     RunBenchmarksFn run_fn) {
+                                     const RunBenchmarksFn& run_fn) {
   int64_t schema_depth = state.range(0);
   int64_t ds_size = state.range(1);
 
@@ -162,7 +188,7 @@ inline void BM_DisjointChainsObjects(benchmark::State& state,
   run_fn(state, root_ds, kObjectSchema, db, {});
 }
 
-inline void BM_DAG(benchmark::State& state, RunBenchmarksFn run_fn) {
+inline void BM_DAG(benchmark::State& state, const RunBenchmarksFn& run_fn) {
   int64_t schema_depth = state.range(0);
   int64_t attr_count = state.range(1);
   int64_t presence_rate = state.range(2);
@@ -191,7 +217,8 @@ inline void BM_DAG(benchmark::State& state, RunBenchmarksFn run_fn) {
   run_fn(state, root_ds, root_schema, db, {});
 }
 
-inline void BM_DAGObjects(benchmark::State& state, RunBenchmarksFn run_fn) {
+inline void BM_DAGObjects(benchmark::State& state,
+                          const RunBenchmarksFn& run_fn) {
   int64_t schema_depth = state.range(0);
   int64_t attr_count = state.range(1);
   int64_t presence_rate = state.range(2);
@@ -231,6 +258,73 @@ inline void BM_DAGObjects(benchmark::State& state, RunBenchmarksFn run_fn) {
   run_fn(state, root_ds, kObjectSchema, db, {});
 }
 
+inline void BM_TreeShapedIntToFloat(benchmark::State& state,
+                                    const RunCastingBenchmarksFn& run_fn) {
+  int64_t schema_depth = state.range(0);
+  int64_t schema_attr_count = state.range(1);
+  int64_t primitive_attr_count = state.range(2);
+  int64_t ds_size = state.range(3);
+  int64_t num_leaves =
+      static_cast<int64_t>(std::pow(schema_attr_count, schema_depth - 1)) *
+      primitive_attr_count * ds_size;
+  state.SetLabel(absl::StrFormat(
+      "ds_size=%d; depth=%d; attrs_per_node=%d+%d; total_leaves=%d", ds_size,
+      schema_depth, schema_attr_count, primitive_attr_count, num_leaves));
+
+  CancellationContext::ScopeGuard cancellation_scope;
+  auto db_a = DataBagImpl::CreateEmptyDatabag();
+  auto db_b = DataBagImpl::CreateEmptyDatabag();
+  auto root_schema_a = DataItem(internal::AllocateExplicitSchema());
+  auto root_schema_b = DataItem(internal::AllocateExplicitSchema());
+  auto root_ds = DataSliceImpl::AllocateEmptyObjects(ds_size);
+  std::vector<DataItem> schemas_a({root_schema_a});
+  std::vector<DataItem> schemas_b({root_schema_b});
+  std::vector<DataSliceImpl> ds({root_ds});
+  int value = 0;
+  for (int64_t i = 0; i < schema_depth; ++i) {
+    std::vector<DataItem> child_schemas_a;
+    std::vector<DataItem> child_schemas_b;
+    std::vector<DataSliceImpl> child_ds;
+    for (uint64_t id = 0; id < schemas_a.size(); ++id) {
+      auto schema_a = schemas_a[id];
+      auto schema_b = schemas_b[id];
+      auto cur_ds = ds[id];
+      for (int64_t j = 0; j < schema_attr_count; ++j)
+      {
+        std::string attr_name = absl::StrCat("layer_", i, "_child_", j);
+        auto child_schema_a = DataItem(internal::AllocateExplicitSchema());
+        auto child_schema_b = DataItem(internal::AllocateExplicitSchema());
+        auto cur_child_ds = DataSliceImpl::AllocateEmptyObjects(ds_size);
+        CHECK_OK(db_a->SetSchemaAttr(schema_a, attr_name, child_schema_a));
+        CHECK_OK(db_b->SetSchemaAttr(schema_b, attr_name, child_schema_b));
+        CHECK_OK(db_a->SetAttr(cur_ds, attr_name, cur_child_ds));
+        child_ds.push_back(std::move(cur_child_ds));
+        child_schemas_a.push_back(std::move(child_schema_a));
+        child_schemas_b.push_back(std::move(child_schema_b));
+      }
+      for (int64_t j = 0; j < primitive_attr_count; ++j)
+      {
+        std::string attr_name = absl::StrCat("layer_", i, "_primitive_", j);
+        auto child_schema_a = DataItem(schema::kInt32);
+        auto child_schema_b = DataItem(schema::kFloat32);
+        std::vector<std::optional<int>> values(ds_size);
+        for (int64_t h = 0; h < ds_size; ++h) {
+          values[h] = ++value;
+        }
+        DataSliceImpl::Create(
+            arolla::CreateDenseArray<int>(values.begin(), values.end()));
+        CHECK_OK(db_a->SetSchemaAttr(schema_a, attr_name, child_schema_a));
+        CHECK_OK(db_b->SetSchemaAttr(schema_b, attr_name, child_schema_b));
+      }
+    }
+    std::swap(ds, child_ds);
+    std::swap(schemas_a, child_schemas_a);
+    std::swap(schemas_b, child_schemas_b);
+  }
+  run_fn(state, root_ds, root_schema_a, db_a, {}, root_schema_b, db_b,
+                          {});
+}
+
 }  // namespace benchmarks_utils
 
 }  // namespace koladata::internal

koladata/internal/op_utils/extract_benchmarks.cc

@@ -14,7 +14,9 @@
 //
 #include "benchmark/benchmark.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "arolla/qtype/base_types.h"
+#include "koladata/internal/casting.h"
 #include "koladata/internal/data_bag.h"
 #include "koladata/internal/data_item.h"
 #include "koladata/internal/data_slice.h"
@@ -42,6 +44,32 @@ void RunBenchmarks(benchmark::State& state, DataSliceImpl& ds, DataItem& schema,
   }
 }
 
+void RunCastingBenchmarks(benchmark::State& state, DataSliceImpl& ds,
+                          DataItem& schema, DataBagImplPtr& databag,
+                          DataBagImpl::FallbackSpan fallbacks,
+                          DataItem& new_schema, DataBagImplPtr& schema_databag,
+                          DataBagImpl::FallbackSpan schema_fallbacks) {
+  auto cast_data_callback =
+      static_cast<absl::StatusOr<internal::DataSliceImpl> (*)(
+          const internal::DataSliceImpl&, const internal::DataItem&)>(
+          schema::CastDataTo);
+  while (state.KeepRunning()) {
+    benchmark::DoNotOptimize(ds);
+    benchmark::DoNotOptimize(schema);
+    benchmark::DoNotOptimize(databag);
+    benchmark::DoNotOptimize(fallbacks);
+    benchmark::DoNotOptimize(new_schema);
+    benchmark::DoNotOptimize(schema_databag);
+    benchmark::DoNotOptimize(schema_fallbacks);
+    auto result_db = DataBagImpl::CreateEmptyDatabag();
+    ExtractOp(result_db.get())(ds, new_schema, *databag, fallbacks,
+                               &*schema_databag, schema_fallbacks,
+                               /*max_depth=*/-1, cast_data_callback)
+        .IgnoreError();
+    benchmark::DoNotOptimize(result_db);
+  }
+}
+
 void BM_DisjointChains(benchmark::State& state) {
   benchmarks_utils::BM_DisjointChains(state, RunBenchmarks);
 }
@@ -62,6 +90,11 @@ void BM_DAGObjects(benchmark::State& state) {
 }
 BENCHMARK(BM_DAGObjects)->Apply(benchmarks_utils::kLayersBenchmarkFn);
 
+void BM_TreeShapedIntToFloat(benchmark::State& state) {
+  benchmarks_utils::BM_TreeShapedIntToFloat(state, RunCastingBenchmarks);
+}
+BENCHMARK(BM_TreeShapedIntToFloat)->Apply(benchmarks_utils::kTreesBenchmarkFn);
+
 inline void BM_ScalarPrimitive(benchmark::State& state) {
   auto input_db = DataBagImpl::CreateEmptyDatabag();
   auto result_db = DataBagImpl::CreateEmptyDatabag();
@@ -77,7 +110,6 @@ inline void BM_ScalarPrimitive(benchmark::State& state) {
     benchmark::DoNotOptimize(result_db);
   }
 }
-
 BENCHMARK(BM_ScalarPrimitive);
 
 }  // namespace