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

Jump to bottom

feat: support sampling table with block and row level simultaneously #16613

Merged
merged 3 commits into from
Oct 16, 2024
Merged

feat: support sampling table with block and row level simultaneously #16613

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
f2a1235
feat: support sampling table with block and row level simultaneously
xudong963 Oct 15, 2024
34dcedf
fix lint
xudong963 Oct 16, 2024
73dc8b1
fix ut
xudong963 Oct 16, 2024
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
96 changes: 62 additions & 34 deletions src/query/ast/src/ast/query.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@ use crate::ast::Identifier;
use crate::ast::Lambda;
use crate::ast::SelectStageOptions;
use crate::ast::WindowDefinition;
use crate::ParseError;
use crate::Result;
use crate::Span;

/// Root node of a query tree
Expand Down Expand Up @@ -623,56 +625,82 @@ impl Display for TemporalClause {
}
}

#[derive(serde::Serialize, serde::Deserialize, Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
pub enum SampleLevel {
ROW,
BLOCK,
}

#[derive(serde::Serialize, serde::Deserialize, Debug, Clone, PartialEq, Drive, DriveMut)]
pub enum SampleConfig {
Probability(f64),
pub enum SampleRowLevel {
RowsNum(f64),
Probability(f64),
}

impl Eq for SampleConfig {}

#[derive(serde::Serialize, serde::Deserialize, Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
pub struct Sample {
pub sample_level: SampleLevel,
pub sample_conf: SampleConfig,
}

impl Sample {
pub fn sample_probability(&self, stats_rows: Option<u64>) -> Option<f64> {
let rand = match &self.sample_conf {
SampleConfig::Probability(probability) => probability / 100.0,
SampleConfig::RowsNum(rows) => {
impl SampleRowLevel {
pub fn sample_probability(&self, stats_rows: Option<u64>) -> Result<Option<f64>> {
let rand = match &self {
SampleRowLevel::Probability(probability) => probability / 100.0,
SampleRowLevel::RowsNum(rows) => {
if let Some(row_num) = stats_rows {
if row_num > 0 {
rows / row_num as f64
} else {
return None;
return Ok(None);
}
} else {
return None;
return Ok(None);
}
}
};
Some(rand)
if rand > 1.0 {
return Err(ParseError(
None,
format!(
"Sample value should be less than or equal to 100, but got {}",
rand * 100.0
),
));
}
Ok(Some(rand))
}
}

impl Eq for SampleRowLevel {}

#[derive(
serde::Serialize, serde::Deserialize, Debug, Clone, PartialEq, Drive, DriveMut, Default,
)]
pub struct SampleConfig {
pub row_level: Option<SampleRowLevel>,
pub block_level: Option<f64>,
}

impl SampleConfig {
pub fn set_row_level_sample(&mut self, value: f64, rows: bool) {
if rows {
self.row_level = Some(SampleRowLevel::RowsNum(value));
} else {
self.row_level = Some(SampleRowLevel::Probability(value));
}
}

pub fn set_block_level_sample(&mut self, probability: f64) {
self.block_level = Some(probability);
}
}

impl Display for Sample {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
impl Eq for SampleConfig {}

impl Display for SampleConfig {
fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
write!(f, "SAMPLE ")?;
match self.sample_level {
SampleLevel::ROW => write!(f, "ROW ")?,
SampleLevel::BLOCK => write!(f, "BLOCK ")?,
if let Some(block_level) = self.block_level {
write!(f, "BLOCK ({}) ", block_level)?;
}
match &self.sample_conf {
SampleConfig::Probability(prob) => write!(f, "({})", prob)?,
SampleConfig::RowsNum(rows) => write!(f, "({} ROWS)", rows)?,
if let Some(row_level) = &self.row_level {
match row_level {
SampleRowLevel::RowsNum(rows) => {
write!(f, "ROW ({} ROWS)", rows)?;
}
SampleRowLevel::Probability(probability) => {
write!(f, "ROW ({})", probability)?;
}
}
}
Ok(())
}
Expand All @@ -692,7 +720,7 @@ pub enum TableReference {
with_options: Option<WithOptions>,
pivot: Option<Box<Pivot>>,
unpivot: Option<Box<Unpivot>>,
sample: Option<Sample>,
sample: Option<SampleConfig>,
},
// `TABLE(expr)[ AS alias ]`
TableFunction {
Expand All @@ -703,7 +731,7 @@ pub enum TableReference {
params: Vec<Expr>,
named_params: Vec<(Identifier, Expr)>,
alias: Option<TableAlias>,
sample: Option<Sample>,
sample: Option<SampleConfig>,
},
// Derived table, which can be a subquery or joined tables or combination of them
Subquery {
Expand Down
53 changes: 20 additions & 33 deletions src/query/ast/src/parser/query.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -702,7 +702,7 @@ pub enum TableReferenceElement {
with_options: Option<WithOptions>,
pivot: Option<Box<Pivot>>,
unpivot: Option<Box<Unpivot>>,
sample: Option<Sample>,
sample: Option<SampleConfig>,
},
// `TABLE(expr)[ AS alias ]`
TableFunction {
Expand All @@ -711,7 +711,7 @@ pub enum TableReferenceElement {
name: Identifier,
params: Vec<TableFunctionParam>,
alias: Option<TableAlias>,
sample: Option<Sample>,
sample: Option<SampleConfig>,
},
// Derived table, which can be a subquery or joined tables or combination of them
Subquery {
Expand Down Expand Up @@ -760,7 +760,7 @@ pub fn table_reference_element(i: Input) -> IResult<WithSpan<TableReferenceEleme
);
let aliased_table = map(
rule! {
#dot_separated_idents_1_to_3 ~ #temporal_clause? ~ #with_options? ~ #table_alias? ~ #pivot? ~ #unpivot? ~ SAMPLE? ~ (ROW | BLOCK)? ~ ("(" ~ #expr ~ ROWS? ~ ")")?
#dot_separated_idents_1_to_3 ~ #temporal_clause? ~ #with_options? ~ #table_alias? ~ #pivot? ~ #unpivot? ~ SAMPLE? ~ (BLOCK ~ "(" ~ #expr ~ ")")? ~ (ROW ~ "(" ~ #expr ~ ROWS? ~ ")")?
},
|(
(catalog, database, table),
Expand All @@ -770,10 +770,10 @@ pub fn table_reference_element(i: Input) -> IResult<WithSpan<TableReferenceEleme
pivot,
unpivot,
sample,
level,
sample_conf,
sample_block_level,
sample_row_level,
)| {
let table_sample = get_table_sample(sample, level, sample_conf);
let table_sample = get_table_sample(sample, sample_block_level, sample_row_level);
TableReferenceElement::Table {
catalog,
database,
Expand Down Expand Up @@ -810,7 +810,7 @@ pub fn table_reference_element(i: Input) -> IResult<WithSpan<TableReferenceEleme
);
let table_function = map(
rule! {
LATERAL? ~ #function_name ~ "(" ~ #comma_separated_list0(table_function_param) ~ ")" ~ #table_alias? ~ SAMPLE? ~ (ROW | BLOCK)? ~ ("(" ~ #expr ~ ROWS? ~ ")")?
LATERAL? ~ #function_name ~ "(" ~ #comma_separated_list0(table_function_param) ~ ")" ~ #table_alias? ~ SAMPLE? ~ (BLOCK ~ "(" ~ #expr ~ ")")? ~ (ROW ~ "(" ~ #expr ~ ROWS? ~ ")")?
},
|(lateral, name, _, params, _, alias, sample, level, sample_conf)| {
let table_sample = get_table_sample(sample, level, sample_conf);
Expand Down Expand Up @@ -871,34 +871,21 @@ pub fn table_reference_element(i: Input) -> IResult<WithSpan<TableReferenceEleme

fn get_table_sample(
sample: Option<&Token>,
level: Option<&Token>,
sample_conf: Option<(&Token, Expr, Option<&Token>, &Token)>,
) -> Option<Sample> {
let mut table_sample = None;
block_level_sample: Option<(&Token, &Token, Expr, &Token)>,
row_level_sample: Option<(&Token, &Token, Expr, Option<&Token>, &Token)>,
) -> Option<SampleConfig> {
let mut default_sample_conf = SampleConfig::default();
if sample.is_some() {
let sample_level = match level {
// If the sample level is not specified, it defaults to ROW
Some(level) => match level.kind {
ROW => SampleLevel::ROW,
BLOCK => SampleLevel::BLOCK,
_ => unreachable!(),
},
None => SampleLevel::ROW,
};
let mut default_sample_conf = SampleConfig::Probability(100.0);
if let Some((_, Expr::Literal { value, .. }, rows, _)) = sample_conf {
default_sample_conf = if rows.is_some() {
SampleConfig::RowsNum(value.as_double().unwrap_or_default())
} else {
SampleConfig::Probability(value.as_double().unwrap_or_default())
};
if let Some((_, _, Expr::Literal { value, .. }, _)) = block_level_sample {
default_sample_conf.set_block_level_sample(value.as_double().unwrap_or_default());
}
if let Some((_, _, Expr::Literal { value, .. }, rows, _)) = row_level_sample {
default_sample_conf
.set_row_level_sample(value.as_double().unwrap_or_default(), rows.is_some());
}
table_sample = Some(Sample {
sample_level,
sample_conf: default_sample_conf,
})
};
table_sample
return Some(default_sample_conf);
}
None
}

struct TableReferenceParser;
Expand Down
3 changes: 2 additions & 1 deletion src/query/ast/tests/it/parser.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,10 +224,11 @@ fn test_statement() {
r#"select * from t sample row (99);"#,
r#"select * from t sample block (99);"#,
r#"select * from t sample row (10 rows);"#,
r#"select * from t sample block (10 rows);"#,
r#"select * from numbers(1000) sample row (99);"#,
r#"select * from numbers(1000) sample block (99);"#,
r#"select * from numbers(1000) sample row (10 rows);"#,
r#"select * from numbers(1000) sample block (99) row (10 rows);"#,
r#"select * from numbers(1000) sample block (99) row (10);"#,
r#"insert into t (c1, c2) values (1, 2), (3, 4);"#,
r#"insert into t (c1, c2) values (1, 2);"#,
r#"insert into table t select * from t2;"#,
Expand Down
Loading
Loading