Merge pull request #33566 from Avogar/fix-avro

Fix segfault in Avro
ClickHouse · Jan 14, 2022 · d54a430 · d54a430
2 parents 1254225 + 1abc899
commit d54a430
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 2 deletions.
diff --git a/src/Processors/Formats/Impl/AvroRowOutputFormat.cpp b/src/Processors/Formats/Impl/AvroRowOutputFormat.cpp
@@ -391,20 +391,27 @@ AvroRowOutputFormat::AvroRowOutputFormat(
 
 AvroRowOutputFormat::~AvroRowOutputFormat() = default;
 
-void AvroRowOutputFormat::writePrefix()
+void AvroRowOutputFormat::createFileWriter()
 {
-    // we have to recreate avro::DataFileWriterBase object due to its interface limitations
     file_writer_ptr = std::make_unique<avro::DataFileWriterBase>(
         std::make_unique<OutputStreamWriteBufferAdapter>(out),
         serializer.getSchema(),
         settings.avro.output_sync_interval,
         getCodec(settings.avro.output_codec));
+}
+
+void AvroRowOutputFormat::writePrefix()
+{
+    // we have to recreate avro::DataFileWriterBase object due to its interface limitations
+    createFileWriter();
 
     file_writer_ptr->syncIfNeeded();
 }
 
 void AvroRowOutputFormat::write(const Columns & columns, size_t row_num)
 {
+    if (!file_writer_ptr)
+        createFileWriter();
     file_writer_ptr->syncIfNeeded();
     serializer.serializeRow(columns, row_num, file_writer_ptr->encoder());
     file_writer_ptr->incr();

diff --git a/src/Processors/Formats/Impl/AvroRowOutputFormat.h b/src/Processors/Formats/Impl/AvroRowOutputFormat.h
@@ -58,6 +58,8 @@ class AvroRowOutputFormat : public IRowOutputFormat
     virtual void writePrefix() override;
     virtual void writeSuffix() override;
 
+    void createFileWriter();
+
     FormatSettings settings;
     AvroSerializer serializer;
     std::unique_ptr<avro::DataFileWriterBase> file_writer_ptr;

diff --git a/tests/queries/0_stateless/02168_avro_bug.reference b/tests/queries/0_stateless/02168_avro_bug.reference
@@ -0,0 +1 @@
+OK
diff --git a/tests/queries/0_stateless/02168_avro_bug.sql b/tests/queries/0_stateless/02168_avro_bug.sql
@@ -0,0 +1,5 @@
+-- Tags: no-fasttest
+insert into table function file('data.avro', 'Avro', 'x UInt64') select * from numbers(10);
+insert into table function file('data.avro', 'Avro', 'x UInt64') select * from numbers(10);
+insert into table function file('data.avro', 'Avro', 'x UInt64') select * from numbers(10);
+select 'OK';