Support selecting a subfield from struct (#294)

Co-authored-by: zhejiangxiaomai <zhenhui.zhao@intel.com>

Author: rui-mo

velox/substrait/SubstraitToVeloxExpr.cpp

@@ -189,6 +189,16 @@ VectorPtr constructFlatVectorForStruct(
   return vector;
 }
 
+core::FieldAccessTypedExprPtr makeFieldAccessExpr(
+    const std::string& name,
+    const TypePtr& type,
+    core::FieldAccessTypedExprPtr input) {
+  if (input) {
+    return std::make_shared<core::FieldAccessTypedExpr>(type, input, name);
+  }
+
+  return std::make_shared<core::FieldAccessTypedExpr>(type, name);
+}
 } // namespace
 
 using facebook::velox::core::variantArrayToVector;
@@ -202,44 +212,26 @@ SubstraitVeloxExprConverter::toVeloxExpr(
   switch (typeCase) {
     case ::substrait::Expression::FieldReference::ReferenceTypeCase::
         kDirectReference: {
-      const auto& dRef = substraitField.direct_reference();
-      VELOX_CHECK(dRef.has_struct_field(), "Struct field expected.");
-      int32_t colIdx = subParser_->parseReferenceSegment(dRef);
-      std::optional<int32_t> childIdx = std::nullopt;
-      if (dRef.struct_field().has_child()) {
-        childIdx =
-            subParser_->parseReferenceSegment(dRef.struct_field().child());
-      }
-
-      const auto& inputTypes = inputType->children();
-      const auto& inputNames = inputType->names();
-      const int64_t inputSize = inputNames.size();
-
-      if (colIdx >= inputSize) {
-        VELOX_FAIL("Missing the column with id '{}' .", colIdx);
-      }
-
-      if (!childIdx.has_value()) {
-        return std::make_shared<core::FieldAccessTypedExpr>(
-            inputTypes[colIdx],
-            std::make_shared<core::InputTypedExpr>(inputTypes[colIdx]),
-            inputNames[colIdx]);
-      } else {
-        // Select a subfield in a struct by name.
-        if (auto inputColumnType = asRowType(inputTypes[colIdx])) {
-          if (childIdx.value() >= inputColumnType->size()) {
-            VELOX_FAIL("Missing the subfield with id '{}' .", childIdx.value());
-          }
-          return std::make_shared<core::FieldAccessTypedExpr>(
-              inputColumnType->childAt(childIdx.value()),
-              std::make_shared<core::FieldAccessTypedExpr>(
-                  inputTypes[colIdx], inputNames[colIdx]),
-              inputColumnType->nameOf(childIdx.value()));
-        } else {
-          VELOX_FAIL("RowType expected.");
+      const auto& directRef = substraitField.direct_reference();
+      core::FieldAccessTypedExprPtr fieldAccess{nullptr};
+      const auto* tmp = &directRef.struct_field();
+
+      auto inputColumnType = inputType;
+      for (;;) {
+        auto idx = tmp->field();
+        fieldAccess = makeFieldAccessExpr(
+            inputColumnType->nameOf(idx),
+            inputColumnType->childAt(idx),
+            fieldAccess);
+
+        if (!tmp->has_child()) {
+          break;
         }
+
+        inputColumnType = asRowType(inputColumnType->childAt(idx));
+        tmp = &tmp->child().struct_field();
       }
-      break;
+      return fieldAccess;
     }
     default:
       VELOX_NYI(

velox/substrait/VeloxToSubstraitExpr.cpp

@@ -368,6 +368,26 @@ void convertVectorValue(
     }
   }
 }
+
+uint32_t getFieldIdForIntermediateNode(
+    const std::string& exprName,
+    const ::substrait::Expression_ReferenceSegment_StructField& structField,
+    const RowTypePtr& inputType) {
+  auto inputColumnType = inputType;
+  std::vector<int32_t> ids;
+  const auto* tmp = &structField;
+  for (;;) {
+    ids.push_back(tmp->field());
+    if (!tmp->has_child()) {
+      break;
+    }
+    tmp = &tmp->child().struct_field();
+  }
+  for (int32_t i = ids.size() - 1; i >= 0; --i) {
+    inputColumnType = asRowType(inputColumnType->childAt(ids[i]));
+  }
+  return inputColumnType->getChildIdx(exprName);
+}
 } // namespace
 
 const ::substrait::Expression& VeloxToSubstraitExprConvertor::toSubstraitExpr(
@@ -448,7 +468,40 @@ VeloxToSubstraitExprConvertor::toSubstraitExpr(
   ::substrait::Expression_ReferenceSegment_StructField* directStruct =
       substraitFieldExpr->mutable_direct_reference()->mutable_struct_field();
 
-  directStruct->set_field(inputType->getChildIdx(exprName));
+  // 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()) {
+    uint32_t idx = inputType->getChildIdx(exprName);
+    directStruct->set_field(idx);
+  } else {
+    auto tmp = toSubstraitExpr(arena, fieldExpr->inputs()[0], inputType)
+                   .selection()
+                   .direct_reference();
+    if (!tmp.has_struct_field()) {
+      uint32_t idx = inputType->getChildIdx(exprName);
+      directStruct->set_field(idx);
+    } else {
+      ::substrait::Expression_ReferenceSegment_StructField* childStruct =
+          google::protobuf::Arena::CreateMessage<
+              ::substrait::Expression_ReferenceSegment_StructField>(&arena);
+      ::substrait::Expression_ReferenceSegment* refSegment =
+          google::protobuf::Arena::CreateMessage<
+              ::substrait::Expression_ReferenceSegment>(&arena);
+      directStruct->MergeFrom(tmp.struct_field());
+      childStruct->set_field(getFieldIdForIntermediateNode(
+          exprName, tmp.struct_field(), inputType));
+      refSegment->set_allocated_struct_field(childStruct);
+      ::substrait::Expression_ReferenceSegment_StructField* innerChild =
+          directStruct;
+      while (innerChild->has_child()) {
+        innerChild = innerChild->mutable_child()->mutable_struct_field();
+      }
+      innerChild->set_allocated_child(refSegment);
+    }
+  }
+
   return *substraitFieldExpr;
 }
 

velox/substrait/tests/Substrait2VeloxPlanConversionTest.cpp

@@ -309,9 +309,9 @@ TEST_F(Substrait2VeloxPlanConversionTest, ifthenTest) {
       "  -- TableScan[table: hive_table, range filters: "
       "[(hd_demo_sk, Filter(IsNotNull, deterministic, null not allowed)),"
       " (hd_vehicle_count, BigintRange: [1, 9223372036854775807] no nulls)], "
-      "remaining filter: (and(or(equalto(ROW[\"hd_buy_potential\"],\">10000\"),"
-      "equalto(ROW[\"hd_buy_potential\"],\"unknown\")),if(greaterthan(ROW[\"hd_vehicle_count\"],0),"
-      "greaterthan(divide(cast ROW[\"hd_dep_count\"] as DOUBLE,cast ROW[\"hd_vehicle_count\"] as DOUBLE),1.2))))]"
+      "remaining filter: (and(or(equalto(\"hd_buy_potential\",\">10000\"),"
+      "equalto(\"hd_buy_potential\",\"unknown\")),if(greaterthan(\"hd_vehicle_count\",0),"
+      "greaterthan(divide(cast \"hd_dep_count\" as DOUBLE,cast \"hd_vehicle_count\" as DOUBLE),1.2))))]"
       " -> n0_0:BIGINT, n0_1:VARCHAR, n0_2:BIGINT, n0_3:BIGINT\n",
       planNode->toString(true, true));
 }

velox/substrait/tests/VeloxSubstraitRoundTripTest.cpp

@@ -23,6 +23,9 @@
 
 #include "velox/functions/sparksql/Register.h"
 #include "velox/substrait/SubstraitToVeloxPlan.h"
+#include "velox/substrait/VeloxToSubstraitPlan.h"
+#include "velox/vector/tests/utils/VectorTestBase.h"
+
 #include "velox/substrait/VariantToVectorConverter.h"
 #include "velox/substrait/VeloxToSubstraitPlan.h"
 
@@ -507,6 +510,37 @@ TEST_F(VeloxSubstraitRoundTripTest, dateType) {
   assertPlanConversion(plan, "SELECT * FROM tmp WHERE c > DATE '1992-01-01'");
 }
 
+TEST_F(VeloxSubstraitRoundTripTest, subField) {
+  RowVectorPtr data = makeRowVector(
+      {"a", "b", "c"},
+      {
+          makeFlatVector<int64_t>({249, 235, 858}),
+          makeFlatVector<int32_t>({581, -708, -133}),
+          makeFlatVector<double>({0.905, 0.968, 0.632}),
+      });
+  createDuckDbTable({data});
+  auto plan =
+      PlanBuilder()
+          .values({data})
+          .project(
+              {"cast(row_constructor(a, b) as row(a bigint, b bigint)) as ab",
+               "a",
+               "b",
+               "c"})
+          .project(
+              {"cast(row_constructor(ab, c) as row(ab row(a bigint, b bigint), c bigint)) as abc",
+               "a",
+               "b"})
+          .project(
+              {"(cast(row_constructor(a, b) as row(a bigint, b bigint))).a",
+               "(abc).ab.a",
+               "(abc).ab.b",
+               "abc.c"})
+          .planNode();
+
+  assertPlanConversion(plan, "SELECT a, a, b, c FROM tmp");
+}
+
 int main(int argc, char** argv) {
   facebook::velox::functions::sparksql::registerFunctions("");
   testing::InitGoogleTest(&argc, argv);