Minor refactoring in the aux-policies in Koladata

 * Renamed the "default" classic policy to "classic" for better clarity
 * Moved policy declarations from py/koladata/types to py/koladata/operators
 * Updated registration logic (particularly, use pybind11 instead of python c api)

PiperOrigin-RevId: 859016491
Change-Id: Ia6d25d89d1c3bdd5e9a4b915d125594ed77a0e59

Author: apronchenkov

docs/api_reference.md

@@ -4864,7 +4864,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] | None = <class 'koladata.expr.view.KodaView'>, repr_fn: Union[Callable[[Expr, NodeTokenView], ReprToken], None] = None, aux_policy: str = 'koladata_default_boxing', 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, view: type[ExprView] | None = <class 'koladata.expr.view.KodaView'>, repr_fn: Union[Callable[[Expr, NodeTokenView], ReprToken], None] = None, 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.
 

koladata/expr/expr_operators.cc

@@ -161,7 +161,7 @@ ToArollaInt64Operator::ToArollaInt64Operator()
     : ToArollaValueOperator(
           "koda_internal.to_arolla_int64",
           arolla::expr::ExprOperatorSignature({{"x"}},
-                                              "koladata_default_boxing"),
+                                              "koladata_classic_aux_policy"),
           "Returns `x` converted into an arolla int64 value.\n"
           "\n"
           "Note that `x` must adhere to the following requirements:\n"
@@ -180,7 +180,7 @@ ToArollaTextOperator::ToArollaTextOperator()
     : ToArollaValueOperator(
           "koda_internal.to_arolla_text",
           arolla::expr::ExprOperatorSignature({{"x"}},
-                                              "koladata_default_boxing"),
+                                              "koladata_classic_aux_policy"),
           "Returns `x` converted into an arolla text value.\n"
           "\n"
           "Note that `x` must adhere to the following requirements:\n"

py/koladata/operators/BUILD

@@ -15,9 +15,10 @@
 # Arolla Expr Operators on Koda abstractions.
 
 load("@com_google_arolla//py/arolla/dynamic_deps:build_defs.bzl", "arolla_py_cc_deps")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
 load("@rules_cc//cc:cc_test.bzl", "cc_test")
 load("@rules_python//python:defs.bzl", "py_library", "py_test")
-load("//py/koladata/dynamic_deps:py_extension.bzl", "koladata_py_extension")
+load("//py/koladata/dynamic_deps:py_extension.bzl", "koladata_pybind_extension")
 load("//py/koladata/operators:optools.bzl", "koladata_cc_operator_package")
 
 package(default_visibility = ["//visibility:private"])
@@ -164,13 +165,21 @@ py_library(
     ],
 )
 
-koladata_py_extension(
-    name = "py_optools_py_ext",
-    srcs = [
-        "py_optools_module.cc",
-        "py_unified_binding_policy.cc",
-        "py_unified_binding_policy.h",
+koladata_pybind_extension(
+    name = "clib",
+    srcs = ["clib.cc"],
+    deps = [
+        ":py_unified_binding_policy",
+        "//koladata/expr:non_determinism",
+        "@com_google_absl//absl/strings:string_view",
+        "@pybind11_abseil//pybind11_abseil:absl_casters",
     ],
+)
+
+cc_library(
+    name = "py_unified_binding_policy",
+    srcs = ["py_unified_binding_policy.cc"],
+    hdrs = ["py_unified_binding_policy.h"],
     deps = [
         "//koladata/expr:expr_operators",
         "//koladata/expr:non_determinism",
@@ -179,7 +188,6 @@ koladata_py_extension(
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/random",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:span",
@@ -189,15 +197,16 @@ koladata_py_extension(
         "@com_google_arolla//arolla/util:status_backport",
         "@com_google_arolla//py/arolla/abc:py_abc",
         "@com_google_arolla//py/arolla/py_utils",
-        "@rules_python//python/cc:current_py_cc_headers",  # buildcleaner: keep
+        "@rules_python//python/cc:current_py_cc_headers",
     ],
 )
 
 py_library(
     name = "unified_binding_policy",
     srcs = ["unified_binding_policy.py"],
     deps = [
-        ":py_optools_py_ext",
+        ":aux_policies",
+        ":clib",
         "//py:python_path",
         "//py/koladata/types:py_boxing",
         "@com_google_arolla//py/arolla",
@@ -208,6 +217,7 @@ py_test(
     name = "unified_binding_policy_test",
     srcs = ["unified_binding_policy_test.py"],
     deps = [
+        ":aux_policies",
         ":op_repr",
         ":unified_binding_policy",
         "//py:python_path",
@@ -221,11 +231,29 @@ py_test(
     ],
 )
 
+py_test(
+    name = "classic_binding_policy_test",
+    srcs = ["classic_binding_policy_test.py"],
+    deps = [
+        ":aux_policies",
+        ":kde_operators",
+        "//py:python_path",
+        "//py/koladata/expr:input_container",
+        "//py/koladata/testing",
+        "//py/koladata/types:data_slice",
+        "//py/koladata/types:ellipsis",
+        "//py/koladata/types:literal_operator",
+        "@com_google_absl_py//absl/testing:absltest",
+        "@com_google_arolla//py/arolla",
+    ],
+)
+
 py_library(
     name = "optools",
     srcs = ["optools.py"],
     visibility = ["//koladata:internal"],
     deps = [
+        ":aux_policies",
         ":op_repr",
         ":qtype_utils",
         ":unified_binding_policy",
@@ -407,6 +435,7 @@ py_library(
     deps = [
         ":arolla_bridge",
         ":assertion",
+        ":aux_policies",
         ":jagged_shape",
         ":masking",
         ":optools",
@@ -429,11 +458,11 @@ py_library(
     srcs = ["tuple.py"],
     deps = [
         ":arolla_bridge",
+        ":aux_policies",
         ":optools",
         ":qtype_utils",
         ":view_overloads",
         "//py:python_path",
-        "//py/koladata/types:py_boxing",
         "//py/koladata/types:schema_constants",
         "@com_google_arolla//py/arolla",
         "@com_google_arolla//py/arolla/jagged_shape",
@@ -446,6 +475,7 @@ py_library(
     deps = [
         ":arolla_bridge",
         ":assertion",
+        ":aux_policies",
         ":bags",
         ":entities",
         ":masking",
@@ -605,6 +635,7 @@ py_library(
     deps = [
         ":arolla_bridge",
         ":assertion",
+        ":aux_policies",
         ":comparison",
         ":jagged_shape",
         ":masking",
@@ -667,6 +698,7 @@ py_library(
     name = "op_repr",
     srcs = ["op_repr.py"],
     deps = [
+        ":aux_policies",
         ":unified_binding_policy",
         "//py:python_path",
         "//py/koladata/types:data_slice",
@@ -795,11 +827,11 @@ py_library(
     name = "bags",
     srcs = ["bags.py"],
     deps = [
+        ":aux_policies",
         ":optools",
         ":qtype_utils",
         ":view_overloads",
         "//py:python_path",
-        "//py/koladata/types:py_boxing",
         "//py/koladata/types:qtypes",
         "//py/koladata/types:schema_constants",
         "@com_google_arolla//py/arolla",
@@ -932,6 +964,17 @@ py_library(
     ],
 )
 
+py_library(
+    name = "aux_policies",
+    srcs = ["aux_policies.py"],
+    deps = [
+        ":clib",
+        "//py:python_path",
+        "//py/koladata/types:py_boxing",
+        "@com_google_arolla//py/arolla",
+    ],
+)
+
 py_test(
     name = "eager_op_utils_test",
     srcs = ["eager_op_utils_test.py"],

py/koladata/operators/aux_policies.py

@@ -0,0 +1,69 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Module with deqring auxiliary policies for Koladata."""
+
+from arolla import arolla
+from koladata.operators import clib
+from koladata.types import py_boxing
+
+# The default auxiliary policy for Koladata operators.
+#
+# This policy implements binding rules that support positional-only,
+# keyword-only, variadic positional, and variadic keyword parameters.
+# It also supports per-parameter boxing control (defaulting to Koladata-specific
+# rules) and allows expressing non-deterministic operator behaviour.
+#
+# NOTE: Non-variadic parameters are effectively compatible with classic
+# Arolla binding rules and are represented as positional-or-keyword parameters.
+# For example, given:
+#
+#   def op(a, *, b, c=default): ...
+#
+#   op(a, b=b)
+#
+# is effectively equivalent to
+#
+#   arolla.abc.bind_op(op, a, b=b)
+#
+# However, variadic parameters are represented differently: positional
+# variadic arguments are packed into a tuple, and keyword variadic
+# arguments are packed into a namedtuple. For example, given:
+#
+#   def op(*args, **kwargs): ...
+#
+#   op(1, 2, x=3, y=4)
+#
+# is equivalent to:
+#
+#   arolla.abc.bind_op(op, args=arolla.tuple(1, 2),
+#                          kwargs=arolla.namedtuple(x=3, y=4))
+#
+# The binding rules implementation can be found in:
+# //py/koladata/operators/py_unified_binding_policy.cc
+UNIFIED_AUX_POLICY = 'koladata_unified_aux_policy'
+
+
+# The classic Arolla binding rules with Koladata-specific boxing: it wraps most
+# values into DataItems (excluding lists and tuples). QValues remain unchanged.
+CLASSIC_AUX_POLICY = 'koladata_classic_aux_policy'
+
+
+clib.register_unified_aux_binding_policy(UNIFIED_AUX_POLICY)
+
+arolla.abc.register_classic_aux_binding_policy_with_custom_boxing(
+    CLASSIC_AUX_POLICY,
+    py_boxing.as_qvalue_or_expr,
+    make_literal_fn=py_boxing.literal,
+)

py/koladata/operators/bags.py

@@ -16,10 +16,10 @@
 
 from arolla import arolla as _arolla
 from arolla.jagged_shape import jagged_shape as _jagged_shape
+from koladata.operators import aux_policies as _aux_policies
 from koladata.operators import optools as _optools
 from koladata.operators import qtype_utils as _qtype_utils
 from koladata.operators import view_overloads as _
-from koladata.types import py_boxing as _py_boxing
 from koladata.types import qtypes as _qtypes
 from koladata.types import schema_constants as _schema_constants
 
@@ -62,7 +62,7 @@ def is_null_bag(bag):  # pylint: disable=unused-argument
         _qtype_utils.expect_data_bag_args(_P.bags),
     ],
     qtype_inference_expr=_qtypes.DATA_BAG,
-    experimental_aux_policy=_py_boxing.DEFAULT_BOXING_POLICY,
+    experimental_aux_policy=_aux_policies.CLASSIC_AUX_POLICY,
 )
 def enriched(*bags):  # pylint: disable=unused-argument
   """Creates a new immutable DataBag enriched by `bags`.
@@ -93,7 +93,7 @@ def enriched(*bags):  # pylint: disable=unused-argument
         _qtype_utils.expect_data_bag_args(_P.bags),
     ],
     qtype_inference_expr=_qtypes.DATA_BAG,
-    experimental_aux_policy=_py_boxing.DEFAULT_BOXING_POLICY,
+    experimental_aux_policy=_aux_policies.CLASSIC_AUX_POLICY,
 )
 def updated(*bags):  # pylint: disable=unused-argument
   """Creates a new immutable DataBag updated by `bags`.

py/koladata/operators/classic_binding_policy_test.py

@@ -0,0 +1,88 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+
+from absl.testing import absltest
+from arolla import arolla
+from koladata.expr import input_container
+from koladata.operators import aux_policies
+from koladata.operators import kde_operators
+from koladata.testing import testing
+from koladata.types import data_slice
+from koladata.types import ellipsis
+from koladata.types import literal_operator
+
+
+ds = data_slice.DataSlice.from_vals
+I = input_container.InputContainer('I')
+kde = kde_operators.kde
+
+
+@arolla.optools.as_lambda_operator(
+    'op_with_classic_binding_policy',
+    experimental_aux_policy=aux_policies.CLASSIC_AUX_POLICY,
+)
+def op_with_classic_binding_policy(x, y):
+  return (x, y)
+
+
+class ClassicBindingPolicyTest(absltest.TestCase):
+
+  def test_basic(self):
+    expr = op_with_classic_binding_policy(1, 2)
+    testing.assert_equal(
+        expr,
+        arolla.abc.bind_op(
+            op_with_classic_binding_policy,
+            literal_operator.literal(ds(1)),
+            literal_operator.literal(ds(2)),
+        ),
+    )
+
+  def test_with_slice(self):
+    expr = op_with_classic_binding_policy(1, slice(1, None, 2))
+    testing.assert_equal(
+        expr,
+        arolla.abc.bind_op(
+            op_with_classic_binding_policy,
+            literal_operator.literal(ds(1)),
+            literal_operator.literal(arolla.types.Slice(ds(1), None, ds(2))),
+        ),
+    )
+
+  def test_with_ellipsis(self):
+    expr = op_with_classic_binding_policy(1, ...)
+    testing.assert_equal(
+        expr,
+        arolla.abc.bind_op(
+            op_with_classic_binding_policy,
+            literal_operator.literal(ds(1)),
+            literal_operator.literal(ellipsis.ellipsis()),
+        ),
+    )
+
+  def test_missing_argument(self):
+    with self.assertRaisesWithLiteralMatch(
+        TypeError, "missing 1 required positional argument: 'y'"
+    ):
+      op_with_classic_binding_policy(1)
+
+  def test_list_unsupported(self):
+    with self.assertRaisesRegex(ValueError, re.escape('list')):
+      op_with_classic_binding_policy(1, [2, 3, 4])
+
+
+if __name__ == '__main__':
+  absltest.main()