Use aux-policies for View assignment in Koladata

Motivation: This change enables the use of ad-hoc operators in Koladata without registration. Previously, unregistered operators defined using `@kd.optools.as_*_operator` defaulted to `arolla.expr.DefaultExprView`, which limited their usability. This change sets KodaView as the default (and allows for further customization).
PiperOrigin-RevId: 862469628
Change-Id: Ib3ac34a4b5642e38edc4e26995e0c90ac2109027

Author: apronchenkov

docs/api_reference.md

@@ -4823,15 +4823,14 @@ Operator definition and registration tooling.
 
 <pre class="no-copy"><code class="lang-text no-auto-prettify">Adds an alias for an operator.</code></pre>
 
-### `kd.optools.add_to_registry(name: str | None = None, *, aliases: Collection[str] = (), unsafe_override: bool = False, view: type[ExprView] = <class 'koladata.expr.view.KodaView'>, repr_fn: Callable[[Expr, NodeTokenView], ReprToken] = <default_op_repr>, via_cc_operator_package: bool = False)` {#kd.optools.add_to_registry}
+### `kd.optools.add_to_registry(name: str | None = None, *, aliases: Collection[str] = (), unsafe_override: bool = False, repr_fn: Callable[[Expr, NodeTokenView], ReprToken] = <default_op_repr>, via_cc_operator_package: bool = False)` {#kd.optools.add_to_registry}
 
 <pre class="no-copy"><code class="lang-text no-auto-prettify">Wrapper around Arolla&#39;s add_to_registry with Koda functionality.
 
 Args:
   name: Optional name of the operator. Otherwise, inferred from the op.
   aliases: Optional aliases for the operator.
   unsafe_override: Whether to override an existing operator.
-  view: View for the operator and its aliases.
   repr_fn: Repr function for the operator and its aliases.
   via_cc_operator_package: If True, the operator will be only registered
     during koladata_cc_operator_package construction, and just looked up in
@@ -4867,7 +4866,7 @@ Returns:
   A decorator that registers an overload for the operator with the
   corresponding name.</code></pre>
 
-### `kd.optools.add_to_registry_as_overloadable(name: str, *, unsafe_override: bool = False, view: type[ExprView] = <class 'koladata.expr.view.KodaView'>, repr_fn: Callable[[Expr, NodeTokenView], ReprToken] = <default_op_repr>, aux_policy: str = 'koladata_classic_aux_policy', via_cc_operator_package: bool = False)` {#kd.optools.add_to_registry_as_overloadable}
+### `kd.optools.add_to_registry_as_overloadable(name: str, *, unsafe_override: bool = False, repr_fn: Callable[[Expr, NodeTokenView], ReprToken] = <default_op_repr>, aux_policy: str = 'koladata_classic_aux_policy', via_cc_operator_package: bool = False)` {#kd.optools.add_to_registry_as_overloadable}
 
 <pre class="no-copy"><code class="lang-text no-auto-prettify">Koda wrapper around Arolla&#39;s add_to_registry_as_overloadable.
 
@@ -4877,7 +4876,6 @@ repr function.
 Args:
   name: Name of the operator.
   unsafe_override: Whether to override an existing operator.
-  view: View for the operator and its aliases.
   repr_fn: Repr function for the operator and its aliases.
   aux_policy: Aux policy for the operator.
   via_cc_operator_package: If True, the operator will be only registered
@@ -4890,7 +4888,7 @@ Args:
 Returns:
   An overloadable registered operator.</code></pre>
 
-### `kd.optools.as_backend_operator(name: str, *, qtype_inference_expr: Expr | QType = DATA_SLICE, qtype_constraints: Iterable[tuple[Expr, str]] = (), deterministic: bool = True, custom_boxing_fn_name_per_parameter: dict[str, str] | None = None) -> Callable[[function], BackendOperator]` {#kd.optools.as_backend_operator}
+### `kd.optools.as_backend_operator(name: str, *, qtype_inference_expr: Expr | QType = DATA_SLICE, qtype_constraints: Iterable[tuple[Expr, str]] = (), deterministic: bool = True, custom_boxing_fn_name_per_parameter: dict[str, str] | None = None, view: str | type[ExprView] = '') -> Callable[[function], BackendOperator]` {#kd.optools.as_backend_operator}
 
 <pre class="no-copy"><code class="lang-text no-auto-prettify">Decorator for Koladata backend operators with a unified binding policy.
 
@@ -4911,12 +4909,14 @@ Args:
   custom_boxing_fn_name_per_parameter: A dictionary specifying a custom boxing
     function per parameter (constants with the boxing functions look like:
     `koladata.types.py_boxing.WITH_*`, e.g. `WITH_PY_FUNCTION_TO_PY_OBJECT`).
+  view: The view for the for the operator, with the default being KodaView
+    (supported values: &#39;&#39;|KodaView, &#39;base&#39;|BaseKodaView, &#39;arolla&#39;|ArollaView).
 
 Returns:
   A decorator that constructs a backend operator based on the provided Python
   function signature.</code></pre>
 
-### `kd.optools.as_lambda_operator(name: str, *, qtype_constraints: Iterable[tuple[Expr, str]] = (), deterministic: bool | None = None, custom_boxing_fn_name_per_parameter: dict[str, str] | None = None, suppress_unused_parameter_warning: bool = False) -> Callable[[function], LambdaOperator | RestrictedLambdaOperator]` {#kd.optools.as_lambda_operator}
+### `kd.optools.as_lambda_operator(name: str, *, qtype_constraints: Iterable[tuple[Expr, str]] = (), deterministic: bool | None = None, custom_boxing_fn_name_per_parameter: dict[str, str] | None = None, suppress_unused_parameter_warning: bool = False, view: str | type[ExprView] = '') -> Callable[[function], LambdaOperator | RestrictedLambdaOperator]` {#kd.optools.as_lambda_operator}
 
 <pre class="no-copy"><code class="lang-text no-auto-prettify">Decorator for Koladata lambda operators with a unified binding policy.
 
@@ -4937,11 +4937,13 @@ Args:
     `koladata.types.py_boxing.WITH_*`, e.g. `WITH_PY_FUNCTION_TO_PY_OBJECT`).
   suppress_unused_parameter_warning: If True, unused parameters will not cause
     a warning.
+  view: The view for the for the operator, with the default being KodaView
+    (supported values: &#39;&#39;|KodaView, &#39;base&#39;|BaseKodaView, &#39;arolla&#39;|ArollaView).
 
 Returns:
   A decorator that constructs a lambda operator by tracing a Python function.</code></pre>
 
-### `kd.optools.as_py_function_operator(name: str, *, qtype_inference_expr: Expr | QType = DATA_SLICE, qtype_constraints: Iterable[tuple[Expr, str]] = (), codec: bytes | None = None, deterministic: bool = True, custom_boxing_fn_name_per_parameter: dict[str, str] | None = None) -> Callable[[function], Operator]` {#kd.optools.as_py_function_operator}
+### `kd.optools.as_py_function_operator(name: str, *, qtype_inference_expr: Expr | QType = DATA_SLICE, qtype_constraints: Iterable[tuple[Expr, str]] = (), codec: bytes | None = None, deterministic: bool = True, custom_boxing_fn_name_per_parameter: dict[str, str] | None = None, view: str | type[ExprView] = '') -> Callable[[function], Operator]` {#kd.optools.as_py_function_operator}
 
 <pre class="no-copy"><code class="lang-text no-auto-prettify">Returns a decorator for defining Koladata-specific py-function operators.
 
@@ -4965,7 +4967,9 @@ Args:
     (i.e., non-deterministic or has side effects).
   custom_boxing_fn_name_per_parameter: A dictionary specifying a custom boxing
     function per parameter (constants with the boxing functions look like:
-    `koladata.types.py_boxing.WITH_*`, e.g. `WITH_PY_FUNCTION_TO_PY_OBJECT`).</code></pre>
+    `koladata.types.py_boxing.WITH_*`, e.g. `WITH_PY_FUNCTION_TO_PY_OBJECT`).
+  view: The view for the for the operator, with the default being KodaView
+    (supported values: &#39;&#39;|KodaView, &#39;base&#39;|BaseKodaView, &#39;arolla&#39;|ArollaView).</code></pre>
 
 ### `kd.optools.as_qvalue(arg: Any) -> QValue` {#kd.optools.as_qvalue}
 

koladata/expr/expr_operators.cc

@@ -80,8 +80,9 @@ class ToArollaValueOperator final
                         arolla::QTypePtr output_qtype)
       : ExprOperatorWithFixedSignature(
             name,
-            arolla::expr::ExprOperatorSignature({{"x"}},
-                                                "koladata_classic_aux_policy"),
+            arolla::expr::ExprOperatorSignature(
+                // Use aux-policy with ArollaView.
+                {{"x"}}, "koladata_unified_aux_policy$arolla:_"),
             doc,
             arolla::FingerprintHasher("::koladata::expr::ToArollaValueOperator")
                 .Combine(name, doc, output_qtype)
@@ -116,8 +117,9 @@ class ToArollaValueOperator final
 InputOperator::InputOperator()
     : arolla::expr::ExprOperatorWithFixedSignature(
           kInternalInput,
-          arolla::expr::ExprOperatorSignature{{"container_name"},
-                                              {"input_key"}},
+          arolla::expr::ExprOperatorSignature(
+              {{"container_name"}, {"input_key"}},
+              "koladata_arolla_classic_aux_policy"),
           "Koda input with DATA_SLICE qtype.\n"
           "\n"
           "Note that this operator cannot be evaluated.\n"
@@ -161,7 +163,10 @@ absl::StatusOr<std::shared_ptr<LiteralOperator>> LiteralOperator::Make(
 LiteralOperator::LiteralOperator(PrivateConstructorTag,
                                  arolla::TypedValue value)
     : arolla::expr::ExprOperatorWithFixedSignature(
-          "koda_internal.literal", arolla::expr::ExprOperatorSignature{},
+          "koda_internal.literal",
+          arolla::expr::ExprOperatorSignature(
+              // Use aux-policy with BaseKodaView.
+              {}, "koladata_unified_aux_policy$base"),
           "Koda literal.",
           arolla::FingerprintHasher("::koladata::expr::LiteralOperator")
               .Combine(value.GetFingerprint())

koladata/expr/init_expr_operators.cc

@@ -43,7 +43,7 @@ AROLLA_INITIALIZER(
                                    arolla::expr::Literal(internal::Ellipsis{})))
                   .status());
           RETURN_IF_ERROR(
-              arolla::expr::RegisterOperator("koda_internal.with_name",
+              arolla::expr::RegisterOperator("kd.annotation.with_name",
                                              MakeNameAnnotationOperator())
                   .status());
           RETURN_IF_ERROR(

py/koladata/expr/view.py

@@ -753,6 +753,12 @@ def uuobj(self, *args, **kwargs):  # pylint: disable=unused-argument
     _raise_eager_only_method('uuobj', 'DataBag')
 
 
+class ArollaView(BaseKodaView, arolla.expr.DefaultExprView):
+  """Arolla view with a few extra methods specific to Koda."""
+
+  _arolla_view_tag = True
+
+
 def has_koda_view(node: arolla.Expr) -> bool:
   """Returns true iff the node has the full koda view."""
   return hasattr(node, '_koda_view_tag')
@@ -763,21 +769,11 @@ def has_base_koda_view(node: arolla.Expr) -> bool:
   return hasattr(node, '_base_koda_view_tag')
 
 
-class ArollaView(BaseKodaView, arolla.expr.DefaultExprView):
-  """Arolla view with a few extra methods specific to Koda."""
-
-  pass
+def has_arolla_view(node: arolla.Expr) -> bool:
+  """Returns true iff the node has the arolla view."""
+  return hasattr(node, '_arolla_view_tag')
 
 
 arolla.abc.set_expr_view_for_qtype(qtypes.DATA_SLICE, KodaView)
 arolla.abc.set_expr_view_for_qtype(qtypes.DATA_BAG, KodaView)
 arolla.abc.set_expr_view_for_qtype(qtypes.JAGGED_SHAPE, KodaView)
-arolla.abc.set_expr_view_for_registered_operator(
-    'koda_internal.input', KodaView
-)
-
-# NOTE: This attaches a BaseKodaView to all literals, including Arolla values.
-# This adds generic Koda methods (such as eval) to the expression.
-arolla.abc.set_expr_view_for_operator_family(
-    '::koladata::expr::LiteralOperator', ArollaView
-)

py/koladata/expr/view_test.py

@@ -45,13 +45,11 @@ class BaseKodaViewTest(parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
-    arolla.abc.set_expr_view_for_registered_operator(
-        'test.op', view.BaseKodaView
-    )
+    arolla.abc.set_expr_view_for_registered_operator(op, view.BaseKodaView)
 
   def tearDown(self):
     # Clear the view.
-    arolla.abc.set_expr_view_for_registered_operator('test.op', None)
+    arolla.abc.set_expr_view_for_registered_operator(op, None)
     super().tearDown()
 
   def test_eval(self):
@@ -83,11 +81,11 @@ class KodaViewTest(parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
-    arolla.abc.set_expr_view_for_registered_operator('test.op', view.KodaView)
+    arolla.abc.set_expr_view_for_registered_operator(op, view.KodaView)
 
   def tearDown(self):
     # Clear the view.
-    arolla.abc.set_expr_view_for_registered_operator('test.op', None)
+    arolla.abc.set_expr_view_for_registered_operator(op, None)
     super().tearDown()
 
   # To be overridden by KodaViewWithTracingTest subclass below.
@@ -646,7 +644,7 @@ def test_make_tuple(*args):
       return arolla.optools.fix_trace_args(args)
 
     arolla.abc.set_expr_view_for_registered_operator(
-        'test_make_tuple', view.KodaView
+        test_make_tuple, view.KodaView
     )
 
     expr = arolla.M.annotation.name(
@@ -896,6 +894,28 @@ def test_annotation_source_location_view(self):
     )
 
 
+class ArollaViewTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    arolla.abc.set_expr_view_for_registered_operator(op, view.ArollaView)
+
+  def tearDown(self):
+    arolla.abc.set_expr_view_for_registered_operator(op, None)
+    super().tearDown()
+
+  def test_has_arolla_view(self):
+    self.assertTrue(view.has_arolla_view(op()))
+    self.assertTrue(view.has_base_koda_view(op()))
+    self.assertFalse(view.has_koda_view(op()))
+
+  def test_less_operator(self):
+    arolla.testing.assert_expr_equal_by_fingerprint(
+        op(C.x) < op(C.y),
+        arolla.abc.unsafe_parse_sexpr((('core.less', (op, C.x), (op, C.y)))),
+    )
+
+
 class KodaViewWithTracingTest(KodaViewTest):
 
   def setUp(self):

py/koladata/operators/BUILD

@@ -205,7 +205,6 @@ py_library(
     name = "unified_binding_policy",
     srcs = ["unified_binding_policy.py"],
     deps = [
-        ":aux_policies",
         ":clib",
         "//py:python_path",
         "//py/koladata/types:py_boxing",
@@ -218,6 +217,7 @@ py_test(
     srcs = ["unified_binding_policy_test.py"],
     deps = [
         ":aux_policies",
+        ":clib",
         ":op_repr",
         ":unified_binding_policy",
         "//py:python_path",
@@ -260,7 +260,6 @@ py_library(
         "//py:python_path",
         "//py/koladata/expr:input_container",
         "//py/koladata/expr:tracing_mode",
-        "//py/koladata/expr:view",
         "//py/koladata/types:py_boxing",
         "//py/koladata/types:qtypes",
         "@com_google_arolla//py/arolla",
@@ -271,12 +270,12 @@ py_test(
     name = "optools_test",
     srcs = ["optools_test.py"],
     deps = [
+        ":aux_policies",
         ":comparison",
         ":jagged_shape",
         ":koda_internal",
         ":math",
         ":optools",
-        ":optools_test_utils",
         ":qtype_utils",
         ":tuple",
         "//py:python_path",
@@ -444,6 +443,7 @@ py_library(
         ":slices",
         "//koladata/operators",
         "//py:python_path",
+        "//py/koladata/expr:view",
         "//py/koladata/fstring",
         "//py/koladata/types:data_slice",
         "//py/koladata/types:py_boxing",
@@ -548,6 +548,7 @@ py_library(
         ":qtype_utils",
         "//koladata/operators",
         "//py:python_path",
+        "//py/koladata/expr:view",
         "//py/koladata/types:data_slice",
         "//py/koladata/types:py_boxing",
         "//py/koladata/types:qtypes",
@@ -647,6 +648,7 @@ py_library(
         ":schema",
         ":view_overloads",
         "//py:python_path",
+        "//py/koladata/expr:view",
         "//py/koladata/types:data_item",
         "//py/koladata/types:data_slice",
         "//py/koladata/types:py_boxing",
@@ -799,6 +801,7 @@ py_library(
         ":slices",
         ":view_overloads",
         "//py:python_path",
+        "//py/koladata/expr:view",
         "//py/koladata/types:data_slice",
         "//py/koladata/types:py_boxing",
         "//py/koladata/types:qtypes",
@@ -972,11 +975,24 @@ py_library(
     deps = [
         ":clib",
         "//py:python_path",
+        "//py/koladata/expr:view",
         "//py/koladata/types:py_boxing",
         "@com_google_arolla//py/arolla",
     ],
 )
 
+py_test(
+    name = "aux_policies_test",
+    srcs = ["aux_policies_test.py"],
+    deps = [
+        ":aux_policies",
+        "//py:python_path",
+        "//py/koladata/expr:view",
+        "@com_google_absl_py//absl/testing:absltest",
+        "@com_google_absl_py//absl/testing:parameterized",
+    ],
+)
+
 py_test(
     name = "eager_op_utils_test",
     srcs = ["eager_op_utils_test.py"],

py/koladata/operators/annotation.py

@@ -17,13 +17,8 @@
 from arolla import arolla
 from koladata.operators import optools
 
-
 # NOTE: Implemented in C++.
 source_location = arolla.abc.lookup_operator('kd.annotation.source_location')
+with_name = arolla.abc.lookup_operator('kd.annotation.with_name')
 
-
-with_name = optools.add_to_registry(
-    name='kd.annotation.with_name',
-    aliases=['kd.with_name'],
-    via_cc_operator_package=True,
-)(arolla.abc.lookup_operator('koda_internal.with_name'))
+optools.add_to_registry('kd.with_name', via_cc_operator_package=True)(with_name)