Merge pull request #14500 from ClickHouse/many-rows-in-totals

Many rows in totals

src/Processors/Formats/IOutputFormat.cpp

@@ -44,6 +44,25 @@ IOutputFormat::Status IOutputFormat::prepare()
     return Status::Finished;
 }
 
+static Chunk prepareTotals(Chunk chunk)
+{
+    if (!chunk.hasRows())
+        return {};
+
+    if (chunk.getNumRows() > 1)
+    {
+        /// This may happen if something like ARRAY JOIN was executed on totals.
+        /// Skip rows except the first one.
+        auto columns = chunk.detachColumns();
+        for (auto & column : columns)
+            column = column->cut(0, 1);
+
+        chunk.setColumns(std::move(columns), 1);
+    }
+
+    return chunk;
+}
+
 void IOutputFormat::work()
 {
     if (!prefix_written)
@@ -70,7 +89,8 @@ void IOutputFormat::work()
             consume(std::move(current_chunk));
             break;
         case Totals:
-            consumeTotals(std::move(current_chunk));
+            if (auto totals = prepareTotals(std::move(current_chunk)))
+                consumeTotals(std::move(totals));
             break;
         case Extremes:
             consumeExtremes(std::move(current_chunk));

src/Processors/Formats/Impl/TemplateBlockOutputFormat.cpp

@@ -186,7 +186,7 @@ void TemplateBlockOutputFormat::finalize()
         switch (static_cast<ResultsetPart>(*format.format_idx_to_column_idx[i]))
         {
             case ResultsetPart::Totals:
-                if (!totals)
+                if (!totals || !totals.hasRows())
                     format.throwInvalidFormat("Cannot print totals for this request", i);
                 writeRow(totals, 0);
                 break;

src/Processors/Transforms/TotalsHavingTransform.cpp

@@ -257,7 +257,8 @@ void TotalsHavingTransform::prepareTotals()
     {
         auto block = finalized_header.cloneWithColumns(totals.detachColumns());
         expression->execute(block);
-        totals = Chunk(block.getColumns(), 1);
+        /// Note: after expression totals may have several rows if `arrayJoin` was used in expression.
+        totals = Chunk(block.getColumns(), block.rows());
     }
 }
 

tests/queries/0_stateless/01472_many_rows_in_totals.reference

@@ -0,0 +1,133 @@
+┏━━━┳━━━━━━━━━━━━━┓
+┃ g ┃ sum(number) ┃
+┡━━━╇━━━━━━━━━━━━━┩
+│ 0 │           2 │
+├───┼─────────────┤
+│ 0 │           2 │
+├───┼─────────────┤
+│ 1 │           4 │
+└───┴─────────────┘
+
+Totals:
+┏━━━┳━━━━━━━━━━━━━┓
+┃ g ┃ sum(number) ┃
+┡━━━╇━━━━━━━━━━━━━┩
+│ 0 │           6 │
+└───┴─────────────┘
+-
+┏━━━┳━━━┓
+┃ g ┃ s ┃
+┡━━━╇━━━┩
+│ 0 │ 2 │
+├───┼───┤
+│ 0 │ 2 │
+├───┼───┤
+│ 1 │ 4 │
+├───┼───┤
+│ 1 │ 4 │
+└───┴───┘
+
+Totals:
+┏━━━┳━━━┓
+┃ g ┃ s ┃
+┡━━━╇━━━┩
+│ 0 │ 6 │
+└───┴───┘
+--
+0	2
+0	2
+1	4
+
+0	6
+-
+0	2
+0	2
+1	4
+1	4
+
+0	6
+--
+{
+	"meta":
+	[
+		{
+			"name": "g",
+			"type": "UInt64"
+		},
+		{
+			"name": "sum(number)",
+			"type": "UInt64"
+		}
+	],
+
+	"data":
+	[
+		{
+			"g": "0",
+			"sum(number)": "2"
+		},
+		{
+			"g": "0",
+			"sum(number)": "2"
+		},
+		{
+			"g": "1",
+			"sum(number)": "4"
+		}
+	],
+
+	"totals":
+	{
+		"g": "0",
+		"sum(number)": "6"
+	},
+
+	"rows": 3,
+
+	"rows_before_limit_at_least": 4
+}
+-
+{
+	"meta":
+	[
+		{
+			"name": "g",
+			"type": "UInt64"
+		},
+		{
+			"name": "s",
+			"type": "UInt64"
+		}
+	],
+
+	"data":
+	[
+		{
+			"g": "0",
+			"s": "2"
+		},
+		{
+			"g": "0",
+			"s": "2"
+		},
+		{
+			"g": "1",
+			"s": "4"
+		},
+		{
+			"g": "1",
+			"s": "4"
+		}
+	],
+
+	"totals":
+	{
+		"g": "0",
+		"s": "6"
+	},
+
+	"rows": 4,
+
+	"rows_before_limit_at_least": 4
+}
+--

tests/queries/0_stateless/01472_many_rows_in_totals.sql

@@ -0,0 +1,15 @@
+set output_format_write_statistics = 0;
+select g, sum(number) from numbers(4) group by bitAnd(number, 1) as g with totals having sum(number) <= arrayJoin([2, 4]) format Pretty;
+select '-';
+select g, s from (select g, sum(number) as s from numbers(4) group by bitAnd(number, 1) as g with totals) array join [1, 2] as a format Pretty;
+select '--';
+
+select g, sum(number) from numbers(4) group by bitAnd(number, 1) as g with totals having sum(number) <= arrayJoin([2, 4]) format TSV;
+select '-';
+select g, s from (select g, sum(number) as s from numbers(4) group by bitAnd(number, 1) as g with totals) array join [1, 2] as a format TSV;
+select '--';
+
+select g, sum(number) from numbers(4) group by bitAnd(number, 1) as g with totals having sum(number) <= arrayJoin([2, 4]) format JSON;
+select '-';
+select g, s from (select g, sum(number) as s from numbers(4) group by bitAnd(number, 1) as g with totals) array join [1, 2] as a format JSON;
+select '--';