sql/vectorized/ColumnarArray.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 package org.apache.spark.sql.vectorized;
0018
0019 import org.apache.spark.annotation.Evolving;
0020 import org.apache.spark.sql.catalyst.expressions.SpecializedGettersReader;
0021 import org.apache.spark.sql.catalyst.expressions.UnsafeArrayData;
0022 import org.apache.spark.sql.catalyst.util.ArrayData;
0023 import org.apache.spark.sql.catalyst.util.GenericArrayData;
0024 import org.apache.spark.sql.types.*;
0025 import org.apache.spark.unsafe.types.CalendarInterval;
0026 import org.apache.spark.unsafe.types.UTF8String;
0027
0028 /**
0029  * Array abstraction in {@link ColumnVector}.
0030  */
0031 @Evolving
0032 public final class ColumnarArray extends ArrayData {
0033   // The data for this array. This array contains elements from
0034   // data[offset] to data[offset + length).
0035   private final ColumnVector data;
0036   private final int offset;
0037   private final int length;
0038
0039   public ColumnarArray(ColumnVector data, int offset, int length) {
0040     this.data = data;
0041     this.offset = offset;
0042     this.length = length;
0043   }
0044
0045   @Override
0046   public int numElements() {
0047     return length;
0048   }
0049
0050   @Override
0051   public ArrayData copy() {
0052     DataType dt = data.dataType();
0053
0054     if (dt instanceof BooleanType) {
0055       return UnsafeArrayData.fromPrimitiveArray(toBooleanArray());
0056     } else if (dt instanceof ByteType) {
0057       return UnsafeArrayData.fromPrimitiveArray(toByteArray());
0058     } else if (dt instanceof ShortType) {
0059       return UnsafeArrayData.fromPrimitiveArray(toShortArray());
0060     } else if (dt instanceof IntegerType || dt instanceof DateType) {
0061       return UnsafeArrayData.fromPrimitiveArray(toIntArray());
0062     } else if (dt instanceof LongType || dt instanceof TimestampType) {
0063       return UnsafeArrayData.fromPrimitiveArray(toLongArray());
0064     } else if (dt instanceof FloatType) {
0065       return UnsafeArrayData.fromPrimitiveArray(toFloatArray());
0066     } else if (dt instanceof DoubleType) {
0067       return UnsafeArrayData.fromPrimitiveArray(toDoubleArray());
0068     } else {
0069       return new GenericArrayData(toObjectArray(dt));
0070     }
0071   }
0072
0073   @Override
0074   public boolean[] toBooleanArray() { return data.getBooleans(offset, length); }
0075
0076   @Override
0077   public byte[] toByteArray() { return data.getBytes(offset, length); }
0078
0079   @Override
0080   public short[] toShortArray() { return data.getShorts(offset, length); }
0081
0082   @Override
0083   public int[] toIntArray() { return data.getInts(offset, length); }
0084
0085   @Override
0086   public long[] toLongArray() { return data.getLongs(offset, length); }
0087
0088   @Override
0089   public float[] toFloatArray() { return data.getFloats(offset, length); }
0090
0091   @Override
0092   public double[] toDoubleArray() { return data.getDoubles(offset, length); }
0093
0094   // TODO: this is extremely expensive.
0095   @Override
0096   public Object[] array() {
0097     DataType dt = data.dataType();
0098     Object[] list = new Object[length];
0099     try {
0100       for (int i = 0; i < length; i++) {
0101         if (!data.isNullAt(offset + i)) {
0102           list[i] = get(i, dt);
0103         }
0104       }
0105       return list;
0106     } catch(Exception e) {
0107       throw new RuntimeException("Could not get the array", e);
0108     }
0109   }
0110
0111   @Override
0112   public boolean isNullAt(int ordinal) { return data.isNullAt(offset + ordinal); }
0113
0114   @Override
0115   public boolean getBoolean(int ordinal) {
0116     return data.getBoolean(offset + ordinal);
0117   }
0118
0119   @Override
0120   public byte getByte(int ordinal) { return data.getByte(offset + ordinal); }
0121
0122   @Override
0123   public short getShort(int ordinal) {
0124     return data.getShort(offset + ordinal);
0125   }
0126
0127   @Override
0128   public int getInt(int ordinal) { return data.getInt(offset + ordinal); }
0129
0130   @Override
0131   public long getLong(int ordinal) { return data.getLong(offset + ordinal); }
0132
0133   @Override
0134   public float getFloat(int ordinal) {
0135     return data.getFloat(offset + ordinal);
0136   }
0137
0138   @Override
0139   public double getDouble(int ordinal) { return data.getDouble(offset + ordinal); }
0140
0141   @Override
0142   public Decimal getDecimal(int ordinal, int precision, int scale) {
0143     return data.getDecimal(offset + ordinal, precision, scale);
0144   }
0145
0146   @Override
0147   public UTF8String getUTF8String(int ordinal) {
0148     return data.getUTF8String(offset + ordinal);
0149   }
0150
0151   @Override
0152   public byte[] getBinary(int ordinal) {
0153     return data.getBinary(offset + ordinal);
0154   }
0155
0156   @Override
0157   public CalendarInterval getInterval(int ordinal) {
0158     return data.getInterval(offset + ordinal);
0159   }
0160
0161   @Override
0162   public ColumnarRow getStruct(int ordinal, int numFields) {
0163     return data.getStruct(offset + ordinal);
0164   }
0165
0166   @Override
0167   public ColumnarArray getArray(int ordinal) {
0168     return data.getArray(offset + ordinal);
0169   }
0170
0171   @Override
0172   public ColumnarMap getMap(int ordinal) {
0173     return data.getMap(offset + ordinal);
0174   }
0175
0176   @Override
0177   public Object get(int ordinal, DataType dataType) {
0178     return SpecializedGettersReader.read(this, ordinal, dataType, false, false);
0179   }
0180
0181   @Override
0182   public void update(int ordinal, Object value) { throw new UnsupportedOperationException(); }
0183
0184   @Override
0185   public void setNullAt(int ordinal) { throw new UnsupportedOperationException(); }
0186 }