1. select from struct.

2. add round-trip test.(nested-struct)

Author: zhejiangxiaomai

velox/substrait/SubstraitToVeloxExpr.cpp

@@ -29,18 +29,37 @@ SubstraitVeloxExprConverter::toVeloxExpr(
         kDirectReference: {
       const auto& directRef = substraitField.direct_reference();
       int32_t colIdx = substraitParser_.parseReferenceSegment(directRef);
+
+      const auto& inputTypes = inputType->children();
       const auto& inputNames = inputType->names();
       const int64_t inputSize = inputNames.size();
-      if (colIdx <= inputSize) {
-        const auto& inputTypes = inputType->children();
-        // Convert type to row.
+      VELOX_CHECK_LT(
+          colIdx, inputSize, "Missing the column with id '{}' .", colIdx);
+      std::optional<int32_t> childIdx;
+      if (directRef.struct_field().has_child()) {
+        childIdx = substraitParser_.parseReferenceSegment(
+            directRef.struct_field().child());
+      }
+
+      if (!childIdx.has_value()) {
         return std::make_shared<core::FieldAccessTypedExpr>(
             inputTypes[colIdx],
             std::make_shared<core::InputTypedExpr>(inputTypes[colIdx]),
             inputNames[colIdx]);
-      } else {
-        VELOX_FAIL("Missing the column with id '{}' .", colIdx);
       }
+      VELOX_CHECK_EQ(inputTypes[colIdx]->kind(), TypeKind::ROW);
+      auto inputColumnType = asRowType(inputTypes[colIdx]);
+      VELOX_CHECK_LT(
+          childIdx.value(),
+          inputColumnType->size(),
+          "Missing the subfield with id '{}' .",
+          childIdx.value());
+      // Select a subfield in a struct by name.
+      return std::make_shared<core::FieldAccessTypedExpr>(
+          inputColumnType->childAt(childIdx.value()),
+          std::make_shared<core::FieldAccessTypedExpr>(
+              inputTypes[colIdx], inputNames[colIdx]),
+          inputColumnType->nameOf(childIdx.value()));
     }
     default:
       VELOX_NYI(

velox/substrait/VeloxToSubstraitExpr.cpp

@@ -114,13 +114,51 @@ VeloxToSubstraitExprConvertor::toSubstraitExpr(
   ::substrait::Expression_FieldReference* substraitFieldExpr =
       google::protobuf::Arena::CreateMessage<
           ::substrait::Expression_FieldReference>(&arena);
-
-  std::string exprName = fieldExpr->name();
-
   ::substrait::Expression_ReferenceSegment_StructField* directStruct =
       substraitFieldExpr->mutable_direct_reference()->mutable_struct_field();
+  std::string exprName = fieldExpr->name();
+  // FieldAccessTypedExpr represents one of two things: a leaf in an expression
+  // or a dereference expression(fieldExpr->isInputColumn() == false)
+  // for a leaf in an expression, find idx from child by exprName.
+  // for a dereference expression, find idx from every child by exprName.
+  if (fieldExpr->isInputColumn()) {
+    auto idx = inputType->getChildIdxIfExists(exprName);
+    if (idx.has_value()) {
+      directStruct->set_field(idx.value());
+    } else {
+      VELOX_USER_FAIL("idx has_value return false.");
+    }
+  } else {
+    int matchCount = 0;
+    uint32_t idxOfExprName = -1;
+    for (auto child : inputType->children()) {
+      auto rowChild = asRowType(child);
+      if (!rowChild) {
+        continue;
+      }
+      auto idxInChild = rowChild->getChildIdxIfExists(exprName);
+      if (idxInChild.has_value()) {
+        matchCount++;
+        idxOfExprName = idxInChild.value();
+      }
+    }
+    if (matchCount == 0) {
+      VELOX_USER_FAIL("exprName :{} no name matched!", exprName);
+    }
+    if (matchCount > 1) {
+      VELOX_USER_FAIL("exprName :{} multiple names matched!", exprName);
+    }
+
+    ::substrait::Expression_ReferenceSegment* refSegment =
+        google::protobuf::Arena::CreateMessage<
+            ::substrait::Expression_ReferenceSegment>(&arena);
+    ::substrait::Expression_ReferenceSegment_StructField* childStruct =
+        refSegment->mutable_struct_field();
+    childStruct->set_field(idxOfExprName);
+    refSegment->set_allocated_struct_field(childStruct);
+    directStruct->set_allocated_child(refSegment);
+  }
 
-  directStruct->set_field(inputType->getChildIdx(exprName));
   return *substraitFieldExpr;
 }
 
@@ -235,20 +273,20 @@ VeloxToSubstraitExprConvertor::toSubstraitNotNullLiteral(
       google::protobuf::Arena::CreateMessage<::substrait::Expression_Literal>(
           &arena);
   switch (variantValue.kind()) {
-    case velox::TypeKind::DOUBLE: {
-      literalExpr->set_fp64(variantValue.value<TypeKind::DOUBLE>());
-      break;
-    }
-    case velox::TypeKind::BIGINT: {
-      literalExpr->set_i64(variantValue.value<TypeKind::BIGINT>());
+    case velox::TypeKind::BOOLEAN: {
+      literalExpr->set_boolean(variantValue.value<TypeKind::BOOLEAN>());
       break;
     }
     case velox::TypeKind::INTEGER: {
       literalExpr->set_i32(variantValue.value<TypeKind::INTEGER>());
       break;
     }
-    case velox::TypeKind::BOOLEAN: {
-      literalExpr->set_boolean(variantValue.value<TypeKind::BOOLEAN>());
+    case velox::TypeKind::BIGINT: {
+      literalExpr->set_i64(variantValue.value<TypeKind::BIGINT>());
+      break;
+    }
+    case velox::TypeKind::DOUBLE: {
+      literalExpr->set_fp64(variantValue.value<TypeKind::DOUBLE>());
       break;
     }
     default:

velox/substrait/tests/VeloxSubstraitRoundTripPlanConverterTest.cpp

@@ -18,10 +18,9 @@
 
 #include "velox/exec/tests/utils/OperatorTestBase.h"
 #include "velox/exec/tests/utils/PlanBuilder.h"
-#include "velox/vector/tests/utils/VectorMaker.h"
-
 #include "velox/substrait/SubstraitToVeloxPlan.h"
 #include "velox/substrait/VeloxToSubstraitPlan.h"
+#include "velox/vector/tests/utils/VectorTestBase.h"
 
 using namespace facebook::velox;
 using namespace facebook::velox::test;
@@ -286,6 +285,32 @@ TEST_F(VeloxSubstraitRoundTripPlanConverterTest, ifThen) {
   assertPlanConversion(plan, "SELECT if (c0=1, c0 + 1, c1 + 2) as x FROM tmp");
 }
 
+TEST_F(VeloxSubstraitRoundTripPlanConverterTest, subField) {
+  RowVectorPtr data = makeRowVector(
+      {"a", "b", "c"},
+      {
+          makeFlatVector<int64_t>(
+              {2499109626526694126, 2342493223442167775, 4077358421272316858}),
+          makeFlatVector<int32_t>({581869302, -708632711, -133711905}),
+          makeFlatVector<double>(
+              {0.90579193414549275, 0.96886777112423139, 0.63235925003444637}),
+      });
+  createDuckDbTable({data});
+  auto plan =
+      PlanBuilder()
+          .values({data})
+          .project(
+              {"cast(row_constructor(a, b) as row(a bigint, b bigint)) as ab",
+               "c"})
+          .project(
+              {"cast(row_constructor(ab, c) as row(ab row(a bigint, b bigint), c bigint)) as abc"})
+          .project({"(abc).ab", "abc.c"})
+          .project({"(ab).a", "(ab).b", "c"})
+          .planNode();
+
+  assertPlanConversion(plan, "SELECT a, b, c FROM tmp");
+}
+
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);
   folly::init(&argc, &argv, false);