ceph/src/arrow/java/gandiva/src/main/java/org/apache/arrow/gandiva/expression/InNode.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.expression;
  19
  20 import java.math.BigDecimal;
  21 import java.nio.charset.Charset;
  22 import java.util.Set;
  23
  24 import org.apache.arrow.gandiva.exceptions.GandivaException;
  25 import org.apache.arrow.gandiva.ipc.GandivaTypes;
  26
  27 import com.google.protobuf.ByteString;
  28
  29 /**
  30  * In Node representation in java.
  31  */
  32 public class InNode implements TreeNode {
  33   private static final Charset charset = Charset.forName("UTF-8");
  34
  35   private final Set<Integer> intValues;
  36   private final Set<Long> longValues;
  37   private final Set<Float> floatValues;
  38   private final Set<Double> doubleValues;
  39   private final Set<BigDecimal> decimalValues;
  40   private final Set<String> stringValues;
  41   private final Set<byte[]> binaryValues;
  42   private final TreeNode input;
  43
  44   private final Integer precision;
  45   private final Integer scale;
  46
  47   private InNode(Set<Integer> values, Set<Long> longValues, Set<String> stringValues, Set<byte[]>
  48           binaryValues, Set<BigDecimal> decimalValues, Integer precision, Integer scale,
  49                  Set<Float> floatValues, Set<Double> doubleValues, TreeNode node) {
  50     this.intValues = values;
  51     this.longValues = longValues;
  52     this.decimalValues = decimalValues;
  53     this.precision = precision;
  54     this.scale = scale;
  55     this.stringValues = stringValues;
  56     this.binaryValues = binaryValues;
  57     this.floatValues = floatValues;
  58     this.doubleValues = doubleValues;
  59     this.input = node;
  60   }
  61
  62   /**
  63    * Makes an IN node for int values.
  64    *
  65    * @param node      Node with the 'IN' clause.
  66    * @param intValues Int values to build the IN node.
  67    * @retur InNode referring to tree node.
  68    */
  69   public static InNode makeIntInExpr(TreeNode node, Set<Integer> intValues) {
  70     return new InNode(intValues,
  71             null, null, null, null, null, null, null,
  72             null, node);
  73   }
  74
  75   /**
  76    * Makes an IN node for long values.
  77    *
  78    * @param node      Node with the 'IN' clause.
  79    * @param longValues Long values to build the IN node.
  80    * @retur InNode referring to tree node.
  81    */
  82   public static InNode makeLongInExpr(TreeNode node, Set<Long> longValues) {
  83     return new InNode(null, longValues,
  84             null, null, null, null, null, null,
  85             null, node);
  86   }
  87
  88   /**
  89    * Makes an IN node for float values.
  90    *
  91    * @param node      Node with the 'IN' clause.
  92    * @param floatValues Float values to build the IN node.
  93    * @retur InNode referring to tree node.
  94    */
  95   public static InNode makeFloatInExpr(TreeNode node, Set<Float> floatValues) {
  96     return new InNode(null, null, null, null, null, null,
  97             null, floatValues, null, node);
  98   }
  99
 100   /**
 101    * Makes an IN node for double values.
 102    *
 103    * @param node      Node with the 'IN' clause.
 104    * @param doubleValues Double values to build the IN node.
 105    * @retur InNode referring to tree node.
 106    */
 107   public static InNode makeDoubleInExpr(TreeNode node, Set<Double> doubleValues) {
 108     return new InNode(null, null, null, null, null,
 109             null, null, null, doubleValues, node);
 110   }
 111
 112   public static InNode makeDecimalInExpr(TreeNode node, Set<BigDecimal> decimalValues,
 113                                          Integer precision, Integer scale) {
 114     return new InNode(null, null, null, null,
 115             decimalValues, precision, scale, null, null, node);
 116   }
 117
 118   public static InNode makeStringInExpr(TreeNode node, Set<String> stringValues) {
 119     return new InNode(null, null, stringValues, null,
 120             null, null, null, null, null, node);
 121   }
 122
 123   public static InNode makeBinaryInExpr(TreeNode node, Set<byte[]> binaryValues) {
 124     return new InNode(null, null, null, binaryValues,
 125             null, null, null, null, null, node);
 126   }
 127
 128   @Override
 129   public GandivaTypes.TreeNode toProtobuf() throws GandivaException {
 130     GandivaTypes.InNode.Builder inNode = GandivaTypes.InNode.newBuilder();
 131
 132     inNode.setNode(input.toProtobuf());
 133
 134     if (intValues != null) {
 135       GandivaTypes.IntConstants.Builder intConstants = GandivaTypes.IntConstants.newBuilder();
 136       intValues.stream().forEach(val -> intConstants.addIntValues(GandivaTypes.IntNode.newBuilder()
 137               .setValue(val).build()));
 138       inNode.setIntValues(intConstants.build());
 139     } else if (longValues != null) {
 140       GandivaTypes.LongConstants.Builder longConstants = GandivaTypes.LongConstants.newBuilder();
 141       longValues.stream().forEach(val -> longConstants.addLongValues(GandivaTypes.LongNode.newBuilder()
 142               .setValue(val).build()));
 143       inNode.setLongValues(longConstants.build());
 144     } else if (floatValues != null) {
 145       GandivaTypes.FloatConstants.Builder floatConstants = GandivaTypes.FloatConstants.newBuilder();
 146       floatValues.stream().forEach(val -> floatConstants.addFloatValues(GandivaTypes.FloatNode.newBuilder()
 147               .setValue(val).build()));
 148       inNode.setFloatValues(floatConstants.build());
 149     } else if (doubleValues != null) {
 150       GandivaTypes.DoubleConstants.Builder doubleConstants = GandivaTypes.DoubleConstants.newBuilder();
 151       doubleValues.stream().forEach(val -> doubleConstants.addDoubleValues(GandivaTypes.DoubleNode.newBuilder()
 152               .setValue(val).build()));
 153       inNode.setDoubleValues(doubleConstants.build());
 154     } else if (decimalValues != null) {
 155       GandivaTypes.DecimalConstants.Builder decimalConstants = GandivaTypes.DecimalConstants.newBuilder();
 156       decimalValues.stream().forEach(val -> decimalConstants.addDecimalValues(GandivaTypes.DecimalNode.newBuilder()
 157               .setValue(val.toPlainString()).setPrecision(precision).setScale(scale).build()));
 158       inNode.setDecimalValues(decimalConstants.build());
 159     } else if (stringValues != null) {
 160       GandivaTypes.StringConstants.Builder stringConstants = GandivaTypes.StringConstants
 161               .newBuilder();
 162       stringValues.stream().forEach(val -> stringConstants.addStringValues(GandivaTypes.StringNode
 163               .newBuilder().setValue(ByteString.copyFrom(val.getBytes(charset))).build()));
 164       inNode.setStringValues(stringConstants.build());
 165     } else if (binaryValues != null) {
 166       GandivaTypes.BinaryConstants.Builder binaryConstants = GandivaTypes.BinaryConstants
 167               .newBuilder();
 168       binaryValues.stream().forEach(val -> binaryConstants.addBinaryValues(GandivaTypes.BinaryNode
 169               .newBuilder().setValue(ByteString.copyFrom(val)).build()));
 170       inNode.setBinaryValues(binaryConstants.build());
 171     }
 172     GandivaTypes.TreeNode.Builder builder = GandivaTypes.TreeNode.newBuilder();
 173     builder.setInNode(inNode.build());
 174     return builder.build();
 175   }
 176 }