dnim/rust-postgres
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Allocate one vec per row, not cell

Browse Source
This commit is contained in:
Steven Fackler 2016-12-18 19:41:23 -08:00
parent ad69145b56
commit 310a4888e9
3 changed files with 93 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

79
src/lib.rs
View File

@ -85,6 +85,7 @@ use std::collections::{VecDeque, HashMap};
use std::fmt;
use std::io;
use std::mem;
use std::ops::Range;
use std::result;
use std::sync::Arc;
use std::time::Duration;
@ -488,7 +489,9 @@ impl InnerConnection {
Ok((param_types, columns))
}
fn read_rows(&mut self, buf: &mut VecDeque<Vec<Option<Vec<u8>>>>) -> Result<bool> {
fn read_rows<F>(&mut self, mut consumer: F) -> Result<bool>
where F: FnMut(RowData)
{
let more_rows;
loop {
match try!(self.read_message()) {
@ -502,8 +505,12 @@ impl InnerConnection {
break;
}
backend::Message::DataRow(body) => {
let row = try!(body.values().map(|v| v.map(ToOwned::to_owned)).collect());
buf.push_back(row);
let mut row = RowData::new();
let mut it = body.values();
while let Some(value) = try!(it.next()) {
row.push(value);
}
consumer(row);
}
backend::Message::ErrorResponse(body) => {
try!(self.wait_for_ready());
@ -694,12 +701,11 @@ impl InnerConnection {
fn read_type(&mut self, oid: Oid) -> Result<Other> {
try!(self.setup_typeinfo_query());
try!(self.raw_execute(TYPEINFO_QUERY, "", 0, &[Type::Oid], &[&oid]));
let mut rows = VecDeque::new();
try!(self.read_rows(&mut rows));
let row = rows.pop_front();
let mut row = None;
try!(self.read_rows(|r| row = Some(r)));
let get_raw = |i: usize| {
row.as_ref().and_then(|r| r.get(i)).and_then(|r| r.as_ref().map(|r| &**r))
row.as_ref().and_then(|r| r.get(i))
};
let (name, type_, elem_oid, rngsubtype, basetype, schema, relid) = {
@ -770,14 +776,13 @@ impl InnerConnection {
fn read_enum_variants(&mut self, oid: Oid) -> Result<Vec<String>> {
try!(self.setup_typeinfo_enum_query());
try!(self.raw_execute(TYPEINFO_ENUM_QUERY, "", 0, &[Type::Oid], &[&oid]));
let mut rows = VecDeque::new();
try!(self.read_rows(&mut rows));
let mut rows = vec![];
try!(self.read_rows(|row| rows.push(row)));
let ctx = SessionInfo::new(&self.parameters);
let mut variants = vec![];
for row in rows {
let raw = row.get(0).and_then(|r| r.as_ref().map(|r| &**r));
variants.push(try!(String::from_sql_nullable(&Type::Name, raw, &ctx)
variants.push(try!(String::from_sql_nullable(&Type::Name, row.get(0), &ctx)
.map_err(Error::Conversion)));
}
@ -804,17 +809,16 @@ impl InnerConnection {
fn read_composite_fields(&mut self, relid: Oid) -> Result<Vec<Field>> {
try!(self.setup_typeinfo_composite_query());
try!(self.raw_execute(TYPEINFO_COMPOSITE_QUERY, "", 0, &[Type::Oid], &[&relid]));
let mut rows = VecDeque::new();
try!(self.read_rows(&mut rows));
let mut rows = vec![];
try!(self.read_rows(|row| rows.push(row)));
let mut fields = vec![];
for row in rows {
let (name, type_) = {
let get_raw = |i: usize| row.get(i).and_then(|r| r.as_ref().map(|r| &**r));
let ctx = SessionInfo::new(&self.parameters);
let name = try!(String::from_sql_nullable(&Type::Name, get_raw(0), &ctx)
let name = try!(String::from_sql_nullable(&Type::Name, row.get(0), &ctx)
.map_err(Error::Conversion));
let type_ = try!(Oid::from_sql_nullable(&Type::Oid, get_raw(1), &ctx)
let type_ = try!(Oid::from_sql_nullable(&Type::Oid, row.get(1), &ctx)
.map_err(Error::Conversion));
(name, type_)
};
@ -1329,6 +1333,43 @@ impl<'a> GenericConnection for Transaction<'a> {
}
}
struct RowData {
buf: Vec<u8>,
indices: Vec<Option<Range<usize>>>,
}
impl RowData {
fn new() -> RowData {
RowData {
buf: vec![],
indices: vec![],
}
}
fn push(&mut self, cell: Option<&[u8]>) {
let index = match cell {
Some(cell) => {
let base = self.buf.len();
self.buf.extend_from_slice(cell);
Some(base..self.buf.len())
}
None => None,
};
self.indices.push(index);
}
fn len(&self) -> usize {
self.indices.len()
}
fn get(&self, index: usize) -> Option<&[u8]> {
match &self.indices[index] {
&Some(ref range) => Some(&self.buf[range.clone()]),
&None => None,
}
}
}
trait OtherNew {
fn new(name: String, oid: Oid, kind: Kind, schema: String) -> Other;
}
@ -1340,13 +1381,13 @@ trait DbErrorNew {
}
trait RowsNew<'a> {
fn new(stmt: &'a Statement<'a>, data: Vec<Vec<Option<Vec<u8>>>>) -> Rows<'a>;
fn new_owned(stmt: Statement<'a>, data: Vec<Vec<Option<Vec<u8>>>>) -> Rows<'a>;
fn new(stmt: &'a Statement<'a>, data: Vec<RowData>) -> Rows<'a>;
fn new_owned(stmt: Statement<'a>, data: Vec<RowData>) -> Rows<'a>;
}
trait LazyRowsNew<'trans, 'stmt> {
fn new(stmt: &'stmt Statement<'stmt>,
data: VecDeque<Vec<Option<Vec<u8>>>>,
data: VecDeque<RowData>,
name: String,
row_limit: i32,
more_rows: bool,

23
src/rows.rs
View File

@ -3,14 +3,13 @@
use fallible_iterator::FallibleIterator;
use postgres_protocol::message::frontend;
use std::ascii::AsciiExt;
use std::borrow::Cow;
use std::collections::VecDeque;
use std::fmt;
use std::io;
use std::ops::Deref;
use std::slice;
use {Result, SessionInfoNew, RowsNew, LazyRowsNew, StatementInternals, WrongTypeNew};
use {Result, SessionInfoNew, RowsNew, LazyRowsNew, StatementInternals, WrongTypeNew, RowData};
use transaction::Transaction;
use types::{FromSql, SessionInfo, WrongType};
use stmt::{Statement, Column};
@ -35,18 +34,18 @@ impl<'a, T> Deref for MaybeOwned<'a, T> {
/// The resulting rows of a query.
pub struct Rows<'stmt> {
stmt: MaybeOwned<'stmt, Statement<'stmt>>,
data: Vec<Vec<Option<Vec<u8>>>>,
data: Vec<RowData>,
}
impl<'a> RowsNew<'a> for Rows<'a> {
fn new(stmt: &'a Statement<'a>, data: Vec<Vec<Option<Vec<u8>>>>) -> Rows<'a> {
fn new(stmt: &'a Statement<'a>, data: Vec<RowData>) -> Rows<'a> {
Rows {
stmt: MaybeOwned::Borrowed(stmt),
data: data,
}
}
fn new_owned(stmt: Statement<'a>, data: Vec<Vec<Option<Vec<u8>>>>) -> Rows<'a> {
fn new_owned(stmt: Statement<'a>, data: Vec<RowData>) -> Rows<'a> {
Rows {
stmt: MaybeOwned::Owned(stmt),
data: data,
@ -112,7 +111,7 @@ impl<'a> IntoIterator for &'a Rows<'a> {
/// An iterator over `Row`s.
pub struct Iter<'a> {
stmt: &'a Statement<'a>,
iter: slice::Iter<'a, Vec<Option<Vec<u8>>>>,
iter: slice::Iter<'a, RowData>,
}
impl<'a> Iterator for Iter<'a> {
@ -148,7 +147,7 @@ impl<'a> ExactSizeIterator for Iter<'a> {}
/// A single result row of a query.
pub struct Row<'a> {
stmt: &'a Statement<'a>,
data: MaybeOwned<'a, Vec<Option<Vec<u8>>>>,
data: MaybeOwned<'a, RowData>,
}
impl<'a> fmt::Debug for Row<'a> {
@ -238,7 +237,7 @@ impl<'a> Row<'a> {
}
let conn = self.stmt.conn().0.borrow();
let value = FromSql::from_sql_nullable(ty,
self.data[idx].as_ref().map(|r| &**r),
self.data.get(idx),
&SessionInfo::new(&conn.parameters));
Some(value.map_err(Error::Conversion))
}
@ -252,7 +251,7 @@ impl<'a> Row<'a> {
where I: RowIndex + fmt::Debug
{
match idx.idx(self.stmt) {
Some(idx) => self.data[idx].as_ref().map(|e| &**e),
Some(idx) => self.data.get(idx),
None => panic!("invalid index {:?}", idx),
}
}
@ -293,7 +292,7 @@ impl<'a> RowIndex for &'a str {
/// A lazily-loaded iterator over the resulting rows of a query.
pub struct LazyRows<'trans, 'stmt> {
stmt: &'stmt Statement<'stmt>,
data: VecDeque<Vec<Option<Vec<u8>>>>,
data: VecDeque<RowData>,
name: String,
row_limit: i32,
more_rows: bool,
@ -303,7 +302,7 @@ pub struct LazyRows<'trans, 'stmt> {
impl<'trans, 'stmt> LazyRowsNew<'trans, 'stmt> for LazyRows<'trans, 'stmt> {
fn new(stmt: &'stmt Statement<'stmt>,
data: VecDeque<Vec<Option<Vec<u8>>>>,
data: VecDeque<RowData>,
name: String,
row_limit: i32,
more_rows: bool,
@ -354,7 +353,7 @@ impl<'trans, 'stmt> LazyRows<'trans, 'stmt> {
try!(conn.stream.write_message(|buf| frontend::execute(&self.name, self.row_limit, buf)));
try!(conn.stream.write_message(|buf| Ok::<(), io::Error>(frontend::sync(buf))));
try!(conn.stream.flush());
conn.read_rows(&mut self.data).map(|more_rows| self.more_rows = more_rows)
conn.read_rows(|row| self.data.push_back(row)).map(|more_rows| self.more_rows = more_rows)
}
/// Returns a slice describing the columns of the `LazyRows`.

37
src/stmt.rs
View File

@ -12,7 +12,7 @@ use error::{Error, DbError};
use types::{SessionInfo, Type, ToSql};
use rows::{Rows, LazyRows};
use transaction::Transaction;
use {bad_response, Connection, StatementInternals, Result, RowsNew, InnerConnection,
use {bad_response, Connection, StatementInternals, Result, RowsNew, InnerConnection, RowData,
SessionInfoNew, LazyRowsNew, DbErrorNew, ColumnNew, StatementInfo, TransactionInternals};
/// A prepared statement.
@ -59,8 +59,9 @@ impl<'conn> StatementInternals<'conn> for Statement<'conn> {
fn into_query(self, params: &[&ToSql]) -> Result<Rows<'conn>> {
check_desync!(self.conn);
self.inner_query("", 0, params)
.map(|(buf, _)| Rows::new_owned(self, buf.into_iter().collect()))
let mut rows = vec![];
try!(self.inner_query("", 0, params, |row| rows.push(row)));
Ok(Rows::new_owned(self, rows))
}
}
@ -77,11 +78,14 @@ impl<'conn> Statement<'conn> {
}
#[allow(type_complexity)]
fn inner_query<'a>(&'a self,
portal_name: &str,
row_limit: i32,
params: &[&ToSql])
-> Result<(VecDeque<Vec<Option<Vec<u8>>>>, bool)> {
fn inner_query<F>(&self,
portal_name: &str,
row_limit: i32,
params: &[&ToSql],
acceptor: F)
-> Result<bool>
where F: FnMut(RowData)
{
let mut conn = self.conn.0.borrow_mut();
try!(conn.raw_execute(&self.info.name,
@ -90,9 +94,7 @@ impl<'conn> Statement<'conn> {
self.param_types(),
params));
let mut buf = VecDeque::new();
let more_rows = try!(conn.read_rows(&mut buf));
Ok((buf, more_rows))
conn.read_rows(acceptor)
}
/// Returns a slice containing the expected parameter types.
@ -200,7 +202,9 @@ impl<'conn> Statement<'conn> {
/// ```
pub fn query<'a>(&'a self, params: &[&ToSql]) -> Result<Rows<'a>> {
check_desync!(self.conn);
self.inner_query("", 0, params).map(|(buf, _)| Rows::new(self, buf.into_iter().collect()))
let mut rows = vec![];
try!(self.inner_query("", 0, params, |row| rows.push(row)));
Ok(Rows::new(self, rows))
}
/// Executes the prepared statement, returning a lazily loaded iterator
@ -239,9 +243,12 @@ impl<'conn> Statement<'conn> {
self.next_portal_id.set(id + 1);
let portal_name = format!("{}p{}", self.info.name, id);
self.inner_query(&portal_name, row_limit, params).map(move |(data, more_rows)| {
LazyRows::new(self, data, portal_name, row_limit, more_rows, false, trans)
})
let mut rows = VecDeque::new();
let more_rows = try!(self.inner_query(&portal_name,
row_limit,
params,
|row| rows.push_back(row)));
Ok(LazyRows::new(self, rows, portal_name, row_limit, more_rows, false, trans))
}
/// Executes a `COPY FROM STDIN` statement, returning the number of rows