OSCL-LXR spark-3.0.0/​sql/​catalyst/​src/​main/​java/​org/​apache/​spark/​sql/​vectorized/​ColumnarBatch.java	Source navigation Diff markup Identifier search General search
	Version: spark-3.0.0 hadoop-3.2.0-src hadoop-3.3.0-src spark-2.4.7 linux-4.15.13 Revert   Change

 /*
  * 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.spark.sql.vectorized;
 
 import java.util.*;
 
 import org.apache.spark.annotation.Evolving;
 import org.apache.spark.sql.catalyst.InternalRow;
 import org.apache.spark.sql.catalyst.expressions.GenericInternalRow;
 import org.apache.spark.sql.types.*;
 import org.apache.spark.unsafe.types.CalendarInterval;
 import org.apache.spark.unsafe.types.UTF8String;
 
 /**
  * This class wraps multiple ColumnVectors as a row-wise table. It provides a row view of this
  * batch so that Spark can access the data row by row. Instance of it is meant to be reused during
  * the entire data loading process.
  */
 @Evolving
 public final class ColumnarBatch implements AutoCloseable {
   private int numRows;
   private final ColumnVector[] columns;
 
   // Staging row returned from `getRow`.
   private final ColumnarBatchRow row;
 
   /**
    * Called to close all the columns in this batch. It is not valid to access the data after
    * calling this. This must be called at the end to clean up memory allocations.
    */
   @Override
   public void close() {
     for (ColumnVector c: columns) {
       c.close();
     }
   }
 
   /**
    * Returns an iterator over the rows in this batch.
    */
   public Iterator<InternalRow> rowIterator() {
     final int maxRows = numRows;
     final ColumnarBatchRow row = new ColumnarBatchRow(columns);
     return new Iterator<InternalRow>() {
       int rowId = 0;
 
       @Override
       public boolean hasNext() {
         return rowId < maxRows;
       }
 
       @Override
       public InternalRow next() {
         if (rowId >= maxRows) {
           throw new NoSuchElementException();
         }
         row.rowId = rowId++;
         return row;
       }
 
       @Override
       public void remove() {
         throw new UnsupportedOperationException();
       }
     };
   }
 
   /**
    * Sets the number of rows in this batch.
    */
   public void setNumRows(int numRows) {
     this.numRows = numRows;
   }
 
   /**
    * Returns the number of columns that make up this batch.
    */
   public int numCols() { return columns.length; }
 
   /**
    * Returns the number of rows for read, including filtered rows.
    */
   public int numRows() { return numRows; }
 
   /**
    * Returns the column at `ordinal`.
    */
   public ColumnVector column(int ordinal) { return columns[ordinal]; }
 
   /**
    * Returns the row in this batch at `rowId`. Returned row is reused across calls.
    */
   public InternalRow getRow(int rowId) {
     assert(rowId >= 0 && rowId < numRows);
     row.rowId = rowId;
     return row;
   }
 
   public ColumnarBatch(ColumnVector[] columns) {
     this(columns, 0);
   }
 
   /**
    * Create a new batch from existing column vectors.
    * @param columns The columns of this batch
    * @param numRows The number of rows in this batch
    */
   public ColumnarBatch(ColumnVector[] columns, int numRows) {
     this.columns = columns;
     this.numRows = numRows;
     this.row = new ColumnarBatchRow(columns);
   }
 }
 
 /**
  * An internal class, which wraps an array of {@link ColumnVector} and provides a row view.
  */
 class ColumnarBatchRow extends InternalRow {
   public int rowId;
   private final ColumnVector[] columns;
 
   ColumnarBatchRow(ColumnVector[] columns) {
     this.columns = columns;
   }
 
   @Override
   public int numFields() { return columns.length; }
 
   @Override
   public InternalRow copy() {
     GenericInternalRow row = new GenericInternalRow(columns.length);
     for (int i = 0; i < numFields(); i++) {
       if (isNullAt(i)) {
         row.setNullAt(i);
       } else {
         DataType dt = columns[i].dataType();
         if (dt instanceof BooleanType) {
           row.setBoolean(i, getBoolean(i));
         } else if (dt instanceof ByteType) {
           row.setByte(i, getByte(i));
         } else if (dt instanceof ShortType) {
           row.setShort(i, getShort(i));
         } else if (dt instanceof IntegerType) {
           row.setInt(i, getInt(i));
         } else if (dt instanceof LongType) {
           row.setLong(i, getLong(i));
         } else if (dt instanceof FloatType) {
           row.setFloat(i, getFloat(i));
         } else if (dt instanceof DoubleType) {
           row.setDouble(i, getDouble(i));
         } else if (dt instanceof StringType) {
           row.update(i, getUTF8String(i).copy());
         } else if (dt instanceof BinaryType) {
           row.update(i, getBinary(i));
         } else if (dt instanceof DecimalType) {
           DecimalType t = (DecimalType)dt;
           row.setDecimal(i, getDecimal(i, t.precision(), t.scale()), t.precision());
         } else if (dt instanceof DateType) {
           row.setInt(i, getInt(i));
         } else if (dt instanceof TimestampType) {
           row.setLong(i, getLong(i));
         } else {
           throw new RuntimeException("Not implemented. " + dt);
         }
       }
     }
     return row;
   }
 
   @Override
   public boolean anyNull() {
     throw new UnsupportedOperationException();
   }
 
   @Override
   public boolean isNullAt(int ordinal) { return columns[ordinal].isNullAt(rowId); }
 
   @Override
   public boolean getBoolean(int ordinal) { return columns[ordinal].getBoolean(rowId); }
 
   @Override
   public byte getByte(int ordinal) { return columns[ordinal].getByte(rowId); }
 
   @Override
   public short getShort(int ordinal) { return columns[ordinal].getShort(rowId); }
 
   @Override
   public int getInt(int ordinal) { return columns[ordinal].getInt(rowId); }
 
   @Override
   public long getLong(int ordinal) { return columns[ordinal].getLong(rowId); }
 
   @Override
   public float getFloat(int ordinal) { return columns[ordinal].getFloat(rowId); }
 
   @Override
   public double getDouble(int ordinal) { return columns[ordinal].getDouble(rowId); }
 
   @Override
   public Decimal getDecimal(int ordinal, int precision, int scale) {
     return columns[ordinal].getDecimal(rowId, precision, scale);
   }
 
   @Override
   public UTF8String getUTF8String(int ordinal) {
     return columns[ordinal].getUTF8String(rowId);
   }
 
   @Override
   public byte[] getBinary(int ordinal) {
     return columns[ordinal].getBinary(rowId);
   }
 
   @Override
   public CalendarInterval getInterval(int ordinal) {
     return columns[ordinal].getInterval(rowId);
   }
 
   @Override
   public ColumnarRow getStruct(int ordinal, int numFields) {
     return columns[ordinal].getStruct(rowId);
   }
 
   @Override
   public ColumnarArray getArray(int ordinal) {
     return columns[ordinal].getArray(rowId);
   }
 
   @Override
   public ColumnarMap getMap(int ordinal) {
     return columns[ordinal].getMap(rowId);
   }
 
   @Override
   public Object get(int ordinal, DataType dataType) {
     if (dataType instanceof BooleanType) {
       return getBoolean(ordinal);
     } else if (dataType instanceof ByteType) {
       return getByte(ordinal);
     } else if (dataType instanceof ShortType) {
       return getShort(ordinal);
     } else if (dataType instanceof IntegerType) {
       return getInt(ordinal);
     } else if (dataType instanceof LongType) {
       return getLong(ordinal);
     } else if (dataType instanceof FloatType) {
       return getFloat(ordinal);
     } else if (dataType instanceof DoubleType) {
       return getDouble(ordinal);
     } else if (dataType instanceof StringType) {
       return getUTF8String(ordinal);
     } else if (dataType instanceof BinaryType) {
       return getBinary(ordinal);
     } else if (dataType instanceof DecimalType) {
       DecimalType t = (DecimalType) dataType;
       return getDecimal(ordinal, t.precision(), t.scale());
     } else if (dataType instanceof DateType) {
       return getInt(ordinal);
     } else if (dataType instanceof TimestampType) {
       return getLong(ordinal);
     } else if (dataType instanceof ArrayType) {
       return getArray(ordinal);
     } else if (dataType instanceof StructType) {
       return getStruct(ordinal, ((StructType)dataType).fields().length);
     } else if (dataType instanceof MapType) {
       return getMap(ordinal);
     } else {
       throw new UnsupportedOperationException("Datatype not supported " + dataType);
     }
   }
 
   @Override
   public void update(int ordinal, Object value) { throw new UnsupportedOperationException(); }
 
   @Override
   public void setNullAt(int ordinal) { throw new UnsupportedOperationException(); }
 }

This page was automatically generated by the 2.1.0 LXR engine. The LXR team