Simple Scala DSL for constructing Apache Arrow schemas.

arrow-schema-dsl.scala

package example

import scala.collection.JavaConverters._
import scala.language.implicitConversions

import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema}

trait FieldBuilder {
  def named(name: String): Field

  def ! : this.type = {
    nullable = false
    this
  }

  def ? : this.type = {
    nullable = true
    this
  }

  protected def isNullable: Boolean = nullable

  private var nullable: Boolean = true
}

trait LowPriorityImplicits {
  implicit class FieldDSL(name: String) {
    def ~ (builder: FieldBuilder): Field = builder.named(name)
  }

  implicit class PrimitiveFieldBuilder(arrowType: ArrowType.PrimitiveType) extends FieldBuilder {
    override def named(name: String): Field = {
      val fieldType = new FieldType(isNullable, arrowType, null)
      new Field(name, fieldType, null)
    }
  }

  class ListFieldBuilder(elementTypeBuilder: FieldBuilder) extends FieldBuilder {
    override def named(name: String): Field = {
      val fieldType = new FieldType(isNullable, ArrowType.List.INSTANCE, null)
      new Field(name, fieldType, Seq(elementTypeBuilder.named(null)).asJava)
    }
  }

  class StructFieldBuilder(fields: Seq[Field]) extends FieldBuilder {
    override def named(name: String): Field = {
      val fieldType = new FieldType(isNullable, ArrowType.Struct.INSTANCE, null)
      new Field(name, fieldType, fields.asJava)
    }
  }
}

object dsl extends LowPriorityImplicits {
  def ArrowSchema(fields: Field*): Schema = new Schema(fields.asJava)

  def ArrowStruct(fields: Field*): StructFieldBuilder = new StructFieldBuilder(fields)

  def ArrowList(elementType: FieldBuilder): ListFieldBuilder = new ListFieldBuilder(elementType)
}

dsl-example.scala

import scala.collection.JavaConverters._

import example.dsl._
import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema}

// Here we are using the address book example schema used in the Twitter Parquet blog post.
// https://blog.twitter.com/engineering/en_us/a/2013/dremel-made-simple-with-parquet.html

// Using the Scala DSL
ArrowSchema(
  // A non-nullable string field.
  "owner" ~ ArrowType.Utf8.INSTANCE.!,
  // A non-nullable list field whose elements are non-nullable strings.
  "ownerPhoneNumbers" ~ ArrowList(ArrowType.Utf8.INSTANCE.!).!,
  // A non-nullable list field whose elements are non-nullable structs.
  "contacts" ~ ArrowList(
    // A non-nullable struct type.
    ArrowStruct(
      // A non-nullable string field.
      "name" ~ ArrowType.Utf8.INSTANCE.!,
      // A nullable string field.
      "phoneNumber" ~ ArrowType.Utf8.INSTANCE.?
    ).!
  ).!
)

// Using the Apache Arrow Java API in Scala
new Schema(
  Seq(
    // A non-nullable string field.
    new Field("owner", new FieldType(false, ArrowType.Utf8.INSTANCE, null), null),
    // A non-nullable list field whose elements are non-nullable strings.
    new Field(
      "ownerPhoneNumbers",
      new FieldType(false, ArrowType.List.INSTANCE, null),
      Seq(new Field(null, new FieldType(false, ArrowType.Utf8.INSTANCE, null), null)).asJava
    ),
    // A non-nullable list field whose elements are non-nullable structs.
    new Field(
      "contacts",
      new FieldType(false, ArrowType.List.INSTANCE, null),
      Seq(
        // A non-nullable struct
        new Field(
          null,
          new FieldType(false, ArrowType.Struct.INSTANCE, null),
          Seq(
            // A non-nullable string field.
            new Field("name", new FieldType(false, ArrowType.Utf8.INSTANCE, null), null),
            // A nullable string field.
            new Field("phoneNumber", FieldType.nullable(ArrowType.Utf8.INSTANCE), null)
          ).asJava
        )
      ).asJava
    )
  ).asJava
)

@advancedxy Refactored it a little bit and removed the `ArrowType->FieldBuilder` implicit conversion. Hopefully, it's easier to reason about this time.