refactor(trace/parquet): Drop dependency on Avro

This change updates the Parquet support library in OpenDC to not rely on
Avro, but instead interface directly with Parquet's reading and writing
functionality, providing less overhead.

1 files changed, 88 insertions, 30 deletions

diff --git a/opendc-trace/opendc-trace-parquet/src/test/kotlin/org/opendc/trace/util/parquet/ParquetTest.kt b/opendc-trace/opendc-trace-parquet/src/test/kotlin/org/opendc/trace/util/parquet/ParquetTest.kt
index 8ef4d1fb..be354319 100644
--- a/opendc-trace/opendc-trace-parquet/src/test/kotlin/org/opendc/trace/util/parquet/ParquetTest.kt
+++ b/opendc-trace/opendc-trace-parquet/src/test/kotlin/org/opendc/trace/util/parquet/ParquetTest.kt
@@ -22,36 +22,81 @@
 
 package org.opendc.trace.util.parquet
 
-import org.apache.avro.SchemaBuilder
-import org.apache.avro.generic.GenericData
-import org.apache.parquet.avro.AvroParquetReader
-import org.apache.parquet.avro.AvroParquetWriter
+import org.apache.hadoop.conf.Configuration
 import org.apache.parquet.hadoop.ParquetFileWriter
+import org.apache.parquet.hadoop.api.ReadSupport
+import org.apache.parquet.hadoop.api.WriteSupport
+import org.apache.parquet.io.api.Converter
+import org.apache.parquet.io.api.GroupConverter
+import org.apache.parquet.io.api.PrimitiveConverter
+import org.apache.parquet.io.api.RecordConsumer
+import org.apache.parquet.io.api.RecordMaterializer
+import org.apache.parquet.schema.MessageType
+import org.apache.parquet.schema.PrimitiveType
+import org.apache.parquet.schema.Type
+import org.apache.parquet.schema.Types
 import org.junit.jupiter.api.*
 import org.junit.jupiter.api.Assertions.assertEquals
-import java.io.File
 import java.nio.file.FileAlreadyExistsException
+import java.nio.file.Files
 import java.nio.file.NoSuchFileException
+import java.nio.file.Path
 
 /**
  * Test suite for the Parquet helper classes.
  */
 internal class ParquetTest {
-    private val schema = SchemaBuilder
-        .record("test")
-        .namespace("org.opendc.format.util")
-        .fields()
-        .name("field").type().intType().noDefault()
-        .endRecord()
+    private lateinit var path: Path
 
-    private lateinit var file: File
+    private val schema = Types.buildMessage()
+        .addField(
+            Types.primitive(PrimitiveType.PrimitiveTypeName.INT32, Type.Repetition.REQUIRED)
+                .named("field")
+        )
+        .named("test")
+    private val writeSupport = object : WriteSupport<Int>() {
+        lateinit var recordConsumer: RecordConsumer
+
+        override fun init(configuration: Configuration): WriteContext {
+            return WriteContext(schema, emptyMap())
+        }
+
+        override fun prepareForWrite(recordConsumer: RecordConsumer) {
+            this.recordConsumer = recordConsumer
+        }
+
+        override fun write(record: Int) {
+            val consumer = recordConsumer
+
+            consumer.startMessage()
+            consumer.startField("field", 0)
+            consumer.addInteger(record)
+            consumer.endField("field", 0)
+            consumer.endMessage()
+        }
+    }
+
+    private val readSupport = object : ReadSupport<Int>() {
+        override fun init(
+            configuration: Configuration,
+            keyValueMetaData: Map<String, String>,
+            fileSchema: MessageType
+        ): ReadContext = ReadContext(fileSchema)
+
+        override fun prepareForRead(
+            configuration: Configuration,
+            keyValueMetaData: Map<String, String>,
+            fileSchema: MessageType,
+            readContext: ReadContext
+        ): RecordMaterializer<Int> = TestRecordMaterializer()
+    }
 
     /**
-     * Setup the test
+     * Set up the test
      */
     @BeforeEach
     fun setUp() {
-        file = File.createTempFile("opendc", "parquet")
+        path = Files.createTempFile("opendc", "parquet")
     }
 
     /**
@@ -59,7 +104,7 @@ internal class ParquetTest {
      */
     @AfterEach
     fun tearDown() {
-        file.delete()
+        Files.deleteIfExists(path)
     }
 
     /**
@@ -68,29 +113,24 @@ internal class ParquetTest {
     @Test
     fun testSmoke() {
         val n = 4
-        val writer = AvroParquetWriter.builder<GenericData.Record>(LocalOutputFile(file))
-            .withSchema(schema)
+        val writer = LocalParquetWriter.builder(path, writeSupport)
             .withWriteMode(ParquetFileWriter.Mode.OVERWRITE)
             .build()
 
         try {
             repeat(n) { i ->
-                val record = GenericData.Record(schema)
-                record.put("field", i)
-                writer.write(record)
+                writer.write(i)
             }
         } finally {
             writer.close()
         }
 
-        val reader = AvroParquetReader.builder<GenericData.Record>(LocalInputFile(file))
-            .build()
-
+        val reader = LocalParquetReader(path, readSupport)
         var counter = 0
         try {
             while (true) {
                 val record = reader.read() ?: break
-                assertEquals(counter++, record.get("field"))
+                assertEquals(counter++, record)
             }
         } finally {
             reader.close()
@@ -105,9 +145,7 @@ internal class ParquetTest {
     @Test
     fun testOverwrite() {
         assertThrows<FileAlreadyExistsException> {
-            AvroParquetWriter.builder<GenericData.Record>(LocalOutputFile(file))
-                .withSchema(schema)
-                .build()
+            LocalParquetWriter.builder(path, writeSupport).build()
         }
     }
 
@@ -116,10 +154,30 @@ internal class ParquetTest {
      */
     @Test
     fun testNonExistent() {
-        file.delete()
+        Files.deleteIfExists(path)
         assertThrows<NoSuchFileException> {
-            AvroParquetReader.builder<GenericData.Record>(LocalInputFile(file))
-                .build()
+            LocalParquetReader(path, readSupport)
+        }
+    }
+
+    private class TestRecordMaterializer : RecordMaterializer<Int>() {
+        private var current: Int = 0
+        private val fieldConverter = object : PrimitiveConverter() {
+            override fun addInt(value: Int) {
+                current = value
+            }
+        }
+        private val root = object : GroupConverter() {
+            override fun getConverter(fieldIndex: Int): Converter {
+                require(fieldIndex == 0)
+                return fieldConverter
+            }
+            override fun start() {}
+            override fun end() {}
         }
+
+        override fun getCurrentRecord(): Int = current
+
+        override fun getRootConverter(): GroupConverter = root
     }
 }