Change the default value kd.map_py(fn, ..., include_missing=False)

Set the default value of `include_missing` to `False`, so that, by default, the function `fn` is applied only for items present in all `args` and `kwargs`.

PiperOrigin-RevId: 714964283
Change-Id: Ic82580783c00c4df2cd03a405a3eaa44885d4110

Author: apronchenkov

py/koladata/ext/py_cloudpickle_test.py

@@ -56,7 +56,7 @@ def test_apply_py(self):
 
   def test_map_py(self):
     def f(x):
-      return x + 1 if x is not None else None
+      return x + 1
 
     x = ds([[1, 2, None, 4], [None, None], [7, 8, 9]])
     res = kd.map_py(py_cloudpickle.py_cloudpickle(f), x)

py/koladata/functor/functor_factories.py

@@ -448,9 +448,9 @@ def map_py_fn(
     schema: The schema to use for resulting DataSlice.
     max_threads: maximum number of threads to use.
     ndim: Dimensionality of items to pass to `f`.
-    include_missing: Specifies whether `f` should be applied to the missing
-      items. By default, the function is applied to all items including the
-      missing. `include_missing=False` can only be used with `ndim=0`.
+    include_missing: Specifies whether `f` applies to all items (`=True`) or
+      only to items present in all `args` and `kwargs` (`=False`, valid only
+      when `ndim=0`); defaults to `False` when `ndim=0`.
     **defaults: Keyword defaults to pass to the function. The values in this map
       may be kde expressions, format strings, or 0-dim DataSlices. See the
       docstring for py_fn for more details.

py/koladata/functor/functor_factories_test.py

@@ -459,16 +459,12 @@ def test_py_fn_explicit_return_type(self):
 
   def test_fstr_fn_simple(self):
     testing.assert_equal(
-        kd.call(
-            functor_factories.fstr_fn(f'{I.x:s} {I.y:s}'), x=1, y=2
-        ),
+        kd.call(functor_factories.fstr_fn(f'{I.x:s} {I.y:s}'), x=1, y=2),
         ds('1 2'),
     )
 
     testing.assert_equal(
-        kd.call(
-            functor_factories.fstr_fn(f'{(I.x + I.y):s}'), x=1, y=2
-        ),
+        kd.call(functor_factories.fstr_fn(f'{(I.x + I.y):s}'), x=1, y=2),
         ds('3'),
     )
 
@@ -483,21 +479,15 @@ def test_fstr_fn_expr(self):
 
   def test_fstr_fn_variable(self):
     testing.assert_equal(
-        kd.call(
-            functor_factories.fstr_fn(f'{V.x:s} {I.y:s}', x=1), y=2
-        ),
+        kd.call(functor_factories.fstr_fn(f'{V.x:s} {I.y:s}', x=1), y=2),
         ds('1 2'),
     )
 
   def test_fstr_fn_no_substitutions(self):
-    with self.assertRaisesRegex(
-        ValueError, 'FString has nothing to format'
-    ):
+    with self.assertRaisesRegex(ValueError, 'FString has nothing to format'):
       _ = kd.call(functor_factories.fstr_fn('abc'))
 
-    with self.assertRaisesRegex(
-        ValueError, 'FString has nothing to format'
-    ):
+    with self.assertRaisesRegex(ValueError, 'FString has nothing to format'):
       # we need to use fstring to avoid the error (f'{I.x}')
       _ = kd.call(functor_factories.fstr_fn('{I.x}'), x=1)
 
@@ -779,7 +769,7 @@ def fn(x):
         kd.call(
             functor_factories.map_py_fn(fn), x=kdi.slice([1, None, None])
         ).to_py(),
-        [2, -1, -1],
+        [2, None, None],
     )
     self.assertEqual(
         kd.call(

py/koladata/operators/py.py

@@ -552,9 +552,9 @@ def my_expansion(x):
     schema: The schema to use for resulting DataSlice.
     max_threads: maximum number of threads to use.
     ndim: Dimensionality of items to pass to `fn`.
-    include_missing: Specifies whether `fn` should be computed to the missing
-      items. By default, the function is applied to all items including the
-      missing. `include_missing=False` can only be used with `ndim=0`.
+    include_missing: Specifies whether `fn` applies to all items (`=True`) or
+      only to items present in all `args` and `kwargs` (`=False`, valid only
+      when `ndim=0`); defaults to `False` when `ndim=0`.
     item_completed_callback: A callback that will be called after each item is
       processed. It will be called in the original thread that called `map_py`
       in case `max_threads` is greater than 1, as we rely on this property for
@@ -573,8 +573,8 @@ def my_expansion(x):
       include_missing, param_name='include_missing'
   )
   if include_missing is None:
-    include_missing = True
-  if not include_missing and ndim != 0:
+    include_missing = ndim != 0
+  elif not include_missing and ndim != 0:
     raise ValueError('`include_missing=False` can only be used with `ndim=0`')
   if not args and not kwargs:
     raise TypeError('expected at least one input DataSlice, got none')

py/koladata/operators/tests/py_map_py_test.py

@@ -42,7 +42,7 @@ class PyMapPyTest(parameterized.TestCase):
 
   def test_map_py_single_arg(self):
     def add_one(x):
-      return x + 1 if x is not None else None
+      return x + 1
 
     x = ds([[1, 2, None, 4], [None, None], [7, 8, 9]])
     res = expr_eval.eval(kde.py.map_py(add_one, x))
@@ -118,8 +118,6 @@ def add_one(x):
 
   def test_map_py_multi_args(self):
     def add_all(x, y, z):
-      if x is None or y is None or z is None:
-        return None
       return x + y + z
 
     val1 = ds([[1, 2, None, 4], [None, None], [7, 8, 9]])
@@ -132,7 +130,7 @@ def add_all(x, y, z):
 
   def test_map_py_texting_output(self):
     def as_string(x):
-      return str(x) if x is not None else None
+      return str(x)
 
     val = ds([[1, 2, None, 4], [None, None], [7, 8, 9]])
     res = expr_eval.eval(kde.py.map_py(as_string, val))
@@ -142,7 +140,7 @@ def as_string(x):
 
   def test_map_py_texting_input(self):
     def as_string(x):
-      return int(x) if x is not None else None
+      return int(x)
 
     val = ds([['1', '2', None, '4'], [None, None], ['7', '8', '9']])
     res = expr_eval.eval(kde.py.map_py(as_string, val))
@@ -152,7 +150,7 @@ def as_string(x):
 
   def test_map_py_with_qtype(self):
     def add_one(x):
-      return x + 1 if x is not None else None
+      return x + 1
 
     val = ds([[1, 2, None, 4], [None, None], [7, 8, 9]])
     res = expr_eval.eval(
@@ -175,17 +173,13 @@ def test_map_py_with_schema(self):
     )
 
     def my_func_dynamic_schema(x):
-      return None if x is None else functions.new(u=x, v=x + 1)
+      return functions.new(u=x, v=x + 1)
 
     def my_func_any_schema(x):
-      return (
-          None
-          if x is None
-          else functions.new(u=x, v=x + 1, schema=schema_constants.ANY)
-      )
+      return functions.new(u=x, v=x + 1, schema=schema_constants.ANY)
 
     def my_func_correct_schema(x):
-      return None if x is None else functions.new(u=x, v=x + 1, schema=schema)
+      return functions.new(u=x, v=x + 1, schema=schema)
 
     res = expr_eval.eval(
         kde.py.map_py(my_func_correct_schema, val, schema=schema)
@@ -246,7 +240,7 @@ def my_func_correct_schema(x):
     )
 
     def my_func_same_bag(schema, x):
-      return None if x is None else db.new(u=x, v=x + 1, schema=schema)
+      return db.new(u=x, v=x + 1, schema=schema)
 
     res = expr_eval.eval(
         kde.py.map_py(
@@ -368,16 +362,14 @@ def my_fn(x):
 
   def test_map_py_scalar_input(self):
     def add_one(x):
-      return x + 1 if x is not None else None
+      return x + 1
 
     val = ds(5)
     res = expr_eval.eval(kde.py.map_py(add_one, val))
     testing.assert_equal(res.no_bag(), ds(6))
 
   def test_map_py_auto_expand(self):
     def my_add(x, y):
-      if x is None or y is None:
-        return None
       return x + y
 
     val1 = ds(1)
@@ -387,8 +379,6 @@ def my_add(x, y):
 
   def test_map_py_raw_input(self):
     def my_add(x, y):
-      if x is None or y is None:
-        return None
       return x + y
 
     res = expr_eval.eval(
@@ -398,7 +388,7 @@ def my_add(x, y):
 
   def test_map_py_dict(self):
     def as_dict(x):
-      return {'x': x, 'y': x + 1 if x is not None else 57}
+      return {'x': x, 'y': x + 1}
 
     val = ds([[1, 2, None, 4], [None, None], [7, 8, 9]])
     res = expr_eval.eval(kde.py.map_py(as_dict, val))
@@ -408,7 +398,9 @@ def as_dict(x):
     )
     testing.assert_equal(
         res['y'].no_bag(),
-        ds([[2, 3, 57, 5], [57, 57], [8, 9, 10]], schema_constants.OBJECT),
+        ds(
+            [[2, 3, None, 5], [None, None], [8, 9, 10]], schema_constants.OBJECT
+        ),
     )
 
   def test_map_py_invalid_qtype(self):
@@ -422,8 +414,6 @@ def as_set(x):
 
   def test_map_py_incompatible_inputs(self):
     def add_x_y(x, y):
-      if x is None or y is None:
-        return None
       return x + y
 
     val1 = ds([[1, 2, None, 4], [None, None], [7, 8, 9]])
@@ -439,8 +429,6 @@ def add_x_y(x, y):
 
   def test_map_py_kwargs(self):
     def my_fn(x, y, z=2, **kwargs):
-      if x is None or y is None or z is None:
-        return None
       return x + y + z + kwargs.get('w', 5)
 
     x = ds([[0, 0, 1], [None, 1, 0]])
@@ -531,13 +519,12 @@ def test_map_py_expanded_results(self):
     with self.subTest('expand_one_dim'):
 
       def ranges(x):
-        return list(range(x or 0))
+        return list(range(x))
 
       res = expr_eval.eval(kde.py.map_py(ranges, val))
       self.assertEqual(res.get_ndim(), 2)
       self.assertEqual(
-          res.to_py(),
-          [[[0], [0, 1], []], [[0, 1, 2, 3], [0, 1, 2, 3, 4]]],
+          res.to_py(), [[[0], [0, 1], None], [[0, 1, 2, 3], [0, 1, 2, 3, 4]]]
       )
 
     with self.subTest('expand_several_dims'):
@@ -549,7 +536,7 @@ def expnd(x):
       self.assertEqual(res.get_ndim(), 2)
       self.assertEqual(
           res.to_py(),
-          [[[[1, -1]], [[2, -1]], [[None, -1]]], [[[4, -1]], [[5, -1]]]],
+          [[[[1, -1]], [[2, -1]], None], [[[4, -1]], [[5, -1]]]],
       )
 
     with self.subTest('agg_and_expand'):