Distinguish missing attribute values and explicitly removed attribute values in DataBagImpl.

PiperOrigin-RevId: 714912235
Change-Id: I2e4af3b5584a1184050aa23a74c544b4aaf461b1

Author: petrmikheev

koladata/internal/BUILD

@@ -572,6 +572,7 @@ cc_library(
         ":object_id",
         ":schema_utils",
         ":sparse_source",
+        ":types_buffer",
         ":uuid_object",
         "//koladata/internal/op_utils:presence_or",
         "@com_google_absl//absl/base:core_headers",

koladata/internal/data_bag.cc

@@ -55,6 +55,7 @@
 #include "koladata/internal/schema_utils.h"
 #include "koladata/internal/slice_builder.h"
 #include "koladata/internal/sparse_source.h"
+#include "koladata/internal/types_buffer.h"
 #include "koladata/internal/uuid_object.h"
 #include "arolla/dense_array/dense_array.h"
 #include "arolla/dense_array/edge.h"
@@ -184,21 +185,20 @@ absl::Status MergeToMutableDenseSourceOnlySparse(
   auto process_source = [&](const SparseSource* source,
                             auto skip_obj_fn) -> absl::Status {
     for (const auto& [key, item] : source->GetAll()) {
-      if (!item.has_value() || skip_obj_fn(key)) {
+      if (skip_obj_fn(key)) {
         continue;
       }
       if (options.data_conflict_policy == MergeOptions::kOverwrite) {
         RETURN_IF_ERROR(result.Set(key, item));
         continue;
       }
-      if (auto this_result = result.Get(key).value_or(DataItem());
-          !this_result.has_value()) {
+      if (auto this_result = result.Get(key); !this_result.has_value()) {
         RETURN_IF_ERROR(result.Set(key, item));
       } else if (options.data_conflict_policy ==
                      MergeOptions::kRaiseOnConflict &&
-                 this_result != item) {
+                 *this_result != item) {
         return absl::FailedPreconditionError(
-            absl::StrCat("conflict ", key, ": ", this_result, " vs ", item));
+            absl::StrCat("conflict ", key, ": ", *this_result, " vs ", item));
       }
     }
     return absl::OkStatus();
@@ -256,6 +256,10 @@ absl::Status MergeToMutableDenseSource(
   ASSIGN_OR_RETURN(
       auto other_items,
       GetAttributeFromSources(objects, {&dense_source}, sparse_sources));
+  // GetAttributeFromSources returns information about removed values in
+  // TypesBuffer withing DataSliceImpl. Ensure that TypesBuffer is there.
+  DCHECK_GT(size, 0);
+  DCHECK(!other_items.types_buffer().id_to_typeidx.empty());
 
   if (options.data_conflict_policy == MergeOptions::kOverwrite) {
     const auto& objects_array = objects.values<ObjectId>();
@@ -271,19 +275,19 @@ absl::Status MergeToMutableDenseSource(
 
   for (int64_t offset = 0; offset < size; ++offset) {
     // TODO: try to use batch `GetAttributeFromSources`.
-    auto other_item = other_items[offset];
-    if (!other_item.has_value()) {
+    if (other_items.types_buffer().id_to_typeidx[offset] ==
+        TypesBuffer::kUnset) {
       continue;
     }
+    auto other_item = other_items[offset];
     auto obj_id = alloc.ObjectByOffset(offset);
-    if (auto this_result = result.Get(obj_id).value_or(DataItem());
-        !this_result.has_value()) {
+    if (auto this_result = result.Get(obj_id); !this_result.has_value()) {
       RETURN_IF_ERROR(result.Set(obj_id, other_item));
     } else {
       if (options.data_conflict_policy == MergeOptions::kRaiseOnConflict &&
-          this_result != other_item) {
+          *this_result != other_item) {
         return absl::FailedPreconditionError(absl::StrCat(
-            "conflict ", obj_id, ": ", this_result, " vs ", other_item));
+            "conflict ", obj_id, ": ", *this_result, " vs ", other_item));
       }
     }
   }
@@ -452,7 +456,6 @@ absl::StatusOr<DataSliceImpl> DataBagImpl::GetAttr(
       }
       source->Get(objs_span, *bldr);
     }
-    bldr->ConvertMaybeRemovedToUnset();
     if (bldr->is_finalized()) {
       break;
     }

koladata/internal/data_bag_merge_test.cc

@@ -37,6 +37,7 @@
 #include "koladata/internal/dtype.h"
 #include "koladata/internal/object_id.h"
 #include "koladata/internal/schema_utils.h"
+#include "koladata/internal/slice_builder.h"
 #include "koladata/internal/uuid_object.h"
 #include "arolla/dense_array/qtype/types.h"
 #include "arolla/qtype/base_types.h"
@@ -442,22 +443,22 @@ TEST(DataBagTest, MergeObjectsOverwriteOnlyDenseSources) {
     ASSERT_OK(db->SetAttr(a, "a", a_value));
     auto db2 = DataBagImpl::CreateEmptyDatabag();
     ASSERT_OK(db2->SetAttr(a[2], "a", DataItem(17.0)));
+    ASSERT_OK(db2->SetAttr(a[3], "a", DataItem()));
     ASSERT_OK(db2->SetAttr(a[5], "a", DataItem(57)));
     ASSERT_OK(db->MergeInplace(
         *db2, MergeOptions{.data_conflict_policy = MergeOptions::kOverwrite}));
     std::vector<DataItem> a_value_expected(a_value.begin(), a_value.end());
     a_value_expected[2] = DataItem(17.0);
+    a_value_expected[3] = DataItem();
     a_value_expected[5] = DataItem(57);
     EXPECT_THAT(db->GetAttr(a, "a"),
                 IsOkAndHolds(ElementsAreArray(a_value_expected)));
   }
   {
-    SCOPED_TRACE("merge dense with sparse overwrite nothing");
+    SCOPED_TRACE("merge dense with empty sparse overwrite nothing");
     auto db = DataBagImpl::CreateEmptyDatabag();
     ASSERT_OK(db->SetAttr(a, "a", a_value));
     auto db2 = DataBagImpl::CreateEmptyDatabag();
-    ASSERT_OK(db2->SetAttr(a[5], "a", a_value[0]));
-    ASSERT_OK(db2->SetAttr(a[5], "a", DataItem()));
     ASSERT_OK(db->MergeInplace(
         *db2, MergeOptions{.data_conflict_policy = MergeOptions::kOverwrite}));
     EXPECT_THAT(db->GetAttr(a, "a"), IsOkAndHolds(ElementsAreArray(a_value)));
@@ -480,8 +481,13 @@ TEST(DataBagTest, MergeObjectsOverwriteOnlyDenseSources) {
     auto db2 = DataBagImpl::CreateEmptyDatabag();
 
     auto b_value = DataSliceImpl::AllocateEmptyObjects(kSize);
-    ASSERT_OK(db2->SetAttr(a, "a", b_value));
-    ASSERT_OK(db2->SetAttr(a[5], "a", DataItem()));
+    SliceBuilder a_but_one(a.size());
+    for (int i = 0; i < a.size(); ++i) {
+      if (i != 5) {
+        a_but_one.InsertIfNotSetAndUpdateAllocIds(i, a[i]);
+      }
+    }
+    ASSERT_OK(db2->SetAttr(std::move(a_but_one).Build(), "a", b_value));
     ASSERT_OK(db->MergeInplace(
         *db2, MergeOptions{.data_conflict_policy = MergeOptions::kOverwrite}));
 
@@ -497,10 +503,15 @@ TEST(DataBagTest, MergeObjectsOverwriteOnlyDenseSources) {
     auto db2 = DataBagImpl::CreateEmptyDatabag();
 
     auto b_value = DataSliceImpl::AllocateEmptyObjects(kSize);
-    ASSERT_OK(db2->SetAttr(a, "a", b_value));
+    SliceBuilder a_but_one(a.size());
+    for (int i = 0; i < a.size(); ++i) {
+      if (i != 5) {
+        a_but_one.InsertIfNotSetAndUpdateAllocIds(i, a[i]);
+      }
+    }
+    ASSERT_OK(db2->SetAttr(std::move(a_but_one).Build(), "a", b_value));
     ASSERT_OK(db2->SetAttr(a[1], "a", DataItem(27.0)));
     ASSERT_OK(db2->SetAttr(a[3], "a", DataItem(57)));
-    ASSERT_OK(db2->SetAttr(a[5], "a", DataItem()));
     ASSERT_OK(db->MergeInplace(
         *db2, MergeOptions{.data_conflict_policy = MergeOptions::kOverwrite}));
 

koladata/internal/data_bag_test.cc

@@ -439,32 +439,47 @@ TEST(DataBagTest, SetGetWithFallbackObjectId) {
     }
 
     // General case
+    std::vector<arolla::OptionalValue<ObjectId>> main_b_ids(size);
     std::vector<arolla::OptionalValue<ObjectId>> main_b_values(size);
+    std::vector<arolla::OptionalValue<ObjectId>> fallback_b_ids(size);
     std::vector<arolla::OptionalValue<ObjectId>> fallback_b_values(size);
     std::vector<arolla::OptionalValue<ObjectId>> merge_b_values(size);
     AllocationId alloc_b1 = Allocate(size);
     AllocationId alloc_b2 = Allocate(size);
     for (size_t i = 0; i < size; ++i) {
       ObjectId id1 = alloc_b1.ObjectByOffset(i);
       ObjectId id2 = alloc_b2.ObjectByOffset(i);
-      if (i % 4 == 0) {
+      if (int group = i % 5; group == 0) {
+        main_b_ids[i] = ds[i].value<ObjectId>();
         main_b_values[i] = id1;
         merge_b_values[i] = id1;
-      } else if (i % 4 == 1) {
+      } else if (group == 1) {
+        fallback_b_ids[i] = ds[i].value<ObjectId>();
         fallback_b_values[i] = id2;
         merge_b_values[i] = id2;
-      } else if (i % 4 == 2) {
+      } else if (group == 2) {
+        main_b_ids[i] = ds[i].value<ObjectId>();
         main_b_values[i] = id1;
+        fallback_b_ids[i] = ds[i].value<ObjectId>();
         fallback_b_values[i] = id2;
         merge_b_values[i] = id1;
+      } else if (group == 3) {
+        // Note: main_b_values contains removed values, fallback not used
+        main_b_ids[i] = ds[i].value<ObjectId>();
+        fallback_b_ids[i] = ds[i].value<ObjectId>();
+        fallback_b_values[i] = id2;
       }
     }
+    auto main_ds_ids_b = DataSliceImpl::Create(
+        arolla::CreateDenseArray<ObjectId>(main_b_ids));
     auto main_ds_b = DataSliceImpl::Create(
         arolla::CreateDenseArray<ObjectId>(main_b_values));
-    ASSERT_OK(db->SetAttr(ds, "b", main_ds_b));
+    ASSERT_OK(db->SetAttr(main_ds_ids_b, "b", main_ds_b));
+    auto fallback_ds_ids_b = DataSliceImpl::Create(
+        arolla::CreateDenseArray<ObjectId>(fallback_b_ids));
     auto fallback_ds_b = DataSliceImpl::Create(
         arolla::CreateDenseArray<ObjectId>(fallback_b_values));
-    ASSERT_OK(db_f->SetAttr(ds, "b", fallback_ds_b));
+    ASSERT_OK(db_f->SetAttr(fallback_ds_ids_b, "b", fallback_ds_b));
     ASSERT_OK_AND_ASSIGN(DataSliceImpl ds_b_get,
                          db->GetAttr(ds, "b", {db_f.get()}));
     auto merge_ds_b = DataSliceImpl::Create(
@@ -526,7 +541,7 @@ TEST(DataBagTest, SetGetWithFallbackPrimitive) {
         merge_b_values[i] = val1;
       } else if (i % 4 == 1) {
         fallback_b_values[i] = val2;
-        merge_b_values[i] = val2;
+        // Note: main_b_values contains removed values, fallback not used
       } else if (i % 4 == 2) {
         main_b_values[i] = val1;
         fallback_b_values[i] = val2;
@@ -559,31 +574,42 @@ TEST(DataBagTest, SetGetWithFallbackPrimitiveMixedType) {
   for (size_t size : {1, 2, 4, 17, 126}) {
     auto ds = DataSliceImpl::AllocateEmptyObjects(size);
 
+    std::vector<arolla::OptionalValue<ObjectId>> main_b_ids(size);
     std::vector<arolla::OptionalValue<int64_t>> main_b_values(size);
+    std::vector<arolla::OptionalValue<ObjectId>> fallback_b_ids(size);
     std::vector<arolla::OptionalValue<double>> fallback_b_values(size);
+
     std::vector<arolla::OptionalValue<int64_t>> merge_b_values_int64(size);
     std::vector<arolla::OptionalValue<double>> merge_b_values_double(size);
     for (size_t i = 0; i < size; ++i) {
       int64_t val1 = i + 1;
       double val2 = -i - 1;
       if (i % 4 == 0) {
+        main_b_ids[i] = ds[i].value<ObjectId>();
         main_b_values[i] = val1;
         merge_b_values_int64[i] = val1;
       } else if (i % 4 == 1) {
+        fallback_b_ids[i] = ds[i].value<ObjectId>();
         fallback_b_values[i] = val2;
         merge_b_values_double[i] = val2;
       } else if (i % 4 == 2) {
+        main_b_ids[i] = ds[i].value<ObjectId>();
+        fallback_b_ids[i] = ds[i].value<ObjectId>();
         main_b_values[i] = val1;
         fallback_b_values[i] = val2;
         merge_b_values_int64[i] = val1;
       }
     }
+    auto main_ds_ids_b =
+        DataSliceImpl::Create(arolla::CreateDenseArray<ObjectId>(main_b_ids));
     auto main_ds_b =
         DataSliceImpl::Create(arolla::CreateDenseArray<int64_t>(main_b_values));
-    ASSERT_OK(db->SetAttr(ds, "b", main_ds_b));
+    ASSERT_OK(db->SetAttr(main_ds_ids_b, "b", main_ds_b));
+    auto fallback_ds_ids_b = DataSliceImpl::Create(
+        arolla::CreateDenseArray<ObjectId>(fallback_b_ids));
     auto fallback_ds_b = DataSliceImpl::Create(
         arolla::CreateDenseArray<double>(fallback_b_values));
-    ASSERT_OK(db_f->SetAttr(ds, "b", fallback_ds_b));
+    ASSERT_OK(db_f->SetAttr(fallback_ds_ids_b, "b", fallback_ds_b));
     ASSERT_OK_AND_ASSIGN(DataSliceImpl ds_b_get,
                          db->GetAttr(ds, "b", {db_f.get()}));
     auto merge_ds_b = DataSliceImpl::Create(