chore(cubestore): Upgrade DF: Implement PanicWorkerNode

Author: srh

rust/cubestore/cubestore/src/queryplanner/panic.rs

@@ -10,6 +10,7 @@ use datafusion::physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, Partitioning, PlanProperties,
     SendableRecordBatchStream,
 };
+use serde::{Deserialize, Serialize};
 use std::any::Any;
 use std::cmp::Ordering;
 use std::fmt::{Formatter, Pointer};
@@ -25,6 +26,16 @@ impl PanicWorkerNode {
             node: Arc::new(self),
         })
     }
+
+    pub fn from_serialized(inputs: &[LogicalPlan], serialized: PanicWorkerSerialized) -> Self {
+        assert_eq!(0, inputs.len());
+        let PanicWorkerSerialized {} = serialized;
+        Self {}
+    }
+
+    pub fn to_serialized(&self) -> PanicWorkerSerialized {
+        PanicWorkerSerialized {}
+    }
 }
 
 lazy_static! {
@@ -81,6 +92,9 @@ impl UserDefinedLogicalNode for PanicWorkerNode {
     }
 }
 
+#[derive(Clone, Serialize, Deserialize, Debug)]
+pub struct PanicWorkerSerialized {}
+
 #[derive(Debug)]
 pub struct PanicWorkerExec {
     properties: PlanProperties,

rust/cubestore/cubestore/src/queryplanner/planning.rs

@@ -38,7 +38,9 @@ use crate::metastore::table::{Table, TablePath};
 use crate::metastore::{
     AggregateFunction, Chunk, Column, IdRow, Index, IndexType, MetaStore, Partition, Schema,
 };
+use crate::queryplanner::metadata_cache::NoopParquetMetadataCache;
 use crate::queryplanner::optimizations::rewrite_plan::{rewrite_plan, PlanRewriter};
+use crate::queryplanner::panic::PanicWorkerSerialized;
 use crate::queryplanner::panic::{plan_panic_worker, PanicWorkerNode};
 use crate::queryplanner::partition_filter::PartitionFilter;
 use crate::queryplanner::providers::InfoSchemaQueryCacheTableProvider;
@@ -50,7 +52,6 @@ use crate::queryplanner::topk::ClusterAggregateTopK;
 use crate::queryplanner::{CubeTableLogical, InfoSchemaTableProvider};
 use crate::table::{cmp_same_types, Row};
 use crate::CubeError;
-use crate::queryplanner::metadata_cache::NoopParquetMetadataCache;
 use datafusion::common;
 use datafusion::common::DFSchemaRef;
 use datafusion::datasource::DefaultTableSource;
@@ -1366,6 +1367,7 @@ pub type Snapshots = Vec<Snapshot>;
 #[derive(Clone, Serialize, Deserialize, Debug)]
 pub enum ExtensionNodeSerialized {
     ClusterSend(ClusterSendSerialized),
+    PanicWorker(PanicWorkerSerialized),
 }
 
 #[derive(Debug, Clone)]

rust/cubestore/cubestore/src/queryplanner/serialized_plan.rs

@@ -9,8 +9,7 @@ use crate::queryplanner::query_executor::{CubeTable, InlineTableId, InlineTableP
 use crate::queryplanner::topk::{ClusterAggregateTopK, SortColumn};
 use crate::queryplanner::udfs::aggregate_udf_by_kind;
 use crate::queryplanner::udfs::{
-    aggregate_kind_by_name, scalar_udf_by_kind, CubeAggregateUDFKind,
-    CubeScalarUDFKind,
+    aggregate_kind_by_name, scalar_udf_by_kind, CubeAggregateUDFKind, CubeScalarUDFKind,
 };
 use crate::queryplanner::{CubeTableLogical, InfoSchemaTableProvider};
 use crate::table::Row;
@@ -1332,23 +1331,29 @@ impl LogicalExtensionCodec for CubeExtensionCodec {
         let serialized = ExtensionNodeSerialized::deserialize(r)
             .map_err(|e| DataFusionError::Execution(format!("try_decode: {}", e)))?;
         Ok(Extension {
-            node: Arc::new(match serialized {
+            node: match serialized {
                 ExtensionNodeSerialized::ClusterSend(serialized) => {
-                    ClusterSendNode::from_serialized(inputs, serialized)
+                    Arc::new(ClusterSendNode::from_serialized(inputs, serialized))
                 }
-            }),
+                ExtensionNodeSerialized::PanicWorker(serialized) => {
+                    Arc::new(PanicWorkerNode::from_serialized(inputs, serialized))
+                }
+            },
         })
     }
 
     fn try_encode(&self, node: &Extension, buf: &mut Vec<u8>) -> datafusion::common::Result<()> {
         use serde::Serialize;
         let mut ser = flexbuffers::FlexbufferSerializer::new();
-        let to_serialize =
-            if let Some(cluster_send) = node.node.as_any().downcast_ref::<ClusterSendNode>() {
-                ExtensionNodeSerialized::ClusterSend(cluster_send.to_serialized())
-            } else {
-                todo!("{:?}", node)
-            };
+        let to_serialize = if let Some(cluster_send) =
+            node.node.as_any().downcast_ref::<ClusterSendNode>()
+        {
+            ExtensionNodeSerialized::ClusterSend(cluster_send.to_serialized())
+        } else if let Some(panic_worker) = node.node.as_any().downcast_ref::<PanicWorkerNode>() {
+            ExtensionNodeSerialized::PanicWorker(panic_worker.to_serialized())
+        } else {
+            todo!("{:?}", node)
+        };
         to_serialize
             .serialize(&mut ser)
             .map_err(|e| DataFusionError::Execution(format!("try_encode: {}", e)))?;