Add DataSlice::CreateWithFlatShape helper.

PiperOrigin-RevId: 715891339
Change-Id: I79d3f4df03a045c97116b7619ee2259ab6e83ced

Author: nbjohnson0

koladata/BUILD

@@ -256,6 +256,7 @@ cc_test(
         "@com_google_absl//absl/strings:str_format",
         "@com_google_arolla//arolla/dense_array",
         "@com_google_arolla//arolla/dense_array/qtype",
+        "@com_google_arolla//arolla/jagged_shape/testing",
         "@com_google_arolla//arolla/qtype",
         "@com_google_arolla//arolla/util",
         "@com_google_arolla//arolla/util/testing",

koladata/data_slice.cc

@@ -848,6 +848,14 @@ absl::StatusOr<DataSlice> DataSlice::CreateWithSchemaFromData(
                 std::move(db), wholeness);
 }
 
+absl::StatusOr<DataSlice> DataSlice::CreateWithFlatShape(
+    internal::DataSliceImpl impl, internal::DataItem schema, DataBagPtr db,
+    Wholeness wholeness) {
+  const auto size = impl.size();
+  return Create(std::move(impl), JaggedShape::FlatFromSize(size),
+                std::move(schema), std::move(db), wholeness);
+}
+
 absl::StatusOr<DataSlice> DataSlice::Create(const internal::DataItem& item,
                                             JaggedShape shape,
                                             internal::DataItem schema,

koladata/data_slice.h

@@ -111,6 +111,11 @@ class DataSlice {
       internal::DataSliceImpl impl, JaggedShape shape, DataBagPtr db = nullptr,
       Wholeness wholeness = Wholeness::kNotWhole);
 
+  // Creates a DataSlice with shape JaggedShape::FlatFromSize(impl.size()).
+  static absl::StatusOr<DataSlice> CreateWithFlatShape(
+      internal::DataSliceImpl impl, internal::DataItem schema,
+      DataBagPtr db = nullptr, Wholeness wholeness = Wholeness::kNotWhole);
+
   // Convenience factory method that accepts JaggedShape, so that we can use
   // implementation-agnostic constructions in visitors passed to VisitImpl.
   static absl::StatusOr<DataSlice> Create(

koladata/data_slice_test.cc

@@ -47,6 +47,7 @@
 #include "koladata/testing/matchers.h"
 #include "arolla/dense_array/dense_array.h"
 #include "arolla/dense_array/qtype/types.h"
+#include "arolla/jagged_shape/testing/matchers.h"
 #include "arolla/qtype/qtype_traits.h"
 #include "arolla/qtype/typed_value.h"
 #include "arolla/util/bytes.h"
@@ -336,6 +337,25 @@ TEST(DataSliceUtils, CreateWithSchemaFromDataError) {
                        HasSubstr("for primitive types")));
 }
 
+TEST(DataSliceTest, CreateWithFlatShape) {
+  {
+    auto slice = *DataSlice::CreateWithFlatShape(
+        internal::DataSliceImpl::Create(std::vector<internal::DataItem>{}),
+        internal::DataItem(schema::kNone));
+    EXPECT_THAT(slice.GetShape(),
+                IsEquivalentTo(DataSlice::JaggedShape::FlatFromSize(0)));
+  }
+  {
+    auto slice = *DataSlice::CreateWithFlatShape(
+        internal::DataSliceImpl::Create(std::vector<internal::DataItem>{
+            internal::DataItem(1), internal::DataItem(2),
+            internal::DataItem(3)}),
+        internal::DataItem(schema::kInt32));
+    EXPECT_THAT(slice.GetShape(),
+                IsEquivalentTo(DataSlice::JaggedShape::FlatFromSize(3)));
+  }
+}
+
 TEST(DataSliceTest, IsWhole) {
   {
     // No DataBag, trivially whole.