Optimize serialization of big allocation with SparseDataSources.

PiperOrigin-RevId: 871792692
Change-Id: I5888c934ab2e9fef64ab8f4b465116c3002122a1

Author: petrmikheev

koladata/internal/data_bag.cc

@@ -3411,11 +3411,6 @@ DataBagImpl::ExtractSmallAllocAttrContent(absl::string_view attr_name) const {
       }
     }
   }
-  std::sort(res.begin(), res.end(),
-            [](const DataBagContent::AttrItemContent& lhs,
-               const DataBagContent::AttrItemContent& rhs) {
-              return lhs.object_id < rhs.object_id;
-            });
   return res;
 }
 
@@ -3433,11 +3428,34 @@ absl::StatusOr<DataBagContent::AttrContent> DataBagImpl::ExtractAttrContent(
     if (size == 0) {
       continue;
     }
-    auto objects = DataSliceImpl::ObjectsFromAllocation(alloc, size);
-    ASSIGN_OR_RETURN(auto values, GetAttributeFromSources(
-                                      objects, dense_sources, sparse_sources));
-    content.allocs.push_back({alloc, values});
+    if (dense_sources.empty()) {
+      if (sparse_sources.size() == 1) {
+        for (const auto& [obj, value] : sparse_sources.front()->GetAll()) {
+          content.items.push_back({obj, value});
+        }
+      } else {
+        absl::flat_hash_set<ObjectId> visited_ids;
+        for (const SparseSource* source : sparse_sources) {
+          for (const auto& [obj, value] : source->GetAll()) {
+            if (bool inserted = visited_ids.insert(obj).second; inserted) {
+              content.items.push_back({obj, value});
+            }
+          }
+        }
+      }
+    } else {
+      auto objects = DataSliceImpl::ObjectsFromAllocation(alloc, size);
+      ASSIGN_OR_RETURN(
+          auto values,
+          GetAttributeFromSources(objects, dense_sources, sparse_sources));
+      content.allocs.push_back({alloc, std::move(values)});
+    }
   }
+  std::sort(content.items.begin(), content.items.end(),
+            [](const DataBagContent::AttrItemContent& lhs,
+               const DataBagContent::AttrItemContent& rhs) {
+              return lhs.object_id < rhs.object_id;
+            });
   return content;
 }
 

koladata/internal/data_bag.h

@@ -210,7 +210,7 @@ class DataBagImpl : public arolla::RefcountedBase {
   // DataSource with lower index will override values of sources with
   // higher indices.
   // Sparse sources always override dense sources.
-  // Returns size of the alloc (can be less then alloc.Capacity() e.g. if
+  // Returns size of the alloc (can be less than alloc.Capacity() e.g. if
   // an initial dense source for this allocation was created from a pre-existing
   // DenseArray).
   int64_t GetAttributeDataSources(AllocationId alloc, absl::string_view attr,

koladata/internal/data_bag_extract_content_test.cc

@@ -72,6 +72,44 @@ TEST(DataBagTest, ExtractSingleItemBigAllocRemovedValues) {
                                     std::bit_ceil(size), DataItem())));
 }
 
+TEST(DataBagTest, ExtractAttrFromSparseSources) {
+  auto db = DataBagImpl::CreateEmptyDatabag();
+  DataSliceImpl objs =
+      DataSliceImpl::ObjectsFromAllocation(Allocate(1000), 1000);
+
+  ASSERT_OK(db->SetAttr(objs[43], "a", DataItem(57)));
+  ASSERT_OK(db->SetAttr(objs[57], "a", DataItem(75)));
+
+  {
+    ASSERT_OK_AND_ASSIGN(auto content, db->ExtractContent());
+    EXPECT_TRUE(content.dicts.empty());
+    EXPECT_TRUE(content.lists.empty());
+    auto a = content.attrs["a"];
+    EXPECT_TRUE(a.allocs.empty());
+    ASSERT_EQ(a.items.size(), 2);
+    EXPECT_EQ(a.items[0].object_id, objs[43].value<ObjectId>());
+    EXPECT_EQ(a.items[0].value, DataItem(57));
+    EXPECT_EQ(a.items[1].object_id, objs[57].value<ObjectId>());
+    EXPECT_EQ(a.items[1].value, DataItem(75));
+  }
+
+  db = db->PartiallyPersistentFork();
+  ASSERT_OK(db->SetAttr(objs[57], "a", DataItem(43)));
+
+  {
+    ASSERT_OK_AND_ASSIGN(auto content, db->ExtractContent());
+    EXPECT_TRUE(content.dicts.empty());
+    EXPECT_TRUE(content.lists.empty());
+    auto a = content.attrs["a"];
+    EXPECT_TRUE(a.allocs.empty());
+    ASSERT_EQ(a.items.size(), 2);
+    EXPECT_EQ(a.items[0].object_id, objs[43].value<ObjectId>());
+    EXPECT_EQ(a.items[0].value, DataItem(57));
+    EXPECT_EQ(a.items[1].object_id, objs[57].value<ObjectId>());
+    EXPECT_EQ(a.items[1].value, DataItem(43));
+  }
+}
+
 TEST(DataBagTest, ExtractRemovedValues) {
   int64_t size = 5;
   auto db = DataBagImpl::CreateEmptyDatabag();

koladata/s11n/decoder.cc

@@ -576,8 +576,14 @@ absl::Status DecodeAttrProto(const KodaV1Proto::AttrProto& attr_proto,
     }
     const TypedValue& tval = input_values[chunk_proto.values_subindex()];
     if (tval.GetType() == arolla::GetQType<internal::DataItem>()) {
-      RETURN_IF_ERROR(db.SetAttr(internal::DataItem(obj), attr_proto.name(),
-                                 tval.UnsafeAs<internal::DataItem>()));
+      if (obj.IsSchema() && attr_proto.name() != schema::kSchemaNameAttr) {
+        RETURN_IF_ERROR(db.SetSchemaAttr(internal::DataItem(obj),
+                                         attr_proto.name(),
+                                         tval.UnsafeAs<internal::DataItem>()));
+      } else {
+        RETURN_IF_ERROR(db.SetAttr(internal::DataItem(obj), attr_proto.name(),
+                                   tval.UnsafeAs<internal::DataItem>()));
+      }
     } else {
       ASSIGN_OR_RETURN(const internal::DataSliceImpl& values,
                        tval.As<internal::DataSliceImpl>());

py/koladata/s11n/data_slice_s11n_test.py

@@ -552,6 +552,17 @@ def test_dtype(self):
       with self.assertRaisesRegex(ValueError, 'unsupported DType: ANY'):
         arolla.s11n.loads(container_proto.SerializeToString())
 
+  def test_memory_usage(self):
+    a = kd.obj_shaped_as(kd.item(1).repeat(int(1e6))).with_attrs(x=1)
+    b = a.S[0].extract()
+    # `b` is a single ObjectId from a big allocation. Here we test that we
+    # serialize the allocation in sparse form (1 ObjectId + value) rather than
+    # in dense form (1 base ObjectId + 1e6 bytes with type and presence
+    # information + value)
+    data = kd.dumps(b)
+    self.assertLess(len(data), 1000)
+    kd.testing.assert_equivalent(kd.loads(data), b)
+
 
 if __name__ == '__main__':
   absltest.main()