[Feature] Support Coalesce for Column Names

pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java

@@ -62,6 +62,7 @@ public enum TransformFunctionType {
 
   IS_NULL("is_null"),
   IS_NOT_NULL("is_not_null"),
+  COALESCE("coalesce"),
 
   IS_DISTINCT_FROM("is_distinct_from"),
   IS_NOT_DISTINCT_FROM("is_not_distinct_from"),

pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/CoalesceTransformFunction.java

@@ -0,0 +1,369 @@
+/**
+ * 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.pinot.core.operator.transform.function;
+
+import com.google.common.base.Preconditions;
+import java.math.BigDecimal;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.function.TransformFunctionType;
+import org.apache.pinot.core.operator.blocks.ProjectionBlock;
+import org.apache.pinot.core.operator.transform.TransformResultMetadata;
+import org.apache.pinot.segment.spi.datasource.DataSource;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.utils.NullValueUtils;
+import org.roaringbitmap.RoaringBitmap;
+
+
+/**
+ * The <code>CoalesceTransformFunction</code> implements the Coalesce operator.
+ *
+ * The results are in String format for first non-null value in the argument list.
+ * If all arguments are null, return a 'null' string.
+ * Note: arguments have to be column names and single type. The type can be either numeric or string.
+ * Number of arguments has to be greater than 0.
+ *
+ * Expected result:
+ * Coalesce(nullColumn, columnA): columnA
+ * Coalesce(columnA, nullColumn): nullColumn
+ * Coalesce(nullColumnA, nullColumnB): "null"
+ *
+ * Note this operator only takes column names for now.
+ * SQL Syntax:
+ *    Coalesce(columnA, columnB)
+ */
+public class CoalesceTransformFunction extends BaseTransformFunction {
+  private TransformFunction[] _transformFunctions;
+  private FieldSpec.DataType _storedType;
+  private TransformResultMetadata _resultMetadata;
+
+  /**
+   * Returns a bit map of corresponding column.
+   * Returns an empty bitmap by default if null option is disabled.
+   */
+  private static RoaringBitmap[] getNullBitMaps(ProjectionBlock projectionBlock,
+      TransformFunction[] transformFunctions) {
+    RoaringBitmap[] roaringBitmaps = new RoaringBitmap[transformFunctions.length];
+    for (int i = 0; i < roaringBitmaps.length; i++) {
+      TransformFunction func = transformFunctions[i];
+      String columnName = ((IdentifierTransformFunction) func).getColumnName();
+      RoaringBitmap nullBitmap = projectionBlock.getBlockValueSet(columnName).getNullBitmap();
+      roaringBitmaps[i] = nullBitmap;
+    }
+    return roaringBitmaps;
+  }
+
+  /**
+   * Get transform int results based on store type.
+   * @param projectionBlock
+   */
+  private int[] getIntTransformResults(ProjectionBlock projectionBlock) {
+    int length = projectionBlock.getNumDocs();
+    int[] results = new int[length];
+    int width = _transformFunctions.length;
+    RoaringBitmap[] nullBitMaps = getNullBitMaps(projectionBlock, _transformFunctions);
+    int[][] data = new int[width][length];
+    RoaringBitmap filledData = new RoaringBitmap();
+    for (int i = 0; i < length; i++) {
+      boolean hasNonNullValue = false;
+      for (int j = 0; j < width; j++) {
+        // Consider value as null only when null option is enabled.
+        if (nullBitMaps[j] != null && nullBitMaps[j].contains(i)) {
+          continue;
+        }
+        if (!filledData.contains(j)) {
+          filledData.add(j);
+          data[j] = _transformFunctions[j].transformToIntValuesSV(projectionBlock);
+        }
+        hasNonNullValue = true;
+        results[i] = data[j][i];
+        break;
+      }
+      if (!hasNonNullValue) {
+        results[i] = (int) NullValueUtils.getDefaultNullValue(_storedType);
+      }
+    }
+    return results;
+  }
+
+  /**
+   * Get transform long results based on store type.
+   * @param projectionBlock
+   */
+  private long[] getLongTransformResults(ProjectionBlock projectionBlock) {
+    int length = projectionBlock.getNumDocs();
+    long[] results = new long[length];
+    int width = _transformFunctions.length;
+    RoaringBitmap[] nullBitMaps = getNullBitMaps(projectionBlock, _transformFunctions);
+    long[][] data = new long[width][length];
+    RoaringBitmap filledData = new RoaringBitmap(); // indicates whether certain column has be filled in data.
+    for (int i = 0; i < length; i++) {
+      boolean hasNonNullValue = false;
+      for (int j = 0; j < width; j++) {
+        // Consider value as null only when null option is enabled.
+        if (nullBitMaps[j] != null && nullBitMaps[j].contains(i)) {
+          continue;
+        }
+        if (!filledData.contains(j)) {
+          filledData.add(j);
+          data[j] = _transformFunctions[j].transformToLongValuesSV(projectionBlock);
+        }
+        hasNonNullValue = true;
+        results[i] = data[j][i];
+        break;
+      }
+      if (!hasNonNullValue) {
+        results[i] = (long) NullValueUtils.getDefaultNullValue(_storedType);
+      }
+    }
+    return results;
+  }
+
+  /**
+   * Get transform float results based on store type.
+   * @param projectionBlock
+   */
+  private float[] getFloatTransformResults(ProjectionBlock projectionBlock) {
+    int length = projectionBlock.getNumDocs();
+    float[] results = new float[length];
+    int width = _transformFunctions.length;
+    RoaringBitmap[] nullBitMaps = getNullBitMaps(projectionBlock, _transformFunctions);
+    float[][] data = new float[width][length];
+    RoaringBitmap filledData = new RoaringBitmap(); // indicates whether certain column has be filled in data.
+    for (int i = 0; i < length; i++) {
+      boolean hasNonNullValue = false;
+      for (int j = 0; j < width; j++) {
+        // Consider value as null only when null option is enabled.
+        if (nullBitMaps[j] != null && nullBitMaps[j].contains(i)) {
+          continue;
+        }
+        if (!filledData.contains(j)) {
+          filledData.add(j);
+          data[j] = _transformFunctions[j].transformToFloatValuesSV(projectionBlock);
+        }
+        hasNonNullValue = true;
+        results[i] = data[j][i];
+        break;
+      }
+      if (!hasNonNullValue) {
+        results[i] = (float) NullValueUtils.getDefaultNullValue(_storedType);
+      }
+    }
+    return results;
+  }
+
+  /**
+   * Get transform double results based on store type.
+   * @param projectionBlock
+   */
+  private double[] getDoublelTransformResults(ProjectionBlock projectionBlock) {
+    int length = projectionBlock.getNumDocs();
+    double[] results = new double[length];
+    int width = _transformFunctions.length;
+    RoaringBitmap[] nullBitMaps = getNullBitMaps(projectionBlock, _transformFunctions);
+    double[][] data = new double[width][length];
+    RoaringBitmap filledData = new RoaringBitmap(); // indicates whether certain column has be filled in data.
+    for (int i = 0; i < length; i++) {
+      boolean hasNonNullValue = false;
+      for (int j = 0; j < width; j++) {
+        // Consider value as null only when null option is enabled.
+        if (nullBitMaps[j] != null && nullBitMaps[j].contains(i)) {
+          continue;
+        }
+        if (!filledData.contains(j)) {
+          filledData.add(j);
+          data[j] = _transformFunctions[j].transformToDoubleValuesSV(projectionBlock);
+        }
+        hasNonNullValue = true;
+        results[i] = data[j][i];
+        break;
+      }
+      if (!hasNonNullValue) {
+        results[i] = (double) NullValueUtils.getDefaultNullValue(_storedType);
+      }
+    }
+    return results;
+  }
+
+  /**
+   * Get transform BigDecimal results based on store type.
+   * @param projectionBlock
+   */
+  private BigDecimal[] getBigDecimalTransformResults(ProjectionBlock projectionBlock) {
+    int length = projectionBlock.getNumDocs();
+    BigDecimal[] results = new BigDecimal[length];
+    int width = _transformFunctions.length;
+    RoaringBitmap[] nullBitMaps = getNullBitMaps(projectionBlock, _transformFunctions);
+    BigDecimal[][] data = new BigDecimal[width][length];
+    RoaringBitmap filledData = new RoaringBitmap(); // indicates whether certain column has be filled in data.
+    for (int i = 0; i < length; i++) {
+      boolean hasNonNullValue = false;
+      for (int j = 0; j < width; j++) {
+        // Consider value as null only when null option is enabled.
+        if (nullBitMaps[j] != null && nullBitMaps[j].contains(i)) {
+          continue;
+        }
+        if (!filledData.contains(j)) {
+          filledData.add(j);
+          data[j] = _transformFunctions[j].transformToBigDecimalValuesSV(projectionBlock);
+        }
+        hasNonNullValue = true;
+        results[i] = data[j][i];
+        break;
+      }
+      if (!hasNonNullValue) {
+        results[i] = (BigDecimal) NullValueUtils.getDefaultNullValue(_storedType);
+      }
+    }
+    return results;
+  }
+
+
+  /**
+   * Get transform String results based on store type.
+   * @param projectionBlock
+   */
+  private String[] getStringTransformResults(ProjectionBlock projectionBlock) {
+    int length = projectionBlock.getNumDocs();
+    String[] results = new String[length];
+    int width = _transformFunctions.length;
+    RoaringBitmap[] nullBitMaps = getNullBitMaps(projectionBlock, _transformFunctions);
+    String[][] data = new String[width][length];
+    RoaringBitmap filledData = new RoaringBitmap(); // indicates whether certain column has be filled in data.
+    for (int i = 0; i < length; i++) {
+      boolean hasNonNullValue = false;
+      for (int j = 0; j < width; j++) {
+        // Consider value as null only when null option is enabled.
+        if (nullBitMaps[j] != null && nullBitMaps[j].contains(i)) {
+          continue;
+        }
+        if (!filledData.contains(j)) {
+          filledData.add(j);
+          data[j] = _transformFunctions[j].transformToStringValuesSV(projectionBlock);
+        }
+        hasNonNullValue = true;
+        results[i] = data[j][i];
+        break;
+      }
+      if (!hasNonNullValue) {
+        results[i] = (String) NullValueUtils.getDefaultNullValue(_storedType);
+      }
+    }
+    return results;
+  }
+
+  @Override
+  public String getName() {
+    return TransformFunctionType.COALESCE.getName();
+  }
+
+  @Override
+  public void init(List<TransformFunction> arguments, Map<String, DataSource> dataSourceMap) {
+    int argSize = arguments.size();
+    Preconditions.checkArgument(argSize > 0, "COALESCE needs to have at least one argument.");
+    _transformFunctions = new TransformFunction[argSize];
+    for (int i = 0; i < argSize; i++) {
+      TransformFunction func = arguments.get(i);
+      Preconditions.checkArgument(func instanceof IdentifierTransformFunction,
+          "Only column names are supported in COALESCE.");
+      FieldSpec.DataType storedType = func.getResultMetadata().getDataType().getStoredType();
+      if (_storedType == null) {
+        _storedType = storedType;
+      } else {
+        Preconditions.checkArgument(storedType.equals(_storedType), "Argument types have to be the same.");
+      }
+      _transformFunctions[i] = func;
+    }
+    switch (_storedType) {
+      case INT:
+        _resultMetadata = INT_SV_NO_DICTIONARY_METADATA;
+        break;
+      case LONG:
+        _resultMetadata = LONG_SV_NO_DICTIONARY_METADATA;
+        break;
+      case FLOAT:
+        _resultMetadata = FLOAT_SV_NO_DICTIONARY_METADATA;
+        break;
+      case DOUBLE:
+        _resultMetadata = DOUBLE_SV_NO_DICTIONARY_METADATA;
+        break;
+      case BIG_DECIMAL:
+        _resultMetadata = BIG_DECIMAL_SV_NO_DICTIONARY_METADATA;
+        break;
+      case STRING:
+        _resultMetadata = STRING_SV_NO_DICTIONARY_METADATA;
+        break;
+      default:
+        throw new UnsupportedOperationException("Coalesce only supports numerical and string data type");
+    }
+  }
+
+  @Override
+  public TransformResultMetadata getResultMetadata() {
+    return _resultMetadata;
+  }
+
+  @Override
+  public String[] transformToStringValuesSV(ProjectionBlock projectionBlock) {
+    if (_storedType != FieldSpec.DataType.STRING) {
+      return super.transformToStringValuesSV(projectionBlock);
+    }
+    return getStringTransformResults(projectionBlock);
+  }
+
+  @Override
+  public int[] transformToIntValuesSV(ProjectionBlock projectionBlock) {
+    if (_storedType != FieldSpec.DataType.INT) {
+      return super.transformToIntValuesSV(projectionBlock);
+    }
+    return getIntTransformResults(projectionBlock);
+  }
+
+  @Override
+  public long[] transformToLongValuesSV(ProjectionBlock projectionBlock) {
+    if (_storedType != FieldSpec.DataType.LONG) {
+      return super.transformToLongValuesSV(projectionBlock);
+    }
+    return getLongTransformResults(projectionBlock);
+  }
+
+  @Override
+  public BigDecimal[] transformToBigDecimalValuesSV(ProjectionBlock projectionBlock) {
+    if (_storedType != FieldSpec.DataType.BIG_DECIMAL) {
+      return super.transformToBigDecimalValuesSV(projectionBlock);
+    }
+    return getBigDecimalTransformResults(projectionBlock);
+  }
+
+  @Override
+  public double[] transformToDoubleValuesSV(ProjectionBlock projectionBlock) {
+    if (_storedType != FieldSpec.DataType.DOUBLE) {
+      return super.transformToDoubleValuesSV(projectionBlock);
+    }
+    return getDoublelTransformResults(projectionBlock);
+  }
+
+  @Override
+  public float[] transformToFloatValuesSV(ProjectionBlock projectionBlock) {
+    if (_storedType != FieldSpec.DataType.FLOAT) {
+      return super.transformToFloatValuesSV(projectionBlock);
+    }
+    return getFloatTransformResults(projectionBlock);
+  }
+}

pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java

@@ -202,6 +202,7 @@ private static Map<String, Class<? extends TransformFunction>> createRegistry()
     typeToImplementation.put(TransformFunctionType.IS_NULL, IsNullTransformFunction.class);
     typeToImplementation.put(TransformFunctionType.IS_NOT_NULL,
         IsNotNullTransformFunction.class);
+    typeToImplementation.put(TransformFunctionType.COALESCE, CoalesceTransformFunction.class);
     typeToImplementation.put(TransformFunctionType.IS_DISTINCT_FROM, IsDistinctFromTransformFunction.class);
     typeToImplementation.put(TransformFunctionType.IS_NOT_DISTINCT_FROM, IsNotDistinctFromTransformFunction.class);