Apache Datafusion Comet and the story of my first contribution to it

Due to RDD legacy, Apache Spark is (mostly) relying on the row-oriented data representation in the memory.

NOTE: This is the second time I've mentioned that Spark "mostly" relies on row-oriented data representation, and it seems appropriate to explain what I mean by this. The explanation is that modern Spark also provides an abstraction for columnar execution, but only a few of the existing Spark operators support this feature. Spark's main abstraction for any operator is the SparkPlan. By implementing doExecuteColumnar and setting supportsColumnar to true, one can make an operator accept both row-oriented and columnar-oriented data.

abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializable {
  def supportsRowBased: Boolean = !supportsColumnar
  def supportsColumnar: Boolean = false
  protected def doExecute(): RDD[InternalRow]
  protected def doExecuteColumnar(): RDD[ColumnarBatch]
}

Row-oriented representation works best when dealing with naturally row-oriented data, particularly unstructured or semi-structured data. For example, processing raw text, CSV files, or XML files is most efficient using a row-oriented format. However, we are now in an era dominated by columnar open formats such as Apache Parquet, Apache ORC, and their associated metadata frameworks (DeltaLake, Apache Iceberg, Apache Hudi, etc.). In this context, Spark not only loses the benefits of working with columns but must also convert columns to rows before applying any operator. While Spark's implementation of the ColumnarToRow operation is very efficient, it still introduces computational overhead. The same principle applies to data processing itself. For instance, Snowflake uses a columnar execution model similar to the well-known "Monet/X100" system. For typical data warehouse workloads, the columnar execution model is significantly more efficient compared to the row-oriented model. The latter is better suited for handling unstructured data (or transactions in transactional databases, though this isn't relevant for data warehouses designed for analytical queries).

NOTE: This topic is a growing area of expertise for me, and I find it challenging to speak confidently about the advantages or disadvantages of virtual function calls compared to code generation. I would like to again reference the article about Databricks Photon, where the authors discuss problems related to the code generation approach.

Spark relies on code generation. When the Catalyst optimizer produces the final Optimized Plan, Spark runs code generation, followed by calls to the Java compiler and JIT compilation. In theory, this approach allows handling very complex cases and leverages modern compilers' branch-prediction techniques. However, in reality, for typical data warehouse analytical workloads, code generation rarely provides advantages over virtual function calls (see the note above). Furthermore, code generation is much harder to maintain, develop, and debug because the final code is only available at runtime. While Spark developer APIs provide a way to access the generated code, developers still need to copy it to a Java file, compile it, analyze the bytecode, and debug it. This process is significantly more complex compared to making calls to kernels that represent operations on columns.

I often hear that Spark is "slow" because it runs on the JVM. While the JVM does interpret bytecode, through JIT compilation and compiler optimizations, Spark can achieve performance comparable to implementations written in C++ or Rust. In fact, I believe choosing the JVM was one of the keys to Spark's success. For instance, extending Spark by adding a new plugin or data source simply requires implementing a few Java/Scala interfaces. If Spark had been written in C++, any extension would require writing C++ code, which is generally more challenging. Therefore, I disagree with the statement that "Spark is slow because of Java," though there might be some merit to this argument in specific cases.

An attempt to create a multi-backend Spark plugin. Possible backends include VeloxDB, ClickHouse, and Apache Arrow, among others. In this case, after optimization, the optimized plan is translated from Spark's list of operators to the corresponding list of Velox or ClickHouse operators.

When a Spark operator cannot be translated to a native one, the system falls back to Spark's JVM execution. This fallback mechanism allows support for features like Spark UDFs and some rarely used operators that no one wants to rewrite into native code. A controversal feature of the Apache Gluten is that it allows fallback to spark and switching back to the native execution in any place of the computational plan. I believe that in some cases it may give an advantage but I also agreed with Databricks developers that explicitly mentioned why they did not implement something like this in their Photon engine:

The last node in a Photon plan is a “transition”node. Unlike the adapter node, the transition node must pivot columnar data to row data so the rowwise legacy Spark SQL engine can operate over it. Since Apache Spark’s scan always produces columnar data when reading columnar formats, we note that one such pivot is required even without Photon. Since we only convert plans to Photon starting at the scan node, adding a single pivot on top of a Photon plan does not cause regressions vs. Spark (both the Spark plan and Photon plan each have a single pivot). However, if we were to eagerly convert arbitrary parts of the plan to use Photon, we could have an arbitrary number of pivots, which could lead to regressions. Today, we elect to be conservative and choose not to do this. In the future, we may investigate weighing the tradeoff of the speedup Photon would provide vs. the slowdown caused by adding an additional columnto-row pivot.

Based on the benchmarks provided by the Gluten core team, it shows a performance advantage of approximately 2x compared to the standard Spark runtime. As I can understand, Apache Gluten (incubating) is an engine behind the Microsft Fabric service (Spark PaaS from MS Azure).

Databricks, the company founded by Spark's original creators, has made an attempt to improve Spark's performance for data warehouse-like use cases. In my opinion, Photon is the most production-ready among all native runtimes for Spark. However, it is a proprietary solution available exclusively within the Databricks Data and AI platform.

This runtime is written from scratch in C++. I was able to test it because I work for a company that is a Databricks customer. In my experience, Photon can provide up to 3x speed improvement and is particularly efficient when combined with Delta Lake (a metadata layer built on top of Apache Parquet, created and maintained by Databricks).

Before we discuss Comet itself, I need to mention the upstream Datafusion project. It is a Rust-based tool focused on creating an embeddable query engine that uses Apache Arrow as its internal memory model. Currently, it is the fastest among all existing engines for single-node Parquet processing.

While Datafusion itself is not a query engine for end users, it is a solid foundation for building tools. For example, datafusion is the engine behind InfluxDB, Databend, ROAPI, GreptimeDB and many others.

Under the hood Datafusion provides a set of columnar operators, query optimizer and integrtion with data sources.

DataFusion Comet is a project that aims to implement something conceptually similar to Databricks Photon, but using DataFusion as the physical execution layer.

Comet works through org.apache.spark.api.plugin.SparkPlugin:

class CometPlugin extends SparkPlugin with Logging {
  override def driverPlugin(): DriverPlugin = new CometDriverPlugin

  override def executorPlugin(): ExecutorPlugin = null
}

abstract class CometExec extends CometPlan {
  def originalPlan: SparkPlan

  override def supportsColumnar: Boolean = true

  override def output: Seq[Attribute] = originalPlan.output

  override def doExecute(): RDD[InternalRow] =
    ColumnarToRowExec(this).doExecute()

  override def executeCollect(): Array[InternalRow] =
    ColumnarToRowExec(this).executeCollect()

  override def outputOrdering: Seq[SortOrder] = originalPlan.outputOrdering

  override def outputPartitioning: Partitioning = originalPlan.outputPartitioning

  def executeColumnarCollectIterator(): (Long, Iterator[ColumnarBatch]) = {
    val countsAndBytes = CometExec.getByteArrayRdd(this).collect()
    val total = countsAndBytes.map(_._1).sum
    val rows = countsAndBytes.iterator
      .flatMap(countAndBytes =>
        CometExec.decodeBatches(countAndBytes._2, this.getClass.getSimpleName))
    (total, rows)
  }

  protected def setSubqueries(planId: Long, sparkPlan: SparkPlan): Unit = {
    sparkPlan.children.foreach(setSubqueries(planId, _))

    sparkPlan.expressions.foreach {
      _.collect { case sub: ScalarSubquery =>
        CometScalarSubquery.setSubquery(planId, sub)
      }
    }
  }

  protected def cleanSubqueries(planId: Long, sparkPlan: SparkPlan): Unit = {
    sparkPlan.children.foreach(cleanSubqueries(planId, _))

    sparkPlan.expressions.foreach {
      _.collect { case sub: ScalarSubquery =>
        CometScalarSubquery.removeSubquery(planId, sub)
      }
    }
  }
}

Protobuf messages

Under the hood, Comet takes a Spark plan and attempts to translate expressions into Datafusion operators where possible. It includes an additional compatibility layer that handles corner cases where Spark results differ from Datafusion. For all supported operators, Comet maintains corresponded protobuf messages:

message Expr {
  oneof expr_struct {
    Literal literal = 2;
    BoundReference bound = 3;
    MathExpr add = 4;
    MathExpr subtract = 5;
    MathExpr multiply = 6;
    MathExpr divide = 7;
    Cast cast = 8;
    BinaryExpr eq = 9;
    BinaryExpr neq = 10;
    ...
    BinaryExpr array_append = 58;
    ArrayInsert array_insert = 59;
  }
}

And in the runtime it looks like this:

def exprToProtoInternal(expr: Expression, inputs: Seq[Attribute]): Option[Expr] = {
  ...
  expr match {
        case a @ Alias(_, _) =>
          val r = exprToProtoInternal(a.child, inputs)
          if (r.isEmpty) {
            withInfo(expr, a.child)
          }

        case cast @ Cast(_: Literal, dataType, _, _) =>
          // This can happen after promoting decimal precisions
          val value = cast.eval()
          exprToProtoInternal(Literal(value, dataType), inputs)

        case UnaryExpression(child) if expr.prettyName == "trycast" =>
          val timeZoneId = SQLConf.get.sessionLocalTimeZone
          handleCast(child, inputs, expr.dataType, Some(timeZoneId), CometEvalMode.TRY)

        case c @ Cast(child, dt, timeZoneId, _) =>
          handleCast(child, inputs, dt, timeZoneId, evalMode(c))

        case add @ Add(left, right, _) if supportedDataType(left.dataType) =>
          createMathExpression(
            left,
            right,
            inputs,
            add.dataType,
            getFailOnError(add),
            (builder, mathExpr) => builder.setAdd(mathExpr))
        ...
    }
}

Using protobuf in this case provides a single location where all messages are defined and allows us to use code generation with protoc to automatically create serializable Java classes and deserializable Rust structs!

Fallback mechanics

Similar to Photon and Gluten, Comet has a fallback mechanism: if it cannot translate the entire plan to native code, it will translate as much as possible. After that, it applies ColumnarToRow (using Comet's own implementation that leverages Apache Arrow) and execution continues using the Spark JVM engine. Comet follows the same policy as Photon and does not switch back to native execution after fallback due to the overhead costs associated with columnar-to-row and row-to-columnar translations.

Performance

On the TPC-H Comet provides about x2-4 speedup over Spark:

On my own benchmarks in the task of wide aggregations (feature engineering for ML cases) speedup was about x3-x5.

The first step was to define the protobuf message and write an additional case for it in Comet's exprToProtoInternal function:

message ArrayInsert {
  Expr src_array_expr = 1;
  Expr pos_expr = 2;
  Expr item_expr = 3;
  bool legacy_negative_index = 4;
}

case expr if expr.prettyName == "array_insert" =>
  val srcExprProto = exprToProto(expr.children(0), inputs, binding)
  val posExprProto = exprToProto(expr.children(1), inputs, binding)
  val itemExprProto = exprToProto(expr.children(2), inputs, binding)
  val legacyNegativeIndex =
    SQLConf.get.getConfString("spark.sql.legacy.negativeIndexInArrayInsert").toBoolean
  if (srcExprProto.isDefined && posExprProto.isDefined && itemExprProto.isDefined) {
    val arrayInsertBuilder = ExprOuterClass.ArrayInsert
      .newBuilder()
      .setSrcArrayExpr(srcExprProto.get)
      .setPosExpr(posExprProto.get)
      .setItemExpr(itemExprProto.get)
      .setLegacyNegativeIndex(legacyNegativeIndex)

    Some(
      ExprOuterClass.Expr
        .newBuilder()
        .setArrayInsert(arrayInsertBuilder)
        .build())
  } else {
  // boring fallback logic is here
  }

NOTE: It is interesting that even new functionality added to Spark 3.4 already has its own "legacy mode." I can imagine how challenging it is to maintain Spark: while you can change literally anything, you must first introduce legacy flags before modifying any existing logic!

The first step was to define a top-level struct that takes all the information from the proto-message and constructs its own state based on this information:

#[derive(Debug, Hash)]
pub struct ArrayInsert {
    src_array_expr: Arc<dyn PhysicalExpr>,
    pos_expr: Arc<dyn PhysicalExpr>,
    item_expr: Arc<dyn PhysicalExpr>,
    legacy_negative_index: bool,
}

impl ArrayInsert {
    pub fn new(
        src_array_expr: Arc<dyn PhysicalExpr>,
        pos_expr: Arc<dyn PhysicalExpr>,
        item_expr: Arc<dyn PhysicalExpr>,
        legacy_negative_index: bool,
    ) -> Self {
        Self {
            src_array_expr,
            pos_expr,
            item_expr,
            legacy_negative_index,
        }
    }
}

To create a Datafusion expression, at minimum the following traits need to be implemented for the defined struct:

• PhysicalExpr
• Display
• PartialEq

impl PhysicalExpr for ArrayInsert {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn data_type(&self, input_schema: &Schema) -> DataFusionResult<DataType> {
      // Datatype checking logic is here
    }

    fn nullable(&self, input_schema: &Schema) -> DataFusionResult<bool> {
        self.src_array_expr.nullable(input_schema)
    }

    fn evaluate(&self, batch: &RecordBatch) -> DataFusionResult<ColumnarValue> {
      // Implementation is here
    }

    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
        vec![&self.src_array_expr, &self.pos_expr, &self.item_expr]
    }

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn PhysicalExpr>>,
    ) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
      // Not important code is here
    }

    fn dyn_hash(&self, _state: &mut dyn Hasher) {
      // Not important code is here
    }
}

impl Display for ArrayInsert {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "ArrayInsert [array: {:?}, pos: {:?}, item: {:?}]",
            self.src_array_expr, self.pos_expr, self.item_expr
        )
    }
}

impl PartialEq<dyn Any> for ArrayInsert {
    fn eq(&self, other: &dyn Any) -> bool {
        down_cast_any_ref(other)
            .downcast_ref::<Self>()
            .map(|x| {
                self.src_array_expr.eq(&x.src_array_expr)
                    && self.pos_expr.eq(&x.pos_expr)
                    && self.item_expr.eq(&x.item_expr)
                    && self.legacy_negative_index.eq(&x.legacy_negative_index)
            })
            .unwrap_or(false)
    }
}

On the native part the first step is to add a rule for the kind of messages I defined:

ExprStruct::ArrayInsert(expr) => {
    let src_array_expr = self.create_expr(
        expr.src_array_expr.as_ref().unwrap(),
        Arc::clone(&input_schema),
    )?;
    let pos_expr =
        self.create_expr(expr.pos_expr.as_ref().unwrap(), Arc::clone(&input_schema))?;
    let item_expr =
        self.create_expr(expr.item_expr.as_ref().unwrap(), Arc::clone(&input_schema))?;
    Ok(Arc::new(ArrayInsert::new(
        src_array_expr,
        pos_expr,
        item_expr,
        expr.legacy_negative_index,
    )))
}

This part was the easiest one from all of my work.

NOTE: My implementation is partially based on the array_append implementation in Datafusion, which is distributed under the same Apache 2.0 license.

Following the pattern established in other parts of Comet's codebase, I decided to implement the logic in a separate small function rather than within the body of PhysicalExpr. Throughout my development process, this decision proved beneficial, as it allowed me to easily debug my work using rust-gdb and simple tests for corner cases. In contrast, debugging the logic within PhysicalExpr would have been significantly more challenoging.

Unlike Spark, Apache Arrow (which underlies DataFusion) has two distinct List implementations: List for integer (i32) indices and LargeList for long (i64) indices. Whether Comet should support LargeList remains an open question, since Spark's maximum array length is limited to java.lang.Integer.MAX_VALUE - 15, which aligns with the List implementation rather than LargeList. However, this presents an interesting topic for future research and contribution, which is why my current array_insert implementation is generic in terms of supported arrays.

fn array_insert<O: OffsetSizeTrait>(
    list_array: &GenericListArray<O>,
    items_array: &ArrayRef,
    pos_array: &ArrayRef,
    legacy_mode: bool,
) -> DataFusionResult<ColumnarValue> {}

Arrays in Apache Arrow

In Apache Arrow arrays are represented in the following way (information from the documentaion):

                               ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─
                                                       ┌ ─ ─ ─ ─ ─ ─ ┐    │
 ┌─────────────┐  ┌───────┐    │     ┌───┐   ┌───┐       ┌───┐ ┌───┐
 │   [A,B,C]   │  │ (0,3) │          │ 1 │   │ 0 │     │ │ 1 │ │ A │ │ 0  │
 ├─────────────┤  ├───────┤    │     ├───┤   ├───┤       ├───┤ ├───┤
 │      []     │  │ (3,3) │          │ 1 │   │ 3 │     │ │ 1 │ │ B │ │ 1  │
 ├─────────────┤  ├───────┤    │     ├───┤   ├───┤       ├───┤ ├───┤
 │    NULL     │  │ (3,4) │          │ 0 │   │ 3 │     │ │ 1 │ │ C │ │ 2  │
 ├─────────────┤  ├───────┤    │     ├───┤   ├───┤       ├───┤ ├───┤
 │     [D]     │  │ (4,5) │          │ 1 │   │ 4 │     │ │ ? │ │ ? │ │ 3  │
 ├─────────────┤  ├───────┤    │     ├───┤   ├───┤       ├───┤ ├───┤
 │  [NULL, F]  │  │ (5,7) │          │ 1 │   │ 5 │     │ │ 1 │ │ D │ │ 4  │
 └─────────────┘  └───────┘    │     └───┘   ├───┤       ├───┤ ├───┤
                                             │ 7 │     │ │ 0 │ │ ? │ │ 5  │
                               │  Validity   └───┘       ├───┤ ├───┤
    Logical       Logical         (nulls)   Offsets    │ │ 1 │ │ F │ │ 6  │
     Values       Offsets      │                         └───┘ └───┘
                                                       │    Values   │    │
                (offsets[i],   │   ListArray               (Array)
               offsets[i+1])                           └ ─ ─ ─ ─ ─ ─ ┘    │
                               └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─

As we can observe, an Arrow array (or List) is composed of three components: an offsets vector, a values vector, and a nulls vector. The values vector is implemented using an Arrow Buffer object, which represents a continuous region of memory. The offsets are used to determine which elements correspond to specific indices in the logical array.

let values = list_array.values();
let offsets = list_array.offsets();
let values_data = values.to_data();
let item_data = items_array.to_data();
let new_capacity = Capacities::Array(values_data.len() + item_data.len());

let mut mutable_values =
    MutableArrayData::with_capacities(vec![&values_data, &item_data], true, new_capacity);

let mut new_offsets = vec![O::usize_as(0)];
let mut new_nulls = Vec::<bool>::with_capacity(list_array.len());

for (row_index, offset_window) in offsets.windows(2).enumerate() {
    let pos = pos_data.values()[row_index];
    let start = offset_window[0].as_usize();
    let end = offset_window[1].as_usize();
    let is_item_null = items_array.is_null(row_index);

    ...
    let new_array_len = std::cmp::max(end - start + 1, corrected_pos);
    if (start + corrected_pos) <= end {
        mutable_values.extend(0, start, start + corrected_pos);
        mutable_values.extend(1, row_index, row_index + 1);
        mutable_values.extend(0, start + corrected_pos, end);
        new_offsets.push(new_offsets[row_index] + O::usize_as(new_array_len));
    } else {
        ...
    }

As primarily a Python developer, I started by defining minimal, simple tests for the native component to address recurring errors like "attempt to subtract with overflow". In Python, I rarely use debugging tools since it's much easier to halt interpretation (for example, in IPython) and analyze the program state - similar to Lisp's REPL-driven development from the 1980s. However, with Rust, this approach isn't possible. Instead, I used rust-gdb for debugging. To my surprise, I found the tool to be remarkably user-friendly and intuitive!

My workflow:

cargo test test_array_insert --no-run
   Compiling datafusion-comet-spark-expr v0.5.0 (/var/home/sem/github/datafusion-comet/native/spark-expr)
   Compiling datafusion-comet v0.5.0 (/var/home/sem/github/datafusion-comet/native/core)
    Finished `test` profile [unoptimized + debuginfo] target(s) in 18.26s
  Executable unittests src/lib.rs (target/debug/deps/comet-db43523bbc882879)
  Executable unittests src/lib.rs (target/debug/deps/datafusion_comet_proto-0a7e2f643944299d)
  Executable unittests src/lib.rs (target/debug/deps/datafusion_comet_spark_expr-0454198feed9fe39)

rust-gdb target/debug/deps/datafusion_comet_spark_expr-0454198feed9fe39

After entering rust-gdb, you can debug your Rust program. For example, typing b list.rs:621 will create a breakpoint on line 621 of the file list.rs. Typing run will execute the binary until it reaches the next breakpoint. At the breakpoint, you can use the print command to view the values of various variables. I strongly recommend anyone interested in Rust to try rust-gdb!

The next step was to define end2end tests of the Comet+Spark:

test("ArrayInsert") {
  assume(isSpark34Plus)
  Seq(true, false).foreach(dictionaryEnabled =>
    withTempDir { dir =>
      val path = new Path(dir.toURI.toString, "test.parquet")
      makeParquetFileAllTypes(path, dictionaryEnabled, 10000)
      val df = spark.read
        .parquet(path.toString)
        .withColumn("arr", array(col("_4"), lit(null), col("_4")))
        .withColumn("arrInsertResult", expr("array_insert(arr, 1, 1)"))
        .withColumn("arrInsertNegativeIndexResult", expr("array_insert(arr, -1, 1)"))
        .withColumn("arrPosGreaterThanSize", expr("array_insert(arr, 8, 1)"))
        .withColumn("arrNegPosGreaterThanSize", expr("array_insert(arr, -8, 1)"))
        .withColumn("arrInsertNone", expr("array_insert(arr, 1, null)"))
      checkSparkAnswerAndOperator(df.select("arrInsertResult"))
      checkSparkAnswerAndOperator(df.select("arrInsertNegativeIndexResult"))
      checkSparkAnswerAndOperator(df.select("arrPosGreaterThanSize"))
      checkSparkAnswerAndOperator(df.select("arrNegPosGreaterThanSize"))
      checkSparkAnswerAndOperator(df.select("arrInsertNone"))
    })
}