PL-406 Allow optional args

Summary:
aadding some changes for optional args

Adding range and kwargs support

Test Plan: n/a

Reviewers: michelle, zasgar, #engineering, oazizi

Reviewed By: #engineering, oazizi

Subscribers: oazizi

Differential Revision: https://phab.corp.pixielabs.ai/D434

GitOrigin-RevId: e6acc2e

.vscode/c_cpp_properties.json

@@ -3,7 +3,7 @@
         {
             "name": "Mac",
             "includePath": [
-                "${workspaceFolder}"
+                "${workspaceFolder}/**"
             ],
             "defines": [],
             "macFrameworkPath": [

src/carnot/compiler/ast_visitor.cc

@@ -122,32 +122,57 @@ StatusOr<std::string> ASTWalker::GetFuncName(const pypa::AstCallPtr& node) {
 StatusOr<ArgMap> ASTWalker::ProcessArgs(const pypa::AstCallPtr& call_ast,
                                         const std::vector<std::string>& expected_args,
                                         bool kwargs_only) {
+  return ProcessArgs(call_ast, expected_args, kwargs_only, {{}});
+}
+StatusOr<ArgMap> ASTWalker::ProcessArgs(
+    const pypa::AstCallPtr& call_ast, const std::vector<std::string>& expected_args,
+    bool kwargs_only, const std::unordered_map<std::string, pypa::AstPtr> default_args) {
   auto arg_ast = call_ast->arglist;
   if (!kwargs_only) {
+    // TODO(philkuz) (PL-406) add support for kwargs
     return error::Unimplemented("Only supporting kwargs for now.");
   }
   ArgMap arg_map;
   // Set to keep track of args that are not yet found.
-  std::unordered_set<std::string> missing_args;
-  missing_args.insert(expected_args.begin(), expected_args.end());
+  std::unordered_set<std::string> missing_or_default_args;
+  missing_or_default_args.insert(expected_args.begin(), expected_args.end());
 
+  std::vector<Status> errors;
   // Iterate through the keywords
   for (auto& k : arg_ast.keywords) {
     pypa::AstKeywordPtr kw_ptr = PYPA_PTR_CAST(Keyword, k);
     std::string key = GetNameID(kw_ptr->name);
-    if (missing_args.find(key) == missing_args.end()) {
-      return CreateAstError(absl::Substitute("Keyword '$0' not expected in function.", key),
-                            call_ast);
+    if (missing_or_default_args.find(key) == missing_or_default_args.end()) {
+      // TODO(philkuz) make a string output version of CreateAstError.
+      errors.push_back(CreateAstError(
+          absl::Substitute("Keyword '$0' not expected in function.", key), call_ast));
+      continue;
     }
-    missing_args.erase(missing_args.find(key));
+    missing_or_default_args.erase(missing_or_default_args.find(key));
     PL_ASSIGN_OR_RETURN(IRNode * value, ProcessData(kw_ptr->value));
     arg_map[key] = value;
   }
-  if (!missing_args.empty()) {
-    return CreateAstError(
-        absl::Substitute("Didn't find keywords '[$0]' in function. Please add them.",
-                         absl::StrJoin(missing_args, ",")),
-        call_ast);
+
+  for (const auto& ma : missing_or_default_args) {
+    auto find_ma = default_args.find(ma);
+    if (find_ma == default_args.end()) {
+      // TODO(philkuz) look for places where ast error might exit prematurely in other parts of the
+      // code.
+      errors.push_back(CreateAstError(
+          absl::Substitute("You must set '$0' directly. No default value found.", ma), call_ast));
+      continue;
+    } else {
+      auto default_value = find_ma->second;
+      PL_ASSIGN_OR_RETURN(IRNode * value, ProcessData(default_value));
+      arg_map[ma] = value;
+    }
+  }
+  if (errors.size() != 0) {
+    std::vector<std::string> msg;
+    for (auto const& e : errors) {
+      msg.push_back(e.ToString());
+    }
+    return error::InvalidArgument(absl::StrJoin(msg, "\n"));
   }
 
   return arg_map;
@@ -225,6 +250,7 @@ StatusOr<IRNode*> ASTWalker::ProcessRangeOp(const pypa::AstCallPtr& node) {
                           node->function);
   }
   PL_ASSIGN_OR_RETURN(RangeIR * ir_node, ir_graph_->MakeNode<RangeIR>());
+
   PL_ASSIGN_OR_RETURN(ArgMap args, ProcessArgs(node, {"time"}, true));
   PL_ASSIGN_OR_RETURN(IRNode * call_result,
                       ProcessAttribute(PYPA_PTR_CAST(Attribute, node->function)));

src/carnot/compiler/ast_visitor.h

@@ -127,9 +127,12 @@ class ASTWalker {
    * @param kwargs_only Whether to only allow keyword args.
    * @return StatusOr<ArgMap>
    */
-  StatusOr<ArgMap> ProcessArgs(const pypa::AstCallPtr& call_ast,
-                               const std::vector<std::string>& expected_args, bool kwargs_only);
+  StatusOr<ArgMap> ProcessArgs(const pypa::AstCallPtr& arg_ast,
+                               const std::vector<std::string>& expected_args, bool kwargs_only,
+                               const std::unordered_map<std::string, pypa::AstPtr> default_args);
 
+  StatusOr<ArgMap> ProcessArgs(const pypa::AstCallPtr& arg_ast,
+                               const std::vector<std::string>& expected_args, bool kwargs_only);
   /**
    * @brief ProcessExprStmtNode handles full lines that are expression statements.
    * ie in the following lines

src/carnot/compiler/ast_visitor_test.cc

@@ -234,6 +234,30 @@ TEST(TimeValueTest, basic) {
                     "\n");
   EXPECT_FALSE(ParseQuery(bad_attribute_value).ok());
 }
+
+TEST(OptionalArgs, DISABLED_group_by_all) {
+  // TODO(philkuz) later diff impl this.
+  std::string agg_query =
+      absl::StrJoin({"queryDF = From(table='cpu', select=['cpu0', 'cpu1'])",
+                     "queryDF.Agg(fn=lambda r : {'sum' : pl.sum(r.cpu0)}).Result(name='agg')"},
+                    "\n");
+  auto status = ParseQuery(agg_query);
+  VLOG(1) << status.ToString();
+  EXPECT_OK(status);
+}
+
+TEST(OptionalArgs, DISABLED_map_copy_relation) {
+  // TODO(philkuz) later diff impl this.
+  // TODO(philkuz) make a relation handler test that confirms the relation is actually copied.
+
+  std::string map_query = absl::StrJoin({"queryDF = From(table='cpu', select=['cpu0', 'cpu1'])",
+                                         "queryDF.Map(fn=lambda r : {'sum' : r.cpu0 + r.cpu1}, "
+                                         "copy_source_cols=True).Result(name='map')"},
+                                        "\n");
+  auto status = ParseQuery(map_query);
+  VLOG(1) << status.ToString();
+  EXPECT_OK(status);
+}
 }  // namespace compiler
 }  // namespace carnot
 }  // namespace pl

src/carnot/compiler/ir_test_utils.h

@@ -12,11 +12,6 @@ void VerifyMemorySource(IRNode* node) {
   EXPECT_EQ(mem_node->table_node()->type(), IRNodeType::StringType);
   EXPECT_EQ(mem_node->select()->type(), IRNodeType::ListType);
   EXPECT_TRUE(mem_node->HasLogicalRepr());
-  // TODO(oazizi or zasgar) - what do you think about doing checks using dependencies of?
-  // I'm not sure what we can accomplish, was wondering if you thoguth it might bet better than the
-  // tests above. They likely overlap in terms of coverage.
-  //
-  // node->graph_ptr()->dag().DependenciesOf(node->id());
 }
 
 void VerifyRange(IRNode* node) {