add GeneratedColumnPlaceholderInputStream (#6796)

close #6801, ref pingcap/tidb#40663

dbms/src/DataStreams/GeneratedColumnPlaceholderBlockInputStream.h

@@ -0,0 +1,97 @@
+// Copyright 2023 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <Columns/ColumnNullable.h>
+#include <Columns/ColumnString.h>
+#include <Columns/ColumnsNumber.h>
+#include <Core/ColumnsWithTypeAndName.h>
+#include <DataStreams/IProfilingBlockInputStream.h>
+#include <DataTypes/DataTypeNullable.h>
+#include <DataTypes/DataTypeString.h>
+#include <Flash/Coprocessor/TiDBTableScan.h>
+
+namespace DB
+{
+class GeneratedColumnPlaceholderBlockInputStream : public IProfilingBlockInputStream
+{
+public:
+    GeneratedColumnPlaceholderBlockInputStream(
+        const BlockInputStreamPtr & input,
+        const std::vector<std::tuple<UInt64, String, DataTypePtr>> & generated_column_infos_,
+        const String & req_id_)
+        : generated_column_infos(generated_column_infos_)
+        , log(Logger::get(req_id_))
+    {
+        children.push_back(input);
+    }
+
+    String getName() const override { return NAME; }
+    Block getHeader() const override
+    {
+        Block block = children.back()->getHeader();
+        insertColumns(block, /*insert_data=*/false);
+        return block;
+    }
+
+    static String getColumnName(UInt64 col_index)
+    {
+        return "generated_column_" + std::to_string(col_index);
+    }
+
+protected:
+    void readPrefix() override
+    {
+        RUNTIME_CHECK(!generated_column_infos.empty());
+        // Validation check.
+        for (size_t i = 1; i < generated_column_infos.size(); ++i)
+        {
+            RUNTIME_CHECK(std::get<0>(generated_column_infos[i]) > std::get<0>(generated_column_infos[i - 1]));
+        }
+    }
+
+    Block readImpl() override
+    {
+        Block block = children.back()->read();
+        insertColumns(block, /*insert_data=*/true);
+        return block;
+    }
+
+private:
+    void insertColumns(Block & block, bool insert_data) const
+    {
+        if (!block)
+            return;
+
+        for (const auto & ele : generated_column_infos)
+        {
+            const auto & col_index = std::get<0>(ele);
+            const auto & col_name = std::get<1>(ele);
+            const auto & data_type = std::get<2>(ele);
+            ColumnPtr column = nullptr;
+            if (insert_data)
+                column = data_type->createColumnConstWithDefaultValue(block.rows());
+            else
+                column = data_type->createColumn();
+            block.insert(col_index, ColumnWithTypeAndName{column, data_type, col_name});
+        }
+    }
+
+    static constexpr auto NAME = "GeneratedColumnPlaceholder";
+    const std::vector<std::tuple<UInt64, String, DataTypePtr>> generated_column_infos;
+    const LoggerPtr log;
+};
+
+} // namespace DB

dbms/src/Flash/Coprocessor/DAGQueryBlockInterpreter.cpp

@@ -185,6 +185,8 @@ void DAGQueryBlockInterpreter::handleMockTableScan(const TiDBTableScan & table_s
         analyzer = std::make_unique<DAGExpressionAnalyzer>(std::move(names_and_types), context);
         pipeline.streams.insert(pipeline.streams.end(), mock_table_scan_streams.begin(), mock_table_scan_streams.end());
     }
+
+    // Ignore handling GeneratedColumnPlaceholderBlockInputStream for now, because we don't support generated column in test framework.
 }
 
 

dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp

@@ -19,6 +19,7 @@
 #include <Common/TiFlashMetrics.h>
 #include <DataStreams/ExpressionBlockInputStream.h>
 #include <DataStreams/FilterBlockInputStream.h>
+#include <DataStreams/GeneratedColumnPlaceholderBlockInputStream.h>
 #include <DataStreams/IProfilingBlockInputStream.h>
 #include <DataStreams/MultiplexInputStream.h>
 #include <DataStreams/NullBlockInputStream.h>
@@ -335,6 +336,8 @@ void DAGStorageInterpreter::executeImpl(DAGPipeline & pipeline)
 
     /// handle timezone/duration cast for local and remote table scan.
     executeCastAfterTableScan(remote_read_streams_start_index, pipeline);
+    /// handle generated column if necessary.
+    executeGeneratedColumnPlaceholder(remote_read_streams_start_index, generated_column_infos, log, pipeline);
     recordProfileStreams(pipeline, table_scan.getTableScanExecutorID());
 
     /// handle filter conditions for local and remote table scan.
@@ -940,6 +943,15 @@ std::tuple<Names, NamesAndTypes, std::vector<ExtraCastAfterTSMode>> DAGStorageIn
         auto const & ci = table_scan.getColumns()[i];
         const ColumnID cid = ci.id;
 
+        if (ci.hasGeneratedColumnFlag())
+        {
+            LOG_DEBUG(log, "got column({}) with generated column flag", i);
+            const auto & data_type = getDataTypeByColumnInfoForComputingLayer(ci);
+            const auto & col_name = GeneratedColumnPlaceholderBlockInputStream::getColumnName(i);
+            generated_column_infos.push_back(std::make_tuple(i, col_name, data_type));
+            source_columns_tmp.emplace_back(NameAndTypePair{col_name, data_type});
+            continue;
+        }
         // Column ID -1 return the handle column
         String name;
         if (cid == TiDBPkColumnID)

dbms/src/Flash/Coprocessor/DAGStorageInterpreter.h

@@ -134,6 +134,8 @@ class DAGStorageInterpreter
     ManageableStoragePtr storage_for_logical_table;
     Names required_columns;
     NamesAndTypes source_columns;
+    // For generated column, just need a placeholder, and TiDB will fill this column.
+    std::vector<std::tuple<UInt64, String, DataTypePtr>> generated_column_infos;
 };
 
 } // namespace DB

dbms/src/Flash/Coprocessor/InterpreterUtils.cpp

@@ -16,6 +16,7 @@
 #include <DataStreams/CreatingSetsBlockInputStream.h>
 #include <DataStreams/ExpressionBlockInputStream.h>
 #include <DataStreams/FilterBlockInputStream.h>
+#include <DataStreams/GeneratedColumnPlaceholderBlockInputStream.h>
 #include <DataStreams/MergeSortingBlockInputStream.h>
 #include <DataStreams/PartialSortingBlockInputStream.h>
 #include <DataStreams/SharedQueryBlockInputStream.h>
@@ -221,4 +222,21 @@ void executePushedDownFilter(
         stream->setExtraInfo("projection after push down filter");
     }
 }
+
+void executeGeneratedColumnPlaceholder(
+    size_t remote_read_streams_start_index,
+    const std::vector<std::tuple<UInt64, String, DataTypePtr>> & generated_column_infos,
+    LoggerPtr log,
+    DAGPipeline & pipeline)
+{
+    if (generated_column_infos.empty())
+        return;
+    assert(remote_read_streams_start_index <= pipeline.streams.size());
+    for (size_t i = 0; i < remote_read_streams_start_index; ++i)
+    {
+        auto & stream = pipeline.streams[i];
+        stream = std::make_shared<GeneratedColumnPlaceholderBlockInputStream>(stream, generated_column_infos, log->identifier());
+        stream->setExtraInfo("generated column placeholder above table scan");
+    }
+}
 } // namespace DB

dbms/src/Flash/Coprocessor/InterpreterUtils.h

@@ -78,4 +78,10 @@ void executePushedDownFilter(
     DAGExpressionAnalyzer & analyzer,
     LoggerPtr log,
     DAGPipeline & pipeline);
+
+void executeGeneratedColumnPlaceholder(
+    size_t remote_read_streams_start_index,
+    const std::vector<std::tuple<UInt64, String, DataTypePtr>> & generated_column_infos,
+    LoggerPtr log,
+    DAGPipeline & pipeline);
 } // namespace DB

dbms/src/Flash/Planner/Plans/PhysicalMockTableScan.cpp

@@ -16,6 +16,7 @@
 #include <Flash/Coprocessor/DAGContext.h>
 #include <Flash/Coprocessor/DAGPipeline.h>
 #include <Flash/Coprocessor/GenSchemaAndColumn.h>
+#include <Flash/Coprocessor/InterpreterUtils.h>
 #include <Flash/Coprocessor/MockSourceStream.h>
 #include <Flash/Pipeline/Exec/PipelineExecBuilder.h>
 #include <Flash/Planner/FinalizeHelper.h>
@@ -67,6 +68,7 @@ std::pair<NamesAndTypes, BlockInputStreams> mockSchemaAndStreams(
     assert(!schema.empty());
     assert(!mock_streams.empty());
 
+    // Ignore handling GeneratedColumnPlaceholderBlockInputStream for now, because we don't support generated column in test framework.
     return {std::move(schema), std::move(mock_streams)};
 }
 } // namespace

dbms/src/Storages/Transaction/TiDB.h

@@ -117,7 +117,8 @@ enum TP
     M(NoDefaultValue, (1 << 12)) \
     M(OnUpdateNow, (1 << 13))    \
     M(PartKey, (1 << 14))        \
-    M(Num, (1 << 15))
+    M(Num, (1 << 15))            \
+    M(GeneratedColumn, (1 << 23))
 
 enum ColumnFlag
 {

tests/fullstack-test/expr/generated_index.test

@@ -0,0 +1,147 @@
+# Copyright 2023 PingCAP, Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+mysql> drop table if exists test.t;
+mysql> create table test.t(c1 varchar(100), c2 varchar(100));
+mysql> insert into test.t values('ABC', 'DEF');
+mysql> alter table test.t set tiflash replica 1;
+func> wait_table test t
+mysql> alter table test.t add index idx2((lower(c2)));
+
+mysql> select /*+ nth_plan(1) */ * from test.t where lower(test.t.c2) = 'def';
++------+------+
+| c1   | c2   |
++------+------+
+| ABC  | DEF  |
++------+------+
+mysql> select /*+ nth_plan(2) */ * from test.t where lower(test.t.c2) = 'def';
++------+------+
+| c1   | c2   |
++------+------+
+| ABC  | DEF  |
++------+------+
+mysql> select /*+ nth_plan(3) */ * from test.t where lower(test.t.c2) = 'def';
++------+------+
+| c1   | c2   |
++------+------+
+| ABC  | DEF  |
++------+------+
+mysql> select /*+ nth_plan(4) */ * from test.t where lower(test.t.c2) = 'def';
++------+------+
+| c1   | c2   |
++------+------+
+| ABC  | DEF  |
++------+------+
+mysql> select /*+ nth_plan(5) */ * from test.t where lower(test.t.c2) = 'def';
++------+------+
+| c1   | c2   |
++------+------+
+| ABC  | DEF  |
++------+------+
+mysql> select /*+ nth_plan(6) */ * from test.t where lower(test.t.c2) = 'def';
++------+------+
+| c1   | c2   |
++------+------+
+| ABC  | DEF  |
++------+------+
+mysql> select /*+ nth_plan(7) */ * from test.t where lower(test.t.c2) = 'def';
++------+------+
+| c1   | c2   |
++------+------+
+| ABC  | DEF  |
++------+------+
+mysql> select /*+ nth_plan(8) */ * from test.t where lower(test.t.c2) = 'def';
++------+------+
+| c1   | c2   |
++------+------+
+| ABC  | DEF  |
++------+------+
+mysql> select /*+ nth_plan(9) */ * from test.t where lower(test.t.c2) = 'def';
++------+------+
+| c1   | c2   |
++------+------+
+| ABC  | DEF  |
++------+------+
+mysql> select /*+ nth_plan(10) */ * from test.t where lower(test.t.c2) = 'def';
++------+------+
+| c1   | c2   |
++------+------+
+| ABC  | DEF  |
++------+------+
+
+mysql> drop table if exists test.t;
+mysql> create table test.t(id int, value int);
+mysql> alter table test.t set tiflash replica 1;
+func> wait_table test t
+mysql> create unique index uk on test.t((tidb_shard(id)), id);
+mysql> select /*+ nth_paln(1) */ max(value) from test.t;
++------------+
+| max(value) |
++------------+
+|       NULL |
++------------+
+mysql> select /*+ nth_paln(2) */ max(value) from test.t;
++------------+
+| max(value) |
++------------+
+|       NULL |
++------------+
+mysql> select /*+ nth_paln(3) */ max(value) from test.t;
++------------+
+| max(value) |
++------------+
+|       NULL |
++------------+
+mysql> select /*+ nth_paln(4) */ max(value) from test.t;
++------------+
+| max(value) |
++------------+
+|       NULL |
++------------+
+mysql> select /*+ nth_paln(5) */ max(value) from test.t;
++------------+
+| max(value) |
++------------+
+|       NULL |
++------------+
+mysql> select /*+ nth_paln(6) */ max(value) from test.t;
++------------+
+| max(value) |
++------------+
+|       NULL |
++------------+
+mysql> select /*+ nth_paln(7) */ max(value) from test.t;
++------------+
+| max(value) |
++------------+
+|       NULL |
++------------+
+mysql> select /*+ nth_paln(8) */ max(value) from test.t;
++------------+
+| max(value) |
++------------+
+|       NULL |
++------------+
+mysql> select /*+ nth_paln(9) */ max(value) from test.t;
++------------+
+| max(value) |
++------------+
+|       NULL |
++------------+
+mysql> select /*+ nth_paln(10) */ max(value) from test.t;
++------------+
+| max(value) |
++------------+
+|       NULL |
++------------+