kosiew commented on code in PR #24925:
URL: https://github.com/apache/datafusion/pull/24925#discussion_r3956912209
##########
datafusion/datasource-avro/src/file_format.rs:
##########
@@ -147,10 +171,173 @@ impl FileFormat for AvroFormat {
Ok(DataSourceExec::from_data_source(conf))
}
+ async fn create_writer_physical_plan(
+ &self,
+ input: Arc<dyn ExecutionPlan>,
+ _state: &dyn Session,
+ conf: FileSinkConfig,
+ order_requirements: Option<LexRequirement>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if conf.insert_op != InsertOp::Append {
+ return not_impl_err!("Overwrites are not implemented yet for Avro
format");
+ }
+
+ let sink = Arc::new(AvroFileSink::new(conf));
+
+ Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _)
+ }
+
fn file_source(
&self,
table_schema: datafusion_datasource::TableSchema,
) -> Arc<dyn FileSource> {
Arc::new(AvroSource::new(table_schema))
}
}
+
+/// Implements [`FileSink`] for Avro Object Container Files
+struct AvroFileSink {
+ config: FileSinkConfig,
+}
+
+impl AvroFileSink {
+ fn new(config: FileSinkConfig) -> Self {
+ Self { config }
+ }
+}
+
+#[async_trait]
+impl FileSink for AvroFileSink {
+ fn config(&self) -> &FileSinkConfig {
+ &self.config
+ }
+
+ async fn spawn_writer_tasks_and_join(
+ &self,
+ context: &Arc<TaskContext>,
+ demux_task: SpawnedTask<Result<()>>,
+ mut file_stream_rx: DemuxedStreamReceiver,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Result<u64> {
+ let mut file_write_tasks: JoinSet<std::result::Result<usize,
DataFusionError>> =
+ JoinSet::new();
+
+ let writer_schema = get_writer_schema(&self.config);
+ while let Some((path, mut rx)) = file_stream_rx.recv().await {
+ let shared_buffer = SharedBuffer::new(INITIAL_BUFFER_BYTES);
+ let mut avro_writer: AvroWriter<SharedBuffer> =
+ WriterBuilder::new(writer_schema.as_ref().clone())
+ .build::<_, AvroOcfFormat>(shared_buffer.clone())
+ .map_err(|e| {
+ internal_datafusion_err!("Failed to create Avro
writer: {e}")
+ })?;
+ let mut object_store_writer = ObjectWriterBuilder::new(
+ FileCompressionType::UNCOMPRESSED,
+ &path,
+ Arc::clone(&object_store),
+ )
+ .with_buffer_size(Some(
+ context
+ .session_config()
+ .options()
+ .execution
+ .objectstore_writer_buffer_size,
+ ))
+ .build()?;
+ file_write_tasks.spawn(async move {
+ let mut row_count = 0;
+ while let Some(batch) = rx.recv().await {
+ row_count += batch.num_rows();
+ avro_writer
+ .write(&batch)
+ .map_err(|e| internal_datafusion_err!("{e}"))?;
+ let mut buff_to_flush =
shared_buffer.buffer.try_lock().unwrap();
+ if buff_to_flush.len() > BUFFER_FLUSH_BYTES {
Review Comment:
Could we add a round-trip test where the encoded Avro output exceeds
`BUFFER_FLUSH_BYTES`, ideally across multiple input batches? The current tests
don't appear to exercise the path where `SharedBuffer` is cleared and later
Avro blocks are appended to the object-store writer. It would be good to cover
this boundary to make sure larger, multi-block files aren't accidentally
truncated or malformed.
--
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]
