perf: use compact pruning for large string IN lists

datafusion/core/tests/parquet/mod.rs

@@ -58,6 +58,7 @@ mod page_pruning;
 mod row_group_pruning;
 mod schema;
 mod schema_coercion;
+mod string_in_list_pruning;
 mod utils;
 
 #[cfg(test)]

datafusion/core/tests/parquet/string_in_list_pruning.rs

@@ -0,0 +1,363 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! End-to-end coverage for compact, large string IN-list pruning. The positive
+//! IN-list cases disable the row and Bloom filters to isolate min/max pruning.
+
+use std::sync::Arc;
+
+use arrow::array::StringArray;
+use arrow::datatypes::{DataType, Field, Schema};
+use arrow::record_batch::RecordBatch;
+use arrow::util::pretty::pretty_format_batches;
+use datafusion::datasource::listing::PartitionedFile;
+use datafusion::datasource::object_store::ObjectStoreUrl;
+use datafusion::datasource::physical_plan::ParquetSource;
+use datafusion::datasource::source::DataSourceExec;
+use datafusion::physical_plan::{ExecutionPlan, collect, displayable};
+use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext};
+use datafusion_common::config::TableParquetOptions;
+use datafusion_common::{ScalarValue, assert_batches_eq};
+use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
+use datafusion_physical_expr::expressions::{col, in_list, lit};
+use datafusion_physical_plan::metrics::{MetricValue, MetricsSet};
+use object_store::path::Path;
+use parquet::arrow::ArrowWriter;
+use parquet::file::properties::{EnabledStatistics, WriterProperties};
+use tempfile::NamedTempFile;
+
+use super::utils::MetricsFinder;
+
+const ROWS_PER_UNIT: usize = 16;
+const UNITS: usize = 4;
+const TOTAL_ROWS: usize = ROWS_PER_UNIT * UNITS;
+const MATCHING_ROWS: usize = ROWS_PER_UNIT * 2;
+
+/// Write either four row groups or four pages in one row group. The second
+/// unit lies in a gap between two members of every test IN list; an enclosing
+/// min/max range for the list cannot prune it.
+fn make_file(page_pruning: bool) -> NamedTempFile {
+    let mut file = tempfile::Builder::new()
+        .prefix("string_in_list_pruning")
+        .suffix(".parquet")
+        .tempfile()
+        .unwrap();
+    let schema = Arc::new(Schema::new(vec![Field::new(
+        "value",
+        DataType::Utf8,
+        false,
+    )]));
+    let values = ["v000000", "v000001", "v000010", "v999999"]
+        .into_iter()
+        .flat_map(|value| std::iter::repeat_n(value, ROWS_PER_UNIT))
+        .collect::<Vec<_>>();
+    let batch = RecordBatch::try_new(
+        Arc::clone(&schema),
+        vec![Arc::new(StringArray::from(values))],
+    )
+    .unwrap();
+    let rows_per_group = if page_pruning {
+        TOTAL_ROWS
+    } else {
+        ROWS_PER_UNIT
+    };
+    let properties = WriterProperties::builder()
+        .set_max_row_group_row_count(Some(rows_per_group))
+        .set_data_page_row_count_limit(ROWS_PER_UNIT)
+        .set_write_batch_size(ROWS_PER_UNIT)
+        .set_dictionary_enabled(false)
+        .set_bloom_filter_enabled(false)
+        .set_statistics_enabled(EnabledStatistics::Page)
+        .build();
+    let mut writer = ArrowWriter::try_new(&mut file, schema, Some(properties)).unwrap();
+    writer.write(&batch).unwrap();
+    let metadata = writer.close().unwrap();
+    assert_eq!(metadata.num_row_groups(), TOTAL_ROWS / rows_per_group);
+    let offsets = metadata.offset_index().unwrap();
+    for row_group in offsets {
+        assert_eq!(
+            row_group[0].page_locations().len(),
+            rows_per_group / ROWS_PER_UNIT
+        );
+    }
+    file
+}
+
+struct ScanOutput {
+    batches: Vec<RecordBatch>,
+    plan: String,
+    metrics: MetricsSet,
+}
+
+impl ScanOutput {
+    fn counter(&self, name: &str) -> usize {
+        self.metrics
+            .sum(|metric| metric.value().name() == name)
+            .unwrap_or_else(|| panic!("missing {name}: {}", self.metrics))
+            .as_usize()
+    }
+
+    fn pruned(&self, name: &str) -> usize {
+        let value = self
+            .metrics
+            .sum(|metric| metric.value().name() == name)
+            .unwrap_or_else(|| panic!("missing {name}: {}", self.metrics));
+        let MetricValue::PruningMetrics {
+            pruning_metrics, ..
+        } = value
+        else {
+            panic!("expected pruning metric {name}: {}", self.metrics);
+        };
+        pruning_metrics.pruned()
+    }
+
+    fn fully_matched(&self, name: &str) -> usize {
+        let value = self
+            .metrics
+            .sum(|metric| metric.value().name() == name)
+            .unwrap_or_else(|| panic!("missing {name}: {}", self.metrics));
+        let MetricValue::PruningMetrics {
+            pruning_metrics, ..
+        } = value
+        else {
+            panic!("expected pruning metric {name}: {}", self.metrics);
+        };
+        pruning_metrics.fully_matched()
+    }
+
+    fn assert_results(&self) {
+        assert_batches_eq!(
+            [
+                "+---------+----+",
+                "| value   | n  |",
+                "+---------+----+",
+                "| v000000 | 16 |",
+                "| v000010 | 16 |",
+                "+---------+----+",
+            ],
+            &self.batches
+        );
+        assert_eq!(self.counter("predicate_evaluation_errors"), 0);
+        assert_eq!(self.counter("pushdown_rows_pruned"), 0);
+        assert_eq!(self.pruned("row_groups_pruned_bloom_filter"), 0);
+    }
+}
+
+async fn scan(
+    file: &NamedTempFile,
+    list_size: usize,
+    max_in_list_size: Option<usize>,
+    page_pruning: bool,
+) -> ScanOutput {
+    let mut config = SessionConfig::new()
+        .with_target_partitions(1)
+        .with_parquet_bloom_filter_pruning(false)
+        .with_parquet_page_index_pruning(page_pruning);
+    config.options_mut().execution.parquet.pushdown_filters = false;
+    if let Some(max_in_list_size) = max_in_list_size {
+        config.options_mut().execution.parquet.max_in_list_size = max_in_list_size;
+    }
+    let ctx = SessionContext::new_with_config(config);
+    ctx.register_parquet(
+        "t",
+        file.path().to_str().unwrap(),
+        ParquetReadOptions::default(),
+    )
+    .await
+    .unwrap();
+    let values = (0..list_size)
+        .map(|index| format!("'v{:06}'", index * 10))
+        .collect::<Vec<_>>()
+        .join(", ");
+    let sql = format!(
+        "SELECT value, count(*) AS n FROM t \
+         WHERE value IN ({values}) GROUP BY value ORDER BY value"
+    );
+    let plan = ctx
+        .sql(&sql)
+        .await
+        .unwrap()
+        .create_physical_plan()
+        .await
+        .unwrap();
+    let plan_text = displayable(plan.as_ref()).indent(true).to_string();
+    let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap();
+    let metrics = MetricsFinder::find_metrics(plan.as_ref()).unwrap();
+    ScanOutput {
+        batches,
+        plan: plan_text,
+        metrics,
+    }
+}
+
+async fn check_string_in_list_pruning(page_pruning: bool) {
+    let file = make_file(page_pruning);
+    for list_size in [20, 21, 256, 1024] {
+        // A zero cap provides a result-equivalence control that cannot use
+        // min/max IN-list pruning at either granularity.
+        let unpruned = scan(&file, list_size, Some(0), page_pruning).await;
+        unpruned.assert_results();
+        assert!(!unpruned.plan.contains("IN_SET_INTERSECTS"));
+        assert_eq!(unpruned.pruned("row_groups_pruned_statistics"), 0);
+        assert_eq!(unpruned.pruned("page_index_rows_pruned"), 0);
+        assert_eq!(unpruned.counter("output_rows"), TOTAL_ROWS);
+
+        let output = scan(&file, list_size, Some(list_size), page_pruning).await;
+        output.assert_results();
+        assert_eq!(
+            pretty_format_batches(&output.batches).unwrap().to_string(),
+            pretty_format_batches(&unpruned.batches)
+                .unwrap()
+                .to_string()
+        );
+        assert_eq!(
+            output.plan.contains("IN_SET_INTERSECTS"),
+            list_size > 20,
+            "list_size={list_size}, plan={}",
+            output.plan
+        );
+        assert_eq!(
+            output.pruned("row_groups_pruned_statistics"),
+            if page_pruning { 0 } else { 2 },
+            "list_size={list_size}, metrics={}",
+            output.metrics
+        );
+        assert_eq!(
+            output.pruned("page_index_rows_pruned"),
+            if page_pruning { MATCHING_ROWS } else { 0 },
+            "list_size={list_size}, metrics={}",
+            output.metrics
+        );
+        assert_eq!(output.counter("output_rows"), MATCHING_ROWS);
+    }
+
+    // The default remains 20: enabling the compact representation must not
+    // silently change the public cap's meaning.
+    let default = scan(&file, 21, None, page_pruning).await;
+    default.assert_results();
+    assert!(!default.plan.contains("IN_SET_INTERSECTS"));
+    assert_eq!(default.pruned("row_groups_pruned_statistics"), 0);
+    assert_eq!(default.pruned("page_index_rows_pruned"), 0);
+    assert_eq!(default.counter("output_rows"), TOTAL_ROWS);
+}
+
+#[tokio::test]
+async fn string_in_list_row_group_pruning() {
+    check_string_in_list_pruning(false).await;
+}
+
+#[tokio::test]
+async fn string_in_list_page_pruning() {
+    check_string_in_list_pruning(true).await;
+}
+
+#[tokio::test]
+async fn string_not_in_list_with_null_does_not_bypass_row_filter() {
+    let mut file = tempfile::Builder::new()
+        .prefix("string_not_in_list_pruning")
+        .suffix(".parquet")
+        .tempfile()
+        .unwrap();
+    let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)]));
+    // The first row group has a known zero null count, and every value lies
+    // in a gap in the IN list. Dropping the NULL list member while inverting
+    // NOT IN would incorrectly prove that this entire row group matches.
+    let values = vec![
+        Some("v000001"),
+        Some("v000001"),
+        Some("v000001"),
+        Some("v000001"),
+        Some("v000000"),
+        Some("v000001"),
+        None,
+        Some("v999999"),
+    ];
+    let batch = RecordBatch::try_new(
+        Arc::clone(&schema),
+        vec![Arc::new(StringArray::from(values))],
+    )
+    .unwrap();
+    let properties = WriterProperties::builder()
+        .set_max_row_group_row_count(Some(4))
+        .set_bloom_filter_enabled(false)
+        .build();
+    let mut writer =
+        ArrowWriter::try_new(&mut file, Arc::clone(&schema), Some(properties)).unwrap();
+    writer.write(&batch).unwrap();
+    assert_eq!(writer.close().unwrap().num_row_groups(), 2);
+
+    // Build the physical source directly so a logical optimizer cannot fold
+    // the SQL NOT IN (..., NULL) filter to an empty relation before the scan.
+    let mut list = (0..21)
+        .map(|index| lit(format!("v{:06}", index * 10)))
+        .collect::<Vec<_>>();
+    list.push(lit(ScalarValue::Utf8(None)));
+    let predicate =
+        in_list(col("value", &schema).unwrap(), list, &true, &schema).unwrap();
+    let location = Path::from_filesystem_path(file.path()).unwrap();
+    let partitioned_file = PartitionedFile::new(
+        location.to_string(),
+        file.as_file().metadata().unwrap().len(),
+    );
+    let ctx =
+        SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1));
+
+    for max_in_list_size in [0, 32] {
+        let mut options = TableParquetOptions::default();
+        options.global.max_in_list_size = max_in_list_size;
+        let source = Arc::new(
+            ParquetSource::new(Arc::clone(&schema))
+                .with_table_parquet_options(options)
+                .with_predicate(Arc::clone(&predicate))
+                .with_pushdown_filters(true)
+                .with_enable_page_index(false)
+                .with_bloom_filter_on_read(false),
+        );
+        let config =
+            FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source)
+                .with_file(partitioned_file.clone())
+                .with_limit(Some(1))
+                .build();
+        let plan: Arc<dyn ExecutionPlan> =
+            Arc::new(DataSourceExec::new(Arc::new(config)));
+        let plan_text = displayable(plan.as_ref()).indent(true).to_string();
+        assert!(plan_text.contains("NOT IN"), "{plan_text}");
+        let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap();
+        let output = ScanOutput {
+            batches,
+            plan: plan_text,
+            metrics: MetricsFinder::find_metrics(plan.as_ref()).unwrap(),
+        };
+
+        assert_eq!(
+            output
+                .batches
+                .iter()
+                .map(RecordBatch::num_rows)
+                .sum::<usize>(),
+            0,
+            "cap={max_in_list_size}, plan={}, metrics={}",
+            output.plan,
+            output.metrics
+        );
+        assert_eq!(output.fully_matched("row_groups_pruned_statistics"), 0);
+        assert_eq!(output.pruned("row_groups_pruned_statistics"), 0);
+        assert_eq!(output.pruned("limit_pruned_row_groups"), 0);
+        assert_eq!(output.counter("pushdown_rows_pruned"), 8);
+        assert_eq!(output.counter("predicate_evaluation_errors"), 0);
+    }
+}