feat: expand operator

proto/batch_plan.proto

@@ -109,6 +109,13 @@ message HashAggNode {
   repeated expr.AggCall agg_calls = 2;
 }
 
+message ExpandNode {
+  message Keys {
+    repeated uint32 keys = 1;
+  }
+  repeated Keys expanded_keys = 1;
+}
+
 message SortAggNode {
   repeated expr.ExprNode group_keys = 1;
   repeated expr.AggCall agg_calls = 2;
@@ -210,6 +217,7 @@ message PlanNode {
     HopWindowNode hop_window = 25;
     TableFunctionNode table_function = 26;
     SysRowSeqScanNode sys_row_seq_scan = 27;
+    ExpandNode expand = 28;
   }
   string identity = 24;
 }

proto/stream_plan.proto

@@ -221,6 +221,13 @@ message LookupUnionNode {
   repeated uint32 order = 1;
 }
 
+message ExpandNode {
+  message Keys {
+    repeated uint32 keys = 1;
+  }
+  repeated Keys expanded_keys = 1;
+}
+
 message StreamNode {
   oneof node_body {
     SourceNode source = 100;
@@ -243,6 +250,7 @@ message StreamNode {
     LookupUnionNode lookup_union = 117;
     UnionNode union = 118;
     DeltaIndexJoinNode delta_index_join = 119;
+    ExpandNode expand = 120;
   }
   // The id for the operator.
   uint64 operator_id = 1;

src/batch/src/executor/expand.rs

@@ -0,0 +1,202 @@
+// Copyright 2022 Singularity Data
+//
+// 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.
+
+use std::sync::Arc;
+
+use futures_async_stream::try_stream;
+use itertools::Itertools;
+use risingwave_common::array::column::Column;
+use risingwave_common::array::{DataChunk, PrimitiveArray, Vis};
+use risingwave_common::catalog::{Field, Schema};
+use risingwave_common::error::ErrorCode::InternalError;
+use risingwave_common::error::{Result, RwError};
+use risingwave_common::types::DataType;
+use risingwave_common::util::chunk_coalesce::{DataChunkBuilder, SlicedDataChunk};
+use risingwave_pb::batch_plan::plan_node::NodeBody;
+
+use super::{BoxedDataChunkStream, BoxedExecutor, BoxedExecutorBuilder, Executor, ExecutorBuilder};
+use crate::task::BatchTaskContext;
+
+pub struct ExpandExecutor {
+    expanded_keys: Vec<Vec<usize>>,
+    child: BoxedExecutor,
+    schema: Schema,
+    identity: String,
+}
+
+impl Executor for ExpandExecutor {
+    fn schema(&self) -> &Schema {
+        &self.schema
+    }
+
+    fn identity(&self) -> &str {
+        &self.identity
+    }
+
+    fn execute(self: Box<Self>) -> BoxedDataChunkStream {
+        self.do_execute()
+    }
+}
+
+impl ExpandExecutor {
+    #[try_stream(boxed, ok = DataChunk, error = RwError)]
+    async fn do_execute(self: Box<Self>) {
+        let mut data_chunk_builder = DataChunkBuilder::with_default_size(self.schema.data_types());
+
+        #[for_await]
+        for data_chunk in self.child.execute() {
+            // TODO: handle dummy chunk.
+            let data_chunk: DataChunk = data_chunk?.compact()?;
+            let cardinality = data_chunk.cardinality();
+            let (columns, vis) = data_chunk.into_parts();
+            assert_eq!(vis, Vis::Compact(cardinality));
+
+            let null_columns: Vec<Column> = columns
+                .iter()
+                .map(|column| {
+                    let array = column.array_ref();
+                    let mut builder = array.create_builder(cardinality)?;
+                    (0..cardinality).try_for_each(|_i| builder.append_null())?;
+                    Ok::<Column, RwError>(Column::new(Arc::new(builder.finish()?)))
+                })
+                .try_collect()?;
+
+            for (i, keys) in self.expanded_keys.iter().enumerate() {
+                let mut new_columns = null_columns.clone();
+                for key in keys {
+                    new_columns[*key] = columns[*key].clone();
+                }
+                let flags = Column::from(PrimitiveArray::<i64>::from_slice(&vec![
+                    Some(i as i64);
+                    cardinality
+                ])?);
+                new_columns.push(flags);
+                let new_data_chunk = DataChunk::new(new_columns, vis.clone());
+                let mut sliced_data_chunk = SlicedDataChunk::new_checked(new_data_chunk)?;
+                loop {
+                    let (left_data, output) = data_chunk_builder.append_chunk(sliced_data_chunk)?;
+                    match (left_data, output) {
+                        (Some(left_data), Some(output)) => {
+                            sliced_data_chunk = left_data;
+                            yield output;
+                        }
+                        (None, Some(output)) => {
+                            yield output;
+                            break;
+                        }
+                        (None, None) => {
+                            break;
+                        }
+                        _ => {
+                            return Err(
+                                InternalError("Data chunk builder error".to_string()).into()
+                            );
+                        }
+                    }
+                }
+            }
+        }
+        if let Some(chunk) = data_chunk_builder.consume_all()? {
+            yield chunk;
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl BoxedExecutorBuilder for ExpandExecutor {
+    async fn new_boxed_executor<C: BatchTaskContext>(
+        source: &ExecutorBuilder<C>,
+        mut inputs: Vec<BoxedExecutor>,
+    ) -> Result<BoxedExecutor> {
+        ensure!(inputs.len() == 1);
+        let expand_node = try_match_expand!(
+            source.plan_node().get_node_body().unwrap(),
+            NodeBody::Expand
+        )?;
+
+        let expanded_keys = expand_node
+            .expanded_keys
+            .iter()
+            .map(|keys| keys.keys.iter().map(|key| *key as usize).collect_vec())
+            .collect_vec();
+
+        let child = inputs.remove(0);
+
+        let mut schema = child.schema().clone();
+        schema
+            .fields
+            .push(Field::with_name(DataType::Int64, "flag"));
+
+        Ok(Box::new(Self {
+            expanded_keys,
+            child,
+            schema,
+            identity: "ExpandExecutor".to_string(),
+        }))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use futures::StreamExt;
+    use risingwave_common::array::{DataChunk, DataChunkTestExt};
+    use risingwave_common::catalog::{Field, Schema};
+    use risingwave_common::types::DataType;
+
+    use super::ExpandExecutor;
+    use crate::executor::test_utils::MockExecutor;
+    use crate::executor::Executor;
+
+    #[tokio::test]
+    async fn test_expand_executor() {
+        let mock_schema = Schema {
+            fields: vec![
+                Field::unnamed(DataType::Int32),
+                Field::unnamed(DataType::Int32),
+                Field::unnamed(DataType::Int32),
+            ],
+        };
+        let expand_schema = Schema {
+            fields: vec![
+                Field::unnamed(DataType::Int32),
+                Field::unnamed(DataType::Int32),
+                Field::unnamed(DataType::Int32),
+                Field::unnamed(DataType::Int64),
+            ],
+        };
+        let mut mock_executor = MockExecutor::new(mock_schema);
+        mock_executor.add(DataChunk::from_pretty(
+            "i i i
+             1 2 3
+             2 3 4",
+        ));
+        let expanded_keys = vec![vec![0, 1], vec![1, 2]];
+        let expand_executor = Box::new(ExpandExecutor {
+            expanded_keys,
+            child: Box::new(mock_executor),
+            schema: expand_schema,
+            identity: "ExpandExecutor".to_string(),
+        });
+        let mut stream = expand_executor.execute();
+        let res = stream.next().await.unwrap().unwrap();
+        let expected_chunk = DataChunk::from_pretty(
+            "i i i I
+             1 2 . 0
+             2 3 . 0
+             . 2 3 1
+             . 3 4 1",
+        );
+        assert_eq!(res, expected_chunk);
+    }
+}

src/batch/src/executor/mod.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 mod delete;
+mod expand;
 mod filter;
 mod generic_exchange;
 mod hash_agg;
@@ -36,6 +37,7 @@ mod values;
 
 use async_recursion::async_recursion;
 pub use delete::*;
+pub use expand::*;
 pub use filter::*;
 use futures::stream::BoxStream;
 pub use generic_exchange::*;
@@ -187,6 +189,7 @@ impl<'a, C: BatchTaskContext> ExecutorBuilder<'a, C> {
             NodeBody::TableFunction => TableFunctionExecutorBuilder,
             NodeBody::HopWindow => HopWindowExecutor,
             NodeBody::SysRowSeqScan => SysRowSeqScanExecutorBuilder,
+            NodeBody::Expand => ExpandExecutor,
         }
         .await?;
         let input_desc = real_executor.identity().to_string();

src/common/src/array/data_chunk.rs

@@ -38,7 +38,7 @@ pub struct DataChunk {
 
 /// `Vis` is a visibility bitmap of rows. When all rows are visible, it is considered compact and
 /// is represented by a single cardinality number rather than that many of ones.
-#[derive(Clone, PartialEq)]
+#[derive(Clone, PartialEq, Debug)]
 pub enum Vis {
     Bitmap(Bitmap),
     Compact(usize), // equivalent to all ones of this size

src/frontend/src/optimizer/plan_node/batch_expand.rs

@@ -0,0 +1,103 @@
+// Copyright 2022 Singularity Data
+//
+// 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.
+
+use std::fmt;
+
+use itertools::Itertools;
+use risingwave_common::error::Result;
+use risingwave_pb::batch_plan::expand_node::Keys;
+use risingwave_pb::batch_plan::plan_node::NodeBody;
+use risingwave_pb::batch_plan::ExpandNode;
+
+use crate::optimizer::plan_node::{
+    LogicalExpand, PlanBase, PlanTreeNodeUnary, ToBatchProst, ToDistributedBatch, ToLocalBatch,
+};
+use crate::optimizer::property::{Distribution, Order};
+use crate::optimizer::PlanRef;
+
+#[derive(Debug, Clone)]
+pub struct BatchExpand {
+    pub base: PlanBase,
+    logical: LogicalExpand,
+}
+
+impl BatchExpand {
+    pub fn new(logical: LogicalExpand) -> Self {
+        let ctx = logical.base.ctx.clone();
+        let base = PlanBase::new_batch(
+            ctx,
+            logical.schema().clone(),
+            Distribution::SomeShard,
+            Order::any(),
+        );
+        BatchExpand { base, logical }
+    }
+
+    pub fn expanded_keys(&self) -> &Vec<Vec<usize>> {
+        self.logical.expanded_keys()
+    }
+}
+
+impl fmt::Display for BatchExpand {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "BatchExpand {{ expanded_keys: {:#?} }}",
+            self.expanded_keys()
+        )
+    }
+}
+
+impl PlanTreeNodeUnary for BatchExpand {
+    fn input(&self) -> PlanRef {
+        self.logical.input()
+    }
+
+    fn clone_with_input(&self, input: PlanRef) -> Self {
+        Self::new(self.logical.clone_with_input(input))
+    }
+}
+
+impl_plan_tree_node_for_unary! { BatchExpand }
+
+impl ToDistributedBatch for BatchExpand {
+    fn to_distributed(&self) -> Result<PlanRef> {
+        let new_input = self.input().to_distributed()?;
+        Ok(self.clone_with_input(new_input).into())
+    }
+}
+
+impl ToBatchProst for BatchExpand {
+    fn to_batch_prost_body(&self) -> NodeBody {
+        NodeBody::Expand(ExpandNode {
+            expanded_keys: self
+                .expanded_keys()
+                .iter()
+                .map(|keys| keys_to_protobuf(keys))
+                .collect_vec(),
+        })
+    }
+}
+
+fn keys_to_protobuf(keys: &[usize]) -> Keys {
+    let keys = keys.iter().map(|key| *key as u32).collect_vec();
+    Keys { keys }
+}
+
+impl ToLocalBatch for BatchExpand {
+    fn to_local(&self) -> Result<PlanRef> {
+        let new_input = self.input().to_local()?;
+        Ok(self.clone_with_input(new_input).into())
+    }
+}