diff --git a/src/lib.rs b/src/lib.rs index 4e6d5f5b..fd3b3cec 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -205,8 +205,10 @@ where &'get self, key: &'get R::Key, ts: Timestamp, + projection_mask: Option>, ) -> Result>, ParquetError> { self.scan(Bound::Included(key), Bound::Unbounded, ts) + .project(projection_mask) .take() .await? .next() @@ -271,16 +273,20 @@ where } } - pub fn project(self, mut projection: Vec) -> Self { - // skip two columns: _null and _ts - for p in &mut projection { - *p += 2; - } + pub fn project(self, projection: Option>) -> Self { + let mask = projection + .map(|mut projection| { + // skip two columns: _null and _ts + for p in &mut projection { + *p += 2; + } - let mask = ProjectionMask::roots( - &arrow_to_parquet_schema(R::arrow_schema()).unwrap(), - projection, - ); + ProjectionMask::roots( + &arrow_to_parquet_schema(R::arrow_schema()).unwrap(), + projection, + ) + }) + .unwrap_or_else(ProjectionMask::all); Self { projection: mask, diff --git a/src/transaction.rs b/src/transaction.rs index 0c1ef44f..02d78fd4 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -52,12 +52,13 @@ where pub async fn get<'get>( &'get self, key: &'get R::Key, + projection_mask: Option>, ) -> Result>, ParquetError> { Ok(match self.local.get(key).and_then(|v| v.as_ref()) { Some(v) => Some(TransactionEntry::Local(v.as_record_ref())), None => self .share - .get(key, self.ts) + .get(key, self.ts, projection_mask) .await? .map(TransactionEntry::Stream), }) @@ -168,7 +169,7 @@ mod tests { txn1.set("foo".to_string()); let txn2 = db.transaction().await; - dbg!(txn2.get(&"foo".to_string()).await.unwrap().is_none()); + dbg!(txn2.get(&"foo".to_string(), None).await.unwrap().is_none()); txn1.commit().await.unwrap(); txn2.commit().await.unwrap(); @@ -176,7 +177,7 @@ mod tests { { let txn3 = db.transaction().await; - dbg!(txn3.get(&"foo".to_string()).await.unwrap().is_none()); + dbg!(txn3.get(&"foo".to_string(), None).await.unwrap().is_none()); txn3.commit().await.unwrap(); } } @@ -235,7 +236,7 @@ mod tests { }); let key = 0.to_string(); - let entry = txn1.get(&key).await.unwrap().unwrap(); + let entry = txn1.get(&key, None).await.unwrap().unwrap(); assert_eq!(entry.get().vstring, 0.to_string()); assert_eq!(entry.get().vu32, Some(0));