feat: add a way to pre-sort data

Cargo.toml

@@ -10,6 +10,7 @@ exclude = [
 
 [dependencies]
 datafusion = "54"
+async-trait = "0.1"
 tokio = {version = "1", features = ["full"] }
 url = "2"
 futures = "0.3"

src/algorithm/connectivity/connected_components.rs

@@ -239,9 +239,7 @@ impl<'a> ConnectedComponentsBuilder<'a> {
             ParquetCheckpointer::new(store_url.clone(), ckpt_base.clone().join("backward"));
 
         // Offload the initial edges; from here on `edges` is disk-backed.
-        let mut edges = edges_ckptr
-            .push(ctx, "initial", prepared_edges, None)
-            .await?;
+        let mut edges = edges_ckptr.push(ctx, "initial", prepared_edges).await?;
         let mut graph_size = edges.clone().count().await?;
         log::info!("after preparation graph has {graph_size} edges");
 
@@ -267,16 +265,14 @@ impl<'a> ConnectedComponentsBuilder<'a> {
             let cc_reps = compute_cc_reps(&edges, r_a, r_b)?;
             // Spill reps to the forward checkpointer; keep the lazy,
             // disk-backed frame for the back pass.
-            let cc_reps = fwd_ckptr
-                .push(ctx, &iteration.to_string(), cc_reps, None)
-                .await?;
+            let cc_reps = fwd_ckptr.push(ctx, &iteration.to_string(), cc_reps).await?;
             forward_reps.push(cc_reps.clone());
 
             // Relabel edges: src -> rep, then dst -> rep, drop self-loops.
             let new_edges = relabel_edges(&edges, &cc_reps)?;
 
             edges = edges_ckptr
-                .push(ctx, &iteration.to_string(), new_edges, None)
+                .push(ctx, &iteration.to_string(), new_edges)
                 .await?;
             // Count the new edge set BEFORE evicting the previous checkpoint.
             // When a contraction iteration empties the edge set (convergence),
@@ -299,7 +295,7 @@ impl<'a> ConnectedComponentsBuilder<'a> {
         } else {
             // Seed the back frontier with the last forward reps.
             let mut frontier = back_ckptr
-                .push(ctx, "seed", forward_reps[n - 1].clone(), None)
+                .push(ctx, "seed", forward_reps[n - 1].clone())
                 .await?;
             // The last forward reps are now consumed; drop them from disk.
             fwd_ckptr.remove_last(ctx, 1).await?;
@@ -318,7 +314,7 @@ impl<'a> ConnectedComponentsBuilder<'a> {
                 let older = forward_reps[t - 1].clone();
                 let composed = back_prop_step(&older, &frontier, acc_a, acc_b)?;
 
-                frontier = back_ckptr.push(ctx, &t.to_string(), composed, None).await?;
+                frontier = back_ckptr.push(ctx, &t.to_string(), composed).await?;
                 // Keep only the latest back checkpoint.
                 back_ckptr.evict(ctx, 1).await?;
                 // `older` (cc_reps_{t}) has been consumed into the new
@@ -362,7 +358,7 @@ impl<'a> ConnectedComponentsBuilder<'a> {
         // hashed representative (or the id itself for isolated vertices).
         let result = if self.use_labels_as_components {
             // spill "raw" components to disk
-            let components = fwd_ckptr.push(ctx, "final_df", final_df, None).await?;
+            let components = fwd_ckptr.push(ctx, "final_df", final_df).await?;
             // Relabel each component to the minimum original id of its members.
             let labels = components.clone().aggregate(
                 vec![col(COMPONENT_COL).alias(CC_COMP_KEY)],

src/algorithm/pregel.rs

@@ -355,15 +355,15 @@ impl PregelBuilder {
         );
 
         let edges_struct = edges_checkpointer
-            .push(
+            .push_pre_sorted(
                 ctx,
                 "edges",
                 self.graph.edges.clone().select(
                     self.edge_columns
                         .iter()
                         .map(|name| col(name).alias(format!("{}_{}", PREGEL_MSG_EDGE, name))),
                 )?,
-                None,
+                &format!("{}_{}", PREGEL_MSG_EDGE, EDGE_SRC),
             )
             .await?;
 
@@ -374,7 +374,7 @@ impl PregelBuilder {
         );
         // offload prepared state to disk
         current_vertices = state_checkpointer
-            .push(ctx, "state-0", current_vertices, None)
+            .push_pre_sorted(ctx, "state-0", current_vertices, VERTEX_ID)
             .await?;
 
         // Main Pregel loop
@@ -473,7 +473,6 @@ impl PregelBuilder {
                     } else {
                         messages_df
                     },
-                    None,
                 )
                 .await?;
 
@@ -491,7 +490,12 @@ impl PregelBuilder {
                 .select(update_columns.clone())?;
 
             current_vertices = state_checkpointer
-                .push(ctx, &format!("state-{}", iteration), new_vertices, None)
+                .push_pre_sorted(
+                    ctx,
+                    &format!("state-{}", iteration),
+                    new_vertices,
+                    VERTEX_ID,
+                )
                 .await?;
 
             // remove the old checkpoints

src/algorithm/subgraph/maximal_independent_set.rs

@@ -143,7 +143,6 @@ impl<'a> MISBuilder<'a> {
                     col(VERTEX_ID).alias(MIS_V),
                     lit(false).alias(MIS_FLAG),
                 ])?,
-                None,
             )
             .await?;
 
@@ -157,7 +156,6 @@ impl<'a> MISBuilder<'a> {
                     col(VERTEX_ID).alias(MIS_V),
                     lit(0.5f64).alias(MIS_PROB),
                 ])?,
-                None,
             )
             .await?;
 
@@ -182,7 +180,6 @@ impl<'a> MISBuilder<'a> {
                         .select_columns(&[EDGE_SRC, EDGE_DST])?,
                     true,
                 )?,
-                None,
             )
             .await?;
 
@@ -210,7 +207,6 @@ impl<'a> MISBuilder<'a> {
                             vec![col(EDGE_SRC)],
                             vec![functions_aggregate::sum::sum(col(MIS_PROB)).alias(MIS_DEG)],
                         )?,
-                    None,
                 )
                 .await?;
 
@@ -243,7 +239,6 @@ impl<'a> MISBuilder<'a> {
                                 .otherwise(lit(0.5))?,
                         )?
                         .select(vec![col(MIS_V), col(MIS_PROB), col(MIS_NOM)])?,
-                    None,
                 )
                 .await?;
 
@@ -305,7 +300,6 @@ impl<'a> MISBuilder<'a> {
                     ctx,
                     &format!("mis_{}", iteration),
                     or_into_mis(current_mis.clone(), new_mis_members)?,
-                    None,
                 )
                 .await?;
 
@@ -325,7 +319,6 @@ impl<'a> MISBuilder<'a> {
                             None,
                         )?
                         .select(vec![col(MIS_V), col(MIS_PROB)])?,
-                    None,
                 )
                 .await?;
 
@@ -350,7 +343,6 @@ impl<'a> MISBuilder<'a> {
                             &[MIS_REM_V],
                             None,
                         )?,
-                    None,
                 )
                 .await?;
 
@@ -378,7 +370,6 @@ impl<'a> MISBuilder<'a> {
                             ctx,
                             &format!("mis_sweep_{}", iteration),
                             or_into_mis(current_mis.clone(), remaining)?,
-                            None,
                         )
                         .await?;
                 }

src/memory.rs

@@ -1,3 +1,4 @@
+mod hash_partitioned;
 mod parquet_checkpointer;
 
 pub use parquet_checkpointer::CheckpointConfig;