Deprecate py_boxing.LIST_TO_SLICE_BOXING_POLICY

PiperOrigin-RevId: 858110336
Change-Id: Idd3dbd6919682b75e930334d9d1a4c1d66967769

Author: apronchenkov

koladata/operators/shapes.cc

@@ -35,6 +35,7 @@
 #include "arolla/qexpr/operators/dense_array/edge_ops.h"
 #include "arolla/qtype/qtype.h"
 #include "arolla/qtype/qtype_traits.h"
+#include "arolla/qtype/tuple_qtype.h"
 #include "arolla/qtype/typed_slot.h"
 #include "arolla/util/repr.h"
 #include "koladata/arolla_utils.h"
@@ -102,25 +103,22 @@ absl::StatusOr<DataSlice::JaggedShape::Edge> GetEdgeFromSizes(
 // `i`. Only rank-0 or rank-1 int DataSlices are supported.
 class JaggedShapeCreateOperator : public arolla::QExprOperator {
  public:
-  explicit JaggedShapeCreateOperator(absl::Span<const arolla::QTypePtr> types)
-      : QExprOperator(types, koladata::GetJaggedShapeQType()) {
-    for (const auto& input_type : types) {
-      DCHECK(input_type == arolla::GetQType<DataSlice>() ||
-             input_type == arolla::GetQType<DataSlice::JaggedShape::Edge>());
-    }
-  }
+  using QExprOperator::QExprOperator;
 
  private:
   absl::StatusOr<std::unique_ptr<arolla::BoundOperator>> DoBind(
       absl::Span<const arolla::TypedSlot> input_slots,
       arolla::TypedSlot output_slot) const override {
-    Slot<DataSlice::JaggedShape> shape_slot =
-        output_slot.UnsafeToSlot<DataSlice::JaggedShape>();
+    int64_t edge_or_slice_count = input_slots[0].SubSlotCount();
+    std::vector<arolla::TypedSlot> edge_or_slice_slots;
+    edge_or_slice_slots.reserve(edge_or_slice_count);
+    for (int64_t i = 0; i < edge_or_slice_count; ++i) {
+      edge_or_slice_slots.emplace_back(input_slots[0].SubSlot(i));
+    }
     return MakeBoundOperator(
         "kd.shapes.new",
-        [edge_or_slice_slots =
-             std::vector(input_slots.begin(), input_slots.end()),
-         shape_slot = std::move(shape_slot)](
+        [edge_or_slice_slots = std::move(edge_or_slice_slots),
+         shape_slot = output_slot.UnsafeToSlot<DataSlice::JaggedShape>()](
             arolla::EvaluationContext* ctx,
             arolla::FramePtr frame) -> absl::Status {
           DataSlice::JaggedShape::EdgeVec edges;
@@ -284,16 +282,27 @@ absl::StatusOr<arolla::OperatorPtr>
 JaggedShapeCreateOperatorFamily::DoGetOperator(
     absl::Span<const arolla::QTypePtr> input_types,
     arolla::QTypePtr output_type) const {
-  for (const auto& input_type : input_types) {
-    if (input_type != arolla::GetQType<DataSlice>() &&
-        input_type != arolla::GetQType<DataSlice::JaggedShape::Edge>()) {
+  if (input_types.size() != 1) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("expected exactly one input, got: ", input_types.size()));
+  }
+  if (!arolla::IsTupleQType(input_types[0])) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("unsupported input type: ", input_types[0]->name()));
+  }
+  for (const auto& type_field : input_types[0]->type_fields()) {
+    if (type_field.GetType() != arolla::GetQType<DataSlice>() &&
+        type_field.GetType() !=
+            arolla::GetQType<DataSlice::JaggedShape::Edge>()) {
       return absl::InvalidArgumentError(
-          absl::StrCat("unsupported input type: ", input_type->name()));
+          absl::StrCat("unsupported input type: ", input_types[0]->name()));
     }
   }
-  return arolla::EnsureOutputQTypeMatches(
-      std::make_shared<JaggedShapeCreateOperator>(input_types), input_types,
-      output_type);
+  if (output_type != koladata::GetJaggedShapeQType()) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "expected output type JaggedShape, got: ", output_type->name()));
+  }
+  return std::make_shared<JaggedShapeCreateOperator>(input_types, output_type);
 }
 
 absl::StatusOr<arolla::OperatorPtr>

py/koladata/expr/BUILD

@@ -65,6 +65,7 @@ py_library(
         "//py:python_path",
         "//py/koladata/fstring",
         "//py/koladata/types:data_slice",
+        "//py/koladata/types:py_boxing",
         "//py/koladata/types:qtypes",
         "//py/koladata/types:schema_constants",
         "//py/koladata/util:kd_functools",

py/koladata/operators/jagged_shape.py

@@ -34,26 +34,31 @@
 _reshape = arolla_bridge._reshape  # pylint: disable=protected-access
 
 
-def _expect_slices_or_edges(value):
-  """Constrains `value` to be a tuple of DataSlices or Edges."""
-  is_slice_or_edge = arolla.LambdaOperator(
+def _expect_data_slice_or_edge_args(param):
+  """Constrains `param` to be a tuple of DataSlices or Edges."""
+  is_data_slice_or_edge = arolla.LambdaOperator(
       'x', (P.x == qtypes.DATA_SLICE) | (P.x == arolla.DENSE_ARRAY_EDGE)
   )
   return (
-      M.seq.all(M.seq.map(is_slice_or_edge, M.qtype.get_field_qtypes(value))),
+      M.qtype.is_tuple_qtype(param)
+      & M.seq.all(
+          M.seq.map(is_data_slice_or_edge, M.qtype.get_field_qtypes(param))
+      ),
       (
-          'all arguments must be DataSlices or Edges, got:'
-          f' {constraints.variadic_name_type_msg(value)}'
+          'expected all arguments to be data slices or edges, got'
+          f' {constraints.variadic_name_type_msg(param)}'
       ),
   )
 
 
 @optools.add_to_registry(via_cc_operator_package=True)
-@arolla.optools.as_backend_operator(
+@optools.as_backend_operator(
     'kd.shapes.new',
-    qtype_constraints=[_expect_slices_or_edges(P.dimensions)],
+    qtype_constraints=[_expect_data_slice_or_edge_args(P.dimensions)],
     qtype_inference_expr=qtypes.JAGGED_SHAPE,
-    experimental_aux_policy=py_boxing.LIST_TO_SLICE_BOXING_POLICY,
+    custom_boxing_fn_name_per_parameter={
+        'dimensions': py_boxing.WITH_LIST_TO_SLICE_SUPPORT,
+    },
 )
 def new(*dimensions):  # pylint: disable=unused-argument
   """Returns a JaggedShape from the provided dimensions.
@@ -88,34 +93,13 @@ def new(*dimensions):  # pylint: disable=unused-argument
   raise NotImplementedError('implemented in the backend')
 
 
-@optools.add_to_registry(via_cc_operator_package=True)
 @arolla.optools.as_backend_operator(
     'kd.shapes._new_with_size',
-    qtype_constraints=[
-        qtype_utils.expect_data_slice(P.result_size),
-        _expect_slices_or_edges(P.dimensions),
-    ],
+    qtype_constraints=[_expect_data_slice_or_edge_args(P.dimensions)],
     qtype_inference_expr=qtypes.JAGGED_SHAPE,
-    experimental_aux_policy=py_boxing.LIST_TO_SLICE_BOXING_POLICY,
 )
 def _new_with_size(result_size, *dimensions):  # pylint: disable=unused-argument
-  """Returns a JaggedShape from the provided dimensions and size.
-
-  It supports a single placeholder dimension argument denoted as `-1`, for which
-  its true value is inferred from the provided `size` argument and remaining
-  `dimensions`. The resulting dimension must be a uniform dimension, i.e. all
-  parent elements must have the same child size.
-
-  Args:
-    result_size: The size of the resulting JaggedShape.
-    *dimensions: A combination of Edges and DataSlices representing the
-      dimensions of the JaggedShape. Edges are used as is, while DataSlices are
-      treated as sizes. DataItems (of ints) are interpreted as uniform
-      dimensions which have the same child size for all parent elements.
-      DataSlices (of ints) are interpreted as a list of sizes, where `ds[i]` is
-      the child size of parent `i`. Only rank-0 or rank-1 int DataSlices are
-      supported.
-  """
+  """(internal) Returns a JaggedShape from the provided dimensions and size."""
   raise NotImplementedError('implemented in the backend')
 
 
@@ -143,7 +127,19 @@ def get_shape(x):  # pylint: disable=unused-argument
 
 
 @optools.add_to_registry(aliases=['kd.reshape'], via_cc_operator_package=True)
-@optools.as_lambda_operator('kd.shapes.reshape')
+@optools.as_lambda_operator(
+    'kd.shapes.reshape',
+    qtype_constraints=[
+        qtype_utils.expect_data_slice(P.x),
+        (
+            M.qtype.is_tuple_qtype(P.shape) | (P.shape == qtypes.JAGGED_SHAPE),
+            (
+                'expected a tuple or a shape, got:'
+                f' {constraints.name_type_msg(P.shape)}'
+            ),
+        ),
+    ],
+)
 def reshape(x, shape):
   """Returns a DataSlice with the provided shape.
 

py/koladata/operators/tests/BUILD

@@ -6640,25 +6640,6 @@ py_test(
     ],
 )
 
-py_test(
-    name = "shapes_new_with_size_test",
-    srcs = ["shapes_new_with_size_test.py"],
-    deps = [
-        "//py:python_path",
-        "//py/koladata/expr:expr_eval",
-        "//py/koladata/expr:input_container",
-        "//py/koladata/expr:view",
-        "//py/koladata/operators:kde_operators",
-        "//py/koladata/testing",
-        "//py/koladata/types:data_slice",
-        "//py/koladata/types:jagged_shape",
-        "//py/koladata/types:literal_operator",
-        "@com_google_absl_py//absl/testing:absltest",
-        "@com_google_absl_py//absl/testing:parameterized",
-        "@com_google_arolla//py/arolla",
-    ],
-)
-
 py_test(
     name = "schema_is_dict_schema_test",
     srcs = ["schema_is_dict_schema_test.py"],

py/koladata/operators/tests/shapes_new_test.py

@@ -68,16 +68,18 @@ def test_boxing(self):
         kde.shapes.new([1]),
         arolla.abc.bind_op(
             kde.shapes.new,
-            literal_operator.literal(data_slice.DataSlice.from_vals([1])),
+            literal_operator.literal(
+                arolla.tuple(data_slice.DataSlice.from_vals([1]))
+            ),
         ),
     )
 
   def test_unsupported_qtype_error(self):
     with self.assertRaisesRegex(
         ValueError,
         re.escape(
-            'all arguments must be DataSlices or Edges, got: *dimensions:'
-            ' (FLOAT32)'
+            'expected all arguments to be data slices or edges, got'
+            ' *dimensions: (FLOAT32)'
         ),
     ):
       kde.shapes.new(arolla.float32(1.0))