Forbid mutable original DataBag and autofreeze mutable argument DataBags in "kd.core.enriched" and "kd.core.updated".

PiperOrigin-RevId: 859044745
Change-Id: I57e654ad711833447a3489161dc1234be6f9f063

Author: petrmikheev

docs/api_reference.md

@@ -708,7 +708,9 @@ Aliases:
 <pre class="no-copy"><code class="lang-text no-auto-prettify">Returns a copy of a DataSlice with a additional fallback DataBag(s).
 
 Values in the original DataBag of `ds` take precedence over the ones in
-`*bag`.
+`*bag`. The original DataBag (if present) must be immutable. If the original
+DataBag is mutable, either freeze `ds` first, or add updates inplace using
+mutable API.
 
 The DataBag attached to the result is a new immutable DataBag that falls back
 to the DataBag of `ds` if present and then to `*bag`.
@@ -1188,7 +1190,9 @@ Aliases:
 <pre class="no-copy"><code class="lang-text no-auto-prettify">Returns a copy of a DataSlice with DataBag(s) of updates applied.
 
 Values in `*bag` take precedence over the ones in the original DataBag of
-`ds`.
+`ds`. The original DataBag (if present) must be immutable. If the original
+DataBag is mutable, either freeze `ds` first, or add updates inplace using
+mutable API.
 
 The DataBag attached to the result is a new immutable DataBag that falls back
 to the DataBag of `ds` if present and then to `*bag`.
@@ -12832,7 +12836,9 @@ Args:
 <pre class="no-copy"><code class="lang-text no-auto-prettify">Returns a copy of a DataSlice with a additional fallback DataBag(s).
 
 Values in the original DataBag of `ds` take precedence over the ones in
-`*bag`.
+`*bag`. The original DataBag (if present) must be immutable. If the original
+DataBag is mutable, either freeze `ds` first, or add updates inplace using
+mutable API.
 
 The DataBag attached to the result is a new immutable DataBag that falls back
 to the DataBag of `ds` if present and then to `*bag`.
@@ -13646,7 +13652,9 @@ Args:
 <pre class="no-copy"><code class="lang-text no-auto-prettify">Returns a copy of a DataSlice with DataBag(s) of updates applied.
 
 Values in `*bag` take precedence over the ones in the original DataBag of
-`ds`.
+`ds`. The original DataBag (if present) must be immutable. If the original
+DataBag is mutable, either freeze `ds` first, or add updates inplace using
+mutable API.
 
 The DataBag attached to the result is a new immutable DataBag that falls back
 to the DataBag of `ds` if present and then to `*bag`.

koladata/operators/core.cc

@@ -18,6 +18,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <string>
 #include <utility>
 #include <vector>
 
@@ -366,27 +367,39 @@ class EnrichedOrUpdatedOperator final : public arolla::QExprOperator {
   absl::StatusOr<std::unique_ptr<arolla::BoundOperator>> DoBind(
       absl::Span<const arolla::TypedSlot> input_slots,
       arolla::TypedSlot output_slot) const override {
+    absl::string_view op_name =
+        is_enriched_operator_ ? "kd.core.enriched" : "kd.core.updated";
     return MakeBoundOperator(
-        is_enriched_operator_ ? "kd.core.enriched" : "kd.core.updated",
+        std::string(op_name),
         [input_slots = std::vector<arolla::TypedSlot>(input_slots.begin(),
                                                       input_slots.end()),
          output_slot = output_slot.UnsafeToSlot<DataSlice>(),
-         is_enriched_operator = is_enriched_operator_](
+         is_enriched_operator = is_enriched_operator_, op_name](
             arolla::EvaluationContext* ctx, arolla::FramePtr frame) {
           const DataSlice& ds =
               frame.Get(input_slots[0].UnsafeToSlot<DataSlice>());
+          if (ds.GetBag() != nullptr && ds.GetBag()->IsMutable()) {
+            return absl::InvalidArgumentError(
+                absl::StrCat(op_name,
+                             " requires the original DataBag to be immutable; "
+                             "either freeze it, or use mutable API"));
+          }
           std::vector<DataBagPtr> db_list;
           db_list.reserve(input_slots.size());
           if (is_enriched_operator) {
             db_list.push_back(ds.GetBag());
             for (size_t i = 1; i < input_slots.size(); ++i) {
+              const auto& bag =
+                  frame.Get(input_slots[i].UnsafeToSlot<DataBagPtr>());
               db_list.push_back(
-                  frame.Get(input_slots[i].UnsafeToSlot<DataBagPtr>()));
+                  bag != nullptr && bag->IsMutable() ? bag->Freeze() : bag);
             }
           } else {
             for (size_t i = input_slots.size() - 1; i >= 1; --i) {
+              const auto& bag =
+                  frame.Get(input_slots[i].UnsafeToSlot<DataBagPtr>());
               db_list.push_back(
-                  frame.Get(input_slots[i].UnsafeToSlot<DataBagPtr>()));
+                  bag != nullptr && bag->IsMutable() ? bag->Freeze() : bag);
             }
             db_list.push_back(ds.GetBag());
           }

py/koladata/benchmarks/set_get_attr_benchmarks.py

@@ -291,7 +291,7 @@ def set_get_multiple_attrs_entity_via_fallback(state):
     update_ds = ds.with_bag(kd.mutable_bag())
     arg_name = 'abc' if state.range(1) else f'abc{i}'
     update_ds.set_attr(arg_name, val)
-    updates.append((arg_name, update_ds))
+    updates.append((arg_name, update_ds.freeze_bag()))
   try:
     _ = ds.missing  # To initialize all lazy initializers and reduce variance.
   except AttributeError:
@@ -410,7 +410,7 @@ def set_get_multiple_attrs_10000_entity_via_fallback(state):
     update_ds = ds.with_bag(kd.mutable_bag())
     arg_name = 'abc' if state.range(1) else f'abc{i}'
     update_ds.set_attr(arg_name, val)
-    updates.append((arg_name, update_ds))
+    updates.append((arg_name, update_ds.freeze_bag()))
   try:
     _ = ds.missing  # To initialize all lazy initializers and reduce variance.
   except AttributeError:

py/koladata/ext/persisted_data/schema_helper.py

@@ -625,7 +625,7 @@ class SchemaHelper:
 
   def __init__(self, schema: kd.types.SchemaItem):
     _check_is_schema_item(schema)
-    self._schema = schema
+    self._schema = schema.freeze_bag()
     artifacts = _analyze_schema(self._schema)
     self._parent_to_child_graph = artifacts.schema_graph
     self._schema_node_name_to_schema = artifacts.schema_node_name_to_schema

py/koladata/functions/tests/set_attr_test.py

@@ -165,16 +165,12 @@ def test_merging(self):
     testing.assert_equal(x.xyz.a, ds([5, 4]).with_bag(db))
 
   def test_merging_with_fallbacks(self):
-    db = kd.mutable_bag()
-    x = db.obj()
-    db2 = kd.mutable_bag()
-    y = db2.new(bar='foo')
-    db3 = kd.mutable_bag()
-    y.with_bag(db3).set_attr('bar', 2)
-    y.with_bag(db3).set_attr('baz', 5)
-    x.foo = y.enriched(db3)
-    testing.assert_equal(x.foo.bar, ds('foo').with_bag(db))
-    testing.assert_equal(x.foo.baz, ds(5).with_bag(db))
+    x = kd.mutable_bag().obj()
+    y = kd.new(bar='foo')
+    z = y.with_bag(kd.bag()).with_attrs(bar=2, baz=5)
+    x.foo = y.enriched(z.get_bag())
+    testing.assert_equal(x.foo.bar.no_bag(), ds('foo'))
+    testing.assert_equal(x.foo.baz.no_bag(), ds(5))
 
   def test_assignment_rhs_error(self):
     x = kd.mutable_bag().obj(x=ds([1, 2]))

py/koladata/operators/core.py

@@ -349,7 +349,9 @@ def updated(ds, *bag):  # pylint: disable=unused-argument
   """Returns a copy of a DataSlice with DataBag(s) of updates applied.
 
   Values in `*bag` take precedence over the ones in the original DataBag of
-  `ds`.
+  `ds`. The original DataBag (if present) must be immutable. If the original
+  DataBag is mutable, either freeze `ds` first, or add updates inplace using
+  mutable API.
 
   The DataBag attached to the result is a new immutable DataBag that falls back
   to the DataBag of `ds` if present and then to `*bag`.
@@ -1270,7 +1272,9 @@ def enriched(ds, *bag):  # pylint: disable=unused-argument
   """Returns a copy of a DataSlice with a additional fallback DataBag(s).
 
   Values in the original DataBag of `ds` take precedence over the ones in
-  `*bag`.
+  `*bag`. The original DataBag (if present) must be immutable. If the original
+  DataBag is mutable, either freeze `ds` first, or add updates inplace using
+  mutable API.
 
   The DataBag attached to the result is a new immutable DataBag that falls back
   to the DataBag of `ds` if present and then to `*bag`.

py/koladata/operators/tests/core_attr_test.py

@@ -85,15 +85,13 @@ def test_attr_update_implicit_casting(self):
 
   def test_attr_update_with_ds_attr(self):
     with self.subTest('object'):
-      db = bag()
-      ds1 = kd.stack(db.obj(x=1), db.obj(y=2))
+      ds1 = kd.stack(kd.obj(x=1), kd.obj(y=2))
       db = kd.core.attr(ds1, ds(['x', 'y']), 42)
       testing.assert_equal(ds1.updated(db).maybe('x').no_bag(), ds([42, None]))
       testing.assert_equal(ds1.updated(db).maybe('y').no_bag(), ds([None, 42]))
 
     with self.subTest('entity'):
-      db = bag()
-      ds1 = db.new(x=ds([1, 3]), y=ds([2, 4]))
+      ds1 = kd.new(x=ds([1, 3]), y=ds([2, 4]))
       db = kd.core.attr(ds1, ds(['x', 'y']), 42)
       testing.assert_equal(ds1.updated(db).x.no_bag(), ds([42, 3]))
       testing.assert_equal(ds1.updated(db).y.no_bag(), ds([2, 42]))

py/koladata/operators/tests/core_attrs_test.py

@@ -83,7 +83,7 @@ def test_multi_attr_overwrite(self):
     )
 
   def test_multi_attr_overwrite_object_schema(self):
-    o = bag().obj(x=1, y=10)
+    o = kd.obj(x=1, y=10)
     db2 = kd.core.attrs(o, x='2', a=1, b='p', c=bag().list([1, 2]))
     self.assertNotEqual(o.get_bag().fingerprint, db2.fingerprint)
 

py/koladata/operators/tests/core_clone_test.py

@@ -116,7 +116,7 @@ def test_entity(self, pass_schema):
   )
   def test_clone_only_reachable(self, pass_schema):
     a_slice = bag().new(b=ds([1, None, 2]), c=ds(['foo', 'bar', 'baz']))
-    o = bag().new(a=a_slice)
+    o = kd.new(a=a_slice)
     fb = bag().new(a=a_slice.no_bag(), c=ds([1, None, 2]))
     o = o.enriched(fb.get_bag())
     if pass_schema:
@@ -155,9 +155,9 @@ def test_uu(self, noise_positioned_in_front, pass_schema):
     fb_noise = bag()
     noise = fb_noise.obj(a=[1, 2, 3])
     if noise_positioned_in_front:
-      o_fb = o.with_bag(noise.enriched(db).get_bag())
+      o_fb = o.with_bag(noise.freeze_bag().enriched(db).get_bag())
     else:
-      o_fb = o.enriched(fb_noise)
+      o_fb = o.freeze_bag().enriched(fb_noise)
 
     if pass_schema:
       result = kd.core.clone(o_fb, schema=o_fb.get_schema())

py/koladata/operators/tests/core_deep_clone_test.py

@@ -153,11 +153,9 @@ def test_named_schema(self, pass_schema):
 
   @parameterized.parameters(True, False)
   def test_clone_only_reachable(self, pass_schema):
-    db = data_bag.DataBag.empty_mutable()
-    fb = data_bag.DataBag.empty_mutable()
-    a_slice = db.new(b=ds([1, None, 2]), c=ds(['foo', 'bar', 'baz']))
-    _ = fb.new(a=a_slice.no_bag(), c=ds([1, None, 2]))
-    o = db.new(a=a_slice)
+    a_slice = kd.new(b=ds([1, None, 2]), c=ds(['foo', 'bar', 'baz']))
+    fb = kd.new(a=a_slice.no_bag(), c=ds([1, None, 2])).get_bag()
+    o = kd.new(a=a_slice)
     merged_bag = o.enriched(fb).get_bag().merge_fallbacks()
     o = o.with_bag(merged_bag)
     if pass_schema:
@@ -257,8 +255,7 @@ def test_metadata_implicit_schema(self):
     kd.core.with_metadata(a.get_obj_schema(), attrs='xy')
 
   def test_metadata_chains(self):
-    db = bag()
-    ds_root = db.new(a=db.obj(x=1, y=2), b=db.obj(foo='bar'))
+    ds_root = kd.new(a=kd.obj(x=1, y=2), b=kd.obj(foo='bar'))
     ds_a = ds_root.a
     for i in range(10):
       schema = kd.core.with_metadata(ds_a.get_obj_schema(), data=f'a_depth_{i}')