Support negative values for dim= in kd.index.

PiperOrigin-RevId: 713646153
Change-Id: I5b442f709b7c6b763c56a7a9ab20816f3835b6c7

Author: odimka

py/koladata/operators/BUILD

@@ -470,6 +470,7 @@ py_library(
     deps = [
         ":arolla_bridge",
         ":assertion",
+        ":comparison",
         ":functor",
         ":jagged_shape",
         ":masking",

py/koladata/operators/slices.py

@@ -20,6 +20,7 @@
 from arolla.jagged_shape import jagged_shape
 from koladata.operators import arolla_bridge
 from koladata.operators import assertion
+from koladata.operators import comparison
 from koladata.operators import functor
 from koladata.operators import jagged_shape as jagged_shape_ops
 from koladata.operators import masking
@@ -762,15 +763,28 @@ def group_by(x, *args):
   return dispatch_op(x, args)
 
 
+@optools.as_lambda_operator(
+    'kde.slices._normalize_dim',
+    qtype_constraints=[
+        qtype_utils.expect_data_slice(P.x),
+        qtype_utils.expect_data_slice(P.dim),
+    ],
+)
+def normalize_dim(x, dim):
+  """Returns dim if dim >= 0, otherwise get_ndim(x) + dim."""
+  # TODO: masking.cond can be slow, optimize
+  return masking.cond(comparison.less(dim, 0), get_ndim(x) + dim, dim)
+
+
 @optools.add_to_registry(aliases=['kde.index'])
 @optools.as_lambda_operator(
     'kde.slices.index',
     qtype_constraints=[
         qtype_utils.expect_data_slice(P.x),
-        qtype_utils.expect_data_slice_or_unspecified(P.dim),
+        qtype_utils.expect_data_slice(P.dim),
     ],
 )
-def index(x, dim=arolla.unspecified()):
+def index(x, dim=-1):
   """Returns the indices of the elements computed over the last dim dimensions.
 
   The resulting slice has the same shape as the input.
@@ -790,24 +804,23 @@ def index(x, dim=arolla.unspecified()):
       # -> kd.slice([[[0, None, 0], [0, 0]], [[None, 1], [1, 1, 1]]])
     kd.index(ds, dim=1)
       # -> kd.slice([[[0, None, 0], [1, 1]], [[None, 0], [1, 1, 1]]])
-    kd.index(ds, dim=2)
+    kd.index(ds, dim=2)  # (same as kd.index(ds, -1) or kd.index(ds))
       # -> kd.slice([[[0, None, 2], [0, 1]], [[None, 1], [0, 1, 2]]])
 
     kd.index(ds) -> kd.index(ds, dim=ds.get_ndim() - 1)
 
   Args:
     x: A DataSlice.
-    dim: The dimension to compute indices over. Requires 0 <= dim < get_ndim(x).
-      If unspecified, it is set to the last dimension of x.
+    dim: The dimension to compute indices over. Requires abs(dim) < get_ndim(x).
+      If dim < 0 then dim = get_ndim(x) + dim.
   """
   x = assertion.with_assertion(
       x,
       get_ndim(x) != 0,
       'kde.slices.index: argument `x` must have non-zero rank',
   )
 
-  dim = M.core.default_if_unspecified(dim, get_ndim(x) - 1)
-
+  dim = normalize_dim(x, dim)
   ndim = get_ndim(x) - dim - 1
   ndim = assertion.with_assertion(
       ndim,

py/koladata/operators/tests/slices_index_test.py

@@ -44,14 +44,12 @@
 QTYPES = frozenset([
     (DATA_SLICE, DATA_SLICE),
     (DATA_SLICE, DATA_SLICE, DATA_SLICE),
-    (DATA_SLICE, arolla.UNSPECIFIED, DATA_SLICE),
 ])
 
 
 class SlicesIndexTest(parameterized.TestCase):
 
   @parameterized.parameters(
-      (ds([5, 6, None, 7]), arolla.unspecified(), ds([0, 1, None, 3], INT64)),
       (
           ds([[5, None, 5], [6, 7], [None, None]]),
           ds(0),
@@ -62,6 +60,11 @@ class SlicesIndexTest(parameterized.TestCase):
           ds(1),
           ds([[0, None, 2], [0, 1], [None, None]], INT64),
       ),
+      (
+          ds([[5, None, 5], [6, 7], [None, None]]),
+          ds(-1),
+          ds([[0, None, 2], [0, 1], [None, None]], INT64),
+      ),
       (
           ds([[['a', 'b', 'c'], ['d', 'e']], [['f', 'g'], ['h', 'i', 'j']]]),
           ds(0),
@@ -72,11 +75,21 @@ class SlicesIndexTest(parameterized.TestCase):
           ds(1),
           ds([[[0, 0, 0], [1, 1]], [[0, 0], [1, 1, 1]]], INT64),
       ),
+      (
+          ds([[['a', 'b', 'c'], ['d', 'e']], [['f', 'g'], ['h', 'i', 'j']]]),
+          ds(-2),
+          ds([[[0, 0, 0], [1, 1]], [[0, 0], [1, 1, 1]]], INT64),
+      ),
       (
           ds([[['a', 'b', 'c'], ['d', 'e']], [['f', 'g'], ['h', 'i', 'j']]]),
           ds(2),
           ds([[[0, 1, 2], [0, 1]], [[0, 1], [0, 1, 2]]], INT64),
       ),
+      (
+          ds([[['a', 'b', 'c'], ['d', 'e']], [['f', 'g'], ['h', 'i', 'j']]]),
+          ds(-1),
+          ds([[[0, 1, 2], [0, 1]], [[0, 1], [0, 1, 2]]], INT64),
+      ),
   )
   def test_eval(self, x, dim, expected_value):
     actual_value = expr_eval.eval(kde.slices.index(x, dim))
@@ -114,7 +127,7 @@ def test_data_item_input_error(self):
     ):
       expr_eval.eval(kde.slices.index(x))
 
-  @parameterized.parameters(-1, 2)
+  @parameterized.parameters(2, -2)
   def test_out_of_bounds_ndim_error(self, ndim):
     x = data_slice.DataSlice.from_vals([1, 2, 3])
     with self.assertRaisesRegex(

py/koladata/operators/tests/slices_range_test.py

@@ -95,7 +95,7 @@ def test_eval(self, *args_and_expected):
   def test_qtype_signatures(self):
     self.assertCountEqual(
         arolla.testing.detect_qtype_signatures(
-            kde.slices.index,
+            kde.slices.range,
             possible_qtypes=test_qtypes.DETECT_SIGNATURES_QTYPES,
         ),
         QTYPES,