Frame writers: Coerce numeric and array types in certain cases. (#16994)

This patch adds "TypeCastSelectors", which is used when writing frames to
perform two coercions:

- When a numeric type is desired and the underlying type is non-numeric or
  unknown, the underlying selector is wrapped, "getObject" is called and the
  result is coerced using "ExprEval.ofType". This differs from the prior
  behavior where the primitive methods like "getLong", "getDouble", etc, would
  be called directly. This fixes an issue where a column would be read as
  all-zeroes when its SQL type is numeric and its physical type is string, which
  can happen when evolving a column's type from string to number.

-  When an array type is desired, the underlying selector is wrapped,
   "getObject" is called, and the result is coerced to Object[]. This coercion
   replaces some earlier logic from #15917.

Author: gianm

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPostShuffleFrameProcessor.java

@@ -109,7 +109,7 @@ public GroupByPostShuffleFrameProcessor(
             RowBasedGrouperHelper.createResultRowBasedColumnSelectorFactory(
                 query,
                 () -> outputRow,
-                RowSignature.Finalization.YES
+                GroupByQueryKit.isFinalize(query) ? RowSignature.Finalization.YES : RowSignature.Finalization.NO
             )
         );
   }

integration-tests-ex/cases/src/test/resources/catalog/incompatibleTypeAssignmentWithValidationDisabled_select.sql

@@ -4,7 +4,7 @@
     "expectedResults": [
       {
         "__time": 1672058096000,
-        "double_col": 0.0
+        "double_col": null
       }
     ]
   }

processing/src/main/java/org/apache/druid/frame/field/FieldWriters.java

@@ -22,6 +22,7 @@
 import org.apache.druid.frame.key.RowKey;
 import org.apache.druid.frame.write.RowBasedFrameWriterFactory;
 import org.apache.druid.frame.write.UnsupportedColumnTypeException;
+import org.apache.druid.frame.write.cast.TypeCastSelectors;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.query.dimension.DefaultDimensionSpec;
 import org.apache.druid.segment.ColumnSelectorFactory;
@@ -101,7 +102,8 @@ private static FieldWriter makeLongWriter(
       final String columnName
   )
   {
-    final ColumnValueSelector<?> selector = selectorFactory.makeColumnValueSelector(columnName);
+    final ColumnValueSelector<?> selector =
+        TypeCastSelectors.makeColumnValueSelector(selectorFactory, columnName, ColumnType.LONG);
     return LongFieldWriter.forPrimitive(selector);
   }
 
@@ -110,7 +112,8 @@ private static FieldWriter makeFloatWriter(
       final String columnName
   )
   {
-    final ColumnValueSelector<?> selector = selectorFactory.makeColumnValueSelector(columnName);
+    final ColumnValueSelector<?> selector =
+        TypeCastSelectors.makeColumnValueSelector(selectorFactory, columnName, ColumnType.FLOAT);
     return FloatFieldWriter.forPrimitive(selector);
   }
 
@@ -119,7 +122,8 @@ private static FieldWriter makeDoubleWriter(
       final String columnName
   )
   {
-    final ColumnValueSelector<?> selector = selectorFactory.makeColumnValueSelector(columnName);
+    final ColumnValueSelector<?> selector =
+        TypeCastSelectors.makeColumnValueSelector(selectorFactory, columnName, ColumnType.DOUBLE);
     return DoubleFieldWriter.forPrimitive(selector);
   }
 
@@ -139,7 +143,8 @@ private static FieldWriter makeStringArrayWriter(
       final boolean removeNullBytes
   )
   {
-    final ColumnValueSelector<?> selector = selectorFactory.makeColumnValueSelector(columnName);
+    final ColumnValueSelector<?> selector =
+        TypeCastSelectors.makeColumnValueSelector(selectorFactory, columnName, ColumnType.STRING_ARRAY);
     return new StringArrayFieldWriter(selector, removeNullBytes);
   }
 
@@ -148,7 +153,8 @@ private static FieldWriter makeLongArrayWriter(
       final String columnName
   )
   {
-    final ColumnValueSelector<?> selector = selectorFactory.makeColumnValueSelector(columnName);
+    final ColumnValueSelector<?> selector =
+        TypeCastSelectors.makeColumnValueSelector(selectorFactory, columnName, ColumnType.LONG_ARRAY);
     return NumericArrayFieldWriter.getLongArrayFieldWriter(selector);
   }
 
@@ -157,7 +163,8 @@ private static FieldWriter makeFloatArrayWriter(
       final String columnName
   )
   {
-    final ColumnValueSelector<?> selector = selectorFactory.makeColumnValueSelector(columnName);
+    final ColumnValueSelector<?> selector =
+        TypeCastSelectors.makeColumnValueSelector(selectorFactory, columnName, ColumnType.FLOAT_ARRAY);
     return NumericArrayFieldWriter.getFloatArrayFieldWriter(selector);
   }
 
@@ -166,7 +173,8 @@ private static FieldWriter makeDoubleArrayWriter(
       final String columnName
   )
   {
-    final ColumnValueSelector<?> selector = selectorFactory.makeColumnValueSelector(columnName);
+    final ColumnValueSelector<?> selector =
+        TypeCastSelectors.makeColumnValueSelector(selectorFactory, columnName, ColumnType.DOUBLE_ARRAY);
     return NumericArrayFieldWriter.getDoubleArrayFieldWriter(selector);
   }
 
@@ -185,7 +193,8 @@ private static FieldWriter makeComplexWriter(
       throw new ISE("No serde for complexTypeName[%s], cannot write column [%s]", columnTypeName, columnName);
     }
 
-    final ColumnValueSelector<?> selector = selectorFactory.makeColumnValueSelector(columnName);
+    final ColumnValueSelector<?> selector =
+        TypeCastSelectors.makeColumnValueSelector(selectorFactory, columnName, ColumnType.ofComplex(columnTypeName));
     return new ComplexFieldWriter(serde, selector);
   }
 }

processing/src/main/java/org/apache/druid/frame/field/NumericArrayFieldSelector.java

@@ -39,7 +39,7 @@
  *
  * @param <ElementType> Type of the individual array elements
  */
-public abstract class NumericArrayFieldSelector<ElementType extends Number> implements ColumnValueSelector
+public abstract class NumericArrayFieldSelector<ElementType extends Number> implements ColumnValueSelector<Object[]>
 {
   /**
    * Memory containing the serialized values of the array
@@ -81,15 +81,15 @@ public void inspectRuntimeShape(RuntimeShapeInspector inspector)
 
   @Nullable
   @Override
-  public Object getObject()
+  public Object[] getObject()
   {
     return computeCurrentArray();
   }
 
   @Override
-  public Class classOfObject()
+  public Class<Object[]> classOfObject()
   {
-    return Object.class;
+    return Object[].class;
   }
 
   @Override
@@ -131,7 +131,7 @@ public boolean isNull()
   public abstract int getIndividualFieldSize();
 
   @Nullable
-  private Number[] computeCurrentArray()
+  private Object[] computeCurrentArray()
   {
     final long fieldPosition = fieldPointer.position();
     final long fieldLength = fieldPointer.length();

processing/src/main/java/org/apache/druid/frame/field/NumericArrayFieldWriter.java

@@ -21,12 +21,10 @@
 
 import org.apache.datasketches.memory.WritableMemory;
 import org.apache.druid.common.config.NullHandling;
-import org.apache.druid.frame.write.FrameWriterUtils;
 import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector;
 import org.apache.druid.segment.ColumnValueSelector;
 
 import javax.annotation.Nullable;
-import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
@@ -128,7 +126,7 @@ public NumericArrayFieldWriter(final ColumnValueSelector selector, NumericFieldW
   @Override
   public long writeTo(WritableMemory memory, long position, long maxSize)
   {
-    Object row = selector.getObject();
+    final Object[] row = (Object[]) selector.getObject();
     if (row == null) {
       int requiredSize = Byte.BYTES;
       if (requiredSize > maxSize) {
@@ -137,18 +135,6 @@ public long writeTo(WritableMemory memory, long position, long maxSize)
       memory.putByte(position, NULL_ROW);
       return requiredSize;
     } else {
-
-      List<? extends Number> list = FrameWriterUtils.getNumericArrayFromObject(row);
-
-      if (list == null) {
-        int requiredSize = Byte.BYTES;
-        if (requiredSize > maxSize) {
-          return -1;
-        }
-        memory.putByte(position, NULL_ROW);
-        return requiredSize;
-      }
-
       // Create a columnValueSelector to write the individual elements re-using the NumericFieldWriter
       AtomicInteger index = new AtomicInteger(0);
       ColumnValueSelector<Number> columnValueSelector = new ColumnValueSelector<Number>()
@@ -199,7 +185,7 @@ public boolean isNull()
         @Override
         public Number getObject()
         {
-          return list.get(index.get());
+          return (Number) row[index.get()];
         }
 
         @Override
@@ -215,7 +201,7 @@ public Class<? extends Number> classOfObject()
       // Next [(1 + Numeric Size) x Number of elements of array] bytes are reserved for the elements of the array and
       //  their null markers
       // Last byte is reserved for array termination
-      int requiredSize = Byte.BYTES + (writer.getNumericSizeBytes() + Byte.BYTES) * list.size() + Byte.BYTES;
+      int requiredSize = Byte.BYTES + (writer.getNumericSizeBytes() + Byte.BYTES) * row.length + Byte.BYTES;
 
       if (requiredSize > maxSize) {
         return -1;
@@ -225,7 +211,7 @@ public Class<? extends Number> classOfObject()
       memory.putByte(position + offset, NON_NULL_ROW);
       offset += Byte.BYTES;
 
-      for (; index.get() < list.size(); index.incrementAndGet()) {
+      for (; index.get() < row.length; index.incrementAndGet()) {
         writer.writeTo(
             memory,
             position + offset,

processing/src/main/java/org/apache/druid/frame/write/FrameWriterUtils.java

@@ -144,60 +144,16 @@ public static List<ByteBuffer> getUtf8ByteBuffersFromStringArraySelector(
       @SuppressWarnings("rawtypes") final BaseObjectColumnValueSelector selector
   )
   {
-    Object row = selector.getObject();
+    final Object[] row = (Object[]) selector.getObject();
     if (row == null) {
       return null;
-    } else if (row instanceof String) {
-      return Collections.singletonList(getUtf8ByteBufferFromString((String) row));
-    }
-
-    final List<ByteBuffer> retVal = new ArrayList<>();
-    if (row instanceof List) {
-      for (int i = 0; i < ((List<?>) row).size(); i++) {
-        retVal.add(getUtf8ByteBufferFromString(((List<String>) row).get(i)));
-      }
-    } else if (row instanceof Object[]) {
-      for (Object value : (Object[]) row) {
-        retVal.add(getUtf8ByteBufferFromString((String) value));
-      }
     } else {
-      throw new ISE("Unexpected type %s found", row.getClass().getName());
-    }
-    return retVal;
-  }
-
-  /**
-   * Retrieves a numeric list from a Java object, given that the object is an instance of something that can be returned
-   * from {@link ColumnValueSelector#getObject()} of valid numeric array selectors representations
-   *
-   * While {@link BaseObjectColumnValueSelector} specifies that only instances of {@code Object[]} can be returned from
-   * the numeric array selectors, this method also handles a few more cases which can be encountered if the selector is
-   * directly implemented on top of the group by stuff
-   */
-  @Nullable
-  public static List<? extends Number> getNumericArrayFromObject(Object row)
-  {
-    if (row == null) {
-      return null;
-    } else if (row instanceof Number) {
-      return Collections.singletonList((Number) row);
-    }
-
-    final List<Number> retVal = new ArrayList<>();
-
-    if (row instanceof List) {
-      for (int i = 0; i < ((List<?>) row).size(); i++) {
-        retVal.add((Number) ((List<?>) row).get(i));
-      }
-    } else if (row instanceof Object[]) {
-      for (Object value : (Object[]) row) {
-        retVal.add((Number) value);
+      final List<ByteBuffer> retVal = new ArrayList<>();
+      for (Object value : row) {
+        retVal.add(getUtf8ByteBufferFromString((String) value));
       }
-    } else {
-      throw new ISE("Unexpected type %s found", row.getClass().getName());
+      return retVal;
     }
-
-    return retVal;
   }
 
   /**
@@ -275,6 +231,7 @@ public static void copyByteBufferToMemoryDisallowingNullBytes(
    * Whenever "allowNullBytes" is true, "removeNullBytes" must be false. Use the methods {@link #copyByteBufferToMemoryAllowingNullBytes}
    * and {@link #copyByteBufferToMemoryDisallowingNullBytes} to copy between the memory
    * <p>
+   *
    * @throws InvalidNullByteException if "allowNullBytes" and "removeNullBytes" is false and a null byte is encountered
    */
   private static void copyByteBufferToMemory(

processing/src/main/java/org/apache/druid/frame/write/cast/ObjectToArrayColumnValueSelector.java

@@ -0,0 +1,101 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.druid.frame.write.cast;
+
+import org.apache.druid.error.DruidException;
+import org.apache.druid.math.expr.ExprEval;
+import org.apache.druid.math.expr.ExpressionType;
+import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector;
+import org.apache.druid.segment.ColumnValueSelector;
+import org.apache.druid.segment.RowIdSupplier;
+
+import javax.annotation.Nullable;
+
+/**
+ * Wraps a {@link ColumnValueSelector}, calls {@link ColumnValueSelector#getObject()}, interprets that value using
+ * {@link ExprEval#ofType}, and casts it using {@link ExprEval#castTo}.
+ */
+public class ObjectToArrayColumnValueSelector implements ColumnValueSelector<Object[]>
+{
+  private final ColumnValueSelector<?> selector;
+  @Nullable
+  private final ExpressionType desiredType;
+  @Nullable
+  private final RowIdSupplier rowIdSupplier;
+
+  public ObjectToArrayColumnValueSelector(
+      final ColumnValueSelector<?> selector,
+      final ExpressionType desiredType,
+      @Nullable final RowIdSupplier rowIdSupplier
+  )
+  {
+    this.selector = selector;
+    this.desiredType = desiredType;
+    this.rowIdSupplier = rowIdSupplier;
+
+    if (!desiredType.isArray() || desiredType.getElementType() == null) {
+      throw DruidException.defensive("Expected array with nonnull element type, got[%s]", desiredType);
+    }
+  }
+
+  @Override
+  public double getDouble()
+  {
+    throw DruidException.defensive("Unexpected call to getDouble on array selector");
+  }
+
+  @Override
+  public float getFloat()
+  {
+    throw DruidException.defensive("Unexpected call to getFloat on array selector");
+  }
+
+  @Override
+  public long getLong()
+  {
+    throw DruidException.defensive("Unexpected call to getLong on array selector");
+  }
+
+  @Override
+  public boolean isNull()
+  {
+    throw DruidException.defensive("Unexpected call to isNull on array selector");
+  }
+
+  @Nullable
+  @Override
+  public Object[] getObject()
+  {
+    return (Object[]) TypeCastSelectors.bestEffortCoerce(selector.getObject(), desiredType);
+  }
+
+  @Override
+  public Class<Object[]> classOfObject()
+  {
+    return Object[].class;
+  }
+
+  @Override
+  public void inspectRuntimeShape(RuntimeShapeInspector inspector)
+  {
+    inspector.visit("selector", selector);
+    inspector.visit("rowIdSupplier", rowIdSupplier);
+  }
+}