Merge pull request #903 from gwenn/series

Sync series with official source
This commit is contained in:
gwenn 2021-03-07 09:32:55 +01:00 committed by GitHub
commit e68d752cf4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -13,7 +13,7 @@ use crate::vtab::{
eponymous_only_module, Context, IndexConstraintOp, IndexInfo, VTab, VTabConnection, VTabCursor, eponymous_only_module, Context, IndexConstraintOp, IndexInfo, VTab, VTabConnection, VTabCursor,
Values, Values,
}; };
use crate::{Connection, Result}; use crate::{Connection, Error, Result};
/// `feature = "series"` Register the "generate_series" module. /// `feature = "series"` Register the "generate_series" module.
pub fn load_module(conn: &Connection) -> Result<()> { pub fn load_module(conn: &Connection) -> Result<()> {
@ -38,6 +38,8 @@ bitflags::bitflags! {
const STEP = 4; const STEP = 4;
// output in descending order // output in descending order
const DESC = 8; const DESC = 8;
// output in descending order
const ASC = 16;
// Both start and stop // Both start and stop
const BOTH = QueryPlanFlags::START.bits | QueryPlanFlags::STOP.bits; const BOTH = QueryPlanFlags::START.bits | QueryPlanFlags::STOP.bits;
} }
@ -71,54 +73,42 @@ unsafe impl<'vtab> VTab<'vtab> for SeriesTab {
fn best_index(&self, info: &mut IndexInfo) -> Result<()> { fn best_index(&self, info: &mut IndexInfo) -> Result<()> {
// The query plan bitmask // The query plan bitmask
let mut idx_num: QueryPlanFlags = QueryPlanFlags::empty(); let mut idx_num: QueryPlanFlags = QueryPlanFlags::empty();
// Index of the start= constraint // Mask of unusable constraints
let mut start_idx = None; let mut unusable_mask: QueryPlanFlags = QueryPlanFlags::empty();
// Index of the stop= constraint // Constraints on start, stop, and step
let mut stop_idx = None; let mut a_idx: [Option<usize>; 3] = [None, None, None];
// Index of the step= constraint
let mut step_idx = None;
for (i, constraint) in info.constraints().enumerate() { for (i, constraint) in info.constraints().enumerate() {
if !constraint.is_usable() { if constraint.column() < SERIES_COLUMN_START {
continue; continue;
} }
if constraint.operator() != IndexConstraintOp::SQLITE_INDEX_CONSTRAINT_EQ { let (i_col, i_mask) = match constraint.column() {
continue; SERIES_COLUMN_START => (0, QueryPlanFlags::START),
SERIES_COLUMN_STOP => (1, QueryPlanFlags::STOP),
SERIES_COLUMN_STEP => (2, QueryPlanFlags::STEP),
_ => {
unreachable!()
} }
match constraint.column() {
SERIES_COLUMN_START => {
start_idx = Some(i);
idx_num |= QueryPlanFlags::START;
}
SERIES_COLUMN_STOP => {
stop_idx = Some(i);
idx_num |= QueryPlanFlags::STOP;
}
SERIES_COLUMN_STEP => {
step_idx = Some(i);
idx_num |= QueryPlanFlags::STEP;
}
_ => {}
}; };
if !constraint.is_usable() {
unusable_mask |= i_mask;
} else if constraint.operator() == IndexConstraintOp::SQLITE_INDEX_CONSTRAINT_EQ {
idx_num |= i_mask;
a_idx[i_col] = Some(i);
} }
}
let mut num_of_arg = 0; // Number of arguments that SeriesTabCursor::filter expects
if let Some(start_idx) = start_idx { let mut n_arg = 0;
num_of_arg += 1; for j in a_idx.iter().flatten() {
let mut constraint_usage = info.constraint_usage(start_idx); n_arg += 1;
constraint_usage.set_argv_index(num_of_arg); let mut constraint_usage = info.constraint_usage(*j);
constraint_usage.set_argv_index(n_arg);
constraint_usage.set_omit(true); constraint_usage.set_omit(true);
} }
if let Some(stop_idx) = stop_idx { if !(unusable_mask & !idx_num).is_empty() {
num_of_arg += 1; return Err(Error::SqliteFailure(
let mut constraint_usage = info.constraint_usage(stop_idx); ffi::Error::new(ffi::SQLITE_CONSTRAINT),
constraint_usage.set_argv_index(num_of_arg); None,
constraint_usage.set_omit(true); ));
}
if let Some(step_idx) = step_idx {
num_of_arg += 1;
let mut constraint_usage = info.constraint_usage(step_idx);
constraint_usage.set_argv_index(num_of_arg);
constraint_usage.set_omit(true);
} }
if idx_num.contains(QueryPlanFlags::BOTH) { if idx_num.contains(QueryPlanFlags::BOTH) {
// Both start= and stop= boundaries are available. // Both start= and stop= boundaries are available.
@ -135,6 +125,8 @@ unsafe impl<'vtab> VTab<'vtab> for SeriesTab {
if let Some(order_by) = order_bys.next() { if let Some(order_by) = order_bys.next() {
if order_by.is_order_by_desc() { if order_by.is_order_by_desc() {
idx_num |= QueryPlanFlags::DESC; idx_num |= QueryPlanFlags::DESC;
} else {
idx_num |= QueryPlanFlags::ASC;
} }
true true
} else { } else {
@ -145,7 +137,9 @@ unsafe impl<'vtab> VTab<'vtab> for SeriesTab {
info.set_order_by_consumed(true); info.set_order_by_consumed(true);
} }
} else { } else {
info.set_estimated_cost(2_147_483_647f64); // If either boundary is missing, we have to generate a huge span
// of numbers. Make this case very expensive so that the query
// planner will work hard to avoid it.
info.set_estimated_rows(2_147_483_647); info.set_estimated_rows(2_147_483_647);
} }
info.set_idx_num(idx_num.bits()); info.set_idx_num(idx_num.bits());
@ -193,7 +187,7 @@ impl SeriesTabCursor<'_> {
} }
unsafe impl VTabCursor for SeriesTabCursor<'_> { unsafe impl VTabCursor for SeriesTabCursor<'_> {
fn filter(&mut self, idx_num: c_int, _idx_str: Option<&str>, args: &Values<'_>) -> Result<()> { fn filter(&mut self, idx_num: c_int, _idx_str: Option<&str>, args: &Values<'_>) -> Result<()> {
let idx_num = QueryPlanFlags::from_bits_truncate(idx_num); let mut idx_num = QueryPlanFlags::from_bits_truncate(idx_num);
let mut i = 0; let mut i = 0;
if idx_num.contains(QueryPlanFlags::START) { if idx_num.contains(QueryPlanFlags::START) {
self.min_value = args.get(i)?; self.min_value = args.get(i)?;
@ -209,8 +203,14 @@ unsafe impl VTabCursor for SeriesTabCursor<'_> {
} }
if idx_num.contains(QueryPlanFlags::STEP) { if idx_num.contains(QueryPlanFlags::STEP) {
self.step = args.get(i)?; self.step = args.get(i)?;
if self.step < 1 { #[allow(clippy::comparison_chain)]
if self.step == 0 {
self.step = 1; self.step = 1;
} else if self.step < 0 {
self.step = -self.step;
if !idx_num.contains(QueryPlanFlags::ASC) {
idx_num |= QueryPlanFlags::DESC;
}
} }
} else { } else {
self.step = 1; self.step = 1;
@ -274,6 +274,7 @@ mod test {
use crate::ffi; use crate::ffi;
use crate::vtab::series; use crate::vtab::series;
use crate::{Connection, Result}; use crate::{Connection, Result};
use fallible_iterator::FallibleIterator;
#[test] #[test]
fn test_series_module() -> Result<()> { fn test_series_module() -> Result<()> {
@ -294,6 +295,18 @@ mod test {
assert_eq!(expected, value?); assert_eq!(expected, value?);
expected += 5; expected += 5;
} }
let mut s =
db.prepare("SELECT * FROM generate_series WHERE start=1 AND stop=9 AND step=2")?;
let series: Vec<i32> = s.query([])?.map(|r| r.get(0)).collect()?;
assert_eq!(vec![1, 3, 5, 7, 9], series);
let mut s = db.prepare("SELECT * FROM generate_series LIMIT 5")?;
let series: Vec<i32> = s.query([])?.map(|r| r.get(0)).collect()?;
assert_eq!(vec![0, 1, 2, 3, 4], series);
let mut s = db.prepare("SELECT * FROM generate_series(0,32,5) ORDER BY value DESC")?;
let series: Vec<i32> = s.query([])?.map(|r| r.get(0)).collect()?;
assert_eq!(vec![30, 25, 20, 15, 10, 5, 0], series);
Ok(()) Ok(())
} }
} }