switch to variant coalesce expression

Author: richardc-db

kernel/kernel-api/src/main/java/io/delta/kernel/Scan.java

@@ -36,6 +36,7 @@
 import io.delta.kernel.internal.util.ColumnMapping;
 import io.delta.kernel.internal.util.PartitionUtils;
 import io.delta.kernel.internal.util.Tuple2;
+import io.delta.kernel.internal.util.VariantUtils;
 
 /**
  * Represents a scan of a Delta table.
@@ -194,6 +195,13 @@ public FilteredColumnarBatch next() {
                     nextDataBatch = nextDataBatch.withDeletedColumnAt(rowIndexOrdinal);
                 }
 
+                // Transform physical variant columns (struct of binaries) into logical variant
+                // columns.
+                nextDataBatch = VariantUtils.withVariantColumns(
+                    tableClient.getExpressionHandler(),
+                    nextDataBatch
+                );
+
                 // Add partition columns
                 nextDataBatch =
                     PartitionUtils.withPartitionColumns(

kernel/kernel-api/src/main/java/io/delta/kernel/internal/util/VariantUtils.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright (2024) The Delta Lake Project Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.delta.kernel.internal.util;
+
+import java.util.Arrays;
+
+import io.delta.kernel.client.ExpressionHandler;
+import io.delta.kernel.data.ColumnVector;
+import io.delta.kernel.data.ColumnarBatch;
+import io.delta.kernel.expressions.*;
+import io.delta.kernel.types.*;
+
+public class VariantUtils {
+    public static ColumnarBatch withVariantColumns(
+            ExpressionHandler expressionHandler,
+            ColumnarBatch dataBatch) {
+        for (int i = 0; i < dataBatch.getSchema().length(); i++) {
+            StructField field = dataBatch.getSchema().at(i);
+            if (!(field.getDataType() instanceof StructType) &&
+                !(field.getDataType() instanceof ArrayType) &&
+                !(field.getDataType() instanceof MapType) &&
+                (field.getDataType() != VariantType.VARIANT ||
+                dataBatch.getColumnVector(i).getDataType() == VariantType.VARIANT)) {
+                continue;
+            }
+
+            ExpressionEvaluator evaluator = expressionHandler.getEvaluator(
+                // Field here is variant type if its actually a variant.
+                // TODO: probably better to pass in the schema as an argument
+                // so the schema is enforced at the expression level. Need to pass in a literal
+                // schema
+                new StructType().add(field),
+                new ScalarExpression(
+                    "variant_coalesce",
+                    Arrays.asList(new Column(field.getName()))
+                ),
+                VariantType.VARIANT
+            );
+
+            // TODO: don't need to pass in the entire batch.
+            ColumnVector variantCol = evaluator.eval(dataBatch);
+            // TODO: make a more efficient way to do this.
+            dataBatch =
+                dataBatch.withDeletedColumnAt(i).withNewColumn(i, field, variantCol);
+        }
+        return dataBatch;
+    }
+
+    private static ColumnVector[] getColumnBatchVectors(ColumnarBatch batch) {
+        ColumnVector[] res = new ColumnVector[batch.getSchema().length()];
+        for (int i = 0; i < batch.getSchema().length(); i++) {
+            res[i] = batch.getColumnVector(i);
+        }
+        return res;
+    }
+}

kernel/kernel-defaults/src/main/java/io/delta/kernel/defaults/internal/data/vector/AbstractColumnVector.java

@@ -79,6 +79,10 @@ public boolean isNullAt(int rowId) {
         return nullability.get()[rowId];
     }
 
+    public Optional<boolean[]> getNullability() {
+        return nullability;
+    }
+
     @Override
     public boolean getBoolean(int rowId) {
         throw unsupportedDataAccessException("boolean");

kernel/kernel-defaults/src/main/java/io/delta/kernel/defaults/internal/data/vector/DefaultArrayVector.java

@@ -88,4 +88,12 @@ public ColumnVector getElements() {
             }
         };
     }
+
+    public ColumnVector getElementVector() {
+        return elementVector;
+    }
+
+    public int[] getOffsets() {
+        return offsets;
+    }
 }

kernel/kernel-defaults/src/main/java/io/delta/kernel/defaults/internal/data/vector/DefaultMapVector.java

@@ -98,4 +98,16 @@ public ColumnVector getValues() {
             }
         };
     }
+
+    public ColumnVector getKeyVector() {
+        return keyVector;
+    }
+
+    public ColumnVector getValueVector() {
+        return valueVector;
+    }
+
+    public int[] getOffsets() {
+        return offsets;
+    }
 }

kernel/kernel-defaults/src/main/java/io/delta/kernel/defaults/internal/expressions/DefaultExpressionEvaluator.java

@@ -15,6 +15,7 @@
  */
 package io.delta.kernel.defaults.internal.expressions;
 
+import java.util.Arrays;
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;
@@ -33,8 +34,7 @@
 import static io.delta.kernel.internal.util.ExpressionUtils.getUnaryChild;
 import static io.delta.kernel.internal.util.Preconditions.checkArgument;
 
-import io.delta.kernel.defaults.internal.data.vector.DefaultBooleanVector;
-import io.delta.kernel.defaults.internal.data.vector.DefaultConstantVector;
+import io.delta.kernel.defaults.internal.data.vector.*;
 import static io.delta.kernel.defaults.internal.expressions.DefaultExpressionUtils.booleanWrapperVector;
 import static io.delta.kernel.defaults.internal.expressions.DefaultExpressionUtils.childAt;
 import static io.delta.kernel.defaults.internal.expressions.DefaultExpressionUtils.compare;
@@ -48,6 +48,7 @@
  */
 public class DefaultExpressionEvaluator implements ExpressionEvaluator {
     private final Expression expression;
+    private final StructType inputSchema;
 
     /**
      * Create a {@link DefaultExpressionEvaluator} instance bound to the given expression and
@@ -68,12 +69,14 @@ public DefaultExpressionEvaluator(
                     "Can not create an expression handler returns result of type %s", outputType);
             throw DeltaErrors.unsupportedExpression(expression, Optional.of(reason));
         }
+        // TODO(richardc-db): Hack to avoid needing to pass the schema into the expression.
+        this.inputSchema = inputSchema;
         this.expression = transformResult.expression;
     }
 
     @Override
     public ColumnVector eval(ColumnarBatch input) {
-        return new ExpressionEvalVisitor(input).visit(expression);
+        return new ExpressionEvalVisitor(input, inputSchema).visit(expression);
     }
 
     @Override
@@ -278,6 +281,21 @@ ExpressionTransformResult visitCoalesce(ScalarExpression coalesce) {
             );
         }
 
+        @Override
+        ExpressionTransformResult visitVariantCoalesce(ScalarExpression variantCoalesce) {
+            checkArgument(
+                variantCoalesce.getChildren().size() == 1,
+                "Expected one input to 'variant_coalesce but received %s",
+                variantCoalesce.getChildren().size()
+            );
+            Expression transformedVariantInput = visit(childAt(variantCoalesce, 0)).expression;
+            return new ExpressionTransformResult(
+                new ScalarExpression(
+                    "VARIANT_COALESCE",
+                    Arrays.asList(transformedVariantInput)),
+                VariantType.VARIANT);
+        }
+
         private Predicate validateIsPredicate(
             Expression baseExpression,
             ExpressionTransformResult result) {
@@ -318,9 +336,11 @@ private Expression transformBinaryComparator(Predicate predicate) {
      */
     private static class ExpressionEvalVisitor extends ExpressionVisitor<ColumnVector> {
         private final ColumnarBatch input;
+        private final StructType inputSchema;
 
-        ExpressionEvalVisitor(ColumnarBatch input) {
+        ExpressionEvalVisitor(ColumnarBatch input, StructType inputSchema) {
             this.input = input;
+            this.inputSchema = inputSchema;
         }
 
         /*
@@ -558,6 +578,119 @@ ColumnVector visitCoalesce(ScalarExpression coalesce) {
             );
         }
 
+        @Override
+        ColumnVector visitVariantCoalesce(ScalarExpression variantCoalesce) {
+            return variantCoalesceImpl(
+                visit(childAt(variantCoalesce, 0)),
+                inputSchema.at(0).getDataType()
+            );
+        }
+
+        private ColumnVector variantCoalesceImpl(ColumnVector inputVec, DataType dt) {
+            if (dt instanceof StructType) {
+                StructType structType = (StructType) dt;
+                DefaultStructVector structVec = (DefaultStructVector) inputVec;
+                ColumnVector[] structColVecs = new ColumnVector[structType.length()];
+                for (int i = 0; i < structType.length(); i++) {
+                    if (structType.at(i).getDataType() instanceof ArrayType ||
+                        structType.at(i).getDataType() instanceof StructType ||
+                        structType.at(i).getDataType() instanceof MapType ||
+                        structType.at(i).getDataType() instanceof VariantType) {
+                        structColVecs[i] = variantCoalesceImpl(
+                            structVec.getChild(i),
+                            structType.at(i).getDataType()
+                        );
+                    } else {
+                        structColVecs[i] = structVec.getChild(i);
+                    }
+                }
+                return new DefaultStructVector(
+                    structVec.getSize(),
+                    structType,
+                    structVec.getNullability(),
+                    structColVecs
+                );
+            }
+
+            if (dt instanceof ArrayType) {
+                ArrayType arrType = (ArrayType) dt;
+                DefaultArrayVector arrVec = (DefaultArrayVector) inputVec;
+
+                if (arrType.getElementType() instanceof ArrayType ||
+                    arrType.getElementType() instanceof StructType ||
+                    arrType.getElementType() instanceof MapType ||
+                    arrType.getElementType() instanceof VariantType) {
+                    ColumnVector elementVec = variantCoalesceImpl(
+                        arrVec.getElementVector(),
+                        arrType.getElementType()
+                    );
+
+                    return new DefaultArrayVector(
+                        arrVec.getSize(),
+                        arrType,
+                        arrVec.getNullability(),
+                        arrVec.getOffsets(),
+                        elementVec
+                    );
+                }
+                return arrVec;
+            }
+
+            if (dt instanceof MapType) {
+                MapType mapType = (MapType) dt;
+                DefaultMapVector mapVec = (DefaultMapVector) inputVec;
+
+                ColumnVector valueVec = mapVec.getValueVector();
+                if (mapType.getValueType() instanceof ArrayType ||
+                    mapType.getValueType() instanceof StructType ||
+                    mapType.getValueType() instanceof MapType ||
+                    mapType.getValueType() instanceof VariantType) {
+                    valueVec = variantCoalesceImpl(
+                        mapVec.getValueVector(),
+                        mapType.getValueType()
+                    );
+                }
+                ColumnVector keyVec = mapVec.getKeyVector();
+                if (mapType.getKeyType() instanceof ArrayType ||
+                    mapType.getKeyType() instanceof StructType ||
+                    mapType.getKeyType() instanceof MapType ||
+                    mapType.getKeyType() instanceof VariantType) {
+                    keyVec = variantCoalesceImpl(
+                        mapVec.getKeyVector(),
+                        mapType.getKeyType()
+                    );
+                }
+                return new DefaultMapVector(
+                    mapVec.getSize(),
+                    mapType,
+                    mapVec.getNullability(),
+                    mapVec.getOffsets(),
+                    keyVec,
+                    valueVec
+                );
+            }
+
+            DefaultStructVector structBackingVariant = (DefaultStructVector) inputVec;
+            checkArgument(
+                structBackingVariant.getChild(0).getDataType() instanceof BinaryType,
+                "Expected struct field 0 backing variant to be binary type. Actual: %s",
+                structBackingVariant.getChild(0).getDataType()
+            );
+            checkArgument(
+                structBackingVariant.getChild(1).getDataType() instanceof BinaryType,
+                "Expected struct field 1 backing variant to be binary type. Actual: %s",
+                structBackingVariant.getChild(1).getDataType()
+            );
+
+            return new DefaultVariantVector(
+                structBackingVariant.getSize(),
+                VariantType.VARIANT,
+                structBackingVariant.getNullability(),
+                structBackingVariant.getChild(0),
+                structBackingVariant.getChild(1)
+            );
+        }
+
         /**
          * Utility method to evaluate inputs to the binary input expression. Also validates the
          * evaluated expression result {@link ColumnVector}s are of the same size.

kernel/kernel-defaults/src/main/java/io/delta/kernel/defaults/internal/expressions/ExpressionVisitor.java

@@ -59,6 +59,8 @@ abstract class ExpressionVisitor<R> {
 
     abstract R visitCoalesce(ScalarExpression ifNull);
 
+    abstract R visitVariantCoalesce(ScalarExpression variantCoalesce);
+
     final R visit(Expression expression) {
         if (expression instanceof PartitionValueExpression) {
             return visitPartitionValue((PartitionValueExpression) expression);
@@ -105,6 +107,8 @@ private R visitScalarExpression(ScalarExpression expression) {
                 return visitIsNull(new Predicate(name, children));
             case "COALESCE":
                 return visitCoalesce(expression);
+            case "VARIANT_COALESCE":
+                return visitVariantCoalesce(expression);
             default:
                 throw new UnsupportedOperationException(
                     String.format("Scalar expression `%s` is not supported.", name));

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

@@ -83,7 +83,12 @@ public static Converter createConverter(
         } else if (typeFromClient instanceof TimestampNTZType) {
             return createTimestampNtzConverter(initialBatchSize, typeFromFile);
         } else if (typeFromClient instanceof VariantType) {
-            return new VariantColumnReader(initialBatchSize);
+            return new RowColumnReader(
+                initialBatchSize,
+                new StructType()
+                    .add("value", BinaryType.BINARY, false)
+                    .add("metadata", BinaryType.BINARY, false),
+                (GroupType) typeFromFile);
         }
 
         throw new UnsupportedOperationException(typeFromClient + " is not supported");