Change BindArguments and CallFunctorWithCompilationCache to use args and knames

apronchenkov · copybara-github · commit 2d38cd30ceff · 2025-01-28T07:01:12.000-08:00
This change simplifies the `kd.functor.map` implementation (and removes the need to copy argument names in `kd.functor.call`).

PiperOrigin-RevId: 720559898
Change-Id: Ie8a49c1cd66681828180514b89d5f8c12bb3bd89
diff --git a/koladata/functor/call.cc b/koladata/functor/call.cc
@@ -107,7 +107,7 @@ absl::StatusOr<std::vector<std::string>> GetVariableEvaluationOrder(
 
 absl::StatusOr<arolla::TypedValue> CallFunctorWithCompilationCache(
     const DataSlice& functor, absl::Span<const arolla::TypedRef> args,
-    absl::Span<const std::pair<std::string, arolla::TypedRef>> kwargs,
+    absl::Span<const std::string> kwnames,
     const expr::EvalOptions& eval_options) {
   ASSIGN_OR_RETURN(bool is_functor, IsFunctor(functor));
   if (!is_functor) {
@@ -117,7 +117,7 @@ absl::StatusOr<arolla::TypedValue> CallFunctorWithCompilationCache(
   ASSIGN_OR_RETURN(auto signature_item, functor.GetAttr(kSignatureAttrName));
   ASSIGN_OR_RETURN(auto signature, KodaSignatureToCppSignature(signature_item));
   ASSIGN_OR_RETURN(auto bound_arguments,
-                   BindArguments(signature, args, kwargs));
+                   BindArguments(signature, args, kwnames));
   ASSIGN_OR_RETURN(auto variable_evaluation_order,
                    GetVariableEvaluationOrder(functor));
   if (variable_evaluation_order.empty() ||
diff --git a/koladata/functor/call.h b/koladata/functor/call.h
@@ -30,20 +30,22 @@ namespace koladata::functor {
 // Calls the given functor with the provided arguments and keyword arguments.
 // The functor would typically be created by the CreateFunctor method,
 // and consists of a returns expression, a signature, and a set of variables.
-// The passed args and kwargs will be bound to the parameters of the signature
-// to produce the values for the inputs (I.foo) in the provided expression and
-// in the variable expressions. The expressions can also refer to the variables
-// via V.foo, in which case the variable expression will be evaluated before
-// evaluating the expression that refers to it. In case of a cycle in variables,
-// an error will be returned.
 //
-// The `eval_options` parameter provides a default buffer factory
-// (typically either the default allocator or an arena allocator)
-// and a cancellation checker, which allows the computation to be
-// interrupted midway if needed.
+// `args` must contain values for positional arguments followed by values for
+// keyword arguments; `kwnames` must contain the names of the keyword arguments,
+// so `kwnames` corresponds to a suffix of `args`. The passed arguments will be
+// bound to the parameters of the signature to produce the values for the inputs
+// (I.foo) in the provided expression and in the variable expressions. The
+// expressions can also refer to the variables via V.foo, in which case the
+// variable expression will be evaluated before evaluating the expression that
+// refers to it. In case of a cycle in variables, an error will be returned.
+//
+// `eval_options` parameter provides a default buffer factory (typically either
+// the default allocator or an arena allocator) and a cancellation context,
+// which allows the computation to be interrupted midway if needed.
 absl::StatusOr<arolla::TypedValue> CallFunctorWithCompilationCache(
     const DataSlice& functor, absl::Span<const arolla::TypedRef> args,
-    absl::Span<const std::pair<std::string, arolla::TypedRef>> kwargs,
+    absl::Span<const std::string> kwnames,
     const expr::EvalOptions& eval_options);
 
 }  // namespace koladata::functor
diff --git a/koladata/functor/call_operator.cc b/koladata/functor/call_operator.cc
@@ -14,6 +14,7 @@
 //
 #include "koladata/functor/call_operator.h"
 
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <utility>
@@ -62,25 +63,21 @@ class CallOperator : public arolla::QExprOperator {
           const auto& fn_data_slice = frame.Get(fn_slot);
 
           std::vector<arolla::TypedRef> arg_refs;
-          arg_refs.reserve(args_slot.SubSlotCount());
-          for (int i = 0; i < args_slot.SubSlotCount(); ++i) {
+          arg_refs.reserve(args_slot.SubSlotCount() +
+                           kwargs_slot.SubSlotCount());
+          for (int64_t i = 0; i < args_slot.SubSlotCount(); ++i) {
             arg_refs.push_back(
                 arolla::TypedRef::FromSlot(args_slot.SubSlot(i), frame));
           }
-
-          auto kwargs_qtype = kwargs_slot.GetType();
-          auto kwarg_names = arolla::GetFieldNames(kwargs_qtype);
-          std::vector<std::pair<std::string, arolla::TypedRef>> kwarg_refs;
-          kwarg_refs.reserve(kwargs_slot.SubSlotCount());
-          for (int i = 0; i < kwargs_slot.SubSlotCount(); ++i) {
-            kwarg_refs.push_back(
-                {kwarg_names[i],
-                 arolla::TypedRef::FromSlot(kwargs_slot.SubSlot(i), frame)});
+          for (int64_t i = 0; i < kwargs_slot.SubSlotCount(); ++i) {
+            arg_refs.push_back(
+                arolla::TypedRef::FromSlot(kwargs_slot.SubSlot(i), frame));
           }
+          auto kwnames = arolla::GetFieldNames(kwargs_slot.GetType());
           ASSIGN_OR_RETURN(
               auto result,
-              functor::CallFunctorWithCompilationCache(
-                  fn_data_slice, arg_refs, kwarg_refs, ctx->options()),
+              functor::CallFunctorWithCompilationCache(fn_data_slice, arg_refs,
+                                                       kwnames, ctx->options()),
               ctx->set_status(std::move(_)));
           if (result.GetType() != output_slot.GetType()) {
             ctx->set_status(absl::InvalidArgumentError(absl::StrFormat(
diff --git a/koladata/functor/call_test.cc b/koladata/functor/call_test.cc
@@ -116,26 +116,20 @@ TEST(CallTest, VariableRhombus) {
       arolla::TypedValue::FromValue(3),
       arolla::TypedValue::FromValue(4),
   };
-  ASSERT_OK_AND_ASSIGN(auto result, CallFunctorWithCompilationCache(
-                                        fn,
-                                        /*args=*/
-                                        {
-                                            inputs[0].AsRef(),
-                                            inputs[1].AsRef(),
-                                            inputs[2].AsRef(),
-                                        },
-                                        /*kwargs=*/{}, /*eval_options=*/{}));
+  ASSERT_OK_AND_ASSIGN(
+      auto result,
+      CallFunctorWithCompilationCache(
+          fn,
+          /*args=*/{inputs[0].AsRef(), inputs[1].AsRef(), inputs[2].AsRef()},
+          /*kwnames=*/{}, /*eval_options=*/{}));
   EXPECT_THAT(result.As<int32_t>(), IsOkAndHolds(2 * ((3 + 5) + (3 + 4))));
 
   ASSERT_OK_AND_ASSIGN(
       result,
       CallFunctorWithCompilationCache(
           fn,
-          /*args=*/
-          {
-              inputs[0].AsRef(),
-          },
-          /*kwargs=*/{{"c", inputs[1].AsRef()}, {"b", inputs[2].AsRef()}},
+          /*args=*/{inputs[0].AsRef(), inputs[1].AsRef(), inputs[2].AsRef()},
+          /*kwnames=*/{"c", "b"},
           /*eval_options=*/{}));
   EXPECT_THAT(result.As<int32_t>(), IsOkAndHolds(2 * ((4 + 5) + (4 + 3))));
 }
@@ -150,7 +144,7 @@ TEST(CallTest, VariableCycle) {
   ASSERT_OK_AND_ASSIGN(auto fn,
                        CreateFunctor(returns_expr, koda_signature,
                                      {{"a", var_a_expr}, {"b", var_b_expr}}));
-  EXPECT_THAT(CallFunctorWithCompilationCache(fn, /*args=*/{}, /*kwargs=*/{},
+  EXPECT_THAT(CallFunctorWithCompilationCache(fn, /*args=*/{}, /*kwnames=*/{},
                                               /*eval_options=*/{}),
               StatusIs(absl::StatusCode::kInvalidArgument,
                        "variable [a] has a dependency cycle"));
@@ -165,7 +159,7 @@ TEST(CallTest, JustLiteral) {
                        CreateFunctor(returns_expr, koda_signature, {}));
   ASSERT_OK_AND_ASSIGN(
       auto result, CallFunctorWithCompilationCache(
-                       fn, /*args=*/{}, /*kwargs=*/{}, /*eval_options=*/{}));
+                       fn, /*args=*/{}, /*kwnames=*/{}, /*eval_options=*/{}));
   EXPECT_THAT(result.As<int32_t>(), IsOkAndHolds(57));
 }
 
@@ -177,7 +171,7 @@ TEST(CallTest, MustBeScalar) {
   ASSERT_OK_AND_ASSIGN(auto fn,
                        CreateFunctor(returns_expr, koda_signature, {}));
   ASSERT_OK_AND_ASSIGN(fn, fn.Reshape(DataSlice::JaggedShape::FlatFromSize(1)));
-  EXPECT_THAT(CallFunctorWithCompilationCache(fn, /*args=*/{}, /*kwargs=*/{},
+  EXPECT_THAT(CallFunctorWithCompilationCache(fn, /*args=*/{}, /*kwnames=*/{},
                                               /*eval_options=*/{}),
               StatusIs(absl::StatusCode::kInvalidArgument,
                        "the first argument of kd.call must be a functor"));
@@ -192,7 +186,7 @@ TEST(CallTest, NoBag) {
                        CreateFunctor(returns_expr, koda_signature, {}));
   EXPECT_THAT(
       CallFunctorWithCompilationCache(fn.WithBag(nullptr), /*args=*/{},
-                                      /*kwargs=*/{}, /*eval_options=*/{}),
+                                      /*kwnames=*/{}, /*eval_options=*/{}),
       StatusIs(absl::StatusCode::kInvalidArgument,
                "the first argument of kd.call must be a functor"));
 }
@@ -209,7 +203,7 @@ TEST(CallTest, DataSliceVariable) {
       auto fn, CreateFunctor(returns_expr, koda_signature, {{"a", var_a}}));
   ASSERT_OK_AND_ASSIGN(auto result, CallFunctorWithCompilationCache(
                                         fn, /*args=*/{},
-                                        /*kwargs=*/{}, /*eval_options=*/{}));
+                                        /*kwnames=*/{}, /*eval_options=*/{}));
   EXPECT_THAT(result.As<DataSlice>(),
               IsOkAndHolds(IsEquivalentTo(var_a.WithBag(fn.GetBag()))));
 }
@@ -240,7 +234,7 @@ TEST(CallTest, EvalError) {
       CallFunctorWithCompilationCache(
           fn,
           /*args=*/{arolla::TypedRef::FromValue(input)},
-          /*kwargs=*/{}, /*eval_options=*/{}),
+          /*kwnames=*/{}, /*eval_options=*/{}),
       StatusIs(absl::StatusCode::kInvalidArgument,
                "expected numerics, got x: DATA_SLICE; while calling math.add "
                "with args {annotation.qtype(L['I.a'], DATA_SLICE), 57}; while "
@@ -277,7 +271,7 @@ TEST(CallTest, Cancellation) {
     expr::EvalOptions eval_options{.cancellation_context = cancel_ctx.get()};
     EXPECT_THAT(CallFunctorWithCompilationCache(
                     fn, /*args=*/{arolla::TypedRef::FromValue(1)},
-                    /*kwargs=*/{}, eval_options),
+                    /*kwnames=*/{}, eval_options),
                 StatusIs(absl::StatusCode::kCancelled));
   }
   {
@@ -292,7 +286,7 @@ TEST(CallTest, Cancellation) {
     ASSERT_OK_AND_ASSIGN(auto result,
                          CallFunctorWithCompilationCache(
                              fn, /*args=*/{arolla::TypedRef::FromValue(1)},
-                             /*kwargs=*/{}, eval_options));
+                             /*kwnames=*/{}, eval_options));
     EXPECT_THAT(result.As<int>(), IsOkAndHolds(3));
   }
 }
diff --git a/koladata/functor/map.cc b/koladata/functor/map.cc
@@ -48,49 +48,35 @@ absl::StatusOr<DataSlice> MapFunctorWithCompilationCache(
   ASSIGN_OR_RETURN(auto aligned_args, shape::Align(std::move(args)));
   DataSlice aligned_functors = std::move(aligned_args.back());
   aligned_args.pop_back();
-  int64_t num_positional = aligned_args.size() - kwnames.size();
   // Pre-allocate the vectors to avoid reallocations in the loop.
   std::vector<arolla::TypedRef> arg_refs;
-  arg_refs.reserve(num_positional);
-  std::vector<std::pair<std::string, arolla::TypedRef>> kwarg_refs;
-  kwarg_refs.reserve(kwnames.size());
+  arg_refs.reserve(aligned_args.size());
   std::vector<DataSlice> result_slices;
   ASSIGN_OR_RETURN(DataSlice missing,
                    DataSlice::Create(internal::DataItem(std::nullopt),
                                      internal::DataItem(schema::kNone)));
 
-  auto call_on_items = [&kwnames, &missing, &arg_refs, &kwarg_refs,
-                        &num_positional, &include_missing,
+  auto call_on_items = [&kwnames, &missing, &arg_refs, &include_missing,
                         &eval_options](const DataSlice& functor,
                                        const std::vector<DataSlice>& arg_slices)
       -> absl::StatusOr<DataSlice> {
     if (!functor.item().has_value()) {
       return missing;
     }
     arg_refs.clear();
-    for (int64_t i = 0; i < num_positional; ++i) {
-      const auto& value = arg_slices[i];
+    for (const auto& value : arg_slices) {
       DCHECK(value.is_item());
       if (!include_missing && !value.item().has_value()) {
         return missing;
       }
       arg_refs.push_back(arolla::TypedRef::FromValue(value));
     }
-    kwarg_refs.clear();
-    for (int64_t i = 0; i < kwnames.size(); ++i) {
-      const auto& value = arg_slices[num_positional + i];
-      DCHECK(value.is_item());
-      if (!include_missing && !value.item().has_value()) {
-        return missing;
-      }
-      kwarg_refs.emplace_back(kwnames[i], arolla::TypedRef::FromValue(value));
-    }
     // We can improve the performance a lot if "functor" is the same for many
     // items, as a lot of the work inside CallFunctorWithCompilationCache could
     // be reused between items then.
     ASSIGN_OR_RETURN(auto result,
-                     CallFunctorWithCompilationCache(functor, arg_refs,
-                                                     kwarg_refs, eval_options));
+                     CallFunctorWithCompilationCache(functor, arg_refs, kwnames,
+                                                     eval_options));
     if (result.GetType() != arolla::GetQType<DataSlice>()) {
       return absl::InvalidArgumentError(absl::StrFormat(
           "the functor is expected to be evaluated to a DataItem"
diff --git a/koladata/functor/map.h b/koladata/functor/map.h
@@ -28,15 +28,17 @@ namespace koladata::functor {
 // Calls the given functors pointwise on items from the provided arguments and
 // keyword arguments, expects them to return DataItems, and stacks
 // those into a single DataSlice.
-// `args` must contain first all positional arguments, then all keyword
-// arguments, so `kwnames` corresponds to a suffix of `args`.
+//
+// `args` must contain values for positional arguments followed by values for
+// keyword arguments; `kwnames` must contain the names of the keyword arguments,
+// so `kwnames` corresponds to a suffix of `args`.
+//
 // `include_missing` controls whether to call the functors on missing items of
 // `args` and `kwargs`.
 //
-// The `eval_options` parameter provides a default buffer factory
-// (typically either the default allocator or an arena allocator)
-// and a cancellation checker, which allows the computation to be
-// interrupted midway if needed.
+// `eval_options` parameter provides a default buffer factory (typically either
+// the default allocator or an arena allocator) and a cancellation context,
+// which allows the computation to be interrupted midway if needed.
 absl::StatusOr<DataSlice> MapFunctorWithCompilationCache(
     const DataSlice& functors, std::vector<DataSlice> args,
     absl::Span<const std::string> kwnames, bool include_missing,
diff --git a/koladata/functor/signature.cc b/koladata/functor/signature.cc
@@ -107,7 +107,12 @@ Signature::Signature(absl::Span<const Parameter> parameters)
 
 absl::StatusOr<std::vector<arolla::TypedValue>> BindArguments(
     const Signature& signature, absl::Span<const arolla::TypedRef> args,
-    absl::Span<const std::pair<std::string, arolla::TypedRef>> kwargs) {
+    absl::Span<const std::string> kwnames) {
+  if (args.size() < kwnames.size()) {
+    return absl::InvalidArgumentError("args.size < kwnames.size()");
+  }
+  const size_t kwargs_offset = args.size() - kwnames.size();
+
   const auto& parameters = signature.parameters();
   const auto& keyword_parameter_index = signature.keyword_parameter_index();
   std::vector<arolla::TypedValue> bound_arguments(
@@ -119,7 +124,7 @@ absl::StatusOr<std::vector<arolla::TypedValue>> BindArguments(
   std::vector<arolla::TypedRef> unknown_kwarg_values;
 
   // Process positional arguments.
-  for (size_t i = 0; i < args.size(); ++i) {
+  for (size_t i = 0; i < kwargs_offset; ++i) {
     if (i >= parameters.size() ||
         (parameters[i].kind != Signature::Parameter::Kind::kPositionalOnly &&
          parameters[i].kind !=
@@ -131,15 +136,17 @@ absl::StatusOr<std::vector<arolla::TypedValue>> BindArguments(
   }
 
   // Process keyword arguments.
-  for (const auto& [name, value] : kwargs) {
+  for (size_t i = kwargs_offset; i < args.size(); ++i) {
+    const auto& name = kwnames[i - kwargs_offset];
+    const auto& value = args[i];
     auto it = keyword_parameter_index.find(name);
     if (it == keyword_parameter_index.end()) {
       unknown_kwarg_names.push_back(name);
       unknown_kwarg_values.push_back(value);
     } else {
       if (bound_arguments[it->second].GetType() != arolla::GetNothingQType()) {
-        return absl::InvalidArgumentError(absl::StrFormat(
-            "parameter [%s] specified twice", parameters[it->second].name));
+        return absl::InvalidArgumentError(
+            absl::StrFormat("parameter [%s] specified twice", name));
       }
       bound_arguments[it->second] = arolla::TypedValue(value);
     }
diff --git a/koladata/functor/signature.h b/koladata/functor/signature.h
@@ -18,7 +18,6 @@
 #include <cstddef>
 #include <optional>
 #include <string>
-#include <utility>
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
@@ -96,7 +95,7 @@ class Signature {
 // must own it.
 absl::StatusOr<std::vector<arolla::TypedValue>> BindArguments(
     const Signature& signature, absl::Span<const arolla::TypedRef> args,
-    absl::Span<const std::pair<std::string, arolla::TypedRef>> kwargs);
+    absl::Span<const std::string> kwnames);
 
 // Makes it possible to use absl::StrCat/absl::StrFormat (with %v) with kinds.
 template <typename Sink>
diff --git a/koladata/functor/signature_test.cc b/koladata/functor/signature_test.cc
