Use FromPy in kd.new.

PiperOrigin-RevId: 859109353
Change-Id: I3c1241ed572bc56bcce7ddd54c2ccc85ccaba0d7

Author: atorero

py/koladata/base/BUILD

@@ -214,6 +214,7 @@ cc_library(
         "//koladata:data_slice",
         "//koladata:data_slice_qtype",
         "//koladata:object_factories",
+        "//py/koladata/base/py_conversions:from_py",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",

py/koladata/base/boxing.cc

@@ -1352,26 +1352,4 @@ class UniversalConverter {
 
 }  // namespace
 
-absl::StatusOr<DataSlice> EntitiesFromPyObject(PyObject* py_obj,
-                                               const DataBagPtr& db,
-                                               AdoptionQueue& adoption_queue) {
-  return EntitiesFromPyObject(py_obj, /*schema=*/std::nullopt,
-                              /*itemid=*/std::nullopt, db, adoption_queue);
-}
-
-absl::StatusOr<DataSlice> EntitiesFromPyObject(
-    PyObject* py_obj, const std::optional<DataSlice>& schema,
-    const std::optional<DataSlice>& itemid, const DataBagPtr& db,
-    AdoptionQueue& adoption_queue) {
-  // NOTE: UniversalConverter does not allow converting multi-dimensional
-  // DataSlices, so we are processing it before invoking the UniversalConverter.
-  if (arolla::python::IsPyQValueInstance(py_obj) && !itemid.has_value()) {
-    ASSIGN_OR_RETURN(auto res,
-                     DataSliceFromPyValue(py_obj, adoption_queue, schema));
-    return EntityCreator::ConvertWithoutAdopt(db, res);
-  }
-  return UniversalConverter<EntityCreator>(db, adoption_queue)
-      .Convert(py_obj, schema, /*from_dim=*/0, itemid);
-}
-
 }  // namespace koladata::python

py/koladata/base/boxing.h

@@ -26,7 +26,6 @@
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "koladata/adoption_utils.h"
-#include "koladata/data_bag.h"
 #include "koladata/data_slice.h"
 #include "koladata/internal/data_item.h"
 
@@ -86,25 +85,6 @@ inline absl::StatusOr<DataSlice> DataSliceFromPyValueNoAdoption(
 absl::StatusOr<DataSlice> DataItemFromPyValue(
     PyObject* py_obj, const std::optional<DataSlice>& schema = std::nullopt);
 
-// Converts Python objects into DataSlices and converts them into appropriate
-// Koda Entities. The conversion is deep, such that all nested structures (e.g.
-// dicts) are also Entities (including Koda Lists and Dicts) or primitive
-// DataSlices.
-//
-// `db` and `adoption_queue` are side outputs. `db` is used to create Koda
-// objects found under `py_obj`, while `adoption_queue` is used to collect
-// DataBag(s) of DataSlice(s) from nested `py_obj`.
-absl::StatusOr<DataSlice> EntitiesFromPyObject(PyObject* py_obj,
-                                               const DataBagPtr& db,
-                                               AdoptionQueue& adoption_queue);
-
-// Same as above, but allows specifying the schemas of Lists / Dicts. When
-// schema == OBJECT, the behavior is the same as `ObjectFromPyObject`.
-absl::StatusOr<DataSlice> EntitiesFromPyObject(
-    PyObject* py_obj, const std::optional<DataSlice>& schema,
-    const std::optional<DataSlice>& itemid, const DataBagPtr& db,
-    AdoptionQueue& adoption_queue);
-
 // Returns an Error from provided status with incompatible schema information
 // during narrow casting.
 absl::Status CreateIncompatibleSchemaErrorFromStatus(

py/koladata/base/py_utils.cc

@@ -33,6 +33,7 @@
 #include "koladata/object_factories.h"
 #include "py/arolla/abc/py_qvalue.h"
 #include "py/koladata/base/boxing.h"
+#include "py/koladata/base/py_conversions/from_py.h"
 #include "arolla/util/status_macros_backport.h"
 
 namespace koladata::python {
@@ -76,7 +77,25 @@ absl::StatusOr<DataSlice> AssignmentRhsFromPyValue(
           " the slice, or kd.dict_like() to create multiple"
           " dictionary instances");
     }
-    return EntitiesFromPyObject(rhs, db, adoption_queue);
+
+    if (arolla::python::IsPyQValueInstance(rhs)) {
+      const auto& typed_value = arolla::python::UnsafeUnwrapPyQValue(rhs);
+      if (typed_value.GetType() == arolla::GetQType<DataSlice>()) {
+        return typed_value.UnsafeAs<DataSlice>();
+      } else {
+        return absl::InvalidArgumentError(absl::StrFormat(
+            "only DataSlice QValues are supported as objects; got: %s.",
+            typed_value.GetType()->name()));
+      }
+    }
+
+    ASSIGN_OR_RETURN(DataSlice res,
+                     FromPy(rhs, std::nullopt, /*from_dim=*/0,
+                            /*dict_as_obj=*/false, /*itemid=*/std::nullopt));
+
+    adoption_queue.Add(res);
+    RETURN_IF_ERROR(adoption_queue.AdoptInto(*db));
+    return res.WithBag(db);
   }
   ASSIGN_OR_RETURN(auto res, DataSliceFromPyValue(rhs, adoption_queue));
   if (res.GetShape().rank() > 0) {

py/koladata/functions/tests/new_test.py

@@ -203,12 +203,9 @@ def test_schema_arg_dict_schema_error(self):
     list_schema = kde.list_schema(item_schema=schema_constants.FLOAT32).eval()
     with self.assertRaisesRegex(
         ValueError,
-        re.escape(
-            'Python Dict can be converted to either Entity or Dict, got schema:'
-            ' DataItem(LIST[FLOAT32]'
-        ),
+        re.escape('schema mismatch: expected list/tuple'),
     ):
-      fns.new({'a': [1, 2, 3], 'b': [4, 5]}, schema=list_schema)
+      _ = fns.new({'a': [1, 2, 3], 'b': [4, 5]}, schema=list_schema)
 
   def test_schema_arg_schema_with_fallback(self):
     schema = kde.schema.new_schema(a=schema_constants.INT32).eval()
@@ -514,12 +511,14 @@ def test_universal_converter_list_of_complex(self):
     )
 
   def test_universal_converter_list_of_different_primitive_lists(self):
-    with self.assertRaisesRegex(ValueError, 'cannot find a common schema'):
-      fns.new([[1, 2], [3.14]])
-    fns.new(
+    x = fns.new([[1, 2], [3.14]])
+    y = fns.new(
         [[1, 2], [3.14]],
-        schema=kde.list_schema(kde.list_schema(schema_constants.FLOAT32)).eval()
+        schema=kde.list_schema(
+            kde.list_schema(schema_constants.FLOAT32)
+        ).eval(),
     )
+    testing.assert_equivalent(x, y)
 
   def test_universal_converter_container_contains_multi_dim_data_slice(self):
     with self.assertRaisesRegex(
@@ -596,12 +595,12 @@ def test_universal_converter_entity(self):
     with self.subTest('item'):
       entity = fns.new(a=42, b='abc')
       new_entity = fns.new(entity)
-      testing.assert_not_equal(entity.get_bag(), new_entity.get_bag())
+      testing.assert_equal(entity.get_bag(), new_entity.get_bag())
       testing.assert_equivalent(new_entity, entity)
     with self.subTest('slice'):
       entity = fns.new(a=ds([1, 2]), b='abc')
       new_entity = fns.new(entity)
-      testing.assert_not_equal(entity.get_bag(), new_entity.get_bag())
+      testing.assert_equal(entity.get_bag(), new_entity.get_bag())
       testing.assert_equivalent(new_entity, entity)
 
   def test_universal_converter_adopt_bag_data(self):
@@ -624,19 +623,29 @@ def test_universal_converter_with_cross_ref_schema_conflict(self):
     d2 = {'b': 37}
     d1['d'] = d2
     d = {'d1': d1, 'd2': d2}
-    with self.assertRaisesRegex(ValueError, 'cannot find a common schema'):
+    with self.assertRaisesRegex(
+        ValueError,
+        'could not parse list of primitives / data items: object with'
+        ' unsupported type: dict',
+    ):
       fns.new(d)
 
   def test_universal_converter_with_itemid(self):
     itemid = kde.uuid_for_list('new').eval()
     res = fns.new([{'a': 42, 'b': 17}, {'c': 12}], itemid=itemid)
-    child_itemid = kde.uuid_for_dict(
-        '__from_py_child__', parent=itemid,
-        list_item_index=ds([0, 1], schema_constants.INT64)
+    child_itemid = kde.uuid(
+        '__from_py_child__',
+        parent=itemid,
+        list_item_index=ds([0, 1], schema_constants.INT64),
+    ).eval()
+    child_dict_item_id = kde.uuid_for_dict(
+        '',
+        base_itemid=child_itemid,
     ).eval()
+
     testing.assert_equal(res.no_bag().get_itemid(), itemid)
     testing.assert_dicts_keys_equal(res[:], ds([['a', 'b'], ['c']]))
-    testing.assert_equal(res[:].no_bag().get_itemid(), child_itemid)
+    testing.assert_equal(res[:].no_bag().get_itemid(), child_dict_item_id)
 
     with self.assertRaisesRegex(
         ValueError, 'itemid argument to list creation, requires List ItemIds'
@@ -650,12 +659,18 @@ def test_universal_converter_with_itemid(self):
 
   def test_universal_converter_recursive_object_error(self):
     d = {'a': 42}
-    d['self'] = d
-    with self.assertRaisesRegex(ValueError, 'recursive .* cannot be converted'):
+    d['a'] = d
+    with self.assertRaisesRegex(
+        ValueError,
+        'recursive Python object cannot be converted',
+    ):
       fns.new(d)
     # Deeper recursion:
     d2 = {'a': {'b': d}}
-    with self.assertRaisesRegex(ValueError, 'recursive .* cannot be converted'):
+    with self.assertRaisesRegex(
+        ValueError,
+        'recursive Python object cannot be converted',
+    ):
       fns.new(d2)
     # Longer cycle:
     d = {'a': 42}
@@ -665,12 +680,17 @@ def test_universal_converter_recursive_object_error(self):
     level_1_d = d
     d = {'top': d}
     bottom_d['cycle'] = level_1_d
-    with self.assertRaisesRegex(ValueError, 'recursive .* cannot be converted'):
+    with self.assertRaisesRegex(
+        ValueError,
+        'recursive Python object cannot be converted',
+    ):
       fns.new(d2)
     # Cycle in list:
-    l = [[1, 2], 3]
+    l = [[1, 2], (3)]
     l[0].append(l)
-    with self.assertRaisesRegex(ValueError, 'recursive .* cannot be converted'):
+    with self.assertRaisesRegex(
+        ValueError, 'schema mismatch: expected list/tuple'
+    ):
       fns.new(l)
 
   def test_universal_converter_with_attrs(self):

py/koladata/types/data_bag.cc

@@ -195,10 +195,12 @@ struct EntityCreatorHelper {
       PyObject* py_obj, const std::optional<DataSlice>& schema_arg,
       const std::optional<DataSlice>& itemid, const DataBagPtr& db,
       AdoptionQueue& adoption_queue) {
-    ASSIGN_OR_RETURN(
-        DataSlice res,
-        EntitiesFromPyObject(py_obj, schema_arg, itemid, db, adoption_queue));
-    return res.WithBag(db);
+    if (arolla::python::IsPyQValueInstance(py_obj) && !itemid.has_value()) {
+      return DataSliceFromPyValue(py_obj, adoption_queue, schema_arg);
+    }
+
+    return FromPy(py_obj, schema_arg, /*from_dim=*/0,
+                  /*dict_as_obj=*/false, itemid);
   }
 };