Support returning row count in prunning aggregate expressions

This lets us effectively prune expressions like IsNotNull

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

Cargo.lock

@@ -1,4 +1,4 @@
 [toolchain]
 channel = "1.90.0"
 components = ["rust-src", "rustfmt", "clippy", "rust-analyzer"]
-profile = "minimal"
+profile = "minimal"

rust-toolchain.toml

@@ -11324,6 +11324,8 @@ impl<K: core::hash::Hash + core::cmp::Eq, V: core::hash::Hash + core::cmp::Eq> c
 
 pub fn vortex_array::expr::pruning::Relation<K, V>::default() -> Self
 
+pub const vortex_array::expr::pruning::ROW_COUNT_FIELD: &str
+
 pub trait vortex_array::expr::pruning::StatsCatalog
 
 pub fn vortex_array::expr::pruning::StatsCatalog::stats_ref(&self, _field_path: &vortex_array::dtype::FieldPath, _stat: vortex_array::expr::stats::Stat) -> core::option::Option<vortex_array::expr::Expression>

vortex-array/public-api.lock

@@ -4,6 +4,7 @@
 pub(crate) mod pruning_expr;
 mod relation;
 
+pub use pruning_expr::ROW_COUNT_FIELD;
 pub use pruning_expr::RequiredStats;
 pub use pruning_expr::checked_pruning_expr;
 pub use pruning_expr::field_path_stat_field_name;
@@ -18,7 +19,8 @@ pub trait StatsCatalog {
     /// Given a field path and statistic, return an expression that when evaluated over the catalog
     /// will return that stat for the referenced field.
     ///
-    /// This is likely to be a column expression, or a literal.
+    /// This is likely to be a column expression, or a literal. Implementations may also expose
+    /// synthetic scope-level values via reserved field paths such as [`ROW_COUNT_FIELD`].
     ///
     /// Returns `None` if the stat is not available for the field path.
     fn stats_ref(&self, _field_path: &FieldPath, _stat: Stat) -> Option<Expression> {

vortex-array/src/expr/pruning/mod.rs

@@ -14,11 +14,17 @@ use crate::dtype::FieldPath;
 use crate::dtype::FieldPathSet;
 use crate::expr::Expression;
 use crate::expr::StatsCatalog;
+use crate::expr::col;
 use crate::expr::get_item;
 use crate::expr::root;
 use crate::expr::stats::Stat;
 
 pub type RequiredStats = Relation<FieldPath, Stat>;
+pub const ROW_COUNT_FIELD: &str = "__vortex_row_count";
+
+fn is_row_count_field(field_path: &FieldPath) -> bool {
+    field_path.parts().len() == 1 && field_path.parts()[0].as_name() == Some(ROW_COUNT_FIELD)
+}
 
 // A catalog that return a stat column whenever it is required, tracking all accessed
 // stats and returning them later.
@@ -43,6 +49,13 @@ struct ScopeStatsCatalog<'a> {
 
 impl StatsCatalog for ScopeStatsCatalog<'_> {
     fn stats_ref(&self, field_path: &FieldPath, stat: Stat) -> Option<Expression> {
+        if is_row_count_field(field_path) {
+            return self
+                .available_stats
+                .contains(field_path)
+                .then(|| col(ROW_COUNT_FIELD));
+        }
+
         let stat_path = field_path.clone().push(stat.name());
 
         if self.available_stats.contains(&stat_path) {
@@ -55,6 +68,10 @@ impl StatsCatalog for ScopeStatsCatalog<'_> {
 
 impl StatsCatalog for TrackingStatsCatalog {
     fn stats_ref(&self, field_path: &FieldPath, stat: Stat) -> Option<Expression> {
+        if is_row_count_field(field_path) {
+            return Some(col(ROW_COUNT_FIELD));
+        }
+
         let mut expr = root();
         let name = field_path_stat_field_name(field_path, stat);
         expr = get_item(name, expr);
@@ -86,6 +103,9 @@ pub fn field_path_stat_field_name(field_path: &FieldPath, stat: Stat) -> FieldNa
 /// cannot hold, and false if it cannot be determined from stats alone whether the positions can
 /// be pruned.
 ///
+/// Include [`ROW_COUNT_FIELD`] in `available_stats` to make the scope row count available to
+/// row-count-aware pruning expressions such as `is_not_null(...)`.
+///
 /// If the falsification logic attempts to access an unknown stat,
 /// this function will return `None`.
 pub fn checked_pruning_expr(

vortex-array/src/expr/pruning/pruning_expr.rs

@@ -11,13 +11,12 @@ use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::arrays::ConstantArray;
 use crate::dtype::DType;
+use crate::dtype::FieldPath;
 use crate::dtype::Nullability;
 use crate::expr::Expression;
 use crate::expr::StatsCatalog;
-use crate::expr::and;
 use crate::expr::eq;
-use crate::expr::gt;
-use crate::expr::lit;
+use crate::expr::pruning::ROW_COUNT_FIELD;
 use crate::expr::stats::Stat;
 use crate::scalar_fn::Arity;
 use crate::scalar_fn::ChildName;
@@ -106,20 +105,12 @@ impl ScalarFnVTable for IsNotNull {
         expr: &Expression,
         catalog: &dyn StatsCatalog,
     ) -> Option<Expression> {
-        // is_not_null is falsified when ALL values are null, i.e. null_count == len.
-        // Since there is no len stat in the zone map, we approximate using IsConstant:
-        // if the zone is constant and has any nulls, then all values must be null.
-        //
-        // TODO(#7187): Add a len stat to enable the more general falsification:
-        //   null_count == len => is_not_null is all false.
-        let null_count_expr = expr.child(0).stat_expression(Stat::NullCount, catalog)?;
-        let is_constant_expr = expr.child(0).stat_expression(Stat::IsConstant, catalog)?;
-        // If the zone is constant (is_constant == true) and has nulls (null_count > 0),
-        // then all values must be null, so is_not_null is all false.
-        Some(and(
-            eq(is_constant_expr, lit(true)),
-            gt(null_count_expr, lit(0u64)),
-        ))
+        // is_not_null is falsified when ALL values are null, i.e. null_count == row_count.
+        let child = expr.child(0);
+        let null_count_expr = child.stat_expression(Stat::NullCount, catalog)?;
+        let row_count_expr =
+            catalog.stats_ref(&FieldPath::from_name(ROW_COUNT_FIELD), Stat::NullCount)?;
+        Some(eq(null_count_expr, row_count_expr))
     }
 }
 
@@ -267,11 +258,9 @@ mod tests {
         use crate::dtype::Field;
         use crate::dtype::FieldPath;
         use crate::dtype::FieldPathSet;
-        use crate::expr::and;
         use crate::expr::col;
         use crate::expr::eq;
-        use crate::expr::gt;
-        use crate::expr::lit;
+        use crate::expr::pruning::ROW_COUNT_FIELD;
         use crate::expr::pruning::checked_pruning_expr;
         use crate::expr::stats::Stat;
 
@@ -281,24 +270,18 @@ mod tests {
             &expr,
             &FieldPathSet::from_iter([
                 FieldPath::from_iter([Field::Name("a".into()), Field::Name("null_count".into())]),
-                FieldPath::from_iter([Field::Name("a".into()), Field::Name("is_constant".into())]),
+                FieldPath::from_name(ROW_COUNT_FIELD),
             ]),
         )
         .unwrap();
 
         assert_eq!(
             &pruning_expr,
-            &and(
-                eq(col("a_is_constant"), lit(true)),
-                gt(col("a_null_count"), lit(0u64)),
-            )
+            &eq(col("a_null_count"), col(ROW_COUNT_FIELD))
         );
         assert_eq!(
             st.map(),
-            &HashMap::from_iter([(
-                FieldPath::from_name("a"),
-                HashSet::from([Stat::NullCount, Stat::IsConstant])
-            )])
+            &HashMap::from_iter([(FieldPath::from_name("a"), HashSet::from([Stat::NullCount]))])
         );
     }
 }

vortex-array/src/scalar_fn/fns/is_not_null.rs

@@ -19,6 +19,7 @@ use vortex_array::dtype::FieldMask;
 use vortex_array::dtype::FieldPath;
 use vortex_array::dtype::FieldPathSet;
 use vortex_array::expr::Expression;
+use vortex_array::expr::pruning::ROW_COUNT_FIELD;
 use vortex_array::expr::pruning::checked_pruning_expr;
 use vortex_error::VortexResult;
 use vortex_layout::LayoutReader;
@@ -139,7 +140,8 @@ impl VortexFile {
                             Field::Name(stat.name().into()),
                         ])
                     })
-                }),
+                })
+                .chain(std::iter::once(FieldPath::from_name(ROW_COUNT_FIELD))),
         );
 
         let Some((predicate, required_stats)) = checked_pruning_expr(filter, &set) else {
@@ -157,6 +159,7 @@ impl VortexFile {
             &required_file_stats,
             stats.stats_sets(),
             fields,
+            self.footer.row_count(),
         )?
         else {
             return Ok(false);

vortex-file/src/file.rs

@@ -9,14 +9,13 @@ use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::StructArray;
 use vortex_array::dtype::Field;
 use vortex_array::dtype::FieldName;
-use vortex_array::dtype::FieldNames;
 use vortex_array::dtype::FieldPath;
 use vortex_array::dtype::StructFields;
+use vortex_array::expr::pruning::ROW_COUNT_FIELD;
 use vortex_array::expr::pruning::field_path_stat_field_name;
 use vortex_array::expr::stats::Stat;
 use vortex_array::expr::stats::StatsProvider;
 use vortex_array::stats::StatsSet;
-use vortex_array::validity::Validity;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
@@ -28,13 +27,14 @@ pub fn extract_relevant_file_stats_as_struct_row(
     access: &HashMap<FieldPath, HashSet<Stat>>,
     stats_sets: &Arc<[StatsSet]>,
     struct_dtype: &StructFields,
+    row_count: u64,
 ) -> VortexResult<Option<ArrayRef>> {
-    if access.is_empty() {
-        return StructArray::try_new(FieldNames::default(), vec![], 1, Validity::NonNullable)
-            .map(|s| Some(s.into_array()));
-    }
+    let mut columns: Vec<(FieldName, ArrayRef)> = Vec::with_capacity(access.len() * 2 + 1);
+    columns.push((
+        FieldName::from(ROW_COUNT_FIELD),
+        ConstantArray::new(row_count, 1).into_array(),
+    ));
 
-    let mut columns: Vec<(FieldName, ArrayRef)> = Vec::with_capacity(access.len() * 2);
     for (field_path, stats) in access.into_iter() {
         if field_path.parts().len() != 1 {
             return Ok(None);
@@ -76,3 +76,38 @@ pub fn extract_relevant_file_stats_as_struct_row(
         StructArray::from_fields(columns.as_slice())?.into_array(),
     ))
 }
+
+#[cfg(test)]
+mod tests {
+    use vortex_array::LEGACY_SESSION;
+    use vortex_array::VortexSessionExecute;
+    use vortex_array::arrays::StructArray;
+    use vortex_array::arrays::struct_::StructArrayExt;
+
+    use super::*;
+
+    #[test]
+    fn includes_row_count_field_when_no_stats_are_needed() {
+        let row = extract_relevant_file_stats_as_struct_row(
+            &HashMap::default(),
+            &Arc::default(),
+            &StructFields::default(),
+            128,
+        )
+        .unwrap()
+        .unwrap()
+        .execute::<StructArray>(&mut LEGACY_SESSION.create_execution_ctx())
+        .unwrap();
+
+        assert_eq!(row.names().as_ref(), &[FieldName::from(ROW_COUNT_FIELD)]);
+        assert_eq!(
+            row.unmasked_field_by_name(ROW_COUNT_FIELD)
+                .unwrap()
+                .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())
+                .unwrap()
+                .as_primitive()
+                .typed_value::<u64>(),
+            Some(128),
+        );
+    }
+}

vortex-file/src/pruning.rs

@@ -23,6 +23,7 @@ use vortex_array::dtype::StructFields;
 use vortex_array::expr::Expression;
 use vortex_array::expr::StatsCatalog;
 use vortex_array::expr::lit;
+use vortex_array::expr::pruning::ROW_COUNT_FIELD;
 use vortex_array::expr::stats::Stat;
 use vortex_array::scalar::Scalar;
 use vortex_array::scalar_fn::fns::literal::Literal;
@@ -115,6 +116,11 @@ impl FileStatsLayoutReader {
 /// Implements [`StatsCatalog`] to provide file-level stats to expressions during pruning evaluation.
 impl StatsCatalog for FileStatsLayoutReader {
     fn stats_ref(&self, field_path: &FieldPath, stat: Stat) -> Option<Expression> {
+        if field_path.parts().len() == 1 && field_path.parts()[0].as_name() == Some(ROW_COUNT_FIELD)
+        {
+            return Some(lit(self.child.row_count()));
+        }
+
         // FileStats currently only holds top-level field statistics.
         if field_path.parts().len() != 1 {
             return None;
@@ -126,7 +132,8 @@ impl StatsCatalog for FileStatsLayoutReader {
 
         let stat_value = field_stats.get(stat)?.as_exact()?;
         let field_dtype = self.struct_fields.field_by_index(field_idx)?;
-        let stat_scalar = Scalar::try_new(field_dtype, Some(stat_value)).ok()?;
+        let stat_dtype = stat.dtype(&field_dtype)?;
+        let stat_scalar = Scalar::try_new(stat_dtype, Some(stat_value)).ok()?;
 
         Some(lit(stat_scalar))
     }
@@ -209,12 +216,14 @@ mod tests {
 
     use vortex_array::ArrayContext;
     use vortex_array::IntoArray as _;
+    use vortex_array::arrays::PrimitiveArray;
     use vortex_array::arrays::StructArray;
     use vortex_array::dtype::DType;
     use vortex_array::dtype::Nullability;
     use vortex_array::dtype::PType;
     use vortex_array::expr::get_item;
     use vortex_array::expr::gt;
+    use vortex_array::expr::is_not_null;
     use vortex_array::expr::lit;
     use vortex_array::expr::root;
     use vortex_array::expr::stats::Precision;
@@ -259,6 +268,18 @@ mod tests {
         )
     }
 
+    fn test_file_null_count_stats(null_count: u64) -> FileStatistics {
+        let mut stats = StatsSet::default();
+        stats.set(
+            Stat::NullCount,
+            Precision::exact(ScalarValue::from(null_count)),
+        );
+        FileStatistics::new(
+            Arc::from([stats]),
+            Arc::from([DType::Primitive(PType::I32, Nullability::Nullable)]),
+        )
+    }
+
     #[test]
     fn pruning_when_filter_out_of_range() -> VortexResult<()> {
         block_on(|handle| async {
@@ -337,4 +358,47 @@ mod tests {
             Ok(())
         })
     }
+
+    #[test]
+    fn pruning_is_not_null_when_file_is_all_null() -> VortexResult<()> {
+        block_on(|handle| async {
+            let session = SESSION.clone().with_handle(handle);
+            let ctx = ArrayContext::empty();
+            let segments = Arc::new(TestSegments::default());
+            let (ptr, eof) = SequenceId::root().split();
+            let struct_array = StructArray::from_fields(
+                [(
+                    "col",
+                    PrimitiveArray::from_option_iter([None::<i32>, None, None, None, None])
+                        .into_array(),
+                )]
+                .as_slice(),
+            )?;
+            let strategy = TableStrategy::new(
+                Arc::new(FlatLayoutStrategy::default()),
+                Arc::new(FlatLayoutStrategy::default()),
+            );
+            let layout = strategy
+                .write_stream(
+                    ctx,
+                    Arc::<TestSegments>::clone(&segments),
+                    struct_array.into_array().to_array_stream().sequenced(ptr),
+                    eof,
+                    &session,
+                )
+                .await?;
+
+            let child = layout.new_reader("".into(), segments, &SESSION)?;
+
+            let reader =
+                FileStatsLayoutReader::new(child, test_file_null_count_stats(5), SESSION.clone());
+
+            let expr = is_not_null(get_item("col", root()));
+            let mask = Mask::new_true(5);
+            let result = reader.pruning_evaluation(&(0..5), &expr, mask)?.await?;
+            assert_eq!(result, Mask::new_false(5));
+
+            Ok(())
+        })
+    }
 }

vortex-file/src/v2/file_stats_reader.rs

@@ -47,6 +47,7 @@ vortex-flatbuffers = { workspace = true, features = ["layout"] }
 vortex-io = { workspace = true }
 vortex-mask = { workspace = true }
 vortex-metrics = { workspace = true }
+vortex-runend = { workspace = true }
 vortex-scan = { workspace = true }
 vortex-sequence = { workspace = true }
 vortex-session = { workspace = true }