Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore: refactor pushdown::TopK. #13095

Merged
merged 3 commits into from
Oct 7, 2023
Merged

chore: refactor pushdown::TopK. #13095

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ddaf1bf
chore: refactor topk.
RinChanNOWWW Oct 7, 2023
74bc09e
Support inner column as topk key.
RinChanNOWWW Oct 7, 2023
a83135d
Merge branch 'main' into topk
RinChanNOWWW Oct 7, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 13 additions & 13 deletions src/query/catalog/src/plan/pushdown.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,9 +107,15 @@ pub struct Filters {
#[derive(Debug, Clone)]
pub struct TopK {
pub limit: usize,
pub order_by: TableField,
/// Record the leaf field of the topk column.
/// - The `name` of `field` will be used to track column in the read block.
/// - The `column_id` of `field` will be used to retrieve column stats from block meta
/// (only used for fuse engine, for parquet table, we will use `leaf_id`).
pub field: TableField,
pub asc: bool,
pub column_id: u32,
/// The index in `table_schema.leaf_fields()`.
/// It's only used for external parquet files reading.
pub leaf_id: usize,
}

pub const TOPK_PUSHDOWN_THRESHOLD: usize = 1000;
Expand All @@ -124,30 +130,24 @@ impl PushDownInfo {
return None;
}

if let RemoteExpr::<String>::ColumnRef { id, .. } = &order.0 {
if let RemoteExpr::<String>::ColumnRef { id, data_type, .. } = &order.0 {
// TODO: support sub column of nested type.
let field = schema.field_with_name(id).ok()?;
if !support(&field.data_type().into()) {
if !support(data_type) {
return None;
}

let leaf_fields = schema.leaf_fields();
let (leaf_id, f) = leaf_fields
.iter()
.enumerate()
.find(|&(_, p)| p == field)
.find(|&(_, p)| p.name() == id)
.unwrap();
// Databend column id is not equal to parquet leaf id when there is nested type.
if f.column_id as usize != leaf_id {
return None;
}
let column_id = f.column_id;

let top_k = TopK {
limit: self.limit.unwrap(),
order_by: field.clone(),
field: f.clone(),
asc: order.1,
column_id,
leaf_id,
};
Some(top_k)
} else {
Expand Down
4 changes: 2 additions & 2 deletions src/query/storages/fuse/src/operations/read/native_data_source_deserializer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,7 @@ impl NativeDeserializeDataTransform {
};

let top_k = top_k.map(|top_k| {
let index = src_schema.index_of(top_k.order_by.name()).unwrap();
let index = src_schema.index_of(top_k.field.name()).unwrap();
let sorter = TopKSorter::new(top_k.limit, top_k.asc);

if !prewhere_columns.contains(&index) {
Expand Down Expand Up @@ -648,7 +648,7 @@ impl Processor for NativeDeserializeDataTransform {
Some(array) => {
let array = array?;
self.read_columns.push(*index);
let data_type = top_k.order_by.data_type().into();
let data_type = top_k.field.data_type().into();
let col = Column::from_arrow(array.as_ref(), &data_type);

arrays.push((*index, array));
Expand Down
14 changes: 7 additions & 7 deletions src/query/storages/fuse/src/operations/read_partitions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -237,7 +237,7 @@ impl FuseTable {
.as_ref()
.filter(|_| self.is_native()) // Only native format supports topk push down.
.and_then(|p| p.top_k(self.schema().as_ref(), RangeIndex::supported_type))
.map(|topk| field_default_value(ctx.clone(), &topk.order_by).map(|d| (topk, d)))
.map(|topk| field_default_value(ctx.clone(), &topk.field).map(|d| (topk, d)))
.transpose()?;

let (mut statistics, parts) =
Expand Down Expand Up @@ -284,23 +284,23 @@ impl FuseTable {
block_metas.sort_by(|a, b| {
let a =
a.1.col_stats
.get(&top_k.column_id)
.get(&top_k.field.column_id)
.unwrap_or(&default_stats);
let b =
b.1.col_stats
.get(&top_k.column_id)
.get(&top_k.field.column_id)
.unwrap_or(&default_stats);
(a.min().as_ref(), a.max().as_ref()).cmp(&(b.min().as_ref(), b.max().as_ref()))
});
} else {
block_metas.sort_by(|a, b| {
let a =
a.1.col_stats
.get(&top_k.column_id)
.get(&top_k.field.column_id)
.unwrap_or(&default_stats);
let b =
b.1.col_stats
.get(&top_k.column_id)
.get(&top_k.field.column_id)
.unwrap_or(&default_stats);
(b.max().as_ref(), b.min().as_ref()).cmp(&(a.max().as_ref(), a.min().as_ref()))
});
Expand Down Expand Up @@ -454,7 +454,7 @@ impl FuseTable {

let sort_min_max = top_k.as_ref().map(|(top_k, default)| {
meta.col_stats
.get(&top_k.column_id)
.get(&top_k.field.column_id)
.map(|stat| (stat.min().clone(), stat.max().clone()))
.unwrap_or((default.clone(), default.clone()))
});
Expand Down Expand Up @@ -494,7 +494,7 @@ impl FuseTable {
let create_on = meta.create_on;

let sort_min_max = top_k.map(|(top_k, default)| {
let stat = meta.col_stats.get(&top_k.column_id);
let stat = meta.col_stats.get(&top_k.field.column_id);
stat.map(|stat| (stat.min().clone(), stat.max().clone()))
.unwrap_or((default.clone(), default))
});
Expand Down
2 changes: 1 addition & 1 deletion src/query/storages/parquet/src/parquet2/parquet_table/partition.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ impl Parquet2Table {
let offset = projected_column_nodes
.column_nodes
.iter()
.position(|node| node.leaf_indices[0] == top_k.column_id as usize)
.position(|node| node.leaf_indices[0] == top_k.leaf_id)
.unwrap();
(top_k, offset)
});
Expand Down
2 changes: 1 addition & 1 deletion src/query/storages/parquet/src/parquet2/pruning.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ impl PartitionPruner {
let min_max = self
.top_k
.as_ref()
.filter(|(tk, _)| tk.column_id as usize == *index)
.filter(|(tk, _)| tk.leaf_id == *index)
.zip(row_group_stats.as_ref())
.map(|((_, offset), stats)| {
let stat = stats[rg_idx].get(&(*offset as u32)).unwrap();
Expand Down
20 changes: 10 additions & 10 deletions src/query/storages/parquet/src/parquet_rs/parquet_reader/read_policy/predicate_and_topk.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ use common_exception::Result;
use common_expression::DataBlock;
use common_expression::DataSchema;
use common_expression::DataSchemaRef;
use common_expression::TableField;
use common_expression::TableSchema;
use common_expression::TopKSorter;
use parquet::arrow::arrow_reader::ParquetRecordBatchReader;
Expand All @@ -37,6 +36,7 @@ use super::utils::evaluate_topk;
use super::utils::read_all;
use crate::parquet_rs::parquet_reader::predicate::ParquetPredicate;
use crate::parquet_rs::parquet_reader::row_group::InMemoryRowGroup;
use crate::parquet_rs::parquet_reader::topk::BuiltTopK;
use crate::parquet_rs::parquet_reader::topk::ParquetTopK;
use crate::parquet_rs::parquet_reader::utils::bitmap_to_boolean_array;
use crate::parquet_rs::parquet_reader::utils::compute_output_field_paths;
Expand All @@ -60,7 +60,7 @@ impl PredicateAndTopkPolicyBuilder {
pub fn create(
schema_desc: &SchemaDescriptor,
predicate: &(Arc<ParquetPredicate>, Vec<usize>),
topk: Option<&(Arc<ParquetTopK>, TableField)>,
topk: Option<&BuiltTopK>,
output_leaves: &[usize],
remain_schema: &TableSchema,
output_schema: &TableSchema,
Expand All @@ -70,8 +70,8 @@ impl PredicateAndTopkPolicyBuilder {

// Compute projections to read columns for each stage (prefetch and remain).
let mut prefetch_leaves = predicate_leaves.iter().cloned().collect::<HashSet<_>>();
if let Some((_, field)) = topk {
prefetch_leaves.insert(field.column_id as usize);
if let Some(topk) = topk {
prefetch_leaves.insert(topk.leaf_id);
}
// Remove prefetch columns
// TODO(parquet): reuse inner columns of a nested type.
Expand All @@ -85,23 +85,23 @@ impl PredicateAndTopkPolicyBuilder {

// Remain fields will not contain predicate fields.
// We just need to remove the topk column if it is contained in remain fields.
let remain_fields = if let Some((_, tf)) = topk {
let remain_fields = if let Some(topk) = topk {
remain_schema
.fields()
.iter()
.cloned()
.filter(|f| f.name() != tf.name())
.filter(|f| f.name() != topk.field.name())
.collect::<Vec<_>>()
} else {
remain_schema.fields().clone()
};

let (prefetch_fields, topk) = if let Some((topk, tf)) = topk {
let fields = vec![tf.clone()]
let (prefetch_fields, topk) = if let Some(topk) = topk {
let fields = vec![topk.field.clone()]
.into_iter()
.chain(predicate.schema().fields().clone().into_iter())
.collect::<Vec<_>>();
(fields, Some(topk.clone()))
(fields, Some(topk.topk.clone()))
} else {
(predicate.schema().fields().clone(), None)
};
Expand Down Expand Up @@ -160,7 +160,7 @@ impl ReadPolicyBuilder for PredicateAndTopkPolicyBuilder {
&row_group,
topk.field_levels(),
selection.clone(),
&None,
topk.field_paths(),
num_rows,
)?;
debug_assert_eq!(block.num_columns(), 1);
Expand Down
18 changes: 11 additions & 7 deletions src/query/storages/parquet/src/parquet_rs/parquet_reader/read_policy/topk_only.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ use common_expression::BlockEntry;
use common_expression::DataBlock;
use common_expression::DataSchema;
use common_expression::DataSchemaRef;
use common_expression::TableField;
use common_expression::TableSchema;
use common_expression::TopKSorter;
use parquet::arrow::arrow_reader::ParquetRecordBatchReader;
Expand All @@ -36,6 +35,7 @@ use super::policy::ReadPolicyImpl;
use super::utils::evaluate_topk;
use super::utils::read_all;
use crate::parquet_rs::parquet_reader::row_group::InMemoryRowGroup;
use crate::parquet_rs::parquet_reader::topk::BuiltTopK;
use crate::parquet_rs::parquet_reader::topk::ParquetTopK;
use crate::parquet_rs::parquet_reader::utils::compute_output_field_paths;
use crate::parquet_rs::parquet_reader::utils::transform_record_batch;
Expand All @@ -55,18 +55,22 @@ pub struct TopkOnlyPolicyBuilder {
impl TopkOnlyPolicyBuilder {
pub fn create(
schema_desc: &SchemaDescriptor,
topk: (Arc<ParquetTopK>, TableField),
topk: &BuiltTopK,
output_schema: &TableSchema,
output_leaves: &[usize],
inner_projection: bool,
) -> Result<Box<dyn ReadPolicyBuilder>> {
let (topk, topk_field) = topk;
let BuiltTopK {
topk,
field: topk_field,
leaf_id,
} = topk;

// Prefetch the topk column. Compute the remain columns.
let remain_leaves = output_leaves
.iter()
.cloned()
.filter(|i| *i != topk_field.column_id as usize)
.filter(|i| i != leaf_id)
.collect::<Vec<_>>();
let remain_projection = ProjectionMask::leaves(schema_desc, remain_leaves);
let remain_fields = output_schema
Expand All @@ -86,12 +90,12 @@ impl TopkOnlyPolicyBuilder {
)?);

let mut src_schema = remain_schema;
src_schema.fields.push(topk_field);
src_schema.fields.push(topk_field.clone());
let src_schema = Arc::new(DataSchema::from(&src_schema));
let dst_schema = Arc::new(DataSchema::from(output_schema));

Ok(Box::new(Self {
topk,
topk: topk.clone(),
remain_projection,
remain_field_levels,
remain_field_paths,
Expand Down Expand Up @@ -125,7 +129,7 @@ impl ReadPolicyBuilder for TopkOnlyPolicyBuilder {
&row_group,
self.topk.field_levels(),
selection.clone(),
&None,
self.topk.field_paths(),
num_rows,
)?;
let prefetched =
Expand Down
7 changes: 3 additions & 4 deletions src/query/storages/parquet/src/parquet_rs/parquet_reader/reader/builder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ use common_catalog::plan::PushDownInfo;
use common_catalog::plan::TopK;
use common_catalog::table_context::TableContext;
use common_exception::Result;
use common_expression::TableField;
use common_expression::TableSchemaRef;
use opendal::Operator;
use parquet::arrow::arrow_to_parquet_schema;
Expand All @@ -37,7 +36,7 @@ use crate::parquet_rs::parquet_reader::policy::POLICY_TOPK_ONLY;
use crate::parquet_rs::parquet_reader::predicate::build_predicate;
use crate::parquet_rs::parquet_reader::predicate::ParquetPredicate;
use crate::parquet_rs::parquet_reader::topk::build_topk;
use crate::parquet_rs::parquet_reader::topk::ParquetTopK;
use crate::parquet_rs::parquet_reader::topk::BuiltTopK;
use crate::parquet_rs::parquet_reader::utils::compute_output_field_paths;
use crate::parquet_rs::parquet_reader::utils::FieldPaths;
use crate::parquet_rs::parquet_reader::NoPretchPolicyBuilder;
Expand All @@ -59,7 +58,7 @@ pub struct ParquetRSReaderBuilder<'a> {

// Can be reused to build multiple readers.
built_predicate: Option<(Arc<ParquetPredicate>, Vec<usize>)>,
built_topk: Option<(Arc<ParquetTopK>, TableField)>,
built_topk: Option<BuiltTopK>,
#[allow(clippy::type_complexity)]
built_output: Option<(
ProjectionMask, // The output projection mask.
Expand Down Expand Up @@ -274,7 +273,7 @@ impl<'a> ParquetRSReaderBuilder<'a> {
let (_, output_leaves, output_schema, paths) = self.built_output.as_ref().unwrap();
TopkOnlyPolicyBuilder::create(
&self.schema_desc,
self.built_topk.clone().unwrap(),
self.built_topk.as_ref().unwrap(),
output_schema,
output_leaves,
paths.is_some(),
Expand Down
Loading
Loading