diff --git a/velox/exec/GroupingSet.cpp b/velox/exec/GroupingSet.cpp
index b0ba52cd9828..9f15ea440c91 100644
--- a/velox/exec/GroupingSet.cpp
+++ b/velox/exec/GroupingSet.cpp
@@ -760,22 +760,23 @@ bool GroupingSet::getOutput(
     }
     return false;
   }
-  extractGroups(folly::Range<char**>(groups, numGroups), result);
+  extractGroups(
+      table_->rows(), folly::Range<char**>(groups, numGroups), result);
   return true;
 }
 
 void GroupingSet::extractGroups(
+    RowContainer* rowContainer,
     folly::Range<char**> groups,
     const RowVectorPtr& result) {
   result->resize(groups.size());
   if (groups.empty()) {
     return;
   }
-  RowContainer& rows = *table_->rows();
-  const auto totalKeys = rows.keyTypes().size();
+  const auto totalKeys = rowContainer->keyTypes().size();
   for (int32_t i = 0; i < totalKeys; ++i) {
     auto& keyVector = result->childAt(i);
-    rows.extractColumn(
+    rowContainer->extractColumn(
         groups.data(),
         groups.size(),
         groupingKeyOutputProjections_[i],
@@ -1252,7 +1253,8 @@ void GroupingSet::extractSpillResult(const RowVectorPtr& result) {
     mergeRows_->listRows(
         &iter, rows.size(), RowContainer::kUnlimited, rows.data());
   }
-  extractGroups(folly::Range<char**>(rows.data(), rows.size()), result);
+  extractGroups(
+      mergeRows_.get(), folly::Range<char**>(rows.data(), rows.size()), result);
   mergeRows_->clear();
 }
 
diff --git a/velox/exec/GroupingSet.h b/velox/exec/GroupingSet.h
index fda8f4eea02f..9edb3cfac34d 100644
--- a/velox/exec/GroupingSet.h
+++ b/velox/exec/GroupingSet.h
@@ -203,7 +203,10 @@ class GroupingSet {
   // Copies the grouping keys and aggregates for 'groups' into 'result' If
   // partial output, extracts the intermediate type for aggregates, final result
   // otherwise.
-  void extractGroups(folly::Range<char**> groups, const RowVectorPtr& result);
+  void extractGroups(
+      RowContainer* container,
+      folly::Range<char**> groups,
+      const RowVectorPtr& result);
 
   // Produces output in if spilling has occurred. First produces data
   // from non-spilled partitions, then merges spill runs and unspilled data