add implementations

Author: adriangb

datafusion/datasource-parquet/src/opener.rs

@@ -100,11 +100,12 @@ impl FileOpener for ParquetOpener {
                 &predicate_creation_errors,
             );
             if let Some(pruning_predicate) = pruning_predicate {
-                let pruning_stats = PrunableStatistics::new(
+                let stats_pruning = PrunableStatistics::new(
                     vec![Arc::clone(stats)],
                     Arc::clone(&self.table_schema),
                 );
-                match pruning_predicate.prune(&pruning_stats) {
+                // TODO: add partition pruning. Need to pipe in the partition columns.
+                match pruning_predicate.prune(&stats_pruning) {
                     Ok(values) => {
                         // We expect a single container -> if all containers are false skip this file
                         if values.into_iter().all(|v| !v) {

datafusion/physical-optimizer/src/pruning.rs

@@ -996,10 +996,113 @@ fn build_statistics_record_batch<S: PruningStatistics>(
     })
 }
 
+pub struct PartitionPruningStatistics {
+    /// Values for each column for each container.
+    /// The outer vectors represent the columns while the inner
+    /// vectors represent the containers.
+    /// The order must match the order of the partition columns in
+    /// [`PartitionPruningStatistics::partition_schema`].
+    partition_values: Vec<Vec<ScalarValue>>,
+    /// The schema of the partition columns.
+    /// This must **not** be the schema of the entire file or table:
+    /// it must only be the schema of the partition columns,
+    /// in the same order as the values in [`PartitionPruningStatistics::partition_values`].
+    partition_schema: SchemaRef,
+}
+
+impl PartitionPruningStatistics {
+    /// Create a new instance of [`PartitionPruningStatistics`].
+    /// 
+    /// Args:
+    /// * `partition_values`: A vector of vectors of [`ScalarValue`]s.
+    ///   The outer vector represents the containers while the inner
+    ///   vector represents the partition values for each column.
+    ///   Note that this is the **opposite** of the order of the
+    ///   partition columns in [`PartitionPruningStatistics::partition_schema`].
+    /// * `partition_schema`: The schema of the partition columns.
+    ///   This must **not** be the schema of the entire file or table:
+    ///   instead it must only be the schema of the partition columns,
+    ///   in the same order as the values in `partition_values`.
+    pub fn new(partition_values: Vec<Vec<ScalarValue>>, partition_schema: SchemaRef) -> Self {
+        let mut partition_valeus_by_column = vec![vec![]; partition_schema.fields().len()];
+        for partition_value in partition_values.iter() {
+            for (i, value) in partition_value.iter().enumerate() {
+                partition_valeus_by_column[i].push(value.clone());
+            }
+        }
+        Self {
+            partition_values: partition_valeus_by_column,
+            partition_schema,
+        }
+    }
+}
+
+impl PruningStatistics for PartitionPruningStatistics {
+    fn min_values(&self, column: &Column) -> Option<ArrayRef> {
+        let index = self.partition_schema.index_of(column.name()).ok()?;
+        let partition_values = self.partition_values.get(index)?;
+        let mut values = Vec::with_capacity(self.partition_values.len());
+        for partition_value in partition_values {
+            match partition_value {
+                ScalarValue::Null => values.push(ScalarValue::Null),
+                _ => values.push(partition_value.clone()),
+            }
+        }
+        match ScalarValue::iter_to_array(values) {
+            Ok(array) => Some(array),
+            Err(_) => {
+                log::warn!(
+                    "Failed to convert min values to array for column {}",
+                    column.name()
+                );
+                None
+            }
+        }
+    }
+
+    fn max_values(&self, column: &Column) -> Option<ArrayRef> {
+        self.min_values(column)
+    }
+
+    fn num_containers(&self) -> usize {
+        self.partition_values.len()
+    }
+
+    fn null_counts(&self, _column: &Column) -> Option<ArrayRef> {
+        None
+    }
+
+    fn row_counts(&self, _column: &Column) -> Option<ArrayRef> {
+        None
+    }
+
+    fn contained(
+        &self,
+        column: &Column,
+        values: &HashSet<ScalarValue>,
+    ) -> Option<BooleanArray> {
+        let index = self.partition_schema.index_of(column.name()).ok()?;
+        let partition_values = self.partition_values.get(index)?;
+        let mut contained = Vec::with_capacity(self.partition_values.len());
+        for partition_value in partition_values {
+            let contained_value = if values.contains(partition_value) {
+                Some(true)
+            } else {
+                Some(false)
+            };
+            contained.push(contained_value);
+        }
+        let array = BooleanArray::from(contained);
+        Some(array)
+    }
+}
+
 /// Prune a set of containers represented by their statistics.
 /// Each [`Statistics`] represents a container (e.g. a file or a partition of files).
 pub struct PrunableStatistics {
+    /// Statistics for each container.
     statistics: Vec<Arc<Statistics>>,
+    /// The schema of the file these statistics are for.
     schema: SchemaRef,
 }
 
@@ -1080,7 +1183,7 @@ impl PruningStatistics for PrunableStatistics {
     /// [`Self::min_values`], [`Self::max_values`], [`Self::null_counts`],
     /// and [`Self::row_counts`].
     fn num_containers(&self) -> usize {
-        1
+        self.statistics.len()
     }
 
     /// Return the number of null values for the named column as an
@@ -1174,6 +1277,72 @@ impl PruningStatistics for PrunableStatistics {
     }
 }
 
+pub struct CompositePruningStatistics {
+    pub statistics: Vec<Box<dyn PruningStatistics>>,
+}
+
+impl CompositePruningStatistics {
+    pub fn new(statistics: Vec<Box<dyn PruningStatistics>>) -> Self {
+        assert!(!statistics.is_empty());
+        Self { statistics }
+    }
+}
+
+impl PruningStatistics for CompositePruningStatistics {
+    fn min_values(&self, column: &Column) -> Option<ArrayRef> {
+        for stats in &self.statistics {
+            if let Some(array) = stats.min_values(column) {
+                return Some(array);
+            }
+        }
+        None
+    }
+
+    fn max_values(&self, column: &Column) -> Option<ArrayRef> {
+        for stats in &self.statistics {
+            if let Some(array) = stats.max_values(column) {
+                return Some(array);
+            }
+        }
+        None
+    }
+
+    fn num_containers(&self) -> usize {
+        self.statistics[0].num_containers()
+    }
+
+    fn null_counts(&self, column: &Column) -> Option<ArrayRef> {
+        for stats in &self.statistics {
+            if let Some(array) = stats.null_counts(column) {
+                return Some(array);
+            }
+        }
+        None
+    }
+
+    fn row_counts(&self, column: &Column) -> Option<ArrayRef> {
+        for stats in &self.statistics {
+            if let Some(array) = stats.row_counts(column) {
+                return Some(array);
+            }
+        }
+        None
+    }
+
+    fn contained(
+        &self,
+        column: &Column,
+        values: &HashSet<ScalarValue>,
+    ) -> Option<BooleanArray> {
+        for stats in &self.statistics {
+            if let Some(array) = stats.contained(column, values) {
+                return Some(array);
+            }
+        }
+        None
+    }
+}
+
 struct PruningExpressionBuilder<'a> {
     column: phys_expr::Column,
     column_expr: Arc<dyn PhysicalExpr>,