From c2478797caaa50dca4013a8cc3ebf65a8c34ed27 Mon Sep 17 00:00:00 2001 From: Dylan Chen Date: Mon, 22 Apr 2024 12:15:27 +0800 Subject: [PATCH 1/2] support column pruning for iceberg --- .../plan_node/logical_iceberg_scan.rs | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/src/frontend/src/optimizer/plan_node/logical_iceberg_scan.rs b/src/frontend/src/optimizer/plan_node/logical_iceberg_scan.rs index f4a14b837f1b2..dc3acb6acfdbf 100644 --- a/src/frontend/src/optimizer/plan_node/logical_iceberg_scan.rs +++ b/src/frontend/src/optimizer/plan_node/logical_iceberg_scan.rs @@ -54,6 +54,15 @@ impl LogicalIcebergScan { pub fn source_catalog(&self) -> Option> { self.core.catalog.clone() } + + pub fn clone_with_required_cols(&self, required_cols: &[usize]) -> Self { + assert!(!required_cols.is_empty()); + let mut core = self.core.clone(); + core.column_catalog = required_cols.iter().map(|idx| core.column_catalog[*idx].clone()).collect(); + let base = PlanBase::new_logical_with_core(&core); + + LogicalIcebergScan { base, core } + } } impl_plan_tree_node_for_leaf! {LogicalIcebergScan} @@ -74,9 +83,13 @@ impl Distill for LogicalIcebergScan { impl ColPrunable for LogicalIcebergScan { fn prune_col(&self, required_cols: &[usize], _ctx: &mut ColumnPruningContext) -> PlanRef { - // TODO: support column pruning for iceberg scan - let mapping = ColIndexMapping::with_remaining_columns(required_cols, self.schema().len()); - LogicalProject::with_mapping(self.clone().into(), mapping).into() + if required_cols.is_empty() { + let mapping = ColIndexMapping::with_remaining_columns(required_cols, self.schema().len()); + // If reuqiured_cols is empty, we use the first column of iceberg to avoid the empty schema. + LogicalProject::with_mapping(self.clone_with_required_cols(&[0]).into() , mapping).into() + } else { + self.clone_with_required_cols(required_cols).into() + } } } From cd05d942869fcd13a0f225ffe6c72ecf8297a497 Mon Sep 17 00:00:00 2001 From: Dylan Chen Date: Mon, 22 Apr 2024 12:20:13 +0800 Subject: [PATCH 2/2] fmt --- .../src/optimizer/plan_node/logical_iceberg_scan.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/frontend/src/optimizer/plan_node/logical_iceberg_scan.rs b/src/frontend/src/optimizer/plan_node/logical_iceberg_scan.rs index dc3acb6acfdbf..b355d6f2057c8 100644 --- a/src/frontend/src/optimizer/plan_node/logical_iceberg_scan.rs +++ b/src/frontend/src/optimizer/plan_node/logical_iceberg_scan.rs @@ -58,7 +58,10 @@ impl LogicalIcebergScan { pub fn clone_with_required_cols(&self, required_cols: &[usize]) -> Self { assert!(!required_cols.is_empty()); let mut core = self.core.clone(); - core.column_catalog = required_cols.iter().map(|idx| core.column_catalog[*idx].clone()).collect(); + core.column_catalog = required_cols + .iter() + .map(|idx| core.column_catalog[*idx].clone()) + .collect(); let base = PlanBase::new_logical_with_core(&core); LogicalIcebergScan { base, core } @@ -84,9 +87,10 @@ impl Distill for LogicalIcebergScan { impl ColPrunable for LogicalIcebergScan { fn prune_col(&self, required_cols: &[usize], _ctx: &mut ColumnPruningContext) -> PlanRef { if required_cols.is_empty() { - let mapping = ColIndexMapping::with_remaining_columns(required_cols, self.schema().len()); + let mapping = + ColIndexMapping::with_remaining_columns(required_cols, self.schema().len()); // If reuqiured_cols is empty, we use the first column of iceberg to avoid the empty schema. - LogicalProject::with_mapping(self.clone_with_required_cols(&[0]).into() , mapping).into() + LogicalProject::with_mapping(self.clone_with_required_cols(&[0]).into(), mapping).into() } else { self.clone_with_required_cols(required_cols).into() }