Fix rollup and cube transforms.

dbms/src/Processors/Transforms/CubeTransform.cpp

@@ -21,10 +21,10 @@ void CubeTransform::consume(Chunk chunk)
 
     current_columns = consumed_chunk.getColumns();
     current_zero_columns.clear();
-    current_zero_columns.reserve(current_columns.size());
+    current_zero_columns.reserve(keys.size());
 
-    for (auto & column : current_columns)
-        current_zero_columns.emplace_back(column->cloneEmpty()->cloneResized(num_rows));
+    for (auto key : keys)
+        current_zero_columns.emplace_back(current_columns[key]->cloneEmpty()->cloneResized(num_rows));
 }
 
 bool CubeTransform::canGenerate()
@@ -41,9 +41,9 @@ Chunk CubeTransform::generate()
         --mask;
 
         auto columns = current_columns;
-        for (size_t i = 0; i < columns.size(); ++i)
+        for (size_t i = 0; i < keys.size(); ++i)
             if (mask & (UInt64(1) << i))
-                columns[i] = current_zero_columns[i];
+                columns[keys[i]] = current_zero_columns[i];
 
         BlocksList cube_blocks = { getInputPort().getHeader().cloneWithColumns(columns) };
         auto cube_block = params->aggregator.mergeBlocks(cube_blocks, false);

dbms/src/Processors/Transforms/RollupTransform.cpp

@@ -15,7 +15,7 @@ RollupTransform::RollupTransform(Block header, AggregatingTransformParamsPtr par
 void RollupTransform::consume(Chunk chunk)
 {
     consumed_chunk = std::move(chunk);
-    last_removed_key = consumed_chunk.getNumColumns();
+    last_removed_key = keys.size();
 }
 
 bool RollupTransform::canGenerate()
@@ -30,10 +30,11 @@ Chunk RollupTransform::generate()
     if (last_removed_key)
     {
         --last_removed_key;
+        auto key = keys[last_removed_key];
 
         auto num_rows = gen_chunk.getNumRows();
         auto columns = gen_chunk.getColumns();
-        columns[last_removed_key] = columns[last_removed_key]->cloneEmpty()->cloneResized(num_rows);
+        columns[key] = columns[key]->cloneEmpty()->cloneResized(num_rows);
 
         BlocksList rollup_blocks = { getInputPort().getHeader().cloneWithColumns(columns) };
         auto rollup_block = params->aggregator.mergeBlocks(rollup_blocks, false);