init

Author: richardc-db

kernel/kernel-defaults/src/main/java/io/delta/kernel/defaults/internal/parquet/ParquetColumnWriters.java

@@ -32,6 +32,7 @@
 import static io.delta.kernel.internal.util.Preconditions.checkArgument;
 
 import io.delta.kernel.defaults.internal.DefaultKernelUtils;
+import io.delta.kernel.defaults.internal.data.vector.DefaultVariantVector;
 import static io.delta.kernel.defaults.internal.parquet.ParquetSchemaUtils.MAX_BYTES_PER_PRECISION;
 
 /**
@@ -142,6 +143,8 @@ private static ColumnWriter createColumnWriter(
             return new MapWriter(colName, fieldIndex, columnVector);
         } else if (dataType instanceof StructType) {
             return new StructWriter(colName, fieldIndex, columnVector);
+        } else if (dataType instanceof VariantType) {
+            return new VariantWriter(colName, fieldIndex, columnVector);
         }
 
         throw new IllegalArgumentException("Unsupported column vector type: " + dataType);
@@ -474,4 +477,31 @@ void writeNonNullRowValue(RecordConsumer recordConsumer, int rowId) {
             recordConsumer.endGroup();
         }
     }
+
+    static class VariantWriter extends ColumnWriter {
+        private ColumnWriter valueWriter;
+        private ColumnWriter metadataWriter;
+
+        VariantWriter(String name, int fieldId, ColumnVector variantColumnVector) {
+            super(name, fieldId, variantColumnVector);
+            valueWriter = new BinaryWriter(
+                "value",
+                0,
+                ((DefaultVariantVector) variantColumnVector).getValueVector()
+            );
+            metadataWriter = new BinaryWriter(
+                "metadata",
+                1,
+                ((DefaultVariantVector) variantColumnVector).getMetadataVector()
+            );
+        }
+
+        @Override
+        void writeNonNullRowValue(RecordConsumer recordConsumer, int rowId) {
+            recordConsumer.startGroup();
+            valueWriter.writeRowValue(recordConsumer, rowId);
+            metadataWriter.writeRowValue(recordConsumer, rowId);
+            recordConsumer.endGroup();
+        }
+    }
 }

kernel/kernel-defaults/src/test/scala/io/delta/kernel/defaults/internal/parquet/ParquetFileWriterSuite.scala

@@ -189,6 +189,30 @@ class ParquetFileWriterSuite extends AnyFunSuite
       }
   }
 
+  test("write variant") {
+    withTempDir { tempPath =>
+      val variantFilePath = goldenTableFile("variantbasic").getAbsolutePath
+      val schema = tableSchema(variantFilePath)
+
+      val physicalSchema = if (hasColumnMappingId(variantFilePath)) {
+        convertToPhysicalSchema(schema, schema, ColumnMapping.COLUMN_MAPPING_MODE_ID)
+      } else {
+        schema
+      }
+      val readData = readParquetUsingKernelAsColumnarBatches(variantFilePath, physicalSchema)
+        // Convert the schema of the data to the physical schema with field ids
+        .map(_.withNewSchema(physicalSchema))
+        // convert the data to filtered columnar batches
+        .map(_.toFiltered(Option.empty[Predicate]))
+
+      val writeOutput = writeToParquetUsingKernel(readData, tempPath.getAbsolutePath)
+
+      val readWrittenData =
+        readParquetUsingKernelAsColumnarBatches(tempPath.getAbsolutePath, physicalSchema)
+      // TODO(r.chen): Finish this and make assertions.
+    }
+  }
+
   test("columnar batches containing different schema") {
     withTempDir { tempPath =>
       val targetDir = tempPath.getAbsolutePath