Use operators for RecordBatch conversion (#5771)

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

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,