Support more complex types when reading into arrow record batch.

[liurenjie1024]

This is a follow up issue of #244, where we have some limitations of reading into arrow record batch:

1. Only primitive type is supported. That's the projected fields count not be either nested filed of structs, or non primitive types such as map, list.
2. Type promotion is missing.
3. Default value is not supported yet.

We should implement sth like ArrowProjectionVisitor to support these.

[sdd]

I'm happy to start looking into this one if you like?

[liurenjie1024]

I'm happy to start looking into this one if you like?

Thanks @sdd !

[sdd]

Hi @liurenjie1024. I have a query regarding flattening. You gave this example in Issue #244:

... For example, when we select (person.address.street, person.name], we have projection mask [1,2], and the returned schema is

schema {
  struct person {
     struct address {
        string street
     }
     string name
  }
}

But what we expect is

schema {
   string person
   string name
}

Originally posted by @liurenjie1024 in #244 (comment)

Why do we expect that? I would have thought it would turn out like this:

schema {
    string street
    string name
}

[liurenjie1024]

@sdd Yes, it was a typo, it should be

schema {
    string street
    string name
}

[sdd]

I'm tackling items 2 and 3 from @liurenjie1024's initial comment on this issue. I've got a design proposal that I'd like to share to get feedback on.

It involves creating a post-processor that can transform RecordBatches that come from the ParquetRecordBatchStream before we send them back to the caller. If we detect that schema migration requires us to promote / insert columns, then we create an instance of this post-processor and append it into the processing pipeline. If it is not needed, the pipeline remains unchanged from its current configuration.

The post-processor could look like this:

use arrow_array::{Array as ArrowArray, ArrayRef as ArrowArrayRef, RecordBatch};
use arrow_schema::{DataType, Schema as ArrowSchema};
use crate::spec::Schema as IcebergSchema;

/// Represents an operation that may need to be performed
/// to transform a RecordBatch coming from a Parquet file record
/// batch stream to match a newer Iceberg schema that has evolved from
/// the one that was used to write the parquet file.
pub(crate) enum EvolutionOp {
    // signifies that a particular column has undergone type promotion,
    // thus the column with the given index needs to be promoted to the
    // specified type
    Promote {
        index: usize,
        target_type: DataType
    },

    // Signifies that a new column has been inserted before the column
    // with index `index`. (we choose "before" rather than "after" so
    // that we can use a usize;  if we insert after, then we need to
    // be able to store -1 here when we want to indicate that the new
    // column is to be added at the front of the list).
    // If multiple columns need to be inserted at a given
    // location, they should all be given the same index, as the index
    // here refers to the original record batch, not the interim state after
    // a preceding operation.
    Add {
        index: usize,
        target_type: DataType,
        value: Option<ArrowArrayRef> // A Scalar
    },

    // signifies that a column has been renamed from one schema to the next.
    // this requires no change to the data within a record batch, only to its
    // schema.
    Rename {
        index: usize,
        target_name: String,
    }

    // The iceberg spec refers to other permissible schema evolution actions
    // (see https://iceberg.apache.org/spec/#schema-evolution):
    // deleting fields and reordering fields.
    // However, these actions can be achieved without needing this
    // slower post-processing step by using the projection mask.
}


pub(crate) struct RecordBatchEvolutionProcessor {
    operations: Vec<EvolutionOp>,

    // Every transformed RecordBatch will have the same schema. We create the
    // target just once and cache it here. Helpfully, Arc<Schema> is needed in
    // the constructor for RecordBatch, so we don't need an expensive copy
    // each time.
    target_schema: OnceCell<Arc<ArrowSchema>>

    // Caching any columns that we need to add is harder as the number of rows
    // in the record batches can vary from batch to batch within the stream,
    // so rather than storing cached added columns here too, we have to
    // generate them on the fly.
}

impl RecordBatchEvolutionProcessor {
    /// returns None if no processor is required
    pub(crate) fn build(source_schema: &ArrowSchemaRef, snapshot_schema: &IcebergSchema, projected_iceberg_field_ids: &[i32]) -> Option<Self> {
        let operations: Vec<_> = Self::generate_operations(
            source_schema,
            snapshot_schema,
            projected_iceberg_field_ids
        );

        if operations.is_empty() {
            None
        } else {
            Some(Self { operations, target_schema: OnceCell::default() })
        }
    }

    pub(crate) fn process_stream(&self, _source_stream: impl Stream<Item=Result<RecordBatch>>) -> Result<ArrowRecordBatchStream> {
        todo!()
    }

    pub(crate) fn process_record_batch(&self, record_batch: RecordBatch) -> Result<RecordBatch> {
        let new_batch_schema = self.target_schema.get_or_init(|| self.create_target_schema());

        RecordBatch::try_new(
            new_batch_schema.clone(),
            self.transform_columns(record_batch.columns())
        )?
    }

    fn generate_operations(source_schema: &ArrowSchemaRef, snapshot_schema: &IcebergSchema, projected_iceberg_field_ids: &[i32]) -> Vec<EvolutionOp> {
        // create the (possibly empty) list of `EvolutionOp`s that we need
        // to apply to the arrays in a record batch with `source_schema` so
        // that it matches the `snapshot_schema`
        todo!();
    }

    fn transform_columns(&self, columns: &[Arc<dyn ArrowArray>]) -> Vec<Arc<dyn ArrowArray>> {
        // iterate over source_columns and self.operations,
        // populating a Vec::with_capacity as we go
        todo!();
    }

    fn create_target_schema(&self) -> Arc<ArrowSchema> {
        todo!();
    }
}

The iceberg ArrowReader would be modified thusly:

impl ArrowReader {
    // ...
    async fn process_file_scan_task(...) {
        // ...

        // create a RecordBatchEvolutionProcessor if our task schema contains columns
        // not present in the parquet file or whose types have been promoted
        let record_batch_evolution_processor = RecordBatchEvolutionProcessor::build(
            record_batch_stream_builder.schema(),
            task.schema(),
            task.project_field_ids(),
        );

        // ...

        // Build the batch stream and send all the RecordBatches that it generates
        // to the requester.
        let mut record_batch_stream: impl Stream<Item=Result<RecordBatch>> = record_batch_stream_builder.build()?;

        if let Some(record_batch_evolution_processor) = record_batch_evolution_processor {
            record_batch_stream = record_batch_evolution_processor.process_stream(record_batch_stream)?;
        }

        while let Some(batch) = record_batch_stream.try_next().await? {
            tx.send(Ok(batch)).await?
        }

        Ok(())
    }
    // ...
}

@liurenjie1024, @Xuanwo: What do you think?

[liurenjie1024]

Hi, @sdd sorry for late reply. I think the design is great, but one missing part is that simple index:usize can't deal with nested data type? For example, the original schema:

schema {
  struct person {
     struct address {
        string street
        int zip
     }
     string name
  }
}

And I want to select person.address, name, and we evolved person.address.zip to long.

One solution I can come up with is that we slightly modify the index to be the post order index of target schema. This way your design would work perfect, what do you think?

[ryzhyk]

@sdd, I'm curious about the status of this. I ran into this issue today and noticed that you're working on a fix, so was wondering if you still plan to do it.

[sdd]

Hi @ryzhyk - I implemented the default value handling and type promotion limitations mentioned in @liurenjie1024's comment at the top of the issue, but the complex type support was not as high a priority to me at the time compared so some other features, so I did not continue with that item.

I don't have a timeframe in mind of when I will pick this up - my main focus has always been on the read side but I wanted to address delete file support first and I'm working on that at the moment.

I'm happy to give someone else the opportunity to work on this if they're available to, and potentially help where I can - otherwise I'll come back to it at some point once I'm done with delete file table scan support.

[ryzhyk]

Thanks for the update @sdd !