Simplify operator implementations in expr_operators.cc.

apronchenkov · copybara-github · commit e5bbf3aa4d4f · 2026-01-23T04:25:32.000-08:00
Notably, instead of using a custom binding policy for the annotation.with_name operator, this changelist configures an existing binding policy to serve the same role.

PiperOrigin-RevId: 860047013
Change-Id: Ib8b12da0241635251efdf662ee8ee2c7e1d622fd
diff --git a/koladata/expr/expr_operators.cc b/koladata/expr/expr_operators.cc
@@ -26,6 +26,7 @@
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
+#include "arolla/expr/annotation_expr_operators.h"
 #include "arolla/expr/basic_expr_operator.h"
 #include "arolla/expr/expr.h"
 #include "arolla/expr/expr_attributes.h"
@@ -68,6 +69,48 @@ absl::Status ValidateTextLiteral(const arolla::expr::ExprAttributes& attr,
   return absl::OkStatus();
 }
 
+// Basic implementation of non-lowerable operators that convert a DataSlice to
+// an Arolla type T. Supports compile-time evaluation if the provided input is
+// a literal.
+class ToArollaValueOperator final
+    : public arolla::expr::BackendExprOperatorTag,
+      public arolla::expr::ExprOperatorWithFixedSignature {
+ public:
+  ToArollaValueOperator(absl::string_view name, absl::string_view doc,
+                        arolla::QTypePtr output_qtype)
+      : ExprOperatorWithFixedSignature(
+            name,
+            arolla::expr::ExprOperatorSignature({{"x"}},
+                                                "koladata_classic_aux_policy"),
+            doc,
+            arolla::FingerprintHasher("::koladata::expr::ToArollaValueOperator")
+                .Combine(name, doc, output_qtype)
+                .Finish()),
+        output_qtype_(output_qtype) {}
+
+  absl::StatusOr<arolla::expr::ExprAttributes> InferAttributes(
+      absl::Span<const arolla::expr::ExprAttributes> inputs) const final {
+    RETURN_IF_ERROR(ValidateOpInputsCount(inputs));
+    if (auto* x_qtype = inputs[0].qtype()) {
+      if (x_qtype != arolla::GetQType<DataSlice>()) {
+        return absl::InvalidArgumentError(
+            absl::StrFormat("expected DATA_SLICE, got x: %s", x_qtype->name()));
+      }
+    } else {
+      return arolla::expr::ExprAttributes{};  // Not ready.
+    }
+    if (const auto& x = inputs[0].qvalue()) {
+      ASSIGN_OR_RETURN(auto result, arolla::InvokeOperator(display_name(), {*x},
+                                                           output_qtype_));
+      return arolla::expr::ExprAttributes{std::move(result)};
+    }
+    return arolla::expr::ExprAttributes{output_qtype_};  // Can not evaluated.
+  }
+
+ private:
+  arolla::QTypePtr output_qtype_;
+};
+
 }  // namespace
 
 InputOperator::InputOperator()
@@ -98,11 +141,14 @@ absl::StatusOr<arolla::expr::ExprAttributes> InputOperator::InferAttributes(
   return arolla::expr::ExprAttributes{};
 }
 
-absl::StatusOr<arolla::expr::ExprNodePtr> MakeLiteral(
-    arolla::TypedValue value) {
-  ASSIGN_OR_RETURN(auto op, expr::LiteralOperator::Make(value));
-  return arolla::expr::ExprNode::UnsafeMakeOperatorNode(
-      std::move(op), {}, arolla::expr::ExprAttributes(std::move(value)));
+bool IsInput(const arolla::expr::ExprNodePtr& node) {
+  if (!node->is_op()) {
+    return false;
+  }
+  return nullptr != arolla::fast_dynamic_downcast_final<const InputOperator*>(
+                        arolla::expr::DecayRegisteredOperator(node->op())
+                            .value_or(nullptr)
+                            .get());
 }
 
 absl::StatusOr<std::shared_ptr<LiteralOperator>> LiteralOperator::Make(
@@ -132,68 +178,18 @@ absl::string_view LiteralOperator::py_qvalue_specialization_key() const {
   return "::koladata::expr::LiteralOperator";
 }
 
-const arolla::TypedValue& LiteralOperator::value() const { return value_; }
-
-absl::StatusOr<arolla::expr::ExprAttributes>
-ToArollaValueOperator::InferAttributes(
-    absl::Span<const arolla::expr::ExprAttributes> inputs) const {
-  RETURN_IF_ERROR(ValidateOpInputsCount(inputs));
-  if (!inputs[0].qtype()) {
-    return arolla::expr::ExprAttributes{};  // Not ready yet.
-  }
-  if (inputs[0].qtype() != arolla::GetQType<DataSlice>()) {
-    return absl::InvalidArgumentError(absl::StrFormat(
-        "expected DATA_SLICE, got x: %s", inputs[0].qtype()->name()));
-  }
-  // Eval if possible.
-  if (inputs[0].qvalue()) {
-    ASSIGN_OR_RETURN(
-        auto casted_value,
-        arolla::InvokeOperator(backend_operator_name_, {*inputs[0].qvalue()},
-                               output_qtype_));
-    return arolla::expr::ExprAttributes{std::move(casted_value)};
-  }
-  // Otherwise, return the output qtype.
-  return arolla::expr::ExprAttributes{output_qtype_};
+absl::StatusOr<arolla::expr::ExprNodePtr> MakeLiteral(
+    arolla::TypedValue value) {
+  ASSIGN_OR_RETURN(auto op, expr::LiteralOperator::Make(value));
+  return arolla::expr::ExprNode::UnsafeMakeOperatorNode(
+      std::move(op), {}, arolla::expr::ExprAttributes(std::move(value)));
 }
 
-ToArollaInt64Operator::ToArollaInt64Operator()
-    : ToArollaValueOperator(
-          "koda_internal.to_arolla_int64",
-          arolla::expr::ExprOperatorSignature({{"x"}},
-                                              "koladata_classic_aux_policy"),
-          "Returns `x` converted into an arolla int64 value.\n"
-          "\n"
-          "Note that `x` must adhere to the following requirements:\n"
-          "* `rank = 0`.\n"
-          "* Have one of the following schemas: NONE, INT32, INT64, OBJECT.\n"
-          "* Have a present value with type INT32 or INT64.\n"
-          "\n"
-          "In all other cases, an exception is raised.\n\n"
-          "Args:\n"
-          "  x: A DataItem to be converted into an arolla int64 value.",
-          arolla::FingerprintHasher("::koladata::expr::ToArollaInt64Operator")
-              .Finish(),
-          "koda_internal.to_arolla_int64", arolla::GetQType<int64_t>()) {}
-
-ToArollaTextOperator::ToArollaTextOperator()
-    : ToArollaValueOperator(
-          "koda_internal.to_arolla_text",
-          arolla::expr::ExprOperatorSignature({{"x"}},
-                                              "koladata_classic_aux_policy"),
-          "Returns `x` converted into an arolla text value.\n"
-          "\n"
-          "Note that `x` must adhere to the following requirements:\n"
-          "* `rank = 0`.\n"
-          "* Have one of the following schemas: NONE, STRING, OBJECT, ANY.\n"
-          "* Have a present value with type TEXT.\n"
-          "\n"
-          "In all other cases, an exception is raised.\n\n"
-          "Args:\n"
-          "  x: A DataItem to be converted into an arolla text value.",
-          arolla::FingerprintHasher("::koladata::expr::ToArollaTextOperator")
-              .Finish(),
-          "koda_internal.to_arolla_text", arolla::GetQType<arolla::Text>()) {}
+bool IsLiteral(const arolla::expr::ExprNodePtr& node) {
+  return node->is_literal() ||
+         (arolla::fast_dynamic_downcast_final<const LiteralOperator*>(
+              node->op().get()) != nullptr);
+}
 
 NonDeterministicOperator::NonDeterministicOperator()
     : arolla::expr::ExprOperatorWithFixedSignature(
@@ -211,20 +207,42 @@ NonDeterministicOperator::InferAttributes(
       arolla::GetQType<internal::NonDeterministicToken>());
 }
 
-bool IsInput(const arolla::expr::ExprNodePtr& node) {
-  if (!node->is_op()) {
-    return false;
-  }
-  return nullptr != arolla::fast_dynamic_downcast_final<const InputOperator*>(
-                        arolla::expr::DecayRegisteredOperator(node->op())
-                            .value_or(nullptr)
-                            .get());
+arolla::expr::ExprOperatorPtr MakeToArollaInt64Operator() {
+  return std::make_shared<ToArollaValueOperator>(
+      "koda_internal.to_arolla_int64",
+      "Returns `x` converted into an arolla int64 value.\n"
+      "\n"
+      "Note that `x` must adhere to the following requirements:\n"
+      "* `rank = 0`.\n"
+      "* Have one of the following schemas: NONE, INT32, INT64, OBJECT.\n"
+      "* Have a present value with type INT32 or INT64.\n"
+      "\n"
+      "In all other cases, an exception is raised.\n\n"
+      "Args:\n"
+      "  x: A DataItem to be converted into an arolla int64 value.",
+      arolla::GetQType<int64_t>());
 }
 
-bool IsLiteral(const arolla::expr::ExprNodePtr& node) {
-  return node->is_literal() ||
-         (arolla::fast_dynamic_downcast_final<const LiteralOperator*>(
-              node->op().get()) != nullptr);
+arolla::expr::ExprOperatorPtr MakeToArollaTextOperator() {
+  return std::make_shared<ToArollaValueOperator>(
+      "koda_internal.to_arolla_text",
+      "Returns `x` converted into an arolla text value.\n"
+      "\n"
+      "Note that `x` must adhere to the following requirements:\n"
+      "* `rank = 0`.\n"
+      "* Have one of the following schemas: NONE, STRING, OBJECT, ANY.\n"
+      "* Have a present value with type TEXT.\n"
+      "\n"
+      "In all other cases, an exception is raised.\n\n"
+      "Args:\n"
+      "  x: A DataItem to be converted into an arolla text value.",
+      arolla::GetQType<arolla::Text>());
+}
+
+arolla::expr::ExprOperatorPtr MakeNameAnnotationOperator() {
+  // Use aux-policy with arolla boxing rules for the second parameter.
+  return std::make_shared<arolla::expr::NameAnnotation>(
+      "koladata_unified_aux_policy:_p:_arolla_as_qvalue_or_expr;01");
 }
 
 absl::StatusOr<InputContainer> InputContainer::Create(
diff --git a/koladata/expr/expr_operators.h b/koladata/expr/expr_operators.h
@@ -28,10 +28,7 @@
 #include "arolla/expr/expr_attributes.h"
 #include "arolla/expr/expr_node.h"
 #include "arolla/expr/expr_operator.h"
-#include "arolla/expr/expr_operator_signature.h"
-#include "arolla/qtype/qtype.h"
 #include "arolla/qtype/typed_value.h"
-#include "arolla/util/fingerprint.h"
 
 namespace koladata::expr {
 
@@ -48,6 +45,9 @@ class InputOperator final
       absl::Span<const arolla::expr::ExprAttributes> inputs) const final;
 };
 
+// Returns true if `node` is an operator node with InputOperator.
+bool IsInput(const arolla::expr::ExprNodePtr& node);
+
 // Non-lowerable stateful operator `koda_internal.literal()` that wraps a
 // TypedValue. This operator allows us to attach a view to non-DataSlice
 // literals.
@@ -66,57 +66,19 @@ class LiteralOperator final
 
   absl::string_view py_qvalue_specialization_key() const final;
 
-  const arolla::TypedValue& value() const;
+  const arolla::TypedValue& value() const { return value_; }
 
  private:
   arolla::TypedValue value_;
 };
 
-// Base class for non-lowerable operators that converts a DataSlice to an Arolla
-// type T. Supports evaluation at operator binding time if the provided input is
-// a literal. Dispatches the actual conversion to a corresponding
-// backend-operator.
-class ToArollaValueOperator
-    : public arolla::expr::BackendExprOperatorTag,
-      public arolla::expr::ExprOperatorWithFixedSignature {
- public:
-  ToArollaValueOperator(absl::string_view name,
-                        arolla::expr::ExprOperatorSignature signature,
-                        absl::string_view doc, arolla::Fingerprint fingerprint,
-                        std::string backend_operator_name,
-                        arolla::QTypePtr output_qtype)
-      : ExprOperatorWithFixedSignature(name, std::move(signature), doc,
-                                       fingerprint),
-        backend_operator_name_(std::move(backend_operator_name)),
-        output_qtype_(output_qtype) {}
-
-  absl::StatusOr<arolla::expr::ExprAttributes> InferAttributes(
-      absl::Span<const arolla::expr::ExprAttributes> inputs) const final;
-
- private:
-  std::string backend_operator_name_;
-  arolla::QTypePtr output_qtype_;
-};
-
-// Non-lowerable operator `koda_internal.to_arolla_int64(x)` that converts
-// DataSlice to int64_t. Supports evaluation at operator binding time if the
-// provided input is a literal.
-class ToArollaInt64Operator final : public ToArollaValueOperator {
- public:
-  ToArollaInt64Operator();
-};
-
-// Non-lowerable operator `koda_internal.to_arolla_text(x)` that converts
-// DataSlice to Text. Supports evaluation at operator binding time if the
-// provided input is a literal.
-class ToArollaTextOperator final : public ToArollaValueOperator {
- public:
-  ToArollaTextOperator();
-};
-
 // Return a literal Expr node.
 absl::StatusOr<arolla::expr::ExprNodePtr> MakeLiteral(arolla::TypedValue value);
 
+// Returns true if `node` is either an arolla literal or an operator node with
+// LiteralOperator.
+bool IsLiteral(const arolla::expr::ExprNodePtr& node);
+
 // Non-lowerable operator
 // ```
 // koda_internal.non_deterministic(
@@ -136,12 +98,18 @@ class NonDeterministicOperator final
       absl::Span<const arolla::expr::ExprAttributes> inputs) const final;
 };
 
-// Returns true if `node` is an operator node with InputOperator.
-bool IsInput(const arolla::expr::ExprNodePtr& node);
+// Non-lowerable operator `koda_internal.to_arolla_int64(x)` that converts
+// DataSlice to int64_t. Supports evaluation at operator binding time if the
+// provided input is a literal.
+arolla::expr::ExprOperatorPtr MakeToArollaInt64Operator();
 
-// Returns true if `node` is either an arolla literal or an operator node with
-// LiteralOperator.
-bool IsLiteral(const arolla::expr::ExprNodePtr& node);
+// Non-lowerable operator `koda_internal.to_arolla_text(x)` that converts
+// DataSlice to Text. Supports evaluation at operator binding time if the
+// provided input is a literal.
+arolla::expr::ExprOperatorPtr MakeToArollaTextOperator();
+
+// Non-lowerable operator `kd.annotation.with_name(x, name)`.
+arolla::expr::ExprOperatorPtr MakeNameAnnotationOperator();
 
 // Helper container to create Koda specific inputs
 class InputContainer {
diff --git a/koladata/expr/init_expr_operators.cc b/koladata/expr/init_expr_operators.cc
@@ -12,8 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 //
-#include <memory>
-
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
@@ -32,6 +30,7 @@ namespace koladata::expr {
 
 AROLLA_INITIALIZER(
         .name = "arolla_operators/koda_internal",
+        .deps = {"arolla_operators/standard"},
         .reverse_deps = {arolla::initializer_dep::kOperators},
         .init_fn = []() -> absl::Status {
           RETURN_IF_ERROR(arolla::expr::RegisterOperator<InputOperator>(
@@ -44,26 +43,25 @@ AROLLA_INITIALIZER(
                                    arolla::expr::Literal(internal::Ellipsis{})))
                   .status());
           RETURN_IF_ERROR(
-              arolla::expr::RegisterOperator(
-                  "koda_internal.with_name",
-                  std::make_shared<arolla::expr::NameAnnotation>(
-                      /*aux_policy=*/"_koladata_annotation_with_name"))
+              arolla::expr::RegisterOperator("koda_internal.with_name",
+                                             MakeNameAnnotationOperator())
+                  .status());
+          RETURN_IF_ERROR(
+              arolla::expr::RegisterOperator("koda_internal.to_arolla_int64",
+                                             MakeToArollaInt64Operator())
+                  .status());
+          RETURN_IF_ERROR(
+              arolla::expr::RegisterOperator("koda_internal.to_arolla_text",
+                                             MakeToArollaTextOperator())
                   .status());
-          RETURN_IF_ERROR(arolla::expr::RegisterOperator(
-                              "koda_internal.to_arolla_int64",
-                              std::make_shared<ToArollaInt64Operator>())
-                              .status());
-          RETURN_IF_ERROR(arolla::expr::RegisterOperator(
-                              "koda_internal.to_arolla_text",
-                              std::make_shared<ToArollaTextOperator>())
-                              .status());
           RETURN_IF_ERROR(
               arolla::expr::RegisterOperator<NonDeterministicOperator>(
-                  "koda_internal.non_deterministic").status());
-          RETURN_IF_ERROR(arolla::expr::RegisterOperator(
-                              "kd.annotation.source_location",
-                              arolla::expr::SourceLocationAnnotation::Make())
-                              .status());
+                  "koda_internal.non_deterministic")
+                  .status());
+          RETURN_IF_ERROR(
+              arolla::expr::RegisterOperatorAlias(
+                  "kd.annotation.source_location", "annotation.source_location")
+                  .status());
           return absl::OkStatus();
         })
 
diff --git a/py/koladata/operators/BUILD b/py/koladata/operators/BUILD
@@ -747,7 +747,6 @@ py_library(
     deps = [
         ":optools",
         "//py:python_path",
-        "//py/koladata/types:py_boxing",
         "@com_google_arolla//py/arolla",
     ],
 )
diff --git a/py/koladata/operators/annotation.py b/py/koladata/operators/annotation.py
diff --git a/py/koladata/operators/tests/annotation_with_name_test.py b/py/koladata/operators/tests/annotation_with_name_test.py

Original file line number	Diff line number	Diff line change
`@@ -747,7 +747,6 @@ py_library(`
`747`	`747`	`deps = [`
`748`	`748`	`":optools",`
`749`	`749`	`"//py:python_path",`
`750`		`- "//py/koladata/types:py_boxing",`
`751`	`750`	`"@com_google_arolla//py/arolla",`
`752`	`751`	`],`
`753`	`752`	`)`