File pruning when set (#3363)

Add whole file pruning to the vortex-file/ffi
Signed-off-by: Joe Isaacs <[email protected]>

---------

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

duckdb-vortex/src/include/vortex_common.hpp

@@ -61,6 +61,10 @@ struct FileReader {
 		return duckdb::make_uniq<FileReader>(file);
 	}
 
+	vx_array_iterator *Scan(const vx_file_scan_options *options) {
+		return Try([&](auto err) { return vx_file_reader_scan(this->file, options, err); });
+	}
+
 	uint64_t FileRowCount() {
 		return Try([&](auto err) { return vx_file_row_count(file, err); });
 	}

duckdb-vortex/src/include/vortex_layout_reader.hpp

@@ -18,8 +18,6 @@
 #include "concurrentqueue.h"
 
 #include "vortex.hpp"
-#include "vortex_extension.hpp"
-#include "vortex_layout_reader.hpp"
 #include "vortex_scan.hpp"
 #include "vortex_common.hpp"
 #include "vortex_expr.hpp"
@@ -42,7 +40,7 @@ struct BindData : public TableFunctionData {
 
 	// Used to read the schema during the bind phase and cached here to
 	// avoid having to open the same file again during the scan phase.
-	unique_ptr<FileReader> initial_file;
+	shared_ptr<FileReader> initial_file;
 
 	// Used to create an arena for protobuf exprs, need a ptr since the bind arg is const.
 	unique_ptr<google::protobuf::Arena> arena;
@@ -107,7 +105,7 @@ struct ScanGlobalState : public GlobalTableFunctionState {
 	// Multi producer, multi consumer lockfree queue.
 	duckdb_moodycamel::ConcurrentQueue<ScanPartition> scan_partitions {8192};
 
-	std::vector<std::shared_ptr<LayoutReader>> layout_readers;
+	std::vector<shared_ptr<FileReader>> file_readers;
 
 	// The column idx that must be returned by the scan.
 	vector<idx_t> column_ids;
@@ -252,9 +250,9 @@ static bool PinFileToThread(ScanGlobalState &global_state) {
 }
 
 static void CreateScanPartitions(ClientContext &context, const BindData &bind, ScanGlobalState &global_state,
-                                 ScanLocalState &local_state, uint64_t file_idx, unique_ptr<FileReader> &file_reader) {
+                                 ScanLocalState &local_state, uint64_t file_idx, FileReader &file_reader) {
 	const auto file_name = global_state.expanded_files[file_idx];
-	const auto row_count = Try([&](auto err) { return vx_file_row_count(file_reader->file, err); });
+	const auto row_count = Try([&](auto err) { return vx_file_row_count(file_reader.file, err); });
 
 	const auto thread_count = std::thread::hardware_concurrency();
 	const auto file_count = global_state.expanded_files.size();
@@ -292,8 +290,7 @@ static void CreateScanPartitions(ClientContext &context, const BindData &bind, S
 	D_ASSERT(global_state.files_partitioned <= global_state.expanded_files.size());
 }
 
-static unique_ptr<ArrayIterator> OpenArrayIter(ScanGlobalState &global_state,
-                                               std::shared_ptr<LayoutReader> &layout_reader,
+static unique_ptr<ArrayIterator> OpenArrayIter(ScanGlobalState &global_state, shared_ptr<FileReader> &file_reader,
                                                ScanPartition row_range_partition) {
 	const auto options = vx_file_scan_options {
 	    .projection = global_state.projected_column_names.data(),
@@ -305,7 +302,7 @@ static unique_ptr<ArrayIterator> OpenArrayIter(ScanGlobalState &global_state,
 	    .row_range_end = row_range_partition.end_row,
 	};
 
-	return make_uniq<ArrayIterator>(layout_reader->Scan(&options));
+	return make_uniq<ArrayIterator>(file_reader->Scan(&options));
 }
 
 // Assigns the next array from the array stream.
@@ -348,8 +345,8 @@ static bool GetNextArray(ClientContext &context, const BindData &bind_data, Scan
 
 		// Layout readers are safe to share across threads for reading. Further, they
 		// are created before pushing partitions of the corresponing files into a queue.
-		auto layout_reader = global_state.layout_readers[partition.file_idx];
-		local_state.array_iterator = OpenArrayIter(global_state, layout_reader, partition);
+		auto file_reader = global_state.file_readers[partition.file_idx];
+		local_state.array_iterator = OpenArrayIter(global_state, file_reader, partition);
 	}
 
 	local_state.currently_scanned_array = local_state.array_iterator->NextArray();
@@ -378,20 +375,25 @@ static void VortexScanFunction(ClientContext &context, TableFunctionInput &data,
 				return;
 			}
 
-			// Free layout readers as long as we pin files to threads.
+			// Free file readers when owned by the thread.
 			if (local_state.scan_partitions.empty() && local_state.thread_local_file_idx.has_value()) {
-				global_state.layout_readers[local_state.thread_local_file_idx.value()] = nullptr;
+				global_state.file_readers[local_state.thread_local_file_idx.value()] = nullptr;
 				local_state.thread_local_file_idx.reset();
 			}
 
 			// Create new scan partitions in case the queue is empty.
 			if (auto file_idx = global_state.next_file_idx.fetch_add(1);
 			    file_idx < global_state.expanded_files.size()) {
-				auto file_name = global_state.expanded_files[file_idx];
-				auto vortex_file =
-				    OpenFileAndVerify(FileSystem::GetFileSystem(context), *bind_data.session, file_name, bind_data);
-				global_state.layout_readers[file_idx] = LayoutReader::CreateFromFile(vortex_file.get());
-				CreateScanPartitions(context, bind_data, global_state, local_state, file_idx, vortex_file);
+				if (file_idx == 0) {
+					global_state.file_readers[0] = bind_data.initial_file;
+				} else {
+					auto file_name = global_state.expanded_files[file_idx];
+					global_state.file_readers[file_idx] =
+					    OpenFileAndVerify(FileSystem::GetFileSystem(context), *bind_data.session, file_name, bind_data);
+				}
+
+				CreateScanPartitions(context, bind_data, global_state, local_state, file_idx,
+				                     *global_state.file_readers[file_idx]);
 			}
 		}
 	}
@@ -491,7 +493,7 @@ void RegisterScanFunction(DatabaseInstance &instance) {
 		}
 
 		// Resizing the empty vector default constructs std::shared pointers at all indices with nullptr.
-		global_state->layout_readers.resize(global_state->expanded_files.size());
+		global_state->file_readers.resize(global_state->expanded_files.size());
 
 		bind.arena->Reset();
 		return std::move(global_state);

duckdb-vortex/src/vortex_scan.cpp

@@ -397,7 +397,7 @@ pub enum FieldOrIdentity {
     Identity,
 }
 
-pub(crate) fn stat_field_name(field: &FieldName, stat: Stat) -> FieldName {
+pub fn stat_field_name(field: &FieldName, stat: Stat) -> FieldName {
     FieldName::from(stat_field_name_string(field, stat))
 }
 

vortex-expr/src/pruning.rs

@@ -119,11 +119,6 @@ typedef struct vx_error vx_error;
  */
 typedef struct vx_file_reader vx_file_reader;
 
-/**
- * A Vortex layout reader.
- */
-typedef struct vx_layout_reader vx_layout_reader;
-
 /**
  * An object that stores registries and caches.
  * This should if possible be reused between queries in ann interactive session.
@@ -404,23 +399,15 @@ struct vx_dtype *vx_file_dtype(const struct vx_file_reader *file);
 /**
  * Build a new `vx_array_iterator` that returns a series of `vx_array`s from a scan over a `vx_layout_reader`.
  */
-struct vx_array_iterator *vx_layout_reader_scan(const struct vx_layout_reader *layout_reader,
-                                                const struct vx_file_scan_options *opts,
-                                                struct vx_error **error);
+struct vx_array_iterator *vx_file_reader_scan(const struct vx_file_reader *file_reader,
+                                              const struct vx_file_scan_options *opts,
+                                              struct vx_error **error);
 
 /**
  * Returns the row count for a given file reader.
  */
 uint64_t vx_file_row_count(struct vx_file_reader *file_reader, struct vx_error **error);
 
-/**
- * Creates a layout reader for a given file.
- */
-struct vx_layout_reader *vx_layout_reader_create(struct vx_file_reader *file_reader,
-                                                 struct vx_error **error);
-
-void vx_layout_reader_free(struct vx_layout_reader *layout_reader);
-
 /**
  * Free the file and all associated resources.
  *

vortex-ffi/cinclude/vortex.h

@@ -29,15 +29,6 @@ pub struct vx_array_iterator {
     pub inner: Option<Box<dyn ArrayIterator>>,
 }
 
-/// Creates a new ArrayIterator wrapper for FFI use.
-pub fn vx_array_iterator<I>(iter: I) -> vortex::error::VortexResult<*mut vx_array_iterator>
-where
-    I: ArrayIterator + 'static,
-{
-    let inner = Some(Box::new(iter) as Box<dyn ArrayIterator>);
-    Ok(Box::into_raw(Box::new(vx_array_iterator { inner })))
-}
-
 /// Attempt to advance the `current` pointer of the iterator.
 ///
 /// A return value of `true` indicates that another element was pulled from the iterator, and a return

vortex-ffi/src/array.rs

@@ -3,7 +3,7 @@
 use std::ffi::{CStr, c_char, c_int, c_uint, c_ulong};
 use std::str::FromStr;
 use std::sync::Arc;
-use std::{ptr, slice};
+use std::{iter, ptr, slice};
 
 use itertools::Itertools;
 use object_store::aws::{AmazonS3Builder, AmazonS3ConfigKey};
@@ -18,7 +18,7 @@ use vortex::error::{VortexError, VortexExpect, VortexResult, vortex_bail, vortex
 use vortex::expr::{ExprRef, Identity, deserialize_expr, select};
 use vortex::file::scan::SplitBy;
 use vortex::file::{VortexFile, VortexOpenOptions, VortexWriteOptions};
-use vortex::layout::LayoutReader;
+use vortex::iter::ArrayIteratorAdapter;
 use vortex::layout::scan::ScanBuilder;
 use vortex::proto::expr::Expr;
 
@@ -33,12 +33,6 @@ pub struct vx_file_reader {
     pub inner: VortexFile,
 }
 
-/// A Vortex layout reader.
-#[allow(non_camel_case_types)]
-pub struct vx_layout_reader {
-    pub inner: Arc<dyn LayoutReader>,
-}
-
 /// Options supplied for opening a file.
 #[repr(C)]
 pub struct vx_file_open_options {
@@ -242,20 +236,33 @@ pub unsafe extern "C-unwind" fn vx_file_dtype(file: *const vx_file_reader) -> *m
 
 /// Build a new `vx_array_iterator` that returns a series of `vx_array`s from a scan over a `vx_layout_reader`.
 #[unsafe(no_mangle)]
-pub unsafe extern "C-unwind" fn vx_layout_reader_scan(
-    layout_reader: *const vx_layout_reader,
+pub unsafe extern "C-unwind" fn vx_file_reader_scan(
+    file_reader: *const vx_file_reader,
     opts: *const vx_file_scan_options,
     error: *mut *mut vx_error,
 ) -> *mut vx_array_iterator {
     try_or(error, ptr::null_mut(), || {
-        let layout_reader = unsafe { layout_reader.as_ref().vortex_expect("null layout reader") };
-        let mut scan_builder = ScanBuilder::new(layout_reader.inner.clone());
+        let file_reader = unsafe { file_reader.as_ref().vortex_expect("null file reader") };
 
         let scan_options = unsafe { opts.as_ref() }.map_or_else(
             || Ok(ScanOptions::default()),
             |options| options.process_scan_options(),
         )?;
 
+        if let Some(expr) = &scan_options.filter_expr {
+            if file_reader.inner.can_prune(expr)? {
+                let dtype = file_reader.inner.dtype().clone();
+                let empty_iter = ArrayIteratorAdapter::new(dtype, iter::empty());
+
+                return Ok(Box::into_raw(Box::new(vx_array_iterator {
+                    inner: Some(Box::new(empty_iter)),
+                })));
+            }
+        };
+
+        let layout_reader = file_reader.inner.layout_reader()?;
+        let mut scan_builder = ScanBuilder::new(layout_reader.clone());
+
         // Apply options if provided.
         if let Some(field_names) = scan_options.field_names {
             // Field names are allowed to be `Some` and empty.
@@ -274,7 +281,9 @@ pub unsafe extern "C-unwind" fn vx_layout_reader_scan(
             scan_builder = scan_builder.with_split_by(split_by_value);
         }
 
-        vx_array_iterator(scan_builder.into_array_iter()?)
+        Ok(Box::into_raw(Box::new(vx_array_iterator {
+            inner: Some(Box::new(scan_builder.into_array_iter()?)),
+        })))
     })
 }
 
@@ -290,27 +299,6 @@ pub extern "C-unwind" fn vx_file_row_count(
     })
 }
 
-/// Creates a layout reader for a given file.
-#[unsafe(no_mangle)]
-pub extern "C-unwind" fn vx_layout_reader_create(
-    file_reader: *mut vx_file_reader,
-    error: *mut *mut vx_error,
-) -> *mut vx_layout_reader {
-    try_or(error, ptr::null_mut(), || {
-        let file_reader = unsafe { file_reader.as_ref().vortex_expect("null file_reader") };
-        let inner = file_reader.inner.layout_reader()?;
-
-        Ok(Box::into_raw(Box::new(vx_layout_reader { inner })))
-    })
-}
-
-#[unsafe(no_mangle)]
-pub extern "C-unwind" fn vx_layout_reader_free(layout_reader: *mut vx_layout_reader) {
-    if !layout_reader.is_null() {
-        drop(unsafe { Box::from_raw(layout_reader) });
-    }
-}
-
 /// Free the file and all associated resources.
 ///
 /// This function will not automatically free any :c:func:`vx_array_iterator` that were built from

vortex-ffi/src/file.rs

@@ -8,12 +8,15 @@ use vortex_array::ArrayRef;
 use vortex_array::stats::StatsSet;
 use vortex_dtype::DType;
 use vortex_error::VortexResult;
+use vortex_expr::ExprRef;
+use vortex_expr::pruning::PruningPredicate;
 use vortex_layout::LayoutReader;
 use vortex_layout::scan::ScanBuilder;
 use vortex_layout::segments::SegmentSource;
 use vortex_metrics::VortexMetrics;
 
 use crate::footer::Footer;
+use crate::pruning::extract_relevant_stat_as_struct_row;
 
 /// Represents a Vortex file, providing access to its metadata and content.
 ///
@@ -80,6 +83,32 @@ impl VortexFile {
     pub fn scan(&self) -> VortexResult<ScanBuilder<ArrayRef>> {
         Ok(ScanBuilder::new(self.layout_reader()?).with_metrics(self.metrics.clone()))
     }
+
+    // Returns true if the expression will never match any rows in the file.
+    pub fn can_prune(&self, filter: &ExprRef) -> VortexResult<bool> {
+        let Some((file_stats, struct_dtype)) = self
+            .footer
+            .statistics()
+            .zip(self.footer.dtype().as_struct())
+        else {
+            return Ok(false);
+        };
+
+        let Some(predicate) = PruningPredicate::try_new(filter) else {
+            return Ok(false);
+        };
+
+        let Some(struct_row) =
+            extract_relevant_stat_as_struct_row(&predicate, file_stats, struct_dtype)?
+        else {
+            return Ok(false);
+        };
+
+        Ok(predicate
+            .evaluate(&struct_row)?
+            .and_then(|p| p.as_constant())
+            .is_some_and(|result| result.as_bool().value() == Some(true)))
+    }
 }
 
 /// A factory for creating segment sources that read data from a Vortex file.

vortex-file/src/file.rs

@@ -96,6 +96,7 @@ mod footer;
 mod generic;
 mod memory;
 mod open;
+mod pruning;
 pub mod segments;
 mod strategy;
 #[cfg(test)]