[wip]

Author: rdettai

datafusion/src/physical_plan/file_format/file_stream.rs

@@ -25,9 +25,15 @@ use crate::{
     datasource::{object_store::ObjectStore, PartitionedFile},
     error::Result as DataFusionResult,
     physical_plan::RecordBatchStream,
+    scalar::ScalarValue,
 };
 use arrow::{
-    datatypes::SchemaRef,
+    array::{
+        make_array, Array, ArrayData, ArrayRef, DictionaryArray, StringArray,
+        UInt8BufferBuilder, UInt8Builder,
+    },
+    buffer::Buffer,
+    datatypes::{DataType, SchemaRef, UInt8Type},
     error::{ArrowError, Result as ArrowResult},
     record_batch::RecordBatch,
 };
@@ -61,8 +67,12 @@ impl<T> FormatReaderOpener for T where
 
 /// A stream that iterates record batch by record batch, file over file.
 pub struct FileStream<F: FormatReaderOpener> {
+    /// Partitioning column names for this stream
+    partition_names: Vec<String>,
     /// An iterator over record batches of the last file returned by file_iter
     batch_iter: BatchIter,
+    /// Partitioning column values for the current batch_iter
+    partition_values: Vec<ScalarValue>,
     /// An iterator over input files
     file_iter: FileIter,
     /// The stream schema (file schema after projection)
@@ -74,6 +84,10 @@ pub struct FileStream<F: FormatReaderOpener> {
     /// is not capable of limiting the number of records in the last batch, the file
     /// stream will take care of truncating it.
     file_reader: F,
+    /// A buffer initialized to zeros that represents the key array of all partition
+    /// columns (partition columns are materialized by dictionary arrays with only one
+    /// value in the dictionary, thus all the keys are equal to zero)
+    key_buffer: Option<Buffer>,
 }
 
 impl<F: FormatReaderOpener> FileStream<F> {
@@ -83,6 +97,7 @@ impl<F: FormatReaderOpener> FileStream<F> {
         file_reader: F,
         schema: SchemaRef,
         limit: Option<usize>,
+        partition_names: Vec<String>,
     ) -> Self {
         let read_iter = files.into_iter().map(move |f| -> DataFusionResult<_> {
             object_store
@@ -93,16 +108,31 @@ impl<F: FormatReaderOpener> FileStream<F> {
         Self {
             file_iter: Box::new(read_iter),
             batch_iter: Box::new(iter::empty()),
+            partition_values: vec![],
             remain: limit,
             schema,
             file_reader,
+            partition_names,
+            key_buffer: None,
         }
     }
 
-    /// Acts as a flat_map of record batches over files.
+    /// Acts as a flat_map of record batches over files. Adds the partitioning
+    /// Columns to the returned record batches.
     fn next_batch(&mut self) -> Option<ArrowResult<RecordBatch>> {
         match self.batch_iter.next() {
-            Some(batch) => Some(batch),
+            Some(Ok(batch)) => {
+                // let mut cols = batch.columns().to_vec();
+                // let partitions = self.partition_names.iter().zip(&self.partition_values);
+                // for p in partitions {
+                //     cols.push(
+                //         Int8Array::from(std::iter::repeat(0u8).take(batch.num_rows()))
+                //             .finish_dict,
+                //     )
+                // }
+                Some(Ok(batch))
+            }
+            Some(Err(e)) => Some(Err(e)),
             None => match self.file_iter.next() {
                 Some(Ok(f)) => {
                     self.batch_iter = (self.file_reader)(f, &self.remain);
@@ -113,6 +143,38 @@ impl<F: FormatReaderOpener> FileStream<F> {
             },
         }
     }
+
+    fn init_and_cache_buf(&mut self, len: usize) -> Buffer {
+        let mut key_buffer_builder = UInt8BufferBuilder::new(len);
+        key_buffer_builder.advance(len); // keys are all 0
+        self.key_buffer = Some(key_buffer_builder.finish());
+        self.key_buffer.as_ref().unwrap().clone()
+    }
+
+    fn create_dict_array(&mut self, val: &ScalarValue, len: usize) -> ArrayRef {
+        let value_datatype = val.get_datatype();
+
+        // build value dictionary
+        let dict_vals = val.to_array();
+
+        // build keys array
+        let key_buffer = match self.key_buffer {
+            None => self.init_and_cache_buf(len),
+            Some(buf) if buf.len() < len => self.init_and_cache_buf(len),
+            Some(buf) => buf.slice(buf.len() - len),
+        };
+
+        // create data type
+        let data_type =
+            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(value_datatype));
+
+        // assemble pieces together
+        let mut builder = ArrayData::builder(data_type)
+            .len(len)
+            .add_buffer(key_buffer);
+        builder = builder.add_child_data(dict_vals.data().clone());
+        Arc::new(DictionaryArray::<UInt8Type>::from(builder.build().unwrap()))
+    }
 }
 
 impl<F: FormatReaderOpener> Stream for FileStream<F> {