Merge branch 'main' into implement_groups_accumulator_count_distinct_primitive_types

Author: coderfender

datafusion-examples/examples/relation_planner/pivot_unpivot.rs

@@ -217,6 +217,25 @@ async fn run_examples(ctx: &SessionContext) -> Result<()> {
     +--------+
     ");
 
+    // Example 7: PIVOT on a quoted mixed-case column
+    // Reuses the parsed column expression so quoted identifiers keep their case.
+    let results = run_example(
+        ctx,
+        "Example 7: PIVOT with quoted mixed-case column",
+        r#"SELECT * FROM point_stats
+           PIVOT (MAX(max_value) FOR "pointNumber" IN ('16951' AS p16951, '16952' AS p16952)) AS p
+           ORDER BY ts"#,
+    )
+    .await?;
+    assert_snapshot!(results, @r"
+    +----------------------+------+--------+--------+
+    | ts                   | port | p16951 | p16952 |
+    +----------------------+------+--------+--------+
+    | 2024-09-01T10:00:00Z | 2411 | 10     | 20     |
+    | 2024-09-01T10:01:00Z | 2411 | 30     | 40     |
+    +----------------------+------+--------+--------+
+    ");
+
     Ok(())
 }
 
@@ -288,6 +307,34 @@ fn register_sample_data(ctx: &SessionContext) -> Result<()> {
         ])?,
     )?;
 
+    // point_stats: grouped data with a quoted mixed-case pivot column.
+    ctx.register_batch(
+        "point_stats",
+        RecordBatch::try_from_iter(vec![
+            (
+                "ts",
+                Arc::new(StringArray::from(vec![
+                    "2024-09-01T10:00:00Z",
+                    "2024-09-01T10:00:00Z",
+                    "2024-09-01T10:01:00Z",
+                    "2024-09-01T10:01:00Z",
+                ])) as ArrayRef,
+            ),
+            (
+                "pointNumber",
+                Arc::new(StringArray::from(vec!["16951", "16952", "16951", "16952"])),
+            ),
+            (
+                "port",
+                Arc::new(StringArray::from(vec!["2411", "2411", "2411", "2411"])),
+            ),
+            (
+                "max_value",
+                Arc::new(Int64Array::from(vec![10, 20, 30, 40])),
+            ),
+        ])?,
+    )?;
+
     Ok(())
 }
 
@@ -415,11 +462,12 @@ fn plan_pivot(
         .collect();
 
     let group_by_cols: Vec<Expr> = schema
-        .fields()
         .iter()
-        .map(|f| f.name().as_str())
-        .filter(|name| *name != pivot_col_name.as_str() && !agg_input_cols.contains(name))
-        .map(col)
+        .filter(|(_, field)| {
+            let name = field.name();
+            name != pivot_col_name.as_str() && !agg_input_cols.contains(&name.as_str())
+        })
+        .map(Expr::from)
         .collect();
 
     // Build CASE expressions for each (aggregate, pivot_value) pair
@@ -434,7 +482,7 @@ fn plan_pivot(
 
         for (value_alias, pivot_value) in &pivot_values {
             // CASE pivot_col WHEN pivot_value THEN agg_input END
-            let case_expr = case(col(&pivot_col_name))
+            let case_expr = case(pivot_col.clone())
                 .when(pivot_value.clone(), agg_input.clone())
                 .end()?;
 

datafusion/core/benches/sql_planner.rs

@@ -130,7 +130,8 @@ fn register_clickbench_hits_table(rt: &Runtime) -> SessionContext {
             format!("{BENCHMARKS_PATH_2}{CLICKBENCH_DATA_PATH}")
         };
 
-    let sql = format!("CREATE EXTERNAL TABLE hits STORED AS PARQUET LOCATION '{path}'");
+    let sql =
+        format!("CREATE EXTERNAL TABLE hits_raw STORED AS PARQUET LOCATION '{path}'");
 
     // ClickBench partitioned dataset was written by an ancient version of pyarrow that
     // that wrote strings with the wrong logical type. To read it correctly, we must
@@ -139,6 +140,17 @@ fn register_clickbench_hits_table(rt: &Runtime) -> SessionContext {
         .unwrap();
     rt.block_on(ctx.sql(&sql)).unwrap();
 
+    // ClickBench stores EventDate as UInt16 (days since 1970-01-01). Create a view
+    // that exposes it as SQL DATE so that queries comparing it with date literals
+    // (e.g. "EventDate >= '2013-07-01'") work correctly during planning.
+    rt.block_on(ctx.sql(
+        "CREATE VIEW hits AS \
+         SELECT * EXCEPT (\"EventDate\"), \
+                CAST(CAST(\"EventDate\" AS INTEGER) AS DATE) AS \"EventDate\" \
+         FROM hits_raw",
+    ))
+    .unwrap();
+
     let count =
         rt.block_on(async { ctx.table("hits").await.unwrap().count().await.unwrap() });
     assert!(count > 0);

datafusion/core/tests/memory_limit/mod.rs

@@ -213,6 +213,7 @@ async fn sort_merge_join_spill() {
         .with_config(config)
         .with_disk_manager_builder(DiskManagerBuilder::default())
         .with_scenario(Scenario::AccessLogStreaming)
+        .with_expected_success()
         .run()
         .await
 }

datafusion/datasource-parquet/src/opener.rs

@@ -920,12 +920,22 @@ impl MetadataLoadedParquetOpen {
         prepared.physical_file_schema = Arc::clone(&physical_file_schema);
 
         // Build predicates for this specific file
-        let (pruning_predicate, page_pruning_predicate) = build_pruning_predicates(
+        let pruning_predicate = build_pruning_predicates(
             prepared.predicate.as_ref(),
             &physical_file_schema,
             &prepared.predicate_creation_errors,
         );
 
+        // Only build page pruning predicate if page index is enabled
+        let page_pruning_predicate = if prepared.enable_page_index {
+            prepared.predicate.as_ref().and_then(|predicate| {
+                let p = build_page_pruning_predicate(predicate, &physical_file_schema);
+                (p.filter_number() > 0).then_some(p)
+            })
+        } else {
+            None
+        };
+
         Ok(FiltersPreparedParquetOpen {
             loaded: MetadataLoadedParquetOpen {
                 prepared,
@@ -945,10 +955,7 @@ impl FiltersPreparedParquetOpen {
         // metadata load above may not have read the page index structures yet.
         // If we need them for reading and they aren't yet loaded, we need to
         // load them now.
-        if should_enable_page_index(
-            self.loaded.prepared.enable_page_index,
-            &self.page_pruning_predicate,
-        ) {
+        if self.page_pruning_predicate.is_some() {
             self.loaded.reader_metadata = load_page_index(
                 self.loaded.reader_metadata,
                 &mut self.loaded.prepared.async_file_reader,
@@ -1661,20 +1668,13 @@ pub(crate) fn build_pruning_predicates(
     predicate: Option<&Arc<dyn PhysicalExpr>>,
     file_schema: &SchemaRef,
     predicate_creation_errors: &Count,
-) -> (
-    Option<Arc<PruningPredicate>>,
-    Option<Arc<PagePruningAccessPlanFilter>>,
-) {
-    let Some(predicate) = predicate.as_ref() else {
-        return (None, None);
-    };
-    let pruning_predicate = build_pruning_predicate(
+) -> Option<Arc<PruningPredicate>> {
+    let predicate = predicate.as_ref()?;
+    build_pruning_predicate(
         Arc::clone(predicate),
         file_schema,
         predicate_creation_errors,
-    );
-    let page_pruning_predicate = build_page_pruning_predicate(predicate, file_schema);
-    (pruning_predicate, Some(page_pruning_predicate))
+    )
 }
 
 /// Returns a `ArrowReaderMetadata` with the page index loaded, loading
@@ -1708,18 +1708,6 @@ async fn load_page_index<T: AsyncFileReader>(
     }
 }
 
-fn should_enable_page_index(
-    enable_page_index: bool,
-    page_pruning_predicate: &Option<Arc<PagePruningAccessPlanFilter>>,
-) -> bool {
-    enable_page_index
-        && page_pruning_predicate.is_some()
-        && page_pruning_predicate
-            .as_ref()
-            .map(|p| p.filter_number() > 0)
-            .unwrap_or(false)
-}
-
 #[cfg(test)]
 mod test {
     use std::sync::Arc;
@@ -1849,6 +1837,12 @@ mod test {
             self
         }
 
+        /// Enable page index.
+        fn with_enable_page_index(mut self, enable: bool) -> Self {
+            self.enable_page_index = enable;
+            self
+        }
+
         /// Set reverse row groups flag.
         fn with_reverse_row_groups(mut self, enable: bool) -> Self {
             self.reverse_row_groups = enable;
@@ -2713,4 +2707,71 @@ mod test {
             "Reverse scan with non-contiguous row groups should correctly map RowSelection"
         );
     }
+
+    /// Test that page pruning predicates are only built and applied when `enable_page_index` is true.
+    ///
+    /// The file has a single row group with 10 pages (10 rows each, values 1..100).
+    /// With page index enabled, pages whose max value <= 90 are pruned, returning only
+    /// the last page (rows 91..100). With page index disabled, all 100 rows are returned
+    /// since neither pushdown nor row-group pruning is active.
+    #[tokio::test]
+    async fn test_page_pruning_predicate_respects_enable_page_index() {
+        use parquet::file::properties::WriterProperties;
+
+        let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+
+        // 100 rows with values 1..=100, written as a single row group with 10 rows per page
+        let values: Vec<i32> = (1..=100).collect();
+        let batch = record_batch!((
+            "a",
+            Int32,
+            values.iter().map(|v| Some(*v)).collect::<Vec<_>>()
+        ))
+        .unwrap();
+        let props = WriterProperties::builder()
+            .set_data_page_row_count_limit(10)
+            .set_write_batch_size(10)
+            .build();
+        let schema = batch.schema();
+        let data_size = write_parquet_batches(
+            Arc::clone(&store),
+            "test.parquet",
+            vec![batch],
+            Some(props),
+        )
+        .await;
+
+        let file = PartitionedFile::new("test.parquet".to_string(), data_size as u64);
+
+        // predicate: a > 90 — should allow page index to prune first 9 pages
+        let predicate = logical2physical(&col("a").gt(lit(90i32)), &schema);
+
+        let make_opener = |enable_page_index| {
+            ParquetOpenerBuilder::new()
+                .with_store(Arc::clone(&store))
+                .with_schema(Arc::clone(&schema))
+                .with_predicate(Arc::clone(&predicate))
+                .with_enable_page_index(enable_page_index)
+                // disable pushdown and row-group pruning so the only pruning path is page index
+                .with_pushdown_filters(false)
+                .with_row_group_stats_pruning(false)
+                .build()
+        };
+        let (_, rows_with_page_index) = count_batches_and_rows(
+            make_opener(true).open(file.clone()).unwrap().await.unwrap(),
+        )
+        .await;
+        let (_, rows_without_page_index) =
+            count_batches_and_rows(make_opener(false).open(file).unwrap().await.unwrap())
+                .await;
+
+        assert_eq!(
+            rows_with_page_index, 10,
+            "page index should prune 9 of 10 pages"
+        );
+        assert_eq!(
+            rows_without_page_index, 100,
+            "without page index all rows are returned"
+        );
+    }
 }

datafusion/datasource-parquet/src/source.rs

@@ -629,17 +629,17 @@ impl FileSource for ParquetSource {
                     write!(f, ", reverse_row_groups=true")?;
                 }
 
-                // Try to build a the pruning predicates.
+                // Try to build the pruning predicates.
                 // These are only generated here because it's useful to have *some*
                 // idea of what pushdown is happening when viewing plans.
-                // However it is important to note that these predicates are *not*
+                // However, it is important to note that these predicates are *not*
                 // necessarily the predicates that are actually evaluated:
                 // the actual predicates are built in reference to the physical schema of
                 // each file, which we do not have at this point and hence cannot use.
-                // Instead we use the logical schema of the file (the table schema without partition columns).
+                // Instead, we use the logical schema of the file (the table schema without partition columns).
                 if let Some(predicate) = &self.predicate {
                     let predicate_creation_errors = Count::new();
-                    if let (Some(pruning_predicate), _) = build_pruning_predicates(
+                    if let Some(pruning_predicate) = build_pruning_predicates(
                         Some(predicate),
                         self.table_schema.table_schema(),
                         &predicate_creation_errors,