pepijnve commented on code in PR #24925:
URL: https://github.com/apache/datafusion/pull/24925#discussion_r3957511905
##########
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:
I've added an SLT that triggers the flush code path multiple times and reads
back the result.
--
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]
