Change from_py and map_py to always assign a bag when schema=OBJECT is requested

Motivation:
 * Removed an edge case in the behaviour of `kd.from_py(..., from_dim=1)` where the result would unexpectedly be assigned a bag when the input list was empty. This edge case was encountered while optimising `map_py_on_present`.
 * Removed discrepancy:
   - `kd.map_py(lambda x: x, [], schema=kd.OBJECT).x` ok
   - `kd.map_py(lambda x: x, [None], schema=kd.OBJECT).x` ok\
   while
   - `kd.map_py(lambda x: x, []).x` ok
   - `kd.map_py(lambda x: x, [None]).x` failure
PiperOrigin-RevId: 714118231
Change-Id: I57aab375f6185a31cc33a9f03a899cca336ee902

Author: apronchenkov

py/koladata/functions/tests/from_py_test.py

@@ -197,8 +197,9 @@ def test_primitive(self):
     self.assertIsNone(item.get_bag())
 
     item = fns.from_py(42, schema=schema_constants.OBJECT)
-    testing.assert_equal(item, ds(42, schema_constants.OBJECT))
-    self.assertIsNone(item.get_bag())
+    testing.assert_equal(item.no_bag(), ds(42, schema_constants.OBJECT))
+    self.assertIsNotNone(item.get_bag())
+    self.assertFalse(item.get_bag().is_mutable())
 
   def test_primitive_casting_error(self):
     with self.assertRaisesRegex(
@@ -225,23 +226,27 @@ def test_primitives_common_schema(self):
 
   def test_primitives_object(self):
     res = fns.from_py([1, 3.14], from_dim=1, schema=schema_constants.OBJECT)
-    testing.assert_equal(res, ds([1, 3.14], schema_constants.OBJECT))
-    self.assertIsNone(res.get_bag())
+    testing.assert_equal(res.no_bag(), ds([1, 3.14], schema_constants.OBJECT))
+    self.assertIsNotNone(res.get_bag())
+    self.assertFalse(res.get_bag().is_mutable())
 
   def test_empty_object(self):
     res = fns.from_py(None, schema=schema_constants.OBJECT)
     testing.assert_equal(res.no_bag(), ds(None, schema_constants.OBJECT))
     self.assertIsNotNone(res.get_bag())
+    self.assertFalse(res.get_bag().is_mutable())
 
     res = fns.from_py([], from_dim=1, schema=schema_constants.OBJECT)
     testing.assert_equal(res.no_bag(), ds([], schema_constants.OBJECT))
     self.assertIsNotNone(res.get_bag())
+    self.assertFalse(res.get_bag().is_mutable())
 
     res = fns.from_py([None, None], from_dim=1, schema=schema_constants.OBJECT)
     testing.assert_equal(
         res.no_bag(), ds([None, None], schema_constants.OBJECT)
     )
     self.assertIsNotNone(res.get_bag())
+    self.assertFalse(res.get_bag().is_mutable())
 
   def test_list_from_dim(self):
     input_list = [[1, 2.0], [3, 4]]
@@ -306,7 +311,9 @@ def test_list_from_dim_with_schema(self):
 
     l3 = fns.from_py(input_list, schema=schema_constants.OBJECT, from_dim=2)
     self.assertEqual(l3.get_ndim(), 2)
-    testing.assert_equal(l3, ds([[1, 2.0], [3, 4]], schema_constants.OBJECT))
+    testing.assert_equal(
+        l3.no_bag(), ds([[1, 2.0], [3, 4]], schema_constants.OBJECT)
+    )
 
   def test_dict_from_dim(self):
     input_dict = [{ds('a'): [1, 2], 'b': [42]}, {ds('c'): [3, 4], 'd': [34]}]

py/koladata/operators/tests/py_map_py_on_cond_test.py

@@ -90,7 +90,7 @@ def test_empty_input(self):
 
     res = expr_eval.eval(kde.py.map_py_on_cond(yes_fn, None, cond, x=val))
     testing.assert_equal(res.no_bag(), ds([], schema_constants.OBJECT))
-    self.assertIsNotNone(res.get_bag())
+    self.assertIsNone(res.get_bag())
 
     res = expr_eval.eval(
         kde.py.map_py_on_cond(

py/koladata/operators/tests/py_map_py_on_present_test.py

@@ -53,7 +53,7 @@ def test_rank_2(self):
     testing.assert_equal(r.no_bag(), ds([[4.5, 5.5], [], []]))
 
   def test_empty_input(self):
-    yes_fn = lambda x: x + 1 if x is not None else 0
+    yes_fn = lambda x: x + 1
 
     val = ds([])
     res = expr_eval.eval(
@@ -64,7 +64,7 @@ def test_empty_input(self):
 
     res = expr_eval.eval(kde.py.map_py_on_present(yes_fn, x=val))
     testing.assert_equal(res.no_bag(), ds([], schema_constants.OBJECT))
-    self.assertIsNotNone(res.get_bag())
+    self.assertIsNone(res.get_bag())
 
     res = expr_eval.eval(
         kde.py.map_py_on_present(yes_fn, x=val, schema=schema_constants.OBJECT)

py/koladata/operators/tests/py_map_py_on_selected_test.py

@@ -77,7 +77,7 @@ def test_empty_input(self):
 
     res = expr_eval.eval(kde.py.map_py_on_selected(fn, cond, x=val))
     testing.assert_equal(res.no_bag(), ds([], schema_constants.OBJECT))
-    self.assertIsNotNone(res.get_bag())
+    self.assertIsNone(res.get_bag())
 
     res = expr_eval.eval(
         kde.py.map_py_on_selected(

py/koladata/operators/tests/py_map_py_test.py

@@ -300,7 +300,7 @@ def my_fn(x):
           kde.py.map_py(my_fn, val, include_missing=include_missing)
       )
       testing.assert_equal(res.no_bag(), ds([[]], schema_constants.OBJECT))
-      self.assertIsNotNone(res.get_bag())
+      self.assertIsNone(res.get_bag())
 
     with self.subTest('schema=FLOAT32'):
       res = expr_eval.eval(
@@ -324,7 +324,7 @@ def my_fn(x):
           )
       )
       testing.assert_equal(res.no_bag(), ds([[]], schema_constants.OBJECT))
-      self.assertIsNotNone(res.get_bag())
+      self.assertFalse(res.get_bag().is_mutable())
 
   @parameterized.parameters(False, True)
   def test_map_py_all_missing_input(self, include_missing):
@@ -364,7 +364,7 @@ def my_fn(x):
           )
       )
       testing.assert_equal(res.no_bag(), ds([[None]], schema_constants.OBJECT))
-      self.assertIsNotNone(res.get_bag())
+      self.assertFalse(res.get_bag().is_mutable())
 
   def test_map_py_scalar_input(self):
     def add_one(x):

py/koladata/types/boxing.cc

@@ -1530,10 +1530,11 @@ absl::StatusOr<DataSlice> GenericFromPyObject(
   DCHECK(res_db == nullptr || res_db->IsMutable());
   if (res_slice.GetBag() == nullptr) {
     ASSIGN_OR_RETURN(res_db, adoption_queue.GetCommonOrMergedDb());
-    // Attach an empty DataBag if the result is empty and has OBJECT schema.
-    if (res_db == nullptr && res_slice.present_count() == 0 &&
-        res_slice.GetSchemaImpl() == schema::kObject) {
+    // If the result has no associated DataBag but an OBJECT schema was
+    // requested, attach an empty DataBag.
+    if (res_db == nullptr && schema && schema->item() == schema::kObject) {
       res_db = DataBag::Empty();
+      res_db->UnsafeMakeImmutable();
     }
     return res_slice.WithBag(std::move(res_db));
   }