[SPARK-27945][SQL] Minimal changes to support columnar processing

sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ArrowColumnVector.java

@@ -48,7 +48,7 @@ public int numNulls() {
   }
 
   @Override
-  public void close() {
+  protected void doClose() {
     if (childColumns != null) {
       for (int i = 0; i < childColumns.length; i++) {
         childColumns[i].close();

sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnVector.java

@@ -50,19 +50,44 @@
 @Evolving
 public abstract class ColumnVector implements AutoCloseable {
 
+  private int refCount = 1;
+
   /**
    * Returns the data type of this column vector.
    */
   public final DataType dataType() { return type; }
 
+  /**
+   * Increment the reference count for this vector. This is an implementation detail and
+   * only BoundReference should call this directly.
+   * @return this for easy chaining.
+   */
+  public final ColumnVector incRefCount() {
+    refCount++;
+    return this;
+  }
+
   /**
    * Cleans up memory for this column vector. The column vector is not usable after this.
+   * In reality it decrements the reference count and if it reaches 0 the resources are released
+   * but this is an implementation detail that most code should just ignore.
    *
    * This overwrites `AutoCloseable.close` to remove the `throws` clause, as column vector is
    * in-memory and we don't expect any exception to happen during closing.
    */
   @Override
-  public abstract void close();
+  public final void close() {
+    refCount--;
+    if (refCount == 0) {
+      doClose();
+    }
+  }
+
+  /**
+   * Actually cleans up memory for this column vector. The column vector is really not usable after
+   * this.
+   */
+  protected abstract void doClose();
 
   /**
    * Returns true if this column vector contains any null values.

sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnarBatch.java

@@ -31,7 +31,7 @@
  * the entire data loading process.
  */
 @Evolving
-public final class ColumnarBatch {
+public final class ColumnarBatch implements AutoCloseable {
   private int numRows;
   private final ColumnVector[] columns;
 
@@ -42,6 +42,7 @@ public final class ColumnarBatch {
    * 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();
@@ -110,7 +111,17 @@ public InternalRow getRow(int rowId) {
   }
 
   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);
   }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/BoundAttribute.scala

@@ -23,6 +23,7 @@ import org.apache.spark.sql.catalyst.errors.attachTree
 import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGenerator, ExprCode, FalseLiteral, JavaCode}
 import org.apache.spark.sql.catalyst.expressions.codegen.Block._
 import org.apache.spark.sql.types._
+import org.apache.spark.sql.vectorized.ColumnarBatch
 
 /**
  * A bound reference points to a specific slot in the input tuple, allowing the actual value
@@ -36,6 +37,15 @@ case class BoundReference(ordinal: Int, dataType: DataType, nullable: Boolean)
 
   private val accessor: (InternalRow, Int) => Any = InternalRow.getAccessor(dataType, nullable)
 
+  override def supportsColumnar: Boolean = true
+
+  override def columnarEval(batch: ColumnarBatch): Any = {
+    // Because of the convention that the returned ColumnVector must be closed by the
+    // caller we increment the reference count for columns taken directly from a batch, so that
+    // the call to close by the caller does not actually release the column's resources.
+    batch.column(ordinal).incRefCount()
+  }
+
   // Use special getter for primitive types (for UnsafeRow)
   override def eval(input: InternalRow): Any = {
     accessor(input, ordinal)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala

@@ -29,6 +29,7 @@ import org.apache.spark.sql.catalyst.trees.TreeNode
 import org.apache.spark.sql.catalyst.util.truncatedString
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
+import org.apache.spark.sql.vectorized.ColumnarBatch
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 // This file defines the basic expression abstract classes in Catalyst.
@@ -80,7 +81,7 @@ import org.apache.spark.sql.types._
  * - [[ComplexTypeMergingExpression]]: to resolve output types of the complex expressions
  *                                     (e.g., [[CaseWhen]]).
  */
-abstract class Expression extends TreeNode[Expression] {
+abstract class Expression extends TreeNode[Expression] with Serializable {
 
   /**
    * Returns true when an expression is a candidate for static evaluation before the query is
@@ -95,6 +96,26 @@ abstract class Expression extends TreeNode[Expression] {
    */
   def foldable: Boolean = false
 
+  /**
+   * Returns true if this expression supports columnar processing through [[columnarEval]].
+   */
+  def supportsColumnar: Boolean = false
+
+  /**
+   * Returns the result of evaluating this expression on the entire [[ColumnarBatch]]. The result of
+   * calling this may be a single [[org.apache.spark.sql.vectorized.ColumnVector]] or a scalar
+   * value. Scalar values typically happen if they are a part of the expression i.e. col("a") + 100.
+   * In this case the 100 is a [[Literal]] that [[Add]] would have to be able to handle.
+   *
+   * By convention any [[org.apache.spark.sql.vectorized.ColumnVector]] returned by [[columnarEval]]
+   * is owned by the caller and will need to be closed by them. This can happen by putting it into
+   * a [[ColumnarBatch]] and closing the batch or by closing the vector directly if it is a
+   * temporary value.
+   */
+  def columnarEval(batch: ColumnarBatch): Any = {
+    throw new IllegalStateException(s"Internal Error ${this.getClass} has column support mismatch")
+  }
+
   /**
    * Returns true when the current expression always return the same result for fixed inputs from
    * children. The non-deterministic expressions should not change in number and order. They should
@@ -287,6 +308,9 @@ trait Unevaluable extends Expression {
   final override def eval(input: InternalRow = null): Any =
     throw new UnsupportedOperationException(s"Cannot evaluate expression: $this")
 
+  final override def columnarEval(batch: ColumnarBatch): Any =
+    throw new UnsupportedOperationException(s"Cannot evaluate expression: $this")
+
   final override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode =
     throw new UnsupportedOperationException(s"Cannot generate code for expression: $this")
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/literals.scala

@@ -45,6 +45,7 @@ import org.apache.spark.sql.catalyst.util._
 import org.apache.spark.sql.catalyst.util.DateTimeUtils.instantToMicros
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
+import org.apache.spark.sql.vectorized.ColumnarBatch
 import org.apache.spark.unsafe.types._
 import org.apache.spark.util.Utils
 
@@ -260,6 +261,7 @@ case class Literal (value: Any, dataType: DataType) extends LeafExpression {
 
   override def foldable: Boolean = true
   override def nullable: Boolean = value == null
+  override def supportsColumnar: Boolean = true
 
   override def toString: String = value match {
     case null => "null"
@@ -300,6 +302,7 @@ case class Literal (value: Any, dataType: DataType) extends LeafExpression {
   }
 
   override def eval(input: InternalRow): Any = value
+  override def columnarEval(batch: ColumnarBatch): Any = value
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
     val javaType = CodeGenerator.javaType(dataType)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala

@@ -25,6 +25,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.plans.logical.EventTimeWatermark
 import org.apache.spark.sql.catalyst.util.quoteIdentifier
 import org.apache.spark.sql.types._
+import org.apache.spark.sql.vectorized.ColumnarBatch
 
 object NamedExpression {
   private val curId = new java.util.concurrent.atomic.AtomicLong()
@@ -152,6 +153,10 @@ case class Alias(child: Expression, name: String)(
   override lazy val resolved =
     childrenResolved && checkInputDataTypes().isSuccess && !child.isInstanceOf[Generator]
 
+  override def supportsColumnar: Boolean = child.supportsColumnar
+
+  override def columnarEval(batch: ColumnarBatch): Any = child.columnarEval(batch)
+
   override def eval(input: InternalRow): Any = child.eval(input)
 
   /** Just a simple passthrough for code generation. */
@@ -237,6 +242,10 @@ case class AttributeReference(
 
   // currently can only handle qualifier of length 2
   require(qualifier.length <= 2)
+
+  override def supportsColumnar: Boolean = true
+  // No columnar eval is needed because this must be bound before it is evaluated
+
   /**
    * Returns true iff the expression id is the same for both attributes.
    */

sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/OffHeapColumnVector.java

@@ -85,8 +85,8 @@ public long valuesNativeAddress() {
   }
 
   @Override
-  public void close() {
-    super.close();
+  protected void doClose() {
+    super.doClose();
     Platform.freeMemory(nulls);
     Platform.freeMemory(data);
     Platform.freeMemory(lengthData);

sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/OnHeapColumnVector.java

@@ -81,8 +81,8 @@ public OnHeapColumnVector(int capacity, DataType type) {
   }
 
   @Override
-  public void close() {
-    super.close();
+  protected void doClose() {
+    super.doClose();
     nulls = null;
     byteData = null;
     shortData = null;

sql/core/src/main/java/org/apache/spark/sql/execution/vectorized/WritableColumnVector.java

@@ -65,7 +65,7 @@ public void reset() {
   }
 
   @Override
-  public void close() {
+  protected void doClose() {
     if (childColumns != null) {
       for (int i = 0; i < childColumns.length; i++) {
         childColumns[i].close();
@@ -604,7 +604,10 @@ public final int appendArray(int length) {
    */
   public final int appendStruct(boolean isNull) {
     if (isNull) {
-      appendNull();
+      // This is the same as appendNull but without the assertion for struct types
+      reserve(elementsAppended + 1);
+      putNull(elementsAppended);
+      elementsAppended++;
       for (WritableColumnVector c: childColumns) {
         if (c.type instanceof StructType) {
           c.appendStruct(true);

sql/core/src/main/scala/org/apache/spark/sql/SparkSessionExtensions.scala

@@ -27,6 +27,7 @@ import org.apache.spark.sql.catalyst.expressions.ExpressionInfo
 import org.apache.spark.sql.catalyst.parser.ParserInterface
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.ColumnarRule
 
 /**
  * :: Experimental ::
@@ -42,6 +43,7 @@ import org.apache.spark.sql.catalyst.rules.Rule
  * <li>Planning Strategies.</li>
  * <li>Customized Parser.</li>
  * <li>(External) Catalog listeners.</li>
+ * <li>Columnar Rules.</li>
  * </ul>
  *
  * The extensions can be used by calling `withExtensions` on the [[SparkSession.Builder]], for
@@ -93,6 +95,23 @@ class SparkSessionExtensions {
   type StrategyBuilder = SparkSession => Strategy
   type ParserBuilder = (SparkSession, ParserInterface) => ParserInterface
   type FunctionDescription = (FunctionIdentifier, ExpressionInfo, FunctionBuilder)
+  type ColumnarRuleBuilder = SparkSession => ColumnarRule
+
+  private[this] val columnarRuleBuilders = mutable.Buffer.empty[ColumnarRuleBuilder]
+
+  /**
+   * Build the override rules for columnar execution.
+   */
+  private[sql] def buildColumnarRules(session: SparkSession): Seq[ColumnarRule] = {
+    columnarRuleBuilders.map(_.apply(session))
+  }
+
+  /**
+   * Inject a rule that can override the columnar execution of an executor.
+   */
+  def injectColumnar(builder: ColumnarRuleBuilder): Unit = {
+    columnarRuleBuilders += builder
+  }
 
   private[this] val resolutionRuleBuilders = mutable.Buffer.empty[RuleBuilder]