Support loading only the requested fields (#31)

Author: seanlinsley

README.md

@@ -94,7 +94,7 @@ async fn example() -> anyhow::Result<()> {
     // Read
     let mut calls = 0;
     let filter = Filter::new(&[database_id], &[granularity], start..=end);
-    for group in CompressedQueryStats::load(db, filter.clone()).await? {
+    for group in CompressedQueryStats::load(db, filter.clone(), ()).await? {
         for stat in group.decompress()? {
             calls += stat.calls;
         }
@@ -106,7 +106,7 @@ async fn example() -> anyhow::Result<()> {
     assert_eq!(2, db.query_one("SELECT count(*) FROM query_stats", &[]).await?.get::<_, i64>(0));
     transaction!(db, {
         let mut stats = Vec::new();
-        for group in CompressedQueryStats::delete(db, filter.clone()).await? {
+        for group in CompressedQueryStats::delete(db, filter.clone(), ()).await? {
             stats.extend(group.decompress()?);
         }
         assert_eq!(0, db.query_one("SELECT count(*) FROM query_stats", &[]).await?.get::<_, i64>(0));
@@ -117,7 +117,7 @@ async fn example() -> anyhow::Result<()> {
         .await?;
     });
     assert_eq!(1, db.query_one("SELECT count(*) FROM query_stats", &[]).await?.get::<_, i64>(0));
-    let group = CompressedQueryStats::load(db, filter).await?.remove(0);
+    let group = CompressedQueryStats::load(db, filter, ()).await?.remove(0);
     assert_eq!(group.start_at, end - Duration::from_secs(120));
     assert_eq!(group.end_at, end - Duration::from_secs(60));
     let stats = group.decompress()?;
@@ -189,6 +189,15 @@ Timestamps support multiple formats:
 - `range_duration` returns the duration of the filter's time range
 - `range_shift` mutably shifts the time range's start and end by a certain amount, e.g. to filter for "today, 7 days ago"
 
+## Loading a subset of fields
+
+Read requests that don't need all fields in a struct can be optimized by declaring which fields they need, allowing pco_store to skip the others. Fields can be specified in several ways:
+- `()` or `Fields::default()`: load all fields
+- `[]` or `Fields::required()`: load only the required fields from `group_by` and `timestamp`
+- `["other"]` or `Fields::new(["other"])`: load extra fields in addition to the required ones
+
+Note that when optional filters are combined with `Fields::required()`, the fields needed by those filters are automatically added to the fields to be loaded.
+
 ## Contributions are welcome to
 
 - support decompression of only the fields requested at runtime

benches/comparison/pco_store.rs

@@ -91,7 +91,7 @@ pub async fn load() -> Result<()> {
     let database_ids: Vec<i64> = db.query_one("SELECT array_agg(DISTINCT database_id) FROM comparison_pco_stores", &[]).await?.get(0);
     let mut stats = Vec::new();
     let filter = Filter::new(&database_ids, SystemTime::UNIX_EPOCH..=SystemTime::now());
-    for group in CompressedQueryStats::load(db, filter).await? {
+    for group in CompressedQueryStats::load(db, filter, ()).await? {
         for stat in group.decompress()? {
             stats.push(stat);
         }

rust-toolchain.toml

@@ -0,0 +1,2 @@
+[toolchain]
+channel = "1.93.0"

src/lib.rs

@@ -73,21 +73,19 @@ pub fn store(args: TokenStream, item: TokenStream) -> TokenStream {
     };
 
     // load and delete
-    let mut fields = vec![quote! { filter: Option<Filter>, }];
+    let mut packed_fields = vec![quote! { filter: Option<Filter>, }];
     let mut load_checks = Vec::new();
     let mut load_where = Vec::new();
     let mut load_params = Vec::new();
-    let mut load_fields = Vec::new();
     let mut bind = 1;
-    let mut index: usize = 0;
     let mut timestamp_ty = None;
     let mut using_chrono = false;
     for field in model.fields.iter() {
         let ident = field.ident.clone().unwrap();
         let ty = field.ty.clone();
         let name = format!("{ident}");
         if group_by.iter().any(|i| *i == ident) {
-            fields.push(quote! { #ident: #ty, });
+            packed_fields.push(quote! { #ident: #ty, });
             load_checks.push(quote! {
                 if filter.#ident.is_empty() {
                     return Err(anyhow::Error::msg(#name.to_string() + " is required"));
@@ -99,7 +97,7 @@ pub fn store(args: TokenStream, item: TokenStream) -> TokenStream {
         } else if timestamp.as_ref().map(|t| *t == ident).unwrap_or(false) {
             using_chrono = !ty.to_token_stream().to_string().contains("SystemTime");
             timestamp_ty = Some(ty.clone());
-            fields.push(quote! { #ident: Vec<u8>, });
+            packed_fields.push(quote! { #ident: Vec<u8>, });
             load_checks.push(quote! {
                 if filter.#ident.is_none() {
                     return Err(anyhow::Error::msg(#name.to_string() + " is required"));
@@ -115,21 +113,13 @@ pub fn store(args: TokenStream, item: TokenStream) -> TokenStream {
                 filter.#ident.as_ref().unwrap().end(),
             });
         } else {
-            fields.push(quote! { #ident: Vec<u8>, });
+            packed_fields.push(quote! { #ident: Vec<u8>, });
         }
-        if timestamp.as_ref().map(|t| *t == ident).unwrap_or(false) {
-            index += 2; // Skip over start_at and end_at
-        }
-        load_fields.push(quote! { #ident: row.get(#index), });
-        index += 1;
     }
-    let fields = tokens(fields);
+    let packed_fields = tokens(packed_fields);
     let load_checks = tokens(load_checks);
     let load_where = if load_where.is_empty() { "true".to_string() } else { load_where.join(" AND ") };
     let load_params = tokens(load_params);
-    let load_fields = tokens(load_fields);
-    let load_sql = format!("SELECT * FROM {table_name} WHERE {load_where}");
-    let delete_sql = format!("DELETE FROM {table_name} WHERE {load_where} RETURNING *");
 
     // decompress
     let mut decompress_fields = Vec::new();
@@ -263,38 +253,53 @@ pub fn store(args: TokenStream, item: TokenStream) -> TokenStream {
     };
     let store_sql = format!("COPY {table_name} ({store_fields}) FROM STDIN BINARY");
 
-    let filter = filter(model, args, using_chrono, &timestamp_ty);
+    let filter = filter(model.clone(), args.clone(), using_chrono, &timestamp_ty);
+    let fields = fields(model, args, packed_name.clone());
     let deserialize_time_range = timestamp_ty.map(|t| deserialize_time_range(&t));
 
     quote! {
+        use serde::Deserialize as _;
+
         #item
 
         #[doc=concat!(" Generated by pco_store to store and load compressed versions of [", stringify!(#name), "]")]
         pub struct #packed_name {
-            #fields
+            #packed_fields
         }
 
         impl #packed_name {
             /// Loads data for the specified filters.
-            pub async fn load(db: &impl ::std::ops::Deref<Target = deadpool_postgres::ClientWrapper>, mut filter: Filter) -> anyhow::Result<Vec<#packed_name>> {
+            pub async fn load(
+                db: &impl ::std::ops::Deref<Target = deadpool_postgres::ClientWrapper>,
+                mut filter: Filter,
+                fields: impl TryInto<Fields>
+            ) -> anyhow::Result<Vec<#packed_name>> {
+                let mut fields = fields.try_into().map_err(|_| anyhow::Error::msg("unknown field"))?;
+                fields.merge_filter(&filter);
                 #load_checks
-                let sql = #load_sql;
+                let sql = "SELECT ".to_string() + fields.select().as_str() + " FROM " + #table_name + " WHERE " + #load_where;
                 let mut results = Vec::new();
                 for row in db.query(&db.prepare_cached(&sql).await?, &[#load_params]).await? {
-                    results.push(#packed_name { filter: Some(filter.clone()), #load_fields });
+                    results.push(fields.load_from_row(row, Some(filter.clone()))?);
                 }
                 Ok(results)
             }
 
             /// Deletes data for the specified filters, returning it to the caller.
             ///
             /// Note that all rows are returned from [decompress][Self::decompress] even if post-decompress filters would normally apply.
-            pub async fn delete(db: &impl ::std::ops::Deref<Target = deadpool_postgres::ClientWrapper>, mut filter: Filter) -> anyhow::Result<Vec<#packed_name>> {
+            pub async fn delete(
+                db: &impl ::std::ops::Deref<Target = deadpool_postgres::ClientWrapper>,
+                mut filter: Filter,
+                fields: impl TryInto<Fields>
+            ) -> anyhow::Result<Vec<#packed_name>> {
+                let mut fields = fields.try_into().map_err(|_| anyhow::Error::msg("unknown field"))?;
+                fields.merge_filter(&filter);
                 #load_checks
-                let sql = #delete_sql;
+                let sql = "DELETE FROM ".to_string() + #table_name + " WHERE " + #load_where + " RETURNING " + fields.select().as_str();
                 let mut results = Vec::new();
                 for row in db.query(&db.prepare_cached(&sql).await?, &[#load_params]).await? {
-                    results.push(#packed_name { filter: None, #load_fields });
+                    results.push(fields.load_from_row(row, None)?);
                 }
                 Ok(results)
             }
@@ -370,6 +375,7 @@ pub fn store(args: TokenStream, item: TokenStream) -> TokenStream {
         }
 
         #filter
+        #fields
         #deserialize_time_range
     }
     .into()
@@ -514,6 +520,162 @@ fn filter(model: ItemStruct, args: Arguments, using_chrono: bool, timestamp_ty:
     }
 }
 
+fn fields(model: ItemStruct, args: Arguments, packed_name: Ident) -> proc_macro2::TokenStream {
+    let name = model.ident.clone();
+    let Arguments { timestamp, group_by, .. } = args;
+    let mut fields = Vec::new();
+    let mut required = Vec::new();
+    let mut merge_filter = Vec::new();
+    let mut select = Vec::new();
+    let mut load = Vec::new();
+    let mut default = Vec::new();
+    let mut from = Vec::new();
+    for field in model.fields.iter() {
+        let ident = field.ident.clone().unwrap();
+        let name = format!("{ident}");
+        let is_timestamp = timestamp.as_ref().map(|t| *t == ident).unwrap_or(false);
+        fields.push(quote! { #ident: bool, });
+        if group_by.iter().any(|i| *i == ident) || is_timestamp {
+            required.push(quote! { #ident: true, });
+        } else {
+            required.push(quote! { #ident: false, });
+            merge_filter.push(quote! {
+                (!filter.#ident.is_empty()).then(|| self.#ident = true);
+            });
+        }
+        select.push(quote! { self.#ident.then(|| fields.push(#name)); });
+        load.push(quote! { #ident: if self.#ident {
+                let v = row.get(index);
+                index += 1;
+                v
+            } else {
+                Default::default()
+            },
+        });
+        default.push(quote! { #ident: true, });
+        from.push(quote! { #name => fields.#ident = true, });
+    }
+    let fields = tokens(fields);
+    let required = tokens(required);
+    let merge_filter = tokens(merge_filter);
+    let select = tokens(select);
+    let load = tokens(load);
+    let default = tokens(default);
+    let from = tokens(from);
+    quote! {
+        #[derive(Clone, Copy, Debug, PartialEq)]
+        #[doc=concat!(" Generated by pco_store to choose which fields to decompress when loading [", stringify!(#name), "]")]
+        pub struct Fields {
+            #fields
+        }
+
+        impl Fields {
+            pub fn new(fields: &[&str]) -> anyhow::Result<Self> {
+                fields.try_into().map_err(|e| anyhow::Error::msg(e))
+            }
+
+            pub fn required() -> Self {
+                Self { #required }
+            }
+
+            fn merge_filter(&mut self, filter: &Filter) {
+                #merge_filter
+            }
+
+            fn select(&self) -> String {
+                let mut fields = Vec::new();
+                #select
+                fields.join(", ")
+            }
+
+            fn load_from_row(&self, row: tokio_postgres::Row, filter: Option<Filter>) -> anyhow::Result<#packed_name> {
+                let mut index = 0;
+                Ok(#packed_name {
+                    filter,
+                    #load
+                })
+            }
+        }
+
+        impl Default for Fields {
+            fn default() -> Self {
+                Self { #default }
+            }
+        }
+
+        impl TryFrom<&[&str]> for Fields {
+            type Error = &'static str;
+            fn try_from(input: &[&str]) -> Result<Self, Self::Error> {
+                let mut fields = Self::required();
+                for s in input {
+                    match *s {
+                        #from
+                        _ => return Err("unknown field"),
+                    }
+                }
+                Ok(fields)
+            }
+        }
+
+        impl<const N: usize> TryFrom<&[&str; N]> for Fields {
+            type Error = &'static str;
+            fn try_from(input: &[&str; N]) -> Result<Self, Self::Error> {
+                Self::try_from(&input[..])
+            }
+        }
+
+        impl TryFrom<Vec<String>> for Fields {
+            type Error = &'static str;
+            fn try_from(input: Vec<String>) -> Result<Self, Self::Error> {
+                let input: Vec<_> = input.iter().map(|s| s.as_str()).collect();
+                Self::try_from(input.as_slice())
+            }
+        }
+
+        impl From<()> for Fields {
+            fn from(_: ()) -> Self {
+                Self::default()
+            }
+        }
+
+        impl<'de> serde::Deserialize<'de> for Fields {
+            fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+            where
+                D: serde::Deserializer<'de>,
+            {
+                deserializer.deserialize_any(FieldsVisitor)
+            }
+        }
+
+        struct FieldsVisitor;
+        impl<'de> serde::de::Visitor<'de> for FieldsVisitor {
+            type Value = Fields;
+
+            fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
+                formatter.write_str("an array of strings matching the struct fields")
+            }
+
+            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
+            where
+                A: serde::de::SeqAccess<'de>,
+            {
+                let mut fields = Vec::new();
+                while let Some(field) = seq.next_element()? {
+                    fields.push(field);
+                }
+                Fields::try_from(fields).map_err(serde::de::Error::custom)
+            }
+
+            fn visit_unit<E>(self) -> Result<Self::Value, E>
+            where
+                E: serde::de::Error,
+            {
+                Ok(Fields::default())
+            }
+        }
+    }
+}
+
 fn deserialize_time_range(timestamp_ty: &Type) -> proc_macro2::TokenStream {
     quote! {
         /// Deserializes many different time range formats:
@@ -527,8 +689,6 @@ fn deserialize_time_range(timestamp_ty: &Type) -> proc_macro2::TokenStream {
             Ok(TimeRange::deserialize(deserializer)?.0)
         }
 
-        use serde::Deserialize as _;
-
         #[derive(Debug, PartialEq)]
         struct TimeRange(Option<std::ops::RangeInclusive<#timestamp_ty>>);
         impl<'de> serde::Deserialize<'de> for TimeRange {

tests/chrono_tests/mod.rs

@@ -47,7 +47,7 @@ async fn systemtime_chrono_interop() {
     {
         let mut actual: Vec<QueryStat> = vec![];
         let filter = Filter::new(&[database_id], start.into()..=end.into());
-        for group in CompressedQueryStats::load(db, filter).await.unwrap() {
+        for group in CompressedQueryStats::load(db, filter, ()).await.unwrap() {
             for stat in group.decompress().unwrap() {
                 actual.push(stat.clone());
             }
@@ -70,7 +70,7 @@ async fn systemtime_chrono_interop() {
         // Read, using chrono::DateTime.
         let mut actual: Vec<QueryStat> = vec![];
         let filter = Filter::new(&[database_id], start..=end);
-        for group in CompressedQueryStats::load(db, filter).await.unwrap() {
+        for group in CompressedQueryStats::load(db, filter, ()).await.unwrap() {
             for stat in group.decompress().unwrap() {
                 actual.push(stat.clone());
             }
@@ -90,7 +90,7 @@ async fn systemtime_chrono_interop() {
     // Read again, using SystemTime.
     let mut actual: Vec<QueryStat> = vec![];
     let filter = Filter::new(&[database_id], start.into()..=end.into());
-    for group in CompressedQueryStats::load(db, filter).await.unwrap() {
+    for group in CompressedQueryStats::load(db, filter, ()).await.unwrap() {
         for stat in group.decompress().unwrap() {
             actual.push(stat.clone());
         }