Add support for kd.slice([kd.obj(...)]) in tracing mode, and surrounding fixes.

The design is roughly similar to kde.fstr. This adds some complexity, but kd.slice and kd.item are very frequently used and I feel it is justified to support this.

PiperOrigin-RevId: 714905688
Change-Id: I438cf9f154401ec3c90638b52c751f4d26752a48

Author: petrmitrichev

py/koladata/functions/BUILD

@@ -35,6 +35,7 @@ py_library(
         ":schema",
         "//py/koladata/fstring",
         "//py/koladata/types:data_bag",
+        "//py/koladata/types:data_item",
         "//py/koladata/types:data_slice",
         "//py/koladata/types:general_eager_ops",
     ],

py/koladata/functions/functions.py

@@ -26,6 +26,7 @@
 from koladata.functions import s11n as _s11n
 from koladata.functions import schema as _schema
 from koladata.types import data_bag as _data_bag
+from koladata.types import data_item as _data_item
 from koladata.types import data_slice as _data_slice
 from koladata.types import general_eager_ops as _general_eager_ops
 
@@ -159,6 +160,14 @@ def new_schema(
 
 py_reference = _py_conversions.py_reference
 
+from_proto = _proto_conversions.from_proto
+to_proto = _proto_conversions.to_proto
+
+dumps = _s11n.dumps
+loads = _s11n.loads
+
+slice = _data_slice.DataSlice.from_vals  # pylint: disable=redefined-builtin
+item = _data_item.DataItem.from_vals
 int32 = _py_conversions.int32
 int64 = _py_conversions.int64
 float32 = _py_conversions.float32
@@ -169,9 +178,20 @@ def new_schema(
 mask = _py_conversions.mask
 expr_quote = _py_conversions.expr_quote
 
-
-from_proto = _proto_conversions.from_proto
-to_proto = _proto_conversions.to_proto
-
-dumps = _s11n.dumps
-loads = _s11n.loads
+slices = _py_types.SimpleNamespace(
+    # We use the top-level functions for `slice` and `item` instead of taking
+    # DataSlice.from_vals again here since for functions implemented in C++
+    # taking them twice from their class returns in different pointers,
+    # so "assertIs" test for aliases fails otherwise.
+    slice=slice,
+    item=item,
+    int32=_py_conversions.int32,
+    int64=_py_conversions.int64,
+    float32=_py_conversions.float32,
+    float64=_py_conversions.float64,
+    str=_py_conversions.str_,
+    bytes=_py_conversions.bytes_,
+    bool=_py_conversions.bool_,
+    mask=_py_conversions.mask,
+    expr_quote=_py_conversions.expr_quote,
+)

py/koladata/functions/tests/BUILD

@@ -850,6 +850,37 @@ py_test(
     ],
 )
 
+py_test(
+    name = "slice_test",
+    srcs = ["slice_test.py"],
+    deps = [
+        "//py/koladata/functions",
+        "//py/koladata/operators:kde_operators",
+        "//py/koladata/testing",
+        "//py/koladata/types:data_slice",
+        "//py/koladata/types:mask_constants",
+        "//py/koladata/types:schema_constants",
+        "@com_google_absl_py//absl/testing:absltest",
+        "@com_google_absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_test(
+    name = "item_test",
+    srcs = ["item_test.py"],
+    deps = [
+        "//py/koladata/functions",
+        "//py/koladata/operators:kde_operators",
+        "//py/koladata/testing",
+        "//py/koladata/types:data_slice",
+        "//py/koladata/types:mask_constants",
+        "//py/koladata/types:schema_constants",
+        "@com_google_absl_py//absl/testing:absltest",
+        "@com_google_absl_py//absl/testing:parameterized",
+        "@com_google_arolla//py/arolla",
+    ],
+)
+
 proto_library(
     name = "test_proto",
     srcs = ["test.proto"],

py/koladata/functions/tests/bool_test.py

@@ -41,6 +41,9 @@ def test_bool_errors(self, x, expected_error_msg):
     with self.assertRaisesRegex(ValueError, re.escape(expected_error_msg)):
       fns.bool(x)
 
+  def test_alias(self):
+    self.assertIs(fns.bool, fns.slices.bool)
+
 
 if __name__ == '__main__':
   absltest.main()

py/koladata/functions/tests/bytes_test.py

@@ -39,6 +39,9 @@ def test_bytes_errors(self, x, expected_error_msg):
     with self.assertRaisesRegex(ValueError, re.escape(expected_error_msg)):
       fns.bytes(x)
 
+  def test_alias(self):
+    self.assertIs(fns.bytes, fns.slices.bytes)
+
 
 if __name__ == '__main__':
   absltest.main()

py/koladata/functions/tests/expr_quote_test.py

@@ -44,6 +44,9 @@ def test_expr_quote_errors(self, x, expected_error_msg):
     with self.assertRaisesRegex(ValueError, re.escape(expected_error_msg)):
       fns.expr_quote(x)
 
+  def test_alias(self):
+    self.assertIs(fns.expr_quote, fns.slices.expr_quote)
+
 
 if __name__ == '__main__':
   absltest.main()

py/koladata/functions/tests/float32_test.py

@@ -49,6 +49,9 @@ def test_float32_errors(self, x, expected_error_msg):
     with self.assertRaisesRegex(ValueError, re.escape(expected_error_msg)):
       fns.float32(x)
 
+  def test_alias(self):
+    self.assertIs(fns.float32, fns.slices.float32)
+
 
 if __name__ == '__main__':
   absltest.main()

py/koladata/functions/tests/float64_test.py

@@ -49,6 +49,9 @@ def test_float64_errors(self, x, expected_error_msg):
     with self.assertRaisesRegex(ValueError, re.escape(expected_error_msg)):
       fns.float64(x)
 
+  def test_alias(self):
+    self.assertIs(fns.float64, fns.slices.float64)
+
 
 if __name__ == '__main__':
   absltest.main()

py/koladata/functions/tests/from_py_test.py

@@ -394,7 +394,7 @@ def test_empty_slice(self):
         a=schema_constants.STRING, b=fns.list_schema(schema_constants.INT32)
     )
     res = fns.from_py([], from_dim=1, schema=schema)
-    testing.assert_equal(res.no_bag(), ds([], schema))
+    testing.assert_equal(res.no_bag(), ds([], schema.no_bag()))
 
   def test_obj_reference(self):
     obj = fns.obj()

py/koladata/functions/tests/int32_test.py

@@ -44,6 +44,9 @@ def test_int32_errors(self, x, expected_error_msg):
     with self.assertRaisesRegex(ValueError, re.escape(expected_error_msg)):
       fns.int32(x)
 
+  def test_alias(self):
+    self.assertIs(fns.int32, fns.slices.int32)
+
 
 if __name__ == '__main__':
   absltest.main()