feat(PartitionedOutput): Add nested RowVector support in the serializer

Author: yingsu00

velox/exec/benchmarks/ExchangeBenchmark.cpp

@@ -170,6 +170,9 @@ enum class ExchangeInputKind {
   kDeep10K,
   kDeep50,
   kStruct1K,
+  kNested1_1K,
+  kNested2_1K,
+  kNested3_1K,
 };
 
 struct ExchangeInputSpec {
@@ -242,6 +245,17 @@ RowTypePtr makeDeepType() {
                 ROW({{"s2_int", INTEGER()}, {"s2_string", VARCHAR()}})))}});
 }
 
+// N levels of single-field ROW wrapping a BIGINT leaf, behind a BIGINT "c0"
+// partition key. Exercises the nested-ROW serialization path; combined with
+// nullPct it produces independent nulls at every ROW level and the leaf.
+RowTypePtr makeNestedType(int level) {
+  TypePtr inner = BIGINT();
+  for (int i = 0; i < level; ++i) {
+    inner = ROW({"r"}, {inner});
+  }
+  return ROW({"c0", "v"}, {BIGINT(), inner});
+}
+
 ExchangeInputSpec makeInputSpec(ExchangeInputKind kind) {
   switch (kind) {
     case ExchangeInputKind::kDeep10K:
@@ -250,6 +264,12 @@ ExchangeInputSpec makeInputSpec(ExchangeInputKind kind) {
       return {"Deep50", makeDeepType(), 2000, 50};
     case ExchangeInputKind::kStruct1K:
       return {"Struct1K", makeStructType(), 100, 1000};
+    case ExchangeInputKind::kNested1_1K:
+      return {"Nested1_1K", makeNestedType(1), 100, 1000};
+    case ExchangeInputKind::kNested2_1K:
+      return {"Nested2_1K", makeNestedType(2), 100, 1000};
+    case ExchangeInputKind::kNested3_1K:
+      return {"Nested3_1K", makeNestedType(3), 100, 1000};
   }
 
   VELOX_UNREACHABLE();
@@ -401,6 +421,21 @@ class ExchangeBenchmark : public VectorTestBase {
             [](auto row) { return static_cast<int128_t>(row); },
             isNull,
             type);
+      case TypeKind::ROW: {
+        // Build each child independently and apply row-level nulls so a ROW
+        // column carries nulls at this level and, recursively, at every nested
+        // ROW level and the leaf.
+        const auto& rowType = type->asRow();
+        std::vector<std::string> names;
+        std::vector<VectorPtr> children;
+        names.reserve(rowType.size());
+        children.reserve(rowType.size());
+        for (auto i = 0; i < rowType.size(); ++i) {
+          names.push_back(rowType.nameOf(i));
+          children.push_back(makeColumn(rowType.childAt(i), numRows, nullPct));
+        }
+        return makeRowVector(names, children, isNull);
+      }
       default:
         VELOX_NYI(
             "makeColumn does not support complex type {} yet",
@@ -718,12 +753,24 @@ EXCHANGE_BENCHMARK_CASE(
     Simple10K_Double_col16,
     makeInputSpec(SimpleColType::kDouble, 16));
 
-// The complex type benchmarks are temporarily disabled.
+// The ARRAY/MAP complex type benchmarks are temporarily disabled because
+// makeColumn does not generate ARRAY/MAP yet.
 // EXCHANGE_BENCHMARK_CASE(Deep10K, makeInputSpec(ExchangeInputKind::kDeep10K));
 // EXCHANGE_BENCHMARK_CASE(Deep50, makeInputSpec(ExchangeInputKind::kDeep50));
 // EXCHANGE_BENCHMARK_CASE(Struct1K,
 // makeInputSpec(ExchangeInputKind::kStruct1K));
 
+// Nested ROW columns (makeColumn supports ROW recursively).
+EXCHANGE_BENCHMARK_CASE(
+    Nested1_1K,
+    makeInputSpec(ExchangeInputKind::kNested1_1K));
+EXCHANGE_BENCHMARK_CASE(
+    Nested2_1K,
+    makeInputSpec(ExchangeInputKind::kNested2_1K));
+EXCHANGE_BENCHMARK_CASE(
+    Nested3_1K,
+    makeInputSpec(ExchangeInputKind::kNested3_1K));
+
 #undef EXCHANGE_BENCHMARK_CASE
 #undef EXCHANGE_BENCHMARK_MODES
 #undef EXCHANGE_BENCHMARK_INPUT

velox/exec/tests/OptimizedPartitionedOutputTest.cpp

@@ -58,12 +58,24 @@ struct TestParam {
 std::vector<TestParam> testParams() {
   std::vector<TestParam> params;
 
-  const std::vector<std::pair<std::string, TypePtr>> types = {
+  // Builds an n-level nested ROW: 1 → ROW(int, bigint); 2 → ROW(int, row1);
+  // 3 → ROW(int, row2); …. The "r" child carries the next nesting level.
+  std::function<TypePtr(int)> nestedRow = [&](int level) -> TypePtr {
+    if (level <= 1) {
+      return ROW({"a", "b"}, {INTEGER(), BIGINT()});
+    }
+    return ROW({"a", "r"}, {INTEGER(), nestedRow(level - 1)});
+  };
+
+  std::vector<std::pair<std::string, TypePtr>> types = {
       {"bool", BOOLEAN()},
       {"tinyint", TINYINT()},
       {"bigint", BIGINT()},
       {"hugeint", HUGEINT()},
   };
+  for (int level = 1; level <= 3; ++level) {
+    types.emplace_back("row" + std::to_string(level), nestedRow(level));
+  }
 
   const std::vector<std::pair<std::string, NullMode>> nullModes = {
       {"no_null", NullMode::kNoNull},
@@ -362,8 +374,56 @@ class OptimizedPartitionedOutputParamTest
     VELOX_UNREACHABLE();
   }
 
-  /// Creates a flat vector of the param's value type with random values and
-  /// nulls applied according to nullMode.
+  /// Recursively creates a vector of the given type with random values. Only
+  /// the top-level row carries row-level nulls (applied by the caller via
+  /// makeRandomValueVector); leaf flat vectors are non-null inside this helper.
+  VectorPtr makeRandomVectorOfType(
+      const TypePtr& type,
+      int numRows,
+      std::mt19937_64& rng) {
+    switch (type->kind()) {
+      case TypeKind::BOOLEAN:
+        return vectorMaker_.flatVector<bool>(
+            numRows, [&](auto /*i*/) -> bool { return rng() % 2 == 0; });
+      case TypeKind::TINYINT:
+        return vectorMaker_.flatVector<int8_t>(
+            numRows,
+            [&](auto /*i*/) -> int8_t { return static_cast<int8_t>(rng()); });
+      case TypeKind::INTEGER:
+        return vectorMaker_.flatVector<int32_t>(
+            numRows,
+            [&](auto /*i*/) -> int32_t { return static_cast<int32_t>(rng()); });
+      case TypeKind::BIGINT:
+        return vectorMaker_.flatVector<int64_t>(
+            numRows,
+            [&](auto /*i*/) -> int64_t { return static_cast<int64_t>(rng()); });
+      case TypeKind::HUGEINT:
+        return vectorMaker_.flatVector<int128_t>(
+            numRows, [&](auto /*i*/) -> int128_t {
+              int64_t hi = static_cast<int64_t>(rng());
+              uint64_t lo = rng();
+              return (static_cast<int128_t>(hi) << 64) |
+                  static_cast<int128_t>(lo);
+            });
+      case TypeKind::ROW: {
+        const auto& rowType = type->asRow();
+        std::vector<VectorPtr> children;
+        children.reserve(rowType.size());
+        for (size_t i = 0; i < rowType.size(); ++i) {
+          children.push_back(
+              makeRandomVectorOfType(rowType.childAt(i), numRows, rng));
+        }
+        return vectorMaker_.rowVector(rowType.names(), children);
+      }
+      default:
+        VELOX_UNREACHABLE("Unsupported value type: {}", type->toString());
+    }
+  }
+
+  /// Creates a vector of the param's value type with random values and nulls
+  /// applied at the top level according to nullMode. For flat scalar types the
+  /// nulls are placed on the leaf. For ROW types the nulls are placed at the
+  /// outermost row level only; child columns themselves remain null-free.
   VectorPtr makeRandomValueVector(int numRows, std::mt19937_64& rng) {
     auto isNullFn = [this](vector_size_t i) -> bool { return isNull(i); };
 
@@ -393,6 +453,16 @@ class OptimizedPartitionedOutputParamTest
                   static_cast<int128_t>(lo);
             },
             isNullFn);
+      case TypeKind::ROW: {
+        auto result = std::dynamic_pointer_cast<RowVector>(
+            makeRandomVectorOfType(param().valueType, numRows, rng));
+        for (vector_size_t i = 0; i < numRows; ++i) {
+          if (isNull(i)) {
+            result->setNull(i, true);
+          }
+        }
+        return result;
+      }
       default:
         VELOX_UNREACHABLE(
             "Unsupported value type: {}", param().valueType->toString());