Expose include_missing parameter in kd.functor.map_py_fn()

PiperOrigin-RevId: 714125675
Change-Id: I7de628b5fbb1a46c1292f6c464ddfeb5f7d62bb0

Author: apronchenkov

py/koladata/functor/functor_factories.py

@@ -435,6 +435,7 @@ def map_py_fn(
     schema: Any = None,
     max_threads: Any = 1,
     ndim: Any = 0,
+    include_missing: Any = None,
     **defaults: Any,
 ) -> data_slice.DataSlice:
   """Returns a Koda functor wrapping a python function for kd.map_py.
@@ -447,6 +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`.
     **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.
@@ -459,7 +463,7 @@ def map_py_fn(
           schema=py_boxing.as_qvalue(schema),
           max_threads=py_boxing.as_qvalue(max_threads),
           ndim=py_boxing.as_qvalue(ndim),
-          include_missing=py_boxing.as_qvalue(True),
+          include_missing=py_boxing.as_qvalue(include_missing),
           item_completed_callback=py_boxing.as_qvalue(None),
           kwargs=I.kwargs,
       ),

py/koladata/functor/functor_factories_test.py

@@ -478,7 +478,7 @@ def test_fstr_fn_expr(self):
             functor_factories.fstr_fn(f'{kde.select(I.x, kdi.present):s}'),
             x=ds([1, None]),
         ),
-        ds(['1', None])
+        ds(['1', None]),
     )
 
   def test_fstr_fn_variable(self):
@@ -773,13 +773,27 @@ def fn(x):
 
   def test_map_py_fn_with_missing_items(self):
     def fn(x):
-      return -1 if x is None else x
+      return -1 if x is None else x + 1
 
     self.assertEqual(
         kd.call(
             functor_factories.map_py_fn(fn), x=kdi.slice([1, None, None])
         ).to_py(),
-        [1, -1, -1],
+        [2, -1, -1],
+    )
+    self.assertEqual(
+        kd.call(
+            functor_factories.map_py_fn(fn, include_missing=True),
+            x=kdi.slice([1, None, None]),
+        ).to_py(),
+        [2, -1, -1],
+    )
+    self.assertEqual(
+        kd.call(
+            functor_factories.map_py_fn(fn, include_missing=False),
+            x=kdi.slice([1, None, None]),
+        ).to_py(),
+        [2, None, None],
     )
 
   def test_map_py_default_arguments(self):