Allocate one vec per row, not cell

sfackler · sfackler · commit 310a4888e90b · 2016-12-18T20:07:54.000-08:00
diff --git a/src/lib.rs b/src/lib.rs
@@ -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,
diff --git a/src/rows.rs b/src/rows.rs
@@ -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`.
diff --git a/src/stmt.rs b/src/stmt.rs
@@ -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
