Merge branch 'develop' into operators-bench

Author: joseph-isaacs

vortex-array/src/arrays/dict/vtable/rules.rs

@@ -15,6 +15,7 @@ use crate::arrays::DictArray;
 use crate::arrays::DictVTable;
 use crate::arrays::ScalarFnArray;
 use crate::builtins::ArrayBuiltins;
+use crate::expr::Pack;
 use crate::optimizer::ArrayOptimizer;
 use crate::optimizer::rules::ArrayParentReduceRule;
 use crate::optimizer::rules::ParentRuleSet;
@@ -44,6 +45,12 @@ impl ArrayParentReduceRule<DictVTable> for DictionaryScalarFnValuesPushDownRule
         // Check that the scalar function can actually be pushed down.
         let sig = parent.scalar_fn().signature();
 
+        // Don't push down pack expressions since we might want to unpack them in exporters
+        // later.
+        if parent.scalar_fn().is::<Pack>() {
+            return Ok(None);
+        }
+
         // If the scalar function is fallible, we cannot push it down since it may fail over a
         // value that isn't referenced by any code.
         if !array.all_values_referenced && sig.is_fallible() {

vortex-duckdb/src/scan.rs

@@ -40,10 +40,8 @@ use vortex::error::vortex_bail;
 use vortex::error::vortex_err;
 use vortex::expr::Expression;
 use vortex::expr::Pack;
-use vortex::expr::and;
 use vortex::expr::and_collect;
 use vortex::expr::col;
-use vortex::expr::lit;
 use vortex::expr::root;
 use vortex::expr::select;
 use vortex::file::OpenOptionsSessionExt;
@@ -53,6 +51,7 @@ use vortex::io::runtime::BlockingRuntime;
 use vortex::io::runtime::current::ThreadSafeIterator;
 use vortex::layout::layouts::USE_VORTEX_OPERATORS;
 use vortex::session::VortexSession;
+use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::RUNTIME;
 use crate::SESSION;
@@ -179,38 +178,32 @@ fn extract_table_filter_expr(
     init: &TableInitInput<VortexTableFunction>,
     column_ids: &[u64],
 ) -> VortexResult<Option<Expression>> {
-    let table_filter_expr = init
-        .table_filter_set()
-        .and_then(|filter| {
-            filter
-                .into_iter()
-                .map(|(idx, ex)| {
-                    let idx_u: usize = idx.as_();
-                    let col_idx: usize = column_ids[idx_u].as_();
-                    let name = init
-                        .bind_data()
-                        .column_names
-                        .get(col_idx)
-                        .vortex_expect("exists");
-                    try_from_table_filter(
-                        &ex,
-                        &col(name.as_str()),
-                        init.bind_data().first_file.dtype(),
-                    )
-                })
-                .reduce(|l, r| l?.zip(r?).map(|(l, r)| Ok(and(l, r))).transpose())
-        })
-        .transpose()?
-        .flatten();
-
-    let complex_filter_expr = and_collect(init.bind_data().filter_exprs.clone());
-    let filter_expr = complex_filter_expr
-        .into_iter()
-        .chain(table_filter_expr)
-        .reduce(and)
-        .unwrap_or_else(|| lit(true));
+    let mut table_filter_exprs: HashSet<Expression> = if let Some(filter) = init.table_filter_set()
+    {
+        filter
+            .into_iter()
+            .map(|(idx, ex)| {
+                let idx_u: usize = idx.as_();
+                let col_idx: usize = column_ids[idx_u].as_();
+                let name = init
+                    .bind_data()
+                    .column_names
+                    .get(col_idx)
+                    .vortex_expect("exists");
+                try_from_table_filter(
+                    &ex,
+                    &col(name.as_str()),
+                    init.bind_data().first_file.dtype(),
+                )
+            })
+            .collect::<VortexResult<Option<HashSet<_>>>>()?
+            .unwrap_or_else(HashSet::new)
+    } else {
+        HashSet::new()
+    };
 
-    Ok(Some(filter_expr))
+    table_filter_exprs.extend(init.bind_data().filter_exprs.clone());
+    Ok(and_collect(table_filter_exprs.into_iter().collect_vec()))
 }
 
 /// Helper function to open a Vortex file from either a local or S3 URL

vortex-scan/src/arrow.rs

@@ -10,9 +10,12 @@ use arrow_schema::SchemaRef;
 use futures::Stream;
 use futures::TryStreamExt;
 use vortex_array::ArrayRef;
+use vortex_array::arrow::ArrowArrayExecutor;
 use vortex_array::arrow::IntoArrowArray;
 use vortex_error::VortexResult;
 use vortex_io::runtime::BlockingRuntime;
+use vortex_layout::layouts::USE_VORTEX_OPERATORS;
+use vortex_session::VortexSession;
 
 use crate::ScanBuilder;
 
@@ -28,9 +31,10 @@ impl ScanBuilder<ArrayRef> {
         runtime: &B,
     ) -> VortexResult<impl RecordBatchReader + 'static> {
         let data_type = DataType::Struct(schema.fields().clone());
+        let session = self.session().clone();
 
         let iter = self
-            .map(move |chunk| to_record_batch(chunk, &data_type))
+            .map(move |chunk| to_record_batch(chunk, &data_type, &session))
             .into_iter(runtime)?
             .map(|result| result.map_err(|e| ArrowError::ExternalError(Box::new(e))));
 
@@ -42,19 +46,29 @@ impl ScanBuilder<ArrayRef> {
         schema: SchemaRef,
     ) -> VortexResult<impl Stream<Item = Result<RecordBatch, ArrowError>> + Send + 'static> {
         let data_type = DataType::Struct(schema.fields().clone());
+        let session = self.session().clone();
 
         let stream = self
-            .map(move |chunk| to_record_batch(chunk, &data_type))
+            .map(move |chunk| to_record_batch(chunk, &data_type, &session))
             .into_stream()?
             .map_err(|e| ArrowError::ExternalError(Box::new(e)));
 
         Ok(stream)
     }
 }
 
-fn to_record_batch(chunk: ArrayRef, data_type: &DataType) -> VortexResult<RecordBatch> {
-    let arrow = chunk.into_arrow(data_type)?;
-    Ok(RecordBatch::from(arrow.as_struct().clone()))
+fn to_record_batch(
+    chunk: ArrayRef,
+    data_type: &DataType,
+    session: &VortexSession,
+) -> VortexResult<RecordBatch> {
+    if *USE_VORTEX_OPERATORS {
+        let arrow = chunk.execute_arrow(data_type, session)?;
+        Ok(RecordBatch::from(arrow.as_struct().clone()))
+    } else {
+        let arrow = chunk.into_arrow(data_type)?;
+        Ok(RecordBatch::from(arrow.as_struct().clone()))
+    }
 }
 
 /// We create an adapter for record batch iterators that supports clone.
@@ -114,6 +128,7 @@ mod tests {
     use vortex_error::VortexResult;
 
     use super::*;
+    use crate::test::SESSION;
 
     fn create_test_struct_array() -> VortexResult<ArrayRef> {
         // Create Arrow arrays
@@ -152,7 +167,7 @@ mod tests {
         let schema = create_arrow_schema();
         let data_type = DataType::Struct(schema.fields().clone());
 
-        let batch = to_record_batch(vortex_array, &data_type)?;
+        let batch = to_record_batch(vortex_array, &data_type, &SESSION)?;
         assert_eq!(batch.num_columns(), 2);
         assert_eq!(batch.num_rows(), 4);
 

vortex-scan/src/scan_builder.rs

@@ -180,6 +180,11 @@ impl<A: 'static + Send> ScanBuilder<A> {
         self.projection.return_dtype(self.layout_reader.dtype())
     }
 
+    /// The session used by the scan.
+    pub fn session(&self) -> &VortexSession {
+        &self.session
+    }
+
     /// Map each split of the scan. The function will be run on the spawned task.
     pub fn map<B: 'static>(
         self,