fix: Correct DataFusion repartitioning (#1554)

While the numbers are worse, I believe this is actually correct.
Also added caching for file statistics in benchmarks, which I think
should make the benchmarks faster overall.

Author: AdamGS

.github/workflows/bench-pr.yml

@@ -30,8 +30,9 @@ jobs:
             name: DataFusion
           - id: random_access
             name: Random Access
-          - id: compress_noci
+          - id: compress
             name: Vortex Compression
+            
     runs-on: self-hosted
     if: ${{ contains(github.event.head_commit.message, '[benchmark]') || github.event.label.name == 'benchmark' && github.event_name == 'pull_request' }}
     steps:

bench-vortex/benches/clickbench.rs

@@ -4,9 +4,8 @@ use std::path::PathBuf;
 use std::process::Command;
 
 use bench_vortex::clickbench::{clickbench_queries, HITS_SCHEMA};
-use bench_vortex::{clickbench, execute_query, idempotent, IdempotentPath};
+use bench_vortex::{clickbench, execute_query, get_session_with_cache, idempotent, IdempotentPath};
 use criterion::{criterion_group, criterion_main, Criterion};
-use datafusion::prelude::SessionContext;
 use tokio::runtime::Builder;
 
 fn benchmark(c: &mut Criterion) {
@@ -42,8 +41,9 @@ fn benchmark(c: &mut Criterion) {
         .unwrap();
     }
 
-    let session_context = SessionContext::new();
+    let session_context = get_session_with_cache();
     let context = session_context.clone();
+
     runtime.block_on(async move {
         clickbench::register_vortex_files(&context, "hits", basepath.as_path(), &HITS_SCHEMA)
             .await

bench-vortex/src/bin/clickbench.rs

@@ -8,11 +8,10 @@ use std::time::{Duration, Instant};
 use bench_vortex::clickbench::{self, clickbench_queries, HITS_SCHEMA};
 use bench_vortex::display::{print_measurements_json, render_table, DisplayFormat};
 use bench_vortex::{
-    execute_query, idempotent, physical_plan, setup_logger, Format, IdempotentPath as _,
-    Measurement,
+    execute_query, get_session_with_cache, idempotent, physical_plan, setup_logger, Format,
+    IdempotentPath as _, Measurement,
 };
 use clap::Parser;
-use datafusion::prelude::SessionContext;
 use indicatif::ProgressBar;
 use itertools::Itertools;
 use log::LevelFilter;
@@ -120,7 +119,7 @@ fn main() {
     let mut all_measurements = Vec::default();
 
     for format in &formats {
-        let session_context = SessionContext::new();
+        let session_context = get_session_with_cache();
         let context = session_context.clone();
         match format {
             Format::Parquet => runtime.block_on(async {

bench-vortex/src/lib.rs

@@ -8,7 +8,10 @@ use std::sync::{Arc, LazyLock};
 use std::time::Duration;
 
 use arrow_array::{RecordBatch, RecordBatchReader};
-use datafusion::prelude::SessionContext;
+use datafusion::execution::cache::cache_manager::CacheManagerConfig;
+use datafusion::execution::cache::cache_unit::{DefaultFileStatisticsCache, DefaultListFilesCache};
+use datafusion::execution::runtime_env::RuntimeEnvBuilder;
+use datafusion::prelude::{SessionConfig, SessionContext};
 use datafusion_physical_plan::{collect, ExecutionPlan};
 use itertools::Itertools;
 use log::LevelFilter;
@@ -297,6 +300,23 @@ impl Measurement {
     }
 }
 
+pub fn get_session_with_cache() -> SessionContext {
+    let cache_config = CacheManagerConfig::default();
+    let file_static_cache = Arc::new(DefaultFileStatisticsCache::default());
+    let list_file_cache = Arc::new(DefaultListFilesCache::default());
+
+    let cache_config = cache_config
+        .with_files_statistics_cache(Some(file_static_cache))
+        .with_list_files_cache(Some(list_file_cache));
+
+    let rt = RuntimeEnvBuilder::new()
+        .with_cache_manager(cache_config)
+        .build_arc()
+        .expect("could not build runtime environment");
+
+    SessionContext::new_with_config_rt(SessionConfig::default(), rt)
+}
+
 #[cfg(test)]
 mod test {
     use std::fs::File;

vortex-datafusion/src/persistent/execution.rs

@@ -2,14 +2,16 @@ use std::fmt;
 use std::sync::Arc;
 
 use datafusion::config::ConfigOptions;
-use datafusion::datasource::physical_plan::{FileGroupPartitioner, FileScanConfig, FileStream};
+use datafusion::datasource::listing::PartitionedFile;
+use datafusion::datasource::physical_plan::{FileScanConfig, FileStream};
 use datafusion_common::{project_schema, Result as DFResult, Statistics};
 use datafusion_execution::{SendableRecordBatchStream, TaskContext};
 use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
 use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
 use datafusion_physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, PlanProperties,
 };
+use itertools::Itertools;
 use vortex_array::Context;
 
 use crate::persistent::opener::VortexFileOpener;
@@ -47,7 +49,7 @@ impl VortexExec {
 
         let plan_properties = PlanProperties::new(
             EquivalenceProperties::new_with_orderings(projected_schema, &orderings),
-            Partitioning::UnknownPartitioning(1),
+            Partitioning::UnknownPartitioning(file_scan_config.file_groups.len()),
             ExecutionMode::Bounded,
         );
 
@@ -60,6 +62,7 @@ impl VortexExec {
             ctx,
         })
     }
+
     pub(crate) fn into_arc(self) -> Arc<dyn ExecutionPlan> {
         Arc::new(self) as _
     }
@@ -103,6 +106,7 @@ impl ExecutionPlan for VortexExec {
         partition: usize,
         context: Arc<TaskContext>,
     ) -> DFResult<SendableRecordBatchStream> {
+        log::debug!("Executing partition {partition}");
         let object_store = context
             .runtime_env()
             .object_store(&self.file_scan_config.object_store_url)?;
@@ -128,21 +132,57 @@ impl ExecutionPlan for VortexExec {
     fn repartitioned(
         &self,
         target_partitions: usize,
-        config: &ConfigOptions,
+        _config: &ConfigOptions,
     ) -> DFResult<Option<Arc<dyn ExecutionPlan>>> {
-        let repartition_file_min_size = config.optimizer.repartition_file_min_size;
-        let repartitioned_file_groups_option = FileGroupPartitioner::new()
-            .with_target_partitions(target_partitions)
-            .with_repartition_file_min_size(repartition_file_min_size)
-            .with_preserve_order_within_groups(self.properties().output_ordering().is_some())
-            .repartition_file_groups(&self.file_scan_config.file_groups);
+        let file_groups = self.file_scan_config.file_groups.clone();
+
+        let repartitioned_file_groups = repartition_by_count(file_groups, target_partitions);
 
         let mut new_plan = self.clone();
-        if let Some(repartitioned_file_groups) = repartitioned_file_groups_option {
-            let mut config = new_plan.file_scan_config;
-            config = config.with_file_groups(repartitioned_file_groups);
-            new_plan.file_scan_config = config;
-        }
+        let mut config = new_plan.file_scan_config;
+        let num_partitions = repartitioned_file_groups.len();
+
+        log::debug!("VortexExec repartitioned to {num_partitions} partitions");
+        config = config.with_file_groups(repartitioned_file_groups);
+        new_plan.file_scan_config = config;
+        new_plan.plan_properties.partitioning = Partitioning::UnknownPartitioning(num_partitions);
+
         Ok(Some(Arc::new(new_plan)))
     }
 }
+
+fn repartition_by_count(
+    file_groups: Vec<Vec<PartitionedFile>>,
+    desired_partitions: usize,
+) -> Vec<Vec<PartitionedFile>> {
+    let all_files = file_groups.into_iter().concat();
+
+    let approx_files_per_partition = all_files.len().div_ceil(desired_partitions);
+    let mut repartitioned_file_groups = Vec::default();
+
+    for chunk in &all_files.into_iter().chunks(approx_files_per_partition) {
+        repartitioned_file_groups.push(chunk.collect::<Vec<_>>());
+    }
+
+    repartitioned_file_groups
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn basic_repartition_test() {
+        let input_file_groups = vec![vec![
+            PartitionedFile::new("a", 0),
+            PartitionedFile::new("b", 0),
+            PartitionedFile::new("c", 0),
+            PartitionedFile::new("d", 0),
+            PartitionedFile::new("e", 0),
+        ]];
+
+        let file_groups = repartition_by_count(input_file_groups, 2);
+
+        assert_eq!(file_groups.len(), 2);
+    }
+}