adriangb commented on code in PR #17861:
URL: https://github.com/apache/datafusion/pull/17861#discussion_r2984333782
##########
datafusion/datasource-avro/src/source.rs:
##########
@@ -56,22 +57,83 @@ impl AvroSource {
}
}
- fn open<R: std::io::Read>(&self, reader: R) -> Result<AvroReader<'static,
R>> {
+ fn open<R: std::io::BufRead>(
+ &self,
+ reader: R,
+ projection: Option<Vec<usize>>,
+ ) -> Result<Reader<R>> {
+ let mut builder = ReaderBuilder::new()
+ .with_batch_size(self.batch_size.expect("Batch size must set
before open"));
+ if let Some(projection) = projection {
+ builder = builder.with_projection(projection);
+ }
+ builder.build(reader).map_err(Into::into)
+ }
+
+ fn projected_file_schema(&self) -> SchemaRef {
let file_schema = self.table_schema.file_schema();
- let projection = Some(
+ if self.projection.file_indices.is_empty() {
+ return Arc::clone(file_schema);
+ }
+
+ Arc::new(Schema::new(
self.projection
.file_indices
.iter()
- .map(|&idx| file_schema.field(idx).name().clone())
+ .map(|idx| file_schema.field(*idx).clone())
.collect::<Vec<_>>(),
- );
- AvroReader::try_new(
- reader,
- &Arc::clone(self.table_schema.file_schema()),
- self.batch_size.expect("Batch size must set before open"),
- projection.as_ref(),
- )
+ ))
+ }
+
+ fn writer_projection_for_schema(
+ &self,
+ writer_schema: &Schema,
+ target_schema: &Schema,
+ ) -> Option<Vec<usize>> {
+ // `arrow-avro` accepts projection ordinals against the file's writer
schema,
+ // while DataFusion plans projection against the logical table schema.
Remap
+ // projected column names to writer ordinals so reader-level pushdown
still
+ // preserves DataFusion's existing name-based projection semantics.
+ let projection = target_schema
+ .fields()
+ .iter()
+ .filter_map(|field| {
+ writer_schema
+ .column_with_name(field.name())
+ .map(|(idx, _)| idx)
+ })
+ .collect::<Vec<_>>();
+
+ let identity_projection = projection.len() ==
writer_schema.fields().len()
+ && projection
+ .iter()
+ .enumerate()
+ .all(|(idx, value)| idx == *value);
+
+ (!identity_projection).then_some(projection)
+ }
+}
+
+fn coerce_batch_to_schema(
+ batch: &RecordBatch,
+ target_schema: SchemaRef,
+) -> Result<RecordBatch> {
Review Comment:
This seems awfully similar to
https://docs.rs/datafusion/latest/datafusion/physical_expr_adapter/struct.BatchAdapter.html
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
