execution/vectorized/MutableColumnarRow.java

0001 /*
0002  * Licensed to the Apache Software Foundation (ASF) under one or more
0003  * contributor license agreements.  See the NOTICE file distributed with
0004  * this work for additional information regarding copyright ownership.
0005  * The ASF licenses this file to You under the Apache License, Version 2.0
0006  * (the "License"); you may not use this file except in compliance with
0007  * the License.  You may obtain a copy of the License at
0008  *
0009  *    http://www.apache.org/licenses/LICENSE-2.0
0010  *
0011  * Unless required by applicable law or agreed to in writing, software
0012  * distributed under the License is distributed on an "AS IS" BASIS,
0013  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
0014  * See the License for the specific language governing permissions and
0015  * limitations under the License.
0016  */
0017
0018 package org.apache.spark.sql.execution.vectorized;
0019
0020 import java.math.BigDecimal;
0021
0022 import org.apache.spark.sql.catalyst.InternalRow;
0023 import org.apache.spark.sql.catalyst.expressions.GenericInternalRow;
0024 import org.apache.spark.sql.types.*;
0025 import org.apache.spark.sql.vectorized.ColumnarArray;
0026 import org.apache.spark.sql.vectorized.ColumnarBatch;
0027 import org.apache.spark.sql.vectorized.ColumnarMap;
0028 import org.apache.spark.sql.vectorized.ColumnarRow;
0029 import org.apache.spark.unsafe.types.CalendarInterval;
0030 import org.apache.spark.unsafe.types.UTF8String;
0031
0032 /**
0033  * A mutable version of {@link ColumnarRow}, which is used in the vectorized hash map for hash
0034  * aggregate, and {@link ColumnarBatch} to save object creation.
0035  *
0036  * Note that this class intentionally has a lot of duplicated code with {@link ColumnarRow}, to
0037  * avoid java polymorphism overhead by keeping {@link ColumnarRow} and this class final classes.
0038  */
0039 public final class MutableColumnarRow extends InternalRow {
0040   public int rowId;
0041   private final WritableColumnVector[] columns;
0042
0043   public MutableColumnarRow(WritableColumnVector[] writableColumns) {
0044     this.columns = writableColumns;
0045   }
0046
0047   @Override
0048   public int numFields() { return columns.length; }
0049
0050   @Override
0051   public InternalRow copy() {
0052     GenericInternalRow row = new GenericInternalRow(columns.length);
0053     for (int i = 0; i < numFields(); i++) {
0054       if (isNullAt(i)) {
0055         row.setNullAt(i);
0056       } else {
0057         DataType dt = columns[i].dataType();
0058         if (dt instanceof BooleanType) {
0059           row.setBoolean(i, getBoolean(i));
0060         } else if (dt instanceof ByteType) {
0061           row.setByte(i, getByte(i));
0062         } else if (dt instanceof ShortType) {
0063           row.setShort(i, getShort(i));
0064         } else if (dt instanceof IntegerType) {
0065           row.setInt(i, getInt(i));
0066         } else if (dt instanceof LongType) {
0067           row.setLong(i, getLong(i));
0068         } else if (dt instanceof FloatType) {
0069           row.setFloat(i, getFloat(i));
0070         } else if (dt instanceof DoubleType) {
0071           row.setDouble(i, getDouble(i));
0072         } else if (dt instanceof StringType) {
0073           row.update(i, getUTF8String(i).copy());
0074         } else if (dt instanceof BinaryType) {
0075           row.update(i, getBinary(i));
0076         } else if (dt instanceof DecimalType) {
0077           DecimalType t = (DecimalType)dt;
0078           row.setDecimal(i, getDecimal(i, t.precision(), t.scale()), t.precision());
0079         } else if (dt instanceof DateType) {
0080           row.setInt(i, getInt(i));
0081         } else if (dt instanceof TimestampType) {
0082           row.setLong(i, getLong(i));
0083         } else {
0084           throw new RuntimeException("Not implemented. " + dt);
0085         }
0086       }
0087     }
0088     return row;
0089   }
0090
0091   @Override
0092   public boolean anyNull() {
0093     throw new UnsupportedOperationException();
0094   }
0095
0096   @Override
0097   public boolean isNullAt(int ordinal) { return columns[ordinal].isNullAt(rowId); }
0098
0099   @Override
0100   public boolean getBoolean(int ordinal) { return columns[ordinal].getBoolean(rowId); }
0101
0102   @Override
0103   public byte getByte(int ordinal) { return columns[ordinal].getByte(rowId); }
0104
0105   @Override
0106   public short getShort(int ordinal) { return columns[ordinal].getShort(rowId); }
0107
0108   @Override
0109   public int getInt(int ordinal) { return columns[ordinal].getInt(rowId); }
0110
0111   @Override
0112   public long getLong(int ordinal) { return columns[ordinal].getLong(rowId); }
0113
0114   @Override
0115   public float getFloat(int ordinal) { return columns[ordinal].getFloat(rowId); }
0116
0117   @Override
0118   public double getDouble(int ordinal) { return columns[ordinal].getDouble(rowId); }
0119
0120   @Override
0121   public Decimal getDecimal(int ordinal, int precision, int scale) {
0122     return columns[ordinal].getDecimal(rowId, precision, scale);
0123   }
0124
0125   @Override
0126   public UTF8String getUTF8String(int ordinal) {
0127     return columns[ordinal].getUTF8String(rowId);
0128   }
0129
0130   @Override
0131   public byte[] getBinary(int ordinal) {
0132     return columns[ordinal].getBinary(rowId);
0133   }
0134
0135   @Override
0136   public CalendarInterval getInterval(int ordinal) {
0137     return columns[ordinal].getInterval(rowId);
0138   }
0139
0140   @Override
0141   public ColumnarRow getStruct(int ordinal, int numFields) {
0142     return columns[ordinal].getStruct(rowId);
0143   }
0144
0145   @Override
0146   public ColumnarArray getArray(int ordinal) {
0147     return columns[ordinal].getArray(rowId);
0148   }
0149
0150   @Override
0151   public ColumnarMap getMap(int ordinal) {
0152     return columns[ordinal].getMap(rowId);
0153   }
0154
0155   @Override
0156   public Object get(int ordinal, DataType dataType) {
0157     if (dataType instanceof BooleanType) {
0158       return getBoolean(ordinal);
0159     } else if (dataType instanceof ByteType) {
0160       return getByte(ordinal);
0161     } else if (dataType instanceof ShortType) {
0162       return getShort(ordinal);
0163     } else if (dataType instanceof IntegerType) {
0164       return getInt(ordinal);
0165     } else if (dataType instanceof LongType) {
0166       return getLong(ordinal);
0167     } else if (dataType instanceof FloatType) {
0168       return getFloat(ordinal);
0169     } else if (dataType instanceof DoubleType) {
0170       return getDouble(ordinal);
0171     } else if (dataType instanceof StringType) {
0172       return getUTF8String(ordinal);
0173     } else if (dataType instanceof BinaryType) {
0174       return getBinary(ordinal);
0175     } else if (dataType instanceof DecimalType) {
0176       DecimalType t = (DecimalType) dataType;
0177       return getDecimal(ordinal, t.precision(), t.scale());
0178     } else if (dataType instanceof DateType) {
0179       return getInt(ordinal);
0180     } else if (dataType instanceof TimestampType) {
0181       return getLong(ordinal);
0182     } else if (dataType instanceof ArrayType) {
0183       return getArray(ordinal);
0184     } else if (dataType instanceof StructType) {
0185       return getStruct(ordinal, ((StructType)dataType).fields().length);
0186     } else if (dataType instanceof MapType) {
0187       return getMap(ordinal);
0188     } else {
0189       throw new UnsupportedOperationException("Datatype not supported " + dataType);
0190     }
0191   }
0192
0193   @Override
0194   public void update(int ordinal, Object value) {
0195     if (value == null) {
0196       setNullAt(ordinal);
0197     } else {
0198       DataType dt = columns[ordinal].dataType();
0199       if (dt instanceof BooleanType) {
0200         setBoolean(ordinal, (boolean) value);
0201       } else if (dt instanceof IntegerType) {
0202         setInt(ordinal, (int) value);
0203       } else if (dt instanceof ShortType) {
0204         setShort(ordinal, (short) value);
0205       } else if (dt instanceof LongType) {
0206         setLong(ordinal, (long) value);
0207       } else if (dt instanceof FloatType) {
0208         setFloat(ordinal, (float) value);
0209       } else if (dt instanceof DoubleType) {
0210         setDouble(ordinal, (double) value);
0211       } else if (dt instanceof DecimalType) {
0212         DecimalType t = (DecimalType) dt;
0213         Decimal d = Decimal.apply((BigDecimal) value, t.precision(), t.scale());
0214         setDecimal(ordinal, d, t.precision());
0215       } else if (dt instanceof CalendarIntervalType) {
0216         setInterval(ordinal, (CalendarInterval) value);
0217       } else {
0218         throw new UnsupportedOperationException("Datatype not supported " + dt);
0219       }
0220     }
0221   }
0222
0223   @Override
0224   public void setNullAt(int ordinal) {
0225     columns[ordinal].putNull(rowId);
0226   }
0227
0228   @Override
0229   public void setBoolean(int ordinal, boolean value) {
0230     columns[ordinal].putNotNull(rowId);
0231     columns[ordinal].putBoolean(rowId, value);
0232   }
0233
0234   @Override
0235   public void setByte(int ordinal, byte value) {
0236     columns[ordinal].putNotNull(rowId);
0237     columns[ordinal].putByte(rowId, value);
0238   }
0239
0240   @Override
0241   public void setShort(int ordinal, short value) {
0242     columns[ordinal].putNotNull(rowId);
0243     columns[ordinal].putShort(rowId, value);
0244   }
0245
0246   @Override
0247   public void setInt(int ordinal, int value) {
0248     columns[ordinal].putNotNull(rowId);
0249     columns[ordinal].putInt(rowId, value);
0250   }
0251
0252   @Override
0253   public void setLong(int ordinal, long value) {
0254     columns[ordinal].putNotNull(rowId);
0255     columns[ordinal].putLong(rowId, value);
0256   }
0257
0258   @Override
0259   public void setFloat(int ordinal, float value) {
0260     columns[ordinal].putNotNull(rowId);
0261     columns[ordinal].putFloat(rowId, value);
0262   }
0263
0264   @Override
0265   public void setDouble(int ordinal, double value) {
0266     columns[ordinal].putNotNull(rowId);
0267     columns[ordinal].putDouble(rowId, value);
0268   }
0269
0270   @Override
0271   public void setDecimal(int ordinal, Decimal value, int precision) {
0272     columns[ordinal].putNotNull(rowId);
0273     columns[ordinal].putDecimal(rowId, value, precision);
0274   }
0275
0276   @Override
0277   public void setInterval(int ordinal, CalendarInterval value) {
0278     columns[ordinal].putNotNull(rowId);
0279     columns[ordinal].putInterval(rowId, value);
0280   }
0281 }