ceph/src/arrow/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/ExpressionRegistry.java

   1 /*
   2  * Licensed to the Apache Software Foundation (ASF) under one or more
   3  * contributor license agreements.  See the NOTICE file distributed with
   4  * this work for additional information regarding copyright ownership.
   5  * The ASF licenses this file to You under the Apache License, Version 2.0
   6  * (the "License"); you may not use this file except in compliance with
   7  * the License.  You may obtain a copy of the License at
   8  *
   9  *    http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 package org.apache.arrow.gandiva.evaluator;
  19
  20 import java.util.List;
  21 import java.util.Set;
  22
  23 import org.apache.arrow.gandiva.exceptions.GandivaException;
  24 import org.apache.arrow.gandiva.ipc.GandivaTypes;
  25 import org.apache.arrow.gandiva.ipc.GandivaTypes.ExtGandivaType;
  26 import org.apache.arrow.gandiva.ipc.GandivaTypes.GandivaDataTypes;
  27 import org.apache.arrow.gandiva.ipc.GandivaTypes.GandivaFunctions;
  28 import org.apache.arrow.gandiva.ipc.GandivaTypes.GandivaType;
  29 import org.apache.arrow.vector.types.DateUnit;
  30 import org.apache.arrow.vector.types.FloatingPointPrecision;
  31 import org.apache.arrow.vector.types.IntervalUnit;
  32 import org.apache.arrow.vector.types.TimeUnit;
  33 import org.apache.arrow.vector.types.pojo.ArrowType;
  34
  35 import com.google.common.collect.Lists;
  36 import com.google.common.collect.Sets;
  37 import com.google.protobuf.InvalidProtocolBufferException;
  38
  39 /**
  40  * Used to get the functions and data types supported by
  41  * Gandiva.
  42  * All types are in Arrow namespace.
  43  */
  44 public class ExpressionRegistry {
  45
  46   private static final int BIT_WIDTH8 = 8;
  47   private static final int BIT_WIDTH_16 = 16;
  48   private static final int BIT_WIDTH_32 = 32;
  49   private static final int BIT_WIDTH_64 = 64;
  50   private static final boolean IS_SIGNED_FALSE = false;
  51   private static final boolean IS_SIGNED_TRUE = true;
  52
  53   private final Set<ArrowType> supportedTypes;
  54   private final Set<FunctionSignature> functionSignatures;
  55
  56   private static volatile ExpressionRegistry INSTANCE;
  57
  58   private ExpressionRegistry(Set<ArrowType> supportedTypes,
  59                              Set<FunctionSignature> functionSignatures) {
  60     this.supportedTypes = supportedTypes;
  61     this.functionSignatures = functionSignatures;
  62   }
  63
  64   /**
  65    * Returns a singleton instance of the class.
  66    * @return singleton instance
  67    * @throws GandivaException if error in Gandiva Library integration.
  68    */
  69   public static ExpressionRegistry getInstance() throws GandivaException {
  70     if (INSTANCE == null) {
  71       synchronized (ExpressionRegistry.class) {
  72         if (INSTANCE == null) {
  73           // ensure library is setup.
  74           JniLoader.getInstance();
  75           Set<ArrowType> typesFromGandiva = getSupportedTypesFromGandiva();
  76           Set<FunctionSignature> functionsFromGandiva = getSupportedFunctionsFromGandiva();
  77           INSTANCE = new ExpressionRegistry(typesFromGandiva, functionsFromGandiva);
  78         }
  79       }
  80     }
  81     return INSTANCE;
  82   }
  83
  84   public Set<FunctionSignature> getSupportedFunctions() {
  85     return functionSignatures;
  86   }
  87
  88   public Set<ArrowType> getSupportedTypes() {
  89     return supportedTypes;
  90   }
  91
  92   private static Set<ArrowType> getSupportedTypesFromGandiva() throws GandivaException {
  93     Set<ArrowType> supportedTypes = Sets.newHashSet();
  94     try {
  95       byte[] gandivaSupportedDataTypes = new ExpressionRegistryJniHelper()
  96               .getGandivaSupportedDataTypes();
  97       GandivaDataTypes gandivaDataTypes = GandivaDataTypes.parseFrom(gandivaSupportedDataTypes);
  98       for (ExtGandivaType type : gandivaDataTypes.getDataTypeList()) {
  99         supportedTypes.add(getArrowType(type));
 100       }
 101     } catch (InvalidProtocolBufferException invalidProtException) {
 102       throw new GandivaException("Could not get supported types.", invalidProtException);
 103     }
 104     return supportedTypes;
 105   }
 106
 107   private static Set<FunctionSignature> getSupportedFunctionsFromGandiva() throws
 108           GandivaException {
 109     Set<FunctionSignature> supportedTypes = Sets.newHashSet();
 110     try {
 111       byte[] gandivaSupportedFunctions = new ExpressionRegistryJniHelper()
 112               .getGandivaSupportedFunctions();
 113       GandivaFunctions gandivaFunctions = GandivaFunctions.parseFrom(gandivaSupportedFunctions);
 114       for (GandivaTypes.FunctionSignature protoFunctionSignature
 115               : gandivaFunctions.getFunctionList()) {
 116
 117         String functionName = protoFunctionSignature.getName();
 118         ArrowType returnType = getArrowType(protoFunctionSignature.getReturnType());
 119         List<ArrowType> paramTypes = Lists.newArrayList();
 120         for (ExtGandivaType type : protoFunctionSignature.getParamTypesList()) {
 121           paramTypes.add(getArrowType(type));
 122         }
 123         FunctionSignature functionSignature = new FunctionSignature(functionName,
 124                                                                     returnType, paramTypes);
 125         supportedTypes.add(functionSignature);
 126       }
 127     } catch (InvalidProtocolBufferException invalidProtException) {
 128       throw new GandivaException("Could not get supported functions.", invalidProtException);
 129     }
 130     return supportedTypes;
 131   }
 132
 133   private static ArrowType getArrowType(ExtGandivaType type) {
 134     switch (type.getType().getNumber()) {
 135       case GandivaType.BOOL_VALUE:
 136         return ArrowType.Bool.INSTANCE;
 137       case GandivaType.UINT8_VALUE:
 138         return new ArrowType.Int(BIT_WIDTH8, IS_SIGNED_FALSE);
 139       case GandivaType.INT8_VALUE:
 140         return new ArrowType.Int(BIT_WIDTH8, IS_SIGNED_TRUE);
 141       case GandivaType.UINT16_VALUE:
 142         return new ArrowType.Int(BIT_WIDTH_16, IS_SIGNED_FALSE);
 143       case GandivaType.INT16_VALUE:
 144         return new ArrowType.Int(BIT_WIDTH_16, IS_SIGNED_TRUE);
 145       case GandivaType.UINT32_VALUE:
 146         return new ArrowType.Int(BIT_WIDTH_32, IS_SIGNED_FALSE);
 147       case GandivaType.INT32_VALUE:
 148         return new ArrowType.Int(BIT_WIDTH_32, IS_SIGNED_TRUE);
 149       case GandivaType.UINT64_VALUE:
 150         return new ArrowType.Int(BIT_WIDTH_64, IS_SIGNED_FALSE);
 151       case GandivaType.INT64_VALUE:
 152         return new ArrowType.Int(BIT_WIDTH_64, IS_SIGNED_TRUE);
 153       case GandivaType.HALF_FLOAT_VALUE:
 154         return new ArrowType.FloatingPoint(FloatingPointPrecision.HALF);
 155       case GandivaType.FLOAT_VALUE:
 156         return new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE);
 157       case GandivaType.DOUBLE_VALUE:
 158         return new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE);
 159       case GandivaType.UTF8_VALUE:
 160         return new ArrowType.Utf8();
 161       case GandivaType.BINARY_VALUE:
 162         return new ArrowType.Binary();
 163       case GandivaType.DATE32_VALUE:
 164         return new ArrowType.Date(DateUnit.DAY);
 165       case GandivaType.DATE64_VALUE:
 166         return new ArrowType.Date(DateUnit.MILLISECOND);
 167       case GandivaType.TIMESTAMP_VALUE:
 168         return new ArrowType.Timestamp(mapArrowTimeUnit(type.getTimeUnit()), null);
 169       case GandivaType.TIME32_VALUE:
 170         return new ArrowType.Time(mapArrowTimeUnit(type.getTimeUnit()),
 171                 BIT_WIDTH_32);
 172       case GandivaType.TIME64_VALUE:
 173         return new ArrowType.Time(mapArrowTimeUnit(type.getTimeUnit()),
 174                 BIT_WIDTH_64);
 175       case GandivaType.NONE_VALUE:
 176         return new ArrowType.Null();
 177       case GandivaType.DECIMAL_VALUE:
 178         return new ArrowType.Decimal(0, 0, 128);
 179       case GandivaType.INTERVAL_VALUE:
 180         return new ArrowType.Interval(mapArrowIntervalUnit(type.getIntervalType()));
 181       case GandivaType.FIXED_SIZE_BINARY_VALUE:
 182       case GandivaType.MAP_VALUE:
 183       case GandivaType.DICTIONARY_VALUE:
 184       case GandivaType.LIST_VALUE:
 185       case GandivaType.STRUCT_VALUE:
 186       case GandivaType.UNION_VALUE:
 187       default:
 188         assert false;
 189     }
 190     return null;
 191   }
 192
 193   private static TimeUnit mapArrowTimeUnit(GandivaTypes.TimeUnit timeUnit) {
 194     switch (timeUnit.getNumber()) {
 195       case GandivaTypes.TimeUnit.MICROSEC_VALUE:
 196         return TimeUnit.MICROSECOND;
 197       case GandivaTypes.TimeUnit.MILLISEC_VALUE:
 198         return TimeUnit.MILLISECOND;
 199       case GandivaTypes.TimeUnit.NANOSEC_VALUE:
 200         return TimeUnit.NANOSECOND;
 201       case GandivaTypes.TimeUnit.SEC_VALUE:
 202         return TimeUnit.SECOND;
 203       default:
 204         return null;
 205     }
 206   }
 207
 208   private static IntervalUnit mapArrowIntervalUnit(GandivaTypes.IntervalType intervalType) {
 209     switch (intervalType.getNumber()) {
 210       case GandivaTypes.IntervalType.YEAR_MONTH_VALUE:
 211         return IntervalUnit.YEAR_MONTH;
 212       case GandivaTypes.IntervalType.DAY_TIME_VALUE:
 213         return IntervalUnit.DAY_TIME;
 214       default:
 215         return null;
 216     }
 217   }
 218
 219 }
 220