Enable parallel transform of kd.switch.

Also minor fixes in repr.

PiperOrigin-RevId: 869902121
Change-Id: I2aebeeab3a330bc17ceef54b69d86ca6c9ea2e6f

Author: timofey-stepanov

docs/api_reference.md

@@ -5338,9 +5338,9 @@ parallel version (DataSlice -> future[DataSlice]), this is work in progress.
 
 Args:
   fn: The functor to transform.
-  allow_runtime_transforms: Whether to allow sub-functors to be not literals,
-    but computed expressions, which will therefore have to be transformed at
-    runtime. This can be slow.
+  allow_runtime_transforms: Whether to allow sub-functors to be not fully
+    defined at transform time (i.e. to depend on functor inputs), which will
+    therefore have to be transformed at runtime. This can be slow.
 
 Returns:
   The transformed functor.</code></pre>

koladata/functor/parallel/BUILD

@@ -85,6 +85,7 @@ cc_library(
         "//koladata/expr:non_determinism",
         "//koladata/functor",
         "//koladata/functor:auto_variables",
+        "//koladata/functor:call",
         "//koladata/functor:signature_utils",
         "//koladata/internal:data_item",
         "//koladata/internal:dtype",

koladata/functor/parallel/transform.cc

@@ -61,14 +61,16 @@
 #include "koladata/expr/expr_operators.h"
 #include "koladata/expr/non_determinism.h"
 #include "koladata/functor/auto_variables.h"
+#include "koladata/functor/call.h"
 #include "koladata/functor/functor.h"
-#include "koladata/functor/parallel/transform_config.pb.h"
 #include "koladata/functor/parallel/transform_config.h"
+#include "koladata/functor/parallel/transform_config.pb.h"
 #include "koladata/functor/signature_utils.h"
 #include "koladata/functor_storage.h"
 #include "koladata/internal/data_item.h"
 #include "koladata/internal/dtype.h"
 #include "koladata/object_factories.h"
+#include "koladata/operators/core.h"
 #include "koladata/operators/slices.h"
 #include "koladata/signature.h"
 #include "koladata/signature_storage.h"
@@ -278,24 +280,15 @@ class InnerTransformManager {
     }
     used_var_names_.insert(result.var_name);
     ASSIGN_OR_RETURN(DataSlice var, functor_.GetAttr(var_name));
-    if (!var.item().holds_value<arolla::expr::ExprQuote>()) {
-      ASSIGN_OR_RETURN(
-          auto transformed_var,
-          TransformEager(arolla::TypedValue::FromValue(std::move(var))));
-      ASSIGN_OR_RETURN(auto transformed_var_slice,
-                       std::move(transformed_var).As<DataSlice>());
-      new_vars_.emplace_back(result.var_name, std::move(transformed_var_slice));
-      results_.emplace(var_name, result);
-      return result;
+    std::optional<arolla::TypedValue> var_qvalue;
+    if (var.item().holds_value<arolla::expr::ExprQuote>()) {
+      ASSIGN_OR_RETURN(var_qvalue, TryComputeLiteral(var));
+    } else {
+      var_qvalue = arolla::TypedValue::FromValue(std::move(var));
     }
-    ASSIGN_OR_RETURN(auto var_expr,
-                     var.item().value<arolla::expr::ExprQuote>().expr());
-    if (expr::IsLiteral(var_expr)) {
-      if (!var_expr->qvalue()) {
-        return absl::InternalError("Literal does not have a value");
-      }
+    if (var_qvalue.has_value()) {
       ASSIGN_OR_RETURN(auto transformed_var,
-                       TransformEager(*var_expr->qvalue()));
+                       TransformEager(*std::move(var_qvalue)));
       auto new_var_expr = arolla::expr::Literal(std::move(transformed_var));
       auto new_var = DataSlice::CreatePrimitive(
           arolla::expr::ExprQuote{std::move(new_var_expr)});
@@ -305,18 +298,17 @@ class InnerTransformManager {
     }
     if (!config_->allow_runtime_transforms()) {
       return absl::InvalidArgumentError(absl::StrCat(
-          "The parallel transformation requires that all sub-functors being"
-          " called are specified as literals, instead of being computed"
-          " dynamically, so that we can transform them recursively. In case you"
-          " are calling a sub-functor that is computed dynamically, but do not"
-          " need to recursively transform it (evaluating this sub-functor"
-          " single-threaded is fine), you can use"
-          " kd.functor.call_fn_normally_when_parallel and its variants to call"
-          " the sub-functor. In case you do need to evaluate a dynamically"
-          " computed sub-functor in a parallel fashion, you can pass"
-          " allow_runtime_transforms=True to kd.parallel.transform, but this"
-          " will be slower and should not be used in production."
-          "\nThe offending sub-functor ",
+          "The parallel transformation requires that all sub-functors being "
+          "called are fully defined at transformation time, so that we can "
+          "transform them recursively. In case you are calling a sub-functor "
+          "that is computed dynamically, but do not need to recursively "
+          "transform it (evaluating this sub-functor single-threaded is fine), "
+          "you can use kd.functor.call_fn_normally_when_parallel and its "
+          "variants to call the sub-functor. In case you do need to evaluate a "
+          "dynamically computed sub-functor in a parallel fashion, you can "
+          "pass allow_runtime_transforms=True to kd.parallel.transform, but "
+          "this will be slower and should not be used in production.\nThe "
+          "offending sub-functor ",
           var_name, ": ", arolla::Repr(var),
           " . The whole functor: ", arolla::Repr(functor_)));
     }
@@ -355,6 +347,38 @@ class InnerTransformManager {
          std::move(sub_functor)});
   }
 
+  // If expr can be literal folded (with respect of its dependencies in
+  // functor_), returns its value. Otherwise returns std::nullopt.
+  absl::StatusOr<std::optional<arolla::TypedValue>> TryComputeLiteral(
+      const DataSlice& expr) const {
+    // TODO: b/477578091 - we are changing the functor in order to reuse
+    // CallFunctorWithCompilationCache. It will be more efficient if we just
+    // evaluate all the variables needed for `expr` explicitly.
+    ASSIGN_OR_RETURN(
+        DataBagPtr returns_db,
+        ops::Attr(functor_,
+                  DataSlice::CreatePrimitive(arolla::Text(kReturnsAttrName)),
+                  expr,
+                  /*override_schema=*/DataSlice::CreatePrimitive(false)));
+    ASSIGN_OR_RETURN(
+        DataBagPtr signature_db,
+        ops::Attr(functor_,
+                  DataSlice::CreatePrimitive(arolla::Text(kSignatureAttrName)),
+                  KodaEmptySignature(),
+                  /*override_schema=*/DataSlice::CreatePrimitive(false)));
+    ASSIGN_OR_RETURN(
+        auto db, DataBag::ImmutableEmptyWithFallbacks({std::move(returns_db),
+                                                       std::move(signature_db),
+                                                       functor_.GetBag()}));
+    auto new_functor = functor_.WithBag(std::move(db));
+    ASSIGN_OR_RETURN(
+        auto res, functor::CallFunctorWithCompilationCache(new_functor, {}, {}),
+        // We consider all functor evaluation errors as caused by being not
+        // actually a literal.
+        std::nullopt);
+    return res;
+  }
+
   const ParallelTransformConfigPtr absl_nonnull& config_;
   const DataSlice& functor_;
   const expr::InputContainer& variable_container_;

koladata/functor/parallel/transform_config.proto

@@ -88,8 +88,8 @@ message ParallelTransformConfigProto {
 
   repeated OperatorReplacement operator_replacements = 1;
 
-  // If true, the sub-functors that are not static but computed will
-  // be transformed to the parallel form at runtime. If false, we will
-  // raise an error in that case.
+  // If true, the sub-functors that are not fully defined at transform time
+  // (i.e. that depend on functor inputs) will be transformed to the parallel
+  // form at runtime. If false, we will raise an error in that case.
   bool allow_runtime_transforms = 2;
 }

koladata/functor/signature_utils.cc

@@ -280,6 +280,12 @@ const Signature& CppArgsKwargsSignature() {
           .value());
   return *val;
 }
+
+const Signature& CppEmptySignature() {
+  static absl::NoDestructor<Signature> val(Signature::Create({}).value());
+  return *val;
+}
+
 }  // namespace
 
 const DataSlice& KodaArgsKwargsSignature() {
@@ -288,4 +294,10 @@ const DataSlice& KodaArgsKwargsSignature() {
   return *val;
 }
 
+const DataSlice& KodaEmptySignature() {
+  static absl::NoDestructor<DataSlice> val{
+      CppSignatureToKodaSignature(CppEmptySignature()).value()};
+  return *val;
+}
+
 }  // namespace koladata::functor

koladata/functor/signature_utils.h

@@ -57,6 +57,8 @@ const DataSlice& NoDefaultValueMarker();
 // Returns functor signature for the *args and **kwargs parameters.
 const DataSlice& KodaArgsKwargsSignature();
 
+// Returns a functor signature with no parameters.
+const DataSlice& KodaEmptySignature();
 
 // Converts a C++ Signature object to a Koda DataItem storing the signature.
 // The returned DataItem will have a new DataBag created to store the triples.

koladata/functor/signature_utils_test.cc

@@ -31,6 +31,7 @@
 #include "koladata/internal/dtype.h"
 #include "koladata/internal/object_id.h"
 #include "koladata/signature.h"
+#include "koladata/signature_storage.h"
 #include "koladata/test_utils.h"
 #include "koladata/testing/matchers.h"
 
@@ -42,6 +43,7 @@ using ::absl_testing::IsOkAndHolds;
 using ::absl_testing::StatusIs;
 using ::koladata::testing::IsEquivalentTo;
 using ::testing::HasSubstr;
+using ::testing::IsEmpty;
 
 TEST(BindArgumentsTest, Basic) {
   Signature::Parameter p1 = {
@@ -380,6 +382,12 @@ TEST(CppSignatureToKodaSignatureTest, Basic) {
   EXPECT_THAT(arolla::Repr(koda_signature), HasSubstr("no_default_value"));
 }
 
+TEST(SignatureUtilsTest, KodaEmptySignature) {
+  ASSERT_OK_AND_ASSIGN(auto signature,
+                       KodaSignatureToCppSignature(KodaEmptySignature()));
+  EXPECT_THAT(signature.parameters(), IsEmpty());
+}
+
 }  // namespace
 
 }  // namespace koladata::functor

py/koladata/functions/parallel.py

@@ -177,13 +177,14 @@ def transform(
 
   Args:
     fn: The functor to transform.
-    allow_runtime_transforms: Whether to allow sub-functors to be not literals,
-      but computed expressions, which will therefore have to be transformed at
-      runtime. This can be slow.
+    allow_runtime_transforms: Whether to allow sub-functors to be not fully
+      defined at transform time (i.e. to depend on functor inputs), which will
+      therefore have to be transformed at runtime. This can be slow.
 
   Returns:
     The transformed functor.
   """
+
   fn = py_boxing.as_qvalue(fn)
   config = koda_internal_parallel.create_transform_config(
       koda_internal_parallel.get_default_transform_config_src().with_attrs(

py/koladata/operators/functor.py

@@ -373,7 +373,13 @@ def _switch_repr(
       # unlike str(k).
       k_repr = repr(k.to_py()) if k != SWITCH_DEFAULT else 'kd.SWITCH_DEFAULT'
       case_reprs[k_repr] = str(f)
-    cases_repr = ',\n'.join(f'{k}: {v}' for k, v in case_reprs.items())
+    cases_repr = (
+        '{\n'
+        + textwrap.indent(
+            '\n'.join(f'{k}: {v},' for k, v in case_reprs.items()), '    '
+        )
+        + '\n}'
+    )
   else:
     cases_repr = f'kd.dict({tokens[case_keys].text}, {tokens[case_fns].text})'
 
@@ -382,9 +388,11 @@ def _switch_repr(
   res = op_repr.ReprToken()
   res.text = (
       f'{node.op.display_name}(\n'
-      f'{textwrap.indent(key_repr, "    ")},\n'
-      f'{textwrap.indent(cases_repr, "    ")},\n'
-      f'{textwrap.indent(args_kwargs, "    ")})'
+      + textwrap.indent(
+          ',\n'.join(filter(None, (key_repr, cases_repr, args_kwargs))),
+          '    ',
+      )
+      + ')'
   )
   return res
 

py/koladata/operators/koda_internal_parallel.py

@@ -2246,6 +2246,69 @@ def _parallel_if(
           cond,
           transformed_yes_fn,
           transformed_no_fn,
+          # async_eval waits until all the future args are ready, but here we
+          # hide the futures inside a tuple in order to propagate to the
+          # downstream operators.
+          (args, return_type_as, kwargs),
+          optools.unified_non_deterministic_arg(),
+      )
+  )
+
+
+@optools.as_lambda_operator(
+    'koda_internal.parallel._parallel_switch_impl',
+)
+def _parallel_switch_impl(
+    executor, key, case_keys, transformed_case_fns, parallel_args
+):
+  """Implementation helper for _parallel_switch."""
+  args = parallel_args[0]
+  return_type_as = parallel_args[1]
+  kwargs = parallel_args[2]
+  return arolla.abc.bind_op(  # pytype: disable=wrong-arg-types
+      functor.switch,
+      key,
+      case_keys,
+      transformed_case_fns,
+      return_type_as=return_type_as,
+      args=M.core.concat_tuples(M.core.make_tuple(executor), args),
+      kwargs=kwargs,
+      non_deterministic=optools.unified_non_deterministic_arg(),
+  )
+
+
+@optools.add_to_registry(via_cc_operator_package=True)
+@optools.as_lambda_operator(
+    'koda_internal.parallel._parallel_switch',
+    qtype_constraints=[
+        qtype_utils.expect_executor(P.executor),
+        qtype_utils.expect_future(P.key),
+        qtype_utils.expect_data_slice(P.transformed_case_fns),
+        qtype_utils.expect_tuple(P.args),
+        qtype_utils.expect_namedtuple(P.kwargs),
+    ],
+)
+def _parallel_switch(
+    executor,
+    key,
+    case_keys,
+    transformed_case_fns,
+    return_type_as,
+    args,
+    kwargs,
+):
+  """The parallel version of kd.switch."""
+  return unwrap_future_to_parallel(
+      async_eval(
+          executor,
+          _parallel_switch_impl,
+          executor,
+          key,
+          case_keys,
+          transformed_case_fns,
+          # async_eval waits until all the future args are ready, but here we
+          # hide the futures inside a tuple in order to propagate to the
+          # downstream operators.
           (args, return_type_as, kwargs),
           optools.unified_non_deterministic_arg(),
       )
@@ -3323,6 +3386,15 @@ def _parallel_with_assertion(executor, x, condition, message_or_fn, args):
       functor_argument_indices: 2
     }
   }
+  operator_replacements {
+    from_op: "kd.functor.switch"
+    to_op: "koda_internal.parallel._parallel_switch"
+    argument_transformation {
+      arguments: EXECUTOR
+      arguments: ORIGINAL_ARGUMENTS
+      functor_argument_indices: 2
+    }
+  }
   operator_replacements {
     from_op: "kd.iterables.make"
     to_op: "koda_internal.parallel._parallel_stream_make"