Optimize reward-amounts query for postgres' (#3819)

* add index for reward-amounts query

* remove where

* include entry

* change query

* created desc for sqlite

* fix tests

Author: imabdulbasit

sequencer/api/migrations/postgres/V1102__reward_merkle_tree_v2_idx_index.sql

@@ -0,0 +1,4 @@
+-- Used by: get_all_reward_accounts() in sequencer/src/api/sql.rs
+
+CREATE INDEX reward_merkle_tree_v2_idx_created
+ON reward_merkle_tree_v2 (idx, created) INCLUDE (entry);

sequencer/api/migrations/sqlite/V902__reward_merkle_tree_v2_idx_index.sql

@@ -0,0 +1,4 @@
+-- Used by: get_all_reward_accounts() in sequencer/src/api/sql.rs
+
+CREATE INDEX reward_merkle_tree_v2_idx_created
+ON reward_merkle_tree_v2 (idx, created DESC);

sequencer/src/api.rs

@@ -6654,7 +6654,8 @@ mod test {
             .iter()
             .map(|(addr, amt)| (*addr, *amt))
             .collect();
-        expected.sort_by_key(|(acct, _)| *acct);
+        // Results are sorted by account address descending
+        expected.sort_by_key(|(acct, _)| std::cmp::Reverse(*acct));
 
         tracing::info!("expected accounts = {expected:?}");
         let limit = expected.len().min(10_000) as u64;
@@ -6808,16 +6809,16 @@ mod test {
         }
 
         // Test pagination
-        // results are sorted by account address
+        // results are sorted by account address descending
         let result_limit_2 = db.get_all_reward_accounts(15, 0, 2).await?;
         assert_eq!(result_limit_2.len(), 2);
-        assert_eq!(result_limit_2[0], (account1, RewardAmount::from(1500u64)));
-        assert_eq!(result_limit_2[1], (account2, RewardAmount::from(2500u64)));
+        assert_eq!(result_limit_2[0], (account4, RewardAmount::from(4000u64)));
+        assert_eq!(result_limit_2[1], (account3, RewardAmount::from(3000u64)));
 
         let result_offset_2 = db.get_all_reward_accounts(15, 2, 2).await?;
         assert_eq!(result_offset_2.len(), 2);
-        assert_eq!(result_offset_2[0], (account3, RewardAmount::from(3000u64)));
-        assert_eq!(result_offset_2[1], (account4, RewardAmount::from(4000u64)));
+        assert_eq!(result_offset_2[0], (account2, RewardAmount::from(2500u64)));
+        assert_eq!(result_offset_2[1], (account1, RewardAmount::from(1500u64)));
 
         Ok(())
     }

sequencer/src/api/sql.rs

@@ -294,27 +294,40 @@ impl CatchupStorage for SqlStorage {
             return Ok(Vec::new());
         }
 
-        // get the latest balance for each account
-        // We use ROW_NUMBER() to rank entries by created height per account,
-        // then select only the most recent entry (rn = 1) for each account.
-        let rows = query_as::<(Value, Value)>(&format!(
+        // get the latest balance for each account.
+        // use DISTINCT ON for Postgres
+        // use ROW_NUMBER() as DISTINCT ON is not supported for SQLite
+        #[cfg(not(feature = "embedded-db"))]
+        let query = format!(
+            "SELECT DISTINCT ON (idx) idx, entry
+               FROM {}
+              WHERE idx IS NOT NULL AND created <= $1
+              ORDER BY idx DESC, created DESC
+              LIMIT $2 OFFSET $3",
+            RewardMerkleTreeV2::state_type()
+        );
+
+        #[cfg(feature = "embedded-db")]
+        let query = format!(
             "SELECT idx, entry FROM (
                  SELECT idx, entry, ROW_NUMBER() OVER (PARTITION BY idx ORDER BY created DESC) as \
              rn
                    FROM {}
-                  WHERE created <= $1 AND idx IS NOT NULL AND entry IS NOT NULL
+                  WHERE created <= $1 AND idx IS NOT NULL
              ) sub
              WHERE rn = 1
-             ORDER BY idx
+             ORDER BY idx DESC
              LIMIT $2 OFFSET $3",
             RewardMerkleTreeV2::state_type()
-        ))
-        .bind(height as i64)
-        .bind(limit as i64)
-        .bind(offset as i64)
-        .fetch_all(tx.as_mut())
-        .await
-        .context("loading reward accounts from storage")?;
+        );
+
+        let rows = query_as::<(Value, Value)>(&query)
+            .bind(height as i64)
+            .bind(limit as i64)
+            .bind(offset as i64)
+            .fetch_all(tx.as_mut())
+            .await
+            .context("loading reward accounts from storage")?;
 
         let mut accounts = Vec::new();
         for (idx, entry) in rows {