Verify alloc ids during DataSliceImpl creation.

PiperOrigin-RevId: 715375786
Change-Id: I7b80ab517989c504f9d8b227962ef1e8a02a82ff

Author: petrmikheev

koladata/internal/data_bag_benchmarks.cc

@@ -209,7 +209,7 @@ void BM_ObjAttributeDataBagManySources(benchmark::State& state) {
     for (int j = i; j < batch_size; j += chain_size) {
       ds_bldr.Set(j, ds_a.values<ObjectId>()[j]);
     }
-    alloc_ids.Insert(ds.allocation_ids());
+    alloc_ids.Insert(ds_a.allocation_ids());
   }
 
   auto ds =

koladata/internal/data_bag_test.cc

@@ -1307,8 +1307,7 @@ TEST(DataBagTest, InternalSetUnitAttrAndReturnMissingObjectsSparseToDense) {
     for (int64_t sz : {kSize / 100 + 1, kSize}) {
       auto objs_bldr = arolla::DenseArrayBuilder<ObjectId>(sz);
       objs_bldr.Set(0, AllocateSingleObject());
-      auto objs = DataSliceImpl::CreateWithAllocIds(
-          ds.allocation_ids(), std::move(objs_bldr).Build());
+      auto objs = DataSliceImpl::Create(std::move(objs_bldr).Build());
       ASSERT_OK_AND_ASSIGN(
           auto _, db->InternalSetUnitAttrAndReturnMissingObjects(objs, "a"));
     }
@@ -1438,8 +1437,8 @@ TEST(DataBagTest, SetGetDataSliceUUid) {
   auto object_array = arolla::CreateFullDenseArray<ObjectId>(
       {alloc1.ObjectByOffset(5), alloc3.ObjectByOffset(0),
        alloc2.ObjectByOffset(9)});
-  ds_b = DataSliceImpl::CreateWithAllocIds(AllocationIdSet({alloc1, alloc2}),
-                                           object_array);
+  ds_b = DataSliceImpl::CreateWithAllocIds(
+      AllocationIdSet({alloc1, alloc2, alloc3}), object_array);
   ASSERT_OK(db->SetAttr(ds, "b", ds_b));
   ASSERT_OK_AND_ASSIGN(ds_b_get, db->GetAttr(ds, "b"));
   ASSERT_EQ(ds_b_get.dtype(), arolla::GetQType<ObjectId>());

koladata/internal/data_slice.h

@@ -325,6 +325,27 @@ bool VerifyNonIntersectingIds(const arolla::DenseArray<T>& main_values,
   return !(process_ids(values) || ...);
 }
 
+template <class T>
+bool VerifyAllocIds(const AllocationIdSet& alloc_ids,
+                    const arolla::DenseArray<T>& values) {
+  if constexpr (std::is_same_v<T, ObjectId>) {
+    bool res = true;
+    values.ForEachPresent([&](int64_t /*id*/, ObjectId obj_id) {
+      if (!obj_id.IsSmallAlloc()) {
+        for (AllocationId alloc_id : alloc_ids.ids()) {
+          if (alloc_id.Contains(obj_id)) return;
+        }
+        res = false;
+      } else if (!alloc_ids.contains_small_allocation_id()) {
+        res = false;
+      }
+    });
+    return res;
+  } else {
+    return true;
+  }
+}
+
 template <class T>
 constexpr bool AreAllTypesDistinct(std::type_identity<T>) {
   return true;
@@ -396,27 +417,30 @@ DataSliceImpl DataSliceImpl::CreateWithAllocIds(
     AllocationIdSet allocation_ids, arolla::DenseArray<T> main_values,
     arolla::DenseArray<Ts>... values) {
   DataSliceImpl res;
-  CreateImpl(res, std::move(main_values), std::move(values)...);
   if constexpr ((std::is_same_v<T, ObjectId> || ... ||
                  std::is_same_v<Ts, ObjectId>)) {
+    DCHECK(data_slice_impl::VerifyAllocIds(allocation_ids, main_values));
+    DCHECK((data_slice_impl::VerifyAllocIds(allocation_ids, values) && ...));
     res.internal_->allocation_ids = std::move(allocation_ids);
   } else {
     (void)allocation_ids;
   }
+  CreateImpl(res, std::move(main_values), std::move(values)...);
   return res;
 }
 
 template <class T>
 DataSliceImpl DataSliceImpl::CreateWithTypesBuffer(TypesBuffer types_buffer,
     AllocationIdSet allocation_ids, arolla::DenseArray<T> values) {
   DataSliceImpl res;
-  CreateImpl(res, std::move(values));
   res.internal_->types_buffer = std::move(types_buffer);
   if constexpr (std::is_same_v<T, ObjectId>) {
+    DCHECK(data_slice_impl::VerifyAllocIds(allocation_ids, values));
     res.internal_->allocation_ids = std::move(allocation_ids);
   } else {
     (void)allocation_ids;
   }
+  CreateImpl(res, std::move(values));
   return res;
 }