support named schemas in traverser and deep_clone

PiperOrigin-RevId: 718827570
Change-Id: I6b085612bbd79212f700cd9dd8c69d172f7b4c19

Author: Koladata Team

koladata/internal/op_utils/BUILD

@@ -207,6 +207,7 @@ cc_test(
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:span",
         "@com_google_arolla//arolla/dense_array",

koladata/internal/op_utils/traverser.h

@@ -285,6 +285,9 @@ class Traverser {
           if (!status.ok()) {
             return;
           }
+          if (attr_name == schema::kSchemaNameAttr) {
+            return;
+          }
           status = PrevisitAttribute(item, attr_name);
         });
     return status;
@@ -339,8 +342,11 @@ class Traverser {
             status = attr_schema_or.status();
             return;
           }
+          auto schema_dtype = (attr_name == schema::kSchemaNameAttr)
+                                  ? schema::kString
+                                  : schema::kSchema;
           status = Previsit(
-              {.item = *attr_schema_or, .schema = DataItem(schema::kSchema)});
+              {.item = *attr_schema_or, .schema = DataItem(schema_dtype)});
         });
     return status;
   }
@@ -459,12 +465,15 @@ class Traverser {
         return;
       }
       auto attr_schema = *attr_schema_or;
-      if (!attr_schema.is_schema()) {
+      DataItem attr_schema_schema = DataItem(schema::kSchema);
+      if (attr_name == schema::kSchemaNameAttr) {
+        attr_schema_schema = DataItem(schema::kString);
+      } else if (!attr_schema.is_schema()) {
         status = absl::InvalidArgumentError(absl::StrFormat(
             "schema %v has unexpected attribute %s", schema, attr_name));
         return;
       }
-      auto attr_value_or = GetValue(attr_schema, DataItem(schema::kSchema));
+      auto attr_value_or = GetValue(attr_schema, attr_schema_schema);
       if (!attr_value_or.ok()) {
         status = attr_value_or.status();
         return;
@@ -538,13 +547,19 @@ class Traverser {
       return VisitDict(item, is_object);
     }
     arolla::DenseArrayBuilder<DataItem> attr_values(attr_names.size());
+    arolla::DenseArrayBuilder<arolla::Text> actual_attr_names(
+        attr_names.size());
     absl::Status status = absl::OkStatus();
+    size_t attr_count = 0;
     attr_names.ForEach([&](int64_t id, bool presence,
                            std::string_view attr_name) {
       DCHECK(presence);
       if (!status.ok()) {
         return;
       }
+      if (attr_name == schema::kSchemaNameAttr) {
+        return;
+      }
       auto attr_item_or = databag_.GetAttr(item.item, attr_name, fallbacks_);
       if (!attr_item_or.ok()) {
         status = attr_item_or.status();
@@ -561,14 +576,17 @@ class Traverser {
         status = value_or.status();
         return;
       }
-      attr_values.Set(id, *value_or);
+      attr_values.Set(attr_count, *value_or);
+      actual_attr_names.Set(attr_count, attr_name);
+      ++attr_count;
     });
     if (!status.ok()) {
       return status;
     }
-    return visitor_->VisitorT::VisitObject(item.item, item.schema, is_object,
-                                            attr_names,
-                                            std::move(attr_values).Build());
+    return visitor_->VisitorT::VisitObject(
+        item.item, item.schema, is_object,
+        std::move(actual_attr_names).Build(attr_count),
+        std::move(attr_values).Build(attr_count));
   }
 
   absl::Status VisitObject(const ItemWithSchema& item) {

koladata/internal/op_utils/traverser_test.cc

@@ -27,6 +27,7 @@
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/check.h"
 #include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/types/span.h"
 #include "koladata/internal/data_bag.h"
@@ -140,6 +141,10 @@ class NoOpVisitor : AbstractVisitor {
       : previsited_(), value_item_(DataItem("get value result")) {}
 
   absl::Status Previsit(const DataItem& item, const DataItem& schema) override {
+    if (!schema.is_schema()) {
+      return absl::InvalidArgumentError(
+          absl::StrFormat("%v is not a schema", schema));
+    }
     previsited_.push_back({item, schema});
     return absl::OkStatus();
   }
@@ -540,6 +545,35 @@ TEST_P(NoOpTraverserTest, SchemaSlice) {
                           {GetFallbackDb(db).get()}));
 }
 
+TEST_P(NoOpTraverserTest, SliceWithNamedSchema) {
+  auto db = DataBagImpl::CreateEmptyDatabag();
+  auto obj_ids = DataSliceImpl::AllocateEmptyObjects(3);
+  auto a0 = obj_ids[0];
+  auto a1 = obj_ids[1];
+  auto a2 = obj_ids[2];
+  auto int_dtype = DataItem(schema::kInt32);
+  auto schema = AllocateSchema();
+  ASSERT_OK_AND_ASSIGN(auto named_schema,
+                       db->CreateUuSchemaFromFields(
+                           absl::StrCat("__named_schema__", "foo"), {}, {}));
+  TriplesT schema_triples = {
+      {schema, {{"x", int_dtype}, {"y", int_dtype}}},
+      {named_schema,
+       {{schema::kSchemaNameAttr, DataItem(arolla::Text("foo"))},
+        {"x", DataItem(schema)},
+        {"y", int_dtype}}}};
+  TriplesT data_triples = {{a0, {{"x", a1}, {"y", DataItem(1)}}},
+                           {a1, {{"x", DataItem(2)}}}};
+  SetDataTriples(*db, data_triples);
+  SetSchemaTriples(*db, schema_triples);
+  SetDataTriples(*db, GenNoiseDataTriples());
+  SetSchemaTriples(*db, GenNoiseSchemaTriples());
+
+  auto ds = DataSliceImpl::Create(CreateDenseArray<DataItem>({a0}));
+  EXPECT_OK(TraverseSlice(ds, DataItem(named_schema), *GetMainDb(db),
+                          {GetFallbackDb(db).get()}));
+}
+
 TEST_P(ObjectsTraverserTest, ObjectsSlice) {
   auto db = DataBagImpl::CreateEmptyDatabag();
   auto obj_ids = DataSliceImpl::AllocateEmptyObjects(10);

py/koladata/operators/tests/core_deep_clone_test.py

@@ -154,6 +154,20 @@ def test_entity(self, pass_schema):
     with self.assertRaisesRegex(AssertionError, 'not equal by fingerprint'):
       testing.assert_equal(result.ref(), o.ref())
 
+  @parameterized.product(
+      pass_schema=[True, False],
+  )
+  def test_named_schema(self, pass_schema):
+    db = data_bag.DataBag.empty()
+    schema = db.named_schema('foo', x=schema_constants.INT32)
+    o = db.new(x=ds([1, 2, 3]), schema=schema)
+    if pass_schema:
+      result = expr_eval.eval(kde.deep_clone(o, schema))
+    else:
+      result = expr_eval.eval(kde.deep_clone(o))
+    testing.assert_equal(result.get_schema().no_bag(), schema.no_bag())
+    testing.assert_equal(result.x.no_bag(), ds([1, 2, 3]))
+
   @parameterized.product(
       pass_schema=[True, False],
   )