pepijnve commented on code in PR #24925:
URL: https://github.com/apache/datafusion/pull/24925#discussion_r3956963030
##########
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:
To be honest, I was kind of lazy here and pretty much copy/pasted from the
`datasource-arrow` crate. I don't think that has a test covering exceeding the
buffer limit either.
Good idea to make sure this code path is actually covered by tests. Just a
test that writes the same batch in a loop enough times so that the threshold is
crossed at least once?
--
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]
