add fuse_block table function

Author: zhyass

src/query/service/src/procedures/systems/fuse_block.rs

@@ -0,0 +1,69 @@
+//  Copyright 2021 Datafuse Labs.
+//
+//  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 common_datablocks::DataBlock;
+use common_datavalues::DataSchema;
+use common_exception::Result;
+
+use crate::procedures::OneBlockProcedure;
+use crate::procedures::Procedure;
+use crate::procedures::ProcedureFeatures;
+use crate::sessions::QueryContext;
+use crate::sessions::TableContext;
+use crate::storages::fuse::table_functions::FuseBlock;
+use crate::storages::fuse::FuseTable;
+
+pub struct FuseBlockProcedure {}
+
+impl FuseBlockProcedure {
+    pub fn try_create() -> Result<Box<dyn Procedure>> {
+        Ok(FuseBlockProcedure {}.into_procedure())
+    }
+}
+
+#[async_trait::async_trait]
+impl OneBlockProcedure for FuseBlockProcedure {
+    fn name(&self) -> &str {
+        "FUSE_BLOCK"
+    }
+
+    fn features(&self) -> ProcedureFeatures {
+        ProcedureFeatures::default().num_arguments(3)
+    }
+
+    async fn all_data(&self, ctx: Arc<QueryContext>, args: Vec<String>) -> Result<DataBlock> {
+        let database_name = args[0].clone();
+        let table_name = args[1].clone();
+        let snapshot_id = args[2].clone();
+        let tenant_id = ctx.get_tenant();
+        let tbl = ctx
+            .get_catalog(&ctx.get_current_catalog())?
+            .get_table(
+                tenant_id.as_str(),
+                database_name.as_str(),
+                table_name.as_str(),
+            )
+            .await?;
+
+        let tbl = FuseTable::try_from_table(tbl.as_ref())?;
+
+        Ok(FuseBlock::new(ctx, tbl, snapshot_id).get_blocks().await?)
+    }
+
+    fn schema(&self) -> Arc<DataSchema> {
+        FuseBlock::schema()
+    }
+}

src/query/service/src/procedures/systems/mod.rs

@@ -13,12 +13,14 @@
 // limitations under the License.
 
 mod clustering_information;
+mod fuse_block;
 mod fuse_segment;
 mod fuse_snapshot;
 mod search_tables;
 mod system;
 
 pub use clustering_information::ClusteringInformationProcedure;
+pub use fuse_block::FuseBlockProcedure;
 pub use fuse_segment::FuseSegmentProcedure;
 pub use fuse_snapshot::FuseSnapshotProcedure;
 pub use search_tables::SearchTablesProcedure;

src/query/service/src/procedures/systems/system.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use crate::procedures::systems::ClusteringInformationProcedure;
+use crate::procedures::systems::FuseBlockProcedure;
 use crate::procedures::systems::FuseSegmentProcedure;
 use crate::procedures::systems::FuseSnapshotProcedure;
 use crate::procedures::systems::SearchTablesProcedure;
@@ -34,6 +35,10 @@ impl SystemProcedure {
             "system$fuse_segment",
             Box::new(FuseSegmentProcedure::try_create),
         );
+        factory.register(
+            "system$fuse_block",
+            Box::new(FuseBlockProcedure::try_create),
+        );
         factory.register(
             "system$search_tables",
             Box::new(SearchTablesProcedure::try_create),

src/query/service/src/table_functions/table_function_factory.rs

@@ -24,6 +24,7 @@ use parking_lot::RwLock;
 use crate::catalogs::SYS_TBL_FUC_ID_END;
 use crate::catalogs::SYS_TBL_FUNC_ID_BEGIN;
 use crate::storages::fuse::table_functions::ClusteringInformationTable;
+use crate::storages::fuse::table_functions::FuseBlockTable;
 use crate::storages::fuse::table_functions::FuseSegmentTable;
 use crate::storages::fuse::table_functions::FuseSnapshotTable;
 use crate::table_functions::async_crash_me::AsyncCrashMeTable;
@@ -105,6 +106,10 @@ impl TableFunctionFactory {
             "fuse_segment".to_string(),
             (next_id(), Arc::new(FuseSegmentTable::create)),
         );
+        creators.insert(
+            "fuse_block".to_string(),
+            (next_id(), Arc::new(FuseBlockTable::create)),
+        );
 
         creators.insert(
             "clustering_information".to_string(),

src/query/storages/fuse/src/table_functions/fuse_blocks/fuse_block.rs

@@ -0,0 +1,110 @@
+//  Copyright 2022 Datafuse Labs.
+//
+//  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 common_datablocks::DataBlock;
+use common_datavalues::prelude::*;
+use common_exception::Result;
+use futures_util::TryStreamExt;
+
+use crate::io::MetaReaders;
+use crate::io::SnapshotHistoryReader;
+use crate::sessions::TableContext;
+use crate::FuseTable;
+
+pub struct FuseBlock<'a> {
+    pub ctx: Arc<dyn TableContext>,
+    pub table: &'a FuseTable,
+    pub snapshot_id: String,
+}
+
+impl<'a> FuseBlock<'a> {
+    pub fn new(ctx: Arc<dyn TableContext>, table: &'a FuseTable, snapshot_id: String) -> Self {
+        Self {
+            ctx,
+            table,
+            snapshot_id,
+        }
+    }
+
+    pub async fn get_blocks(&self) -> Result<DataBlock> {
+        let tbl = self.table;
+        let maybe_snapshot = tbl.read_table_snapshot(self.ctx.clone()).await?;
+        if let Some(snapshot) = maybe_snapshot {
+            // prepare the stream of snapshot
+            let snapshot_version = tbl.snapshot_format_version();
+            let snapshot_location = tbl
+                .meta_location_generator
+                .snapshot_location_from_uuid(&snapshot.snapshot_id, snapshot_version)?;
+            let reader = MetaReaders::table_snapshot_reader(self.ctx.clone());
+            let mut snapshot_stream = reader.snapshot_history(
+                snapshot_location,
+                snapshot_version,
+                tbl.meta_location_generator().clone(),
+            );
+
+            // find the element by snapshot_id in stream
+            while let Some(snapshot) = snapshot_stream.try_next().await? {
+                if snapshot.snapshot_id.simple().to_string() == self.snapshot_id {
+                    let len = snapshot.summary.block_count as usize;
+                    let snapshot_id = vec![self.snapshot_id.clone().into_bytes()];
+                    let timestamp =
+                        vec![snapshot.timestamp.map(|dt| (dt.timestamp_micros()) as i64)];
+                    let mut block_location: Vec<Vec<u8>> = Vec::with_capacity(len);
+                    let mut block_size: Vec<u64> = Vec::with_capacity(len);
+                    let mut bloom_filter_location: Vec<Option<Vec<u8>>> = Vec::with_capacity(len);
+                    let mut bloom_filter_size: Vec<u64> = Vec::with_capacity(len);
+
+                    let reader = MetaReaders::segment_info_reader(self.ctx.as_ref());
+                    for (x, ver) in &snapshot.segments {
+                        let segment = reader.read(x, None, *ver).await?;
+                        segment.blocks.clone().into_iter().for_each(|block| {
+                            block_location.push(block.location.0.into_bytes());
+                            block_size.push(block.block_size);
+                            bloom_filter_location.push(
+                                block
+                                    .bloom_filter_index_location
+                                    .map(|(s, _)| s.into_bytes()),
+                            );
+                            bloom_filter_size.push(block.bloom_filter_index_size);
+                        });
+                    }
+
+                    return Ok(DataBlock::create(FuseBlock::schema(), vec![
+                        Arc::new(ConstColumn::new(Series::from_data(snapshot_id), len)),
+                        Arc::new(ConstColumn::new(Series::from_data(timestamp), len)),
+                        Series::from_data(block_location),
+                        Series::from_data(block_size),
+                        Series::from_data(bloom_filter_location),
+                        Series::from_data(bloom_filter_size),
+                    ]));
+                }
+            }
+        }
+
+        Ok(DataBlock::empty_with_schema(Self::schema()))
+    }
+
+    pub fn schema() -> Arc<DataSchema> {
+        DataSchemaRefExt::create(vec![
+            DataField::new("snapshot_id", Vu8::to_data_type()),
+            DataField::new_nullable("timestamp", TimestampType::new_impl(6)),
+            DataField::new("block_location", Vu8::to_data_type()),
+            DataField::new("block_size", u64::to_data_type()),
+            DataField::new_nullable("bloom_filter_location", Vu8::to_data_type()),
+            DataField::new("bloom_filter_size", u64::to_data_type()),
+        ])
+    }
+}