perf(datafusion): push filters into Paimon Parquet readers (#574)

Author: JingsongLi

benchmarks/tpcds/src/cli.rs

@@ -113,6 +113,9 @@ pub struct RuntimeArgs {
     /// DataFusion execution partitions. Defaults to available CPUs.
     #[arg(long)]
     pub target_partitions: Option<usize>,
+    /// Record batch size used by both DataFusion and Paimon file readers.
+    #[arg(long, default_value_t = 8192)]
+    pub batch_size: usize,
     /// Evaluate pushed filters during Parquet scans, in addition to statistics pruning.
     #[arg(long)]
     pub parquet_pushdown_filters: bool,
@@ -129,12 +132,16 @@ pub struct RuntimeArgs {
 
 impl RuntimeArgs {
     pub fn to_config(&self) -> Result<BenchmarkRuntimeConfig, String> {
+        if self.batch_size == 0 {
+            return Err("--batch-size must be greater than zero".to_string());
+        }
         let defaults = BenchmarkRuntimeConfig::default();
         Ok(BenchmarkRuntimeConfig {
             target_partitions: self
                 .target_partitions
                 .unwrap_or(defaults.target_partitions)
                 .max(1),
+            batch_size: self.batch_size,
             parquet_pushdown_filters: self.parquet_pushdown_filters,
             memory_limit_bytes: self.memory_limit_gib.map(gib_to_usize).transpose()?,
             spill_dir: self.spill_dir.clone(),

benchmarks/tpcds/src/context.rs

@@ -28,6 +28,8 @@ use serde::{Deserialize, Serialize};
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct BenchmarkRuntimeConfig {
     pub target_partitions: usize,
+    #[serde(default = "default_batch_size")]
+    pub batch_size: usize,
     #[serde(default)]
     pub parquet_pushdown_filters: bool,
     pub memory_limit_bytes: Option<usize>,
@@ -41,6 +43,7 @@ impl Default for BenchmarkRuntimeConfig {
             target_partitions: std::thread::available_parallelism()
                 .map(usize::from)
                 .unwrap_or(1),
+            batch_size: default_batch_size(),
             parquet_pushdown_filters: false,
             memory_limit_bytes: None,
             spill_dir: None,
@@ -49,6 +52,10 @@ impl Default for BenchmarkRuntimeConfig {
     }
 }
 
+fn default_batch_size() -> usize {
+    8192
+}
+
 pub fn build_sql_context(config: &BenchmarkRuntimeConfig) -> DataFusionResult<SQLContext> {
     let mut runtime = RuntimeEnvBuilder::new();
     if let Some(memory_limit) = config.memory_limit_bytes {
@@ -66,7 +73,8 @@ pub fn build_sql_context(config: &BenchmarkRuntimeConfig) -> DataFusionResult<SQ
     let mut session_config = current_state
         .config()
         .clone()
-        .with_target_partitions(config.target_partitions.max(1));
+        .with_target_partitions(config.target_partitions.max(1))
+        .with_batch_size(config.batch_size.max(1));
     session_config
         .options_mut()
         .execution
@@ -100,6 +108,13 @@ pub async fn open_catalog_session(
     let mut sql = build_sql_context(runtime_config)?;
     sql.register_catalog_with_default_db("paimon", catalog.clone(), Some(database))
         .await?;
+    // Keep decoder batch boundaries identical for Parquet and Paimon. Paimon
+    // retains Java's table default, so use its supported session override.
+    sql.sql(&format!(
+        "SET 'paimon.read.batch-size' = '{}'",
+        runtime_config.batch_size
+    ))
+    .await?;
     Ok(CatalogSession {
         sql,
         catalog,

benchmarks/tpcds/src/lib.rs

@@ -183,6 +183,7 @@ mod tests {
         let spill_dir = TempDir::new().unwrap();
         let ctx = build_sql_context(&BenchmarkRuntimeConfig {
             target_partitions: 3,
+            batch_size: 4096,
             parquet_pushdown_filters: true,
             memory_limit_bytes: Some(32 * 1024 * 1024),
             spill_dir: Some(spill_dir.path().to_path_buf()),
@@ -198,6 +199,10 @@ mod tests {
                 .target_partitions,
             3
         );
+        assert_eq!(
+            ctx.ctx().state().config_options().execution.batch_size,
+            4096
+        );
         assert!(
             ctx.ctx()
                 .state()

benchmarks/tpcds/tests/smoke.rs

@@ -225,6 +225,7 @@ async fn command_orchestration_loads_and_writes_a_run_report() {
     .unwrap();
     let runtime = RuntimeArgs {
         target_partitions: Some(2),
+        batch_size: 8192,
         parquet_pushdown_filters: false,
         memory_limit_gib: None,
         spill_dir: None,

crates/integrations/datafusion/src/config.rs

@@ -40,7 +40,6 @@ mod blob_descriptor_functions;
 mod blob_reader;
 mod blob_view;
 mod catalog;
-pub mod config;
 mod delete;
 mod error;
 mod filter_pushdown;