refactor: reuse arena result tuple across execution pipeline

Author: KKould

examples/hello_world.rs

@@ -14,7 +14,7 @@
 
 #[cfg(all(not(target_arch = "wasm32"), feature = "orm"))]
 mod app {
-    use kite_sql::db::{DataBaseBuilder, Database, ResultIter};
+    use kite_sql::db::{DataBaseBuilder, Database};
     use kite_sql::errors::DatabaseError;
     use kite_sql::storage::Storage;
     use kite_sql::Model;

examples/transaction.rs

@@ -14,7 +14,7 @@
 
 #[cfg(not(target_arch = "wasm32"))]
 mod app {
-    use kite_sql::db::{DataBaseBuilder, ResultIter};
+    use kite_sql::db::DataBaseBuilder;
     use kite_sql::errors::DatabaseError;
     use kite_sql::types::tuple::Tuple;
     use kite_sql::types::value::DataValue;

src/bin/server.rs

@@ -15,7 +15,7 @@
 use async_trait::async_trait;
 use clap::Parser;
 use futures::stream;
-use kite_sql::db::{DBTransaction, DataBaseBuilder, Database, ResultIter};
+use kite_sql::db::{BorrowResultIter, DBTransaction, DataBaseBuilder, Database};
 use kite_sql::errors::DatabaseError;
 use kite_sql::storage::rocksdb::RocksStorage;
 use kite_sql::types::tuple::{SchemaRef, Tuple};
@@ -28,6 +28,7 @@ use pgwire::api::query::{PlaceholderExtendedQueryHandler, SimpleQueryHandler};
 use pgwire::api::results::{DataRowEncoder, FieldFormat, FieldInfo, QueryResponse, Response, Tag};
 use pgwire::api::{ClientInfo, NoopErrorHandler, PgWireServerHandlers, Type};
 use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
+use pgwire::messages::data::DataRow;
 use pgwire::tokio::process_socket;
 use std::fmt::Debug;
 use std::io;
@@ -213,43 +214,48 @@ impl SimpleQueryHandler for SessionBackend {
             _ => {
                 let mut guard = self.tx.lock();
 
-                let mut tuples = Vec::new();
                 let response = if let Some(transaction) = guard.as_mut() {
                     let mut iter = unsafe { transaction.as_mut().run(query) }
                         .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
-                    for tuple in iter.by_ref() {
-                        tuples.push(tuple.map_err(|e| PgWireError::ApiError(Box::new(e)))?);
-                    }
-                    let schema = iter.schema().clone();
+                    let response = encode_query_result(&mut iter)?;
                     iter.done()
                         .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
-                    encode_tuples(&schema, tuples)?
+                    response
                 } else {
                     let mut iter = self
                         .inner
                         .run(query)
                         .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
-                    for tuple in iter.by_ref() {
-                        tuples.push(tuple.map_err(|e| PgWireError::ApiError(Box::new(e)))?);
-                    }
-                    let schema = iter.schema().clone();
+                    let response = encode_query_result(&mut iter)?;
                     iter.done()
                         .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
-                    encode_tuples(&schema, tuples)?
+                    response
                 };
                 Ok(vec![Response::Query(response)])
             }
         }
     }
 }
 
-fn encode_tuples<'a>(schema: &SchemaRef, tuples: Vec<Tuple>) -> PgWireResult<QueryResponse<'a>> {
-    if tuples.is_empty() {
-        return Ok(QueryResponse::new(Arc::new(vec![]), stream::empty()));
+fn encode_query_result<'a, I>(iter: &mut I) -> PgWireResult<QueryResponse<'a>>
+where
+    I: BorrowResultIter,
+{
+    let fields = encode_fields(iter.schema())?;
+    let mut results = Vec::new();
+
+    while let Some(tuple) = iter
+        .next_borrowed_tuple()
+        .map_err(|e| PgWireError::ApiError(Box::new(e)))?
+    {
+        results.push(encode_tuple(fields.clone(), tuple));
     }
 
-    let mut results = Vec::with_capacity(tuples.len());
-    let schema = Arc::new(
+    Ok(QueryResponse::new(fields, stream::iter(results)))
+}
+
+fn encode_fields(schema: &SchemaRef) -> PgWireResult<Arc<Vec<FieldInfo>>> {
+    Ok(Arc::new(
         schema
             .iter()
             .map(|column| {
@@ -264,41 +270,37 @@ fn encode_tuples<'a>(schema: &SchemaRef, tuples: Vec<Tuple>) -> PgWireResult<Que
                 ))
             })
             .collect::<PgWireResult<Vec<FieldInfo>>>()?,
-    );
-
-    for tuple in tuples {
-        let mut encoder = DataRowEncoder::new(schema.clone());
-        for value in tuple.values {
-            match value.logical_type() {
-                LogicalType::SqlNull => encoder.encode_field(&None::<i8>),
-                LogicalType::Boolean => encoder.encode_field(&value.bool()),
-                LogicalType::Tinyint => encoder.encode_field(&value.i8()),
-                LogicalType::UTinyint => encoder.encode_field(&value.u8().map(|v| v as i8)),
-                LogicalType::Smallint => encoder.encode_field(&value.i16()),
-                LogicalType::USmallint => encoder.encode_field(&value.u16().map(|v| v as i16)),
-                LogicalType::Integer => encoder.encode_field(&value.i32()),
-                LogicalType::UInteger => encoder.encode_field(&value.u32()),
-                LogicalType::Bigint => encoder.encode_field(&value.i64()),
-                LogicalType::UBigint => encoder.encode_field(&value.u64().map(|v| v as i64)),
-                LogicalType::Float => encoder.encode_field(&value.float()),
-                LogicalType::Double => encoder.encode_field(&value.double()),
-                LogicalType::Char(..) | LogicalType::Varchar(..) => {
-                    encoder.encode_field(&value.utf8())
-                }
-                LogicalType::Date => encoder.encode_field(&value.date()),
-                LogicalType::DateTime => encoder.encode_field(&value.datetime()),
-                LogicalType::Time(_) => encoder.encode_field(&value.time()),
-                LogicalType::Decimal(_, _) => {
-                    encoder.encode_field(&value.decimal().map(|decimal| decimal.to_string()))
-                }
-                _ => unreachable!(),
-            }?;
-        }
+    ))
+}
 
-        results.push(encoder.finish());
+fn encode_tuple(schema: Arc<Vec<FieldInfo>>, tuple: &Tuple) -> PgWireResult<DataRow> {
+    let mut encoder = DataRowEncoder::new(schema);
+    for value in &tuple.values {
+        match value.logical_type() {
+            LogicalType::SqlNull => encoder.encode_field(&None::<i8>),
+            LogicalType::Boolean => encoder.encode_field(&value.bool()),
+            LogicalType::Tinyint => encoder.encode_field(&value.i8()),
+            LogicalType::UTinyint => encoder.encode_field(&value.u8().map(|v| v as i8)),
+            LogicalType::Smallint => encoder.encode_field(&value.i16()),
+            LogicalType::USmallint => encoder.encode_field(&value.u16().map(|v| v as i16)),
+            LogicalType::Integer => encoder.encode_field(&value.i32()),
+            LogicalType::UInteger => encoder.encode_field(&value.u32()),
+            LogicalType::Bigint => encoder.encode_field(&value.i64()),
+            LogicalType::UBigint => encoder.encode_field(&value.u64().map(|v| v as i64)),
+            LogicalType::Float => encoder.encode_field(&value.float()),
+            LogicalType::Double => encoder.encode_field(&value.double()),
+            LogicalType::Char(..) | LogicalType::Varchar(..) => encoder.encode_field(&value.utf8()),
+            LogicalType::Date => encoder.encode_field(&value.date()),
+            LogicalType::DateTime => encoder.encode_field(&value.datetime()),
+            LogicalType::Time(_) => encoder.encode_field(&value.time()),
+            LogicalType::Decimal(_, _) => {
+                encoder.encode_field(&value.decimal().map(|decimal| decimal.to_string()))
+            }
+            _ => unreachable!(),
+        }?;
     }
 
-    Ok(QueryResponse::new(schema, stream::iter(results)))
+    encoder.finish()
 }
 
 fn into_pg_type(data_type: &LogicalType) -> PgWireResult<Type> {

src/db.rs

@@ -666,14 +666,38 @@ impl<S: Storage> Database<S> {
     }
 }
 
-/// Common interface for result iterators returned by database execution APIs.
-///
-/// A result iterator streams [`Tuple`] values and exposes the output schema of
-/// the current statement.
-pub trait ResultIter: Iterator<Item = Result<Tuple, DatabaseError>> {
+/// Borrowing interface for result iterators returned by database execution APIs.
+pub trait BorrowResultIter {
     /// Returns the output schema for the current result set.
     fn schema(&self) -> &SchemaRef;
 
+    /// Returns the next row as a borrowed tuple.
+    fn next_borrowed_tuple(&mut self) -> Result<Option<&Tuple>, DatabaseError>;
+
+    /// Creates a mapped iterator that transforms borrowed tuples into owned output values.
+    fn map_result<F, O>(self, mapper: F) -> MappedResultIter<Self, F, O>
+    where
+        Self: Sized,
+        F: for<'a> FnMut(&'a SchemaRef, &'a Tuple) -> Result<O, DatabaseError>,
+    {
+        let schema = self.schema().clone();
+        MappedResultIter {
+            inner: self,
+            mapper,
+            schema,
+            _marker: PhantomData,
+        }
+    }
+
+    /// Finishes consuming the iterator and flushes any remaining work.
+    fn done(self) -> Result<(), DatabaseError>;
+}
+
+/// Common interface for owned-tuple result iterators.
+///
+/// This remains for compatibility with existing callers that expect
+/// `Iterator<Item = Result<Tuple, DatabaseError>>`.
+pub trait ResultIter: BorrowResultIter + Iterator<Item = Result<Tuple, DatabaseError>> {
     #[cfg(feature = "orm")]
     /// Converts this iterator into a typed ORM iterator.
     ///
@@ -687,9 +711,46 @@ pub trait ResultIter: Iterator<Item = Result<Tuple, DatabaseError>> {
     {
         OrmIter::new(self)
     }
+}
 
-    /// Finishes consuming the iterator and flushes any remaining work.
-    fn done(self) -> Result<(), DatabaseError>;
+impl<I> ResultIter for I where I: BorrowResultIter + Iterator<Item = Result<Tuple, DatabaseError>> {}
+
+/// Typed adapter over a borrowing result iterator.
+pub struct MappedResultIter<I, F, O> {
+    inner: I,
+    mapper: F,
+    schema: SchemaRef,
+    _marker: PhantomData<O>,
+}
+
+impl<I, F, O> MappedResultIter<I, F, O>
+where
+    I: BorrowResultIter,
+    F: for<'a> FnMut(&'a SchemaRef, &'a Tuple) -> Result<O, DatabaseError>,
+{
+    pub fn schema(&self) -> &SchemaRef {
+        &self.schema
+    }
+
+    pub fn done(self) -> Result<(), DatabaseError> {
+        self.inner.done()
+    }
+}
+
+impl<I, F, O> Iterator for MappedResultIter<I, F, O>
+where
+    I: BorrowResultIter,
+    F: for<'a> FnMut(&'a SchemaRef, &'a Tuple) -> Result<O, DatabaseError>,
+{
+    type Item = Result<O, DatabaseError>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.inner.next_borrowed_tuple() {
+            Ok(Some(tuple)) => Some((self.mapper)(&self.schema, tuple)),
+            Ok(None) => None,
+            Err(err) => Some(Err(err)),
+        }
+    }
 }
 
 #[cfg(feature = "orm")]
@@ -760,6 +821,33 @@ impl<S: Storage> Drop for DatabaseIter<'_, S> {
     }
 }
 
+impl<S: Storage> DatabaseIter<'_, S> {
+    #[inline]
+    pub fn schema(&self) -> &SchemaRef {
+        unsafe { (*self.inner).schema() }
+    }
+
+    #[inline]
+    pub fn next_borrowed_tuple(&mut self) -> Result<Option<&Tuple>, DatabaseError> {
+        let result = unsafe { (*self.inner).next_borrowed_tuple() };
+        if result.as_ref().is_ok_and(Option::is_none) {
+            self._guard = None;
+        }
+        result
+    }
+
+    #[inline]
+    pub fn done(mut self) -> Result<(), DatabaseError> {
+        unsafe {
+            Box::from_raw(mem::replace(&mut self.inner, std::ptr::null_mut())).done()?;
+        }
+        unsafe {
+            Box::from_raw(mem::replace(&mut self.transaction, std::ptr::null_mut())).commit()?;
+        }
+        Ok(())
+    }
+}
+
 impl<S: Storage> Iterator for DatabaseIter<'_, S> {
     type Item = Result<Tuple, DatabaseError>;
 
@@ -772,19 +860,17 @@ impl<S: Storage> Iterator for DatabaseIter<'_, S> {
     }
 }
 
-impl<S: Storage> ResultIter for DatabaseIter<'_, S> {
+impl<S: Storage> BorrowResultIter for DatabaseIter<'_, S> {
     fn schema(&self) -> &SchemaRef {
-        unsafe { (*self.inner).schema() }
+        DatabaseIter::schema(self)
     }
 
-    fn done(mut self) -> Result<(), DatabaseError> {
-        unsafe {
-            Box::from_raw(mem::replace(&mut self.inner, std::ptr::null_mut())).done()?;
-        }
-        unsafe {
-            Box::from_raw(mem::replace(&mut self.transaction, std::ptr::null_mut())).commit()?;
-        }
-        Ok(())
+    fn next_borrowed_tuple(&mut self) -> Result<Option<&Tuple>, DatabaseError> {
+        DatabaseIter::next_borrowed_tuple(self)
+    }
+
+    fn done(self) -> Result<(), DatabaseError> {
+        DatabaseIter::done(self)
     }
 }
 
@@ -851,34 +937,55 @@ impl<'a, T: Transaction + 'a> TransactionIter<'a, T> {
     fn new(schema: SchemaRef, executor: Executor<'a, T>) -> Self {
         Self { executor, schema }
     }
+
+    #[inline]
+    pub fn schema(&self) -> &SchemaRef {
+        &self.schema
+    }
+
+    #[inline]
+    pub fn next_borrowed_tuple(&mut self) -> Result<Option<&Tuple>, DatabaseError> {
+        self.executor.next_tuple()
+    }
+
+    #[inline]
+    pub fn done(mut self) -> Result<(), DatabaseError> {
+        while self.next_borrowed_tuple()?.is_some() {}
+        Ok(())
+    }
 }
 
 impl<T: Transaction> Iterator for TransactionIter<'_, T> {
     type Item = Result<Tuple, DatabaseError>;
 
     fn next(&mut self) -> Option<Self::Item> {
-        self.executor.next()
+        match self.executor.next_tuple() {
+            Ok(Some(tuple)) => Some(Ok(tuple.clone())),
+            Ok(None) => None,
+            Err(err) => Some(Err(err)),
+        }
     }
 }
 
-impl<T: Transaction> ResultIter for TransactionIter<'_, T> {
+impl<T: Transaction> BorrowResultIter for TransactionIter<'_, T> {
     fn schema(&self) -> &SchemaRef {
-        &self.schema
+        TransactionIter::schema(self)
     }
 
-    fn done(mut self) -> Result<(), DatabaseError> {
-        for result in self.by_ref() {
-            let _ = result?;
-        }
-        Ok(())
+    fn next_borrowed_tuple(&mut self) -> Result<Option<&Tuple>, DatabaseError> {
+        TransactionIter::next_borrowed_tuple(self)
+    }
+
+    fn done(self) -> Result<(), DatabaseError> {
+        TransactionIter::done(self)
     }
 }
 
 #[cfg(all(test, not(target_arch = "wasm32")))]
 pub(crate) mod test {
     use crate::binder::{Binder, BinderContext};
     use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef};
-    use crate::db::{DataBaseBuilder, DatabaseError, ResultIter};
+    use crate::db::{DataBaseBuilder, DatabaseError};
     use crate::expression::ScalarExpression;
     use crate::planner::operator::join::JoinCondition;
     use crate::planner::operator::Operator;