address comments

apache · Mar 5, 2024 · 465baf4 · 465baf4
1 parent 43cbcd8
commit 465baf4
Showing 1 changed file with 37 additions and 55 deletions.
diff --git a/cpp/velox/substrait/VeloxSubstraitSignature.cc b/cpp/velox/substrait/VeloxSubstraitSignature.cc
@@ -121,29 +121,20 @@ TypePtr VeloxSubstraitSignature::fromSubstraitSignature(const std::string& signa
     return str.size() >= prefix.size() && str.substr(0, prefix.size()) == prefix;
   };
 
-  if (startWith(signature, "dec")) {
-    // Decimal type name is in the format of dec<precision,scale>.
-    auto precisionStart = signature.find_first_of('<');
-    auto tokenIndex = signature.find_first_of(',');
-    auto scaleEnd = signature.find_first_of('>');
-    auto precision = stoi(signature.substr(precisionStart + 1, (tokenIndex - precisionStart - 1)));
-    auto scale = stoi(signature.substr(tokenIndex + 1, (scaleEnd - tokenIndex - 1)));
-    return DECIMAL(precision, scale);
-  }
-
-  if (startWith(signature, "struct")) {
-    // Struct type name is in the format of struct<T1,T2,...,Tn>.
-    auto structStart = signature.find_first_of('<');
-    auto structEnd = signature.find_last_of('>');
+  auto parseNestedTypeSignature = [&](const std::string& signature) -> std::vector<TypePtr> {
+    auto start = signature.find_first_of('<');
+    auto end = signature.find_last_of('>');
     VELOX_CHECK(
-        structEnd - structStart > 1, "Native validation failed due to: more information is needed to create RowType");
-    std::string childrenTypes = signature.substr(structStart + 1, structEnd - structStart - 1);
+        end - start > 1,
+        "Native validation failed due to: more information is needed to create nested type for {}",
+        signature);
+
+    std::string childrenTypes = signature.substr(start + 1, end - start - 1);
 
     // Split the types with delimiter.
     std::string delimiter = ",";
     std::size_t pos;
     std::vector<TypePtr> types;
-    std::vector<std::string> names;
     while ((pos = childrenTypes.find(delimiter)) != std::string::npos) {
       auto typeStr = childrenTypes.substr(0, pos);
       std::size_t endPos = pos;
@@ -160,16 +151,44 @@ TypePtr VeloxSubstraitSignature::fromSubstraitSignature(const std::string& signa
         }
       }
       types.emplace_back(fromSubstraitSignature(typeStr));
-      names.emplace_back("");
       childrenTypes.erase(0, endPos + delimiter.length());
     }
     if (childrenTypes.size() > 0 && !startWith(childrenTypes, ">")) {
       types.emplace_back(fromSubstraitSignature(childrenTypes));
+    }
+    return types;
+  };
+
+  if (startWith(signature, "dec")) {
+    // Decimal type name is in the format of dec<precision,scale>.
+    auto precisionStart = signature.find_first_of('<');
+    auto tokenIndex = signature.find_first_of(',');
+    auto scaleEnd = signature.find_first_of('>');
+    auto precision = stoi(signature.substr(precisionStart + 1, (tokenIndex - precisionStart - 1)));
+    auto scale = stoi(signature.substr(tokenIndex + 1, (scaleEnd - tokenIndex - 1)));
+    return DECIMAL(precision, scale);
+  }
+
+  if (startWith(signature, "struct")) {
+    // Struct type name is in the format of struct<T1,T2,...,Tn>.
+    auto types = parseNestedTypeSignature(signature);
+    std::vector<std::string> names;
+    names.resize(types.size());
+    for (int i = 0; i < types.size(); i++) {
       names.emplace_back("");
     }
     return std::make_shared<RowType>(std::move(names), std::move(types));
   }
 
+  if (startWith(signature, "map")) {
+    // Map type name is in the format of map<T1,T2>.
+    auto types = parseNestedTypeSignature(signature);
+    if (types.size() != 2) {
+      VELOX_UNSUPPORTED("Substrait type signature conversion to Velox type not supported for {}.", signature);
+    }
+    return MAP(std::move(types)[0], std::move(types)[1]);
+  }
+
   if (startWith(signature, "list")) {
     auto listStart = signature.find_first_of('<');
     auto listEnd = signature.find_last_of('>');
@@ -183,43 +202,6 @@ TypePtr VeloxSubstraitSignature::fromSubstraitSignature(const std::string& signa
     return ARRAY(elementType);
   }
 
-  if (startWith(signature, "map")) {
-    // Map type name is in the format of map<T1,T2>.
-    auto mapStart = signature.find_first_of('<');
-    auto mapEnd = signature.find_last_of('>');
-    VELOX_CHECK(
-        mapEnd - mapStart > 1,
-        "Native validation failed due to: more information is needed to create MapType: {}",
-        signature);
-    std::string childrenTypes = signature.substr(mapStart + 1, mapEnd - mapStart - 1);
-    std::string delimiter = ",";
-    std::size_t pos;
-    std::vector<TypePtr> types;
-    while ((pos = childrenTypes.find(delimiter)) != std::string::npos) {
-      auto typeStr = childrenTypes.substr(0, pos);
-      std::size_t endPos = pos;
-      if (startWith(typeStr, "dec") || startWith(typeStr, "struct") || startWith(typeStr, "map") ||
-          startWith(typeStr, "list")) {
-        endPos = childrenTypes.find(">") + 1;
-        if (endPos > pos) {
-          typeStr += childrenTypes.substr(pos, endPos - pos);
-        } else {
-          typeStr += childrenTypes.substr(pos);
-          endPos = childrenTypes.size();
-        }
-      }
-      types.emplace_back(fromSubstraitSignature(typeStr));
-      childrenTypes.erase(0, endPos + delimiter.length());
-    }
-    if (childrenTypes.size() > 0 && !startWith(childrenTypes, ">")) {
-      types.emplace_back(fromSubstraitSignature(childrenTypes));
-    }
-    if (types.size() != 2) {
-      VELOX_UNSUPPORTED("Substrait type signature conversion to Velox type not supported for {}.", signature);
-    }
-    return MAP(std::move(types)[0], std::move(types)[1]);
-  }
-
   VELOX_UNSUPPORTED("Substrait type signature conversion to Velox type not supported for {}.", signature);
 }