Add DataSliceImpl::CreateWithTypesBuffer (will be used in DenseSource, needed for removed values support)

PiperOrigin-RevId: 714900264
Change-Id: Ie9ce652512c4054fb75cd3902f9586fe1983405d

Author: petrmikheev

koladata/internal/BUILD

@@ -419,6 +419,7 @@ cc_test(
         ":dtype",
         ":object_id",
         ":types",
+        ":types_buffer",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/strings",

koladata/internal/data_slice.h

@@ -86,6 +86,14 @@ class DataSliceImpl {
                                           arolla::DenseArray<T> main_values,
                                           arolla::DenseArray<Ts>... values);
 
+  // Creates a single-type DataSliceImpl with provided TypesBuffer. TypesBuffer
+  // in single-type slice is optional, used in DenseSource to distinguish
+  // missing and removed values.
+  template <class T>
+  static DataSliceImpl CreateWithTypesBuffer(TypesBuffer types_buffer,
+                                             AllocationIdSet allocation_ids,
+                                             arolla::DenseArray<T> values);
+
   static DataSliceImpl Create(const arolla::DenseArray<DataItem>& items);
 
   // Returns 0 dimension DataSliceImpl with specified values. In case of
@@ -398,6 +406,21 @@ DataSliceImpl DataSliceImpl::CreateWithAllocIds(
   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>) {
+    res.internal_->allocation_ids = std::move(allocation_ids);
+  } else {
+    (void)allocation_ids;
+  }
+  return res;
+}
+
+
 using DataSlicePtr = std::unique_ptr<DataSliceImpl>;
 
 }  // namespace koladata::internal

koladata/internal/data_slice_test.cc

@@ -31,6 +31,7 @@
 #include "koladata/internal/object_id.h"
 #include "koladata/internal/slice_builder.h"
 #include "koladata/internal/types.h"
+#include "koladata/internal/types_buffer.h"
 #include "arolla/dense_array/dense_array.h"
 #include "arolla/dense_array/qtype/types.h"
 #include "arolla/memory/optional_value.h"
@@ -150,6 +151,27 @@ TEST(DataSliceImpl, Create) {
   EXPECT_THAT(ds.values<int64_t>(), ElementsAre(57, 75, 19));
 }
 
+TEST(DataSliceImpl, CreateWithTypesBuffer) {
+  constexpr int64_t kSize = 3;
+  auto array =
+      arolla::CreateDenseArray<int64_t>({std::nullopt, 75, std::nullopt});
+  TypesBuffer types_buffer;
+  types_buffer.InitAllUnset(3);
+  types_buffer.id_to_typeidx[1] =
+      types_buffer.get_or_add_typeidx(ScalarTypeId<int64_t>());
+  types_buffer.id_to_typeidx[2] = TypesBuffer::kRemoved;
+  DataSliceImpl ds = DataSliceImpl::CreateWithTypesBuffer(
+      std::move(types_buffer), AllocationIdSet(), array);
+  EXPECT_EQ(ds.size(), kSize);
+  EXPECT_EQ(ds.dtype(), arolla::GetQType<int64_t>());
+  EXPECT_FALSE(ds.is_empty_and_unknown());
+  EXPECT_THAT(ds.allocation_ids(), IsEmpty());
+  EXPECT_THAT(ds, ElementsAre(DataItem(), 75, DataItem()));
+  EXPECT_EQ(ds.types_buffer().id_to_typeidx[0], TypesBuffer::kUnset);
+  EXPECT_EQ(ds.types_buffer().id_to_typeidx[1], 0);
+  EXPECT_EQ(ds.types_buffer().id_to_typeidx[2], TypesBuffer::kRemoved);
+}
+
 TEST(DataSliceImpl, CreateFromDataItemSpan) {
   {
     auto array = std::vector<DataItem>{};