[query] Add Hail <-> Arrow schema conversion

Any proposed open columnar format that we intend to evolve hail to make
use of can be described with an Arrow schema, making Arrow's types
a common denominator for future work in this direction. This change adds
the mapping between Hail physical types and that schema, in both
directions, in a new `is.hail.types.arrow` package.

`ArrowSchemaConversion` converts a PStruct to an Arrow Schema and back.
There's a 1:1 mapping between Arrow's nullability and PType
requiredness, except at the root struct, which Arrow gives no
nullability and which always reads back as required. Reading takes the
reference genomes in scope: it resolves each locus's genome by name and
checks that the stored contigs match it.

Mapping:
- bool, int32, int64, float32, float64: Arrow primitives (signed ints)
- str, binary: utf8, binary (32-bit offsets)
- array<T>: list<element: T>
- set<T>: list<element: T> with `hail.set`
- dict<K, V>: map<entries: struct<key: K, value: V>>, keysSorted=true,
  and keys are always non-nullable
- struct: struct
- tuple: struct with fields "0", "1", ... plus `hail.tuple`
- call: int32 with `hail.call`, metadata {"encoding":1}
- locus: int64 packed as (contig_index << 32) | position, with
  `hail.locus` storing the reference genome name and ordered contigs.
  Every locus type reads back as PPackedLocus, whose representation is
  this storage
- interval<T>: struct<start, end, includes_start, includes_end> with
  `hail.interval`
- ndarray: Arrow's canonical `arrow.variable_shape_tensor`
  (struct<data: list<T>, shape: fixed_size_list<int32>[ndim]>), which
  arrow-java does not provide

Extension metadata is a JSON object. Arrow -> Hail accepts only what
Hail -> Arrow produces and rejects anything else with a clear error.
That includes unregistered extension types, dictionary encoding (on any
field, including structural children such as map entries), child fields
on primitive types and duplicate field names. Types that cannot be
stored (RNG state, void) are also rejected. Extension types do not
create Arrow vectors, since the columnar writer will drive parquet-java
directly.

Declare arrow-vector 18.3.0 (Spark 4.1.2's version) as a compile-only
dependency now that Hail uses it directly.

Author: chrisvittal

hail/hail/package.mill

@@ -73,6 +73,7 @@ object `package` extends HailScalaModule:
       `breeze` :: "2.1.0",
       `blas` :: "3.0.4",
       `arpack` :: "3.0.4",
+      `arrow-vector` :: "18.3.0",
       `avro` :: "1.12.1",
       `commons-codec` :: "1.19.0",
       `commons-io` :: "2.21.0",

hail/hail/src/is/hail/types/arrow/ArrowSchemaConversion.scala

@@ -0,0 +1,251 @@
+package is.hail.types.arrow
+
+import is.hail.types.physical._
+import is.hail.utils.fatal
+import is.hail.variant.ReferenceGenomeEnv
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.arrow.vector.types.FloatingPointPrecision
+import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema}
+
+/** Conversion between Hail physical types and the Arrow schema of the columnar format.
+  *
+  * Arrow nullability carries Hail requiredness at every level below the root, with one exception:
+  * dict keys are always non-nullable, so a dict whose key type is optional reads back with a
+  * required key type. The root struct has no Arrow nullability, so it always reads back as
+  * required.
+  *
+  * Every locus type reads back as a [[PPackedLocus]], whose representation is the Arrow storage.
+  *
+  * Arrow -> Hail accepts exactly the Arrow types that Hail -> Arrow produces.
+  */
+object ArrowSchemaConversion {
+  ExtensionTypes.register()
+
+  val ElementFieldName: String = "element"
+  val EntriesFieldName: String = "entries"
+  val KeyFieldName: String = "key"
+  val ValueFieldName: String = "value"
+
+  private val Int32 = new ArrowType.Int(32, true)
+  private val Int64 = new ArrowType.Int(64, true)
+  private val Float32 = new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)
+  private val Float64 = new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)
+
+  private val IntervalFieldNames =
+    IndexedSeq("start", "end", "includes_start", "includes_end")
+
+  private val TensorFieldNames = IndexedSeq("data", "shape")
+
+  def toArrowSchema(t: PStruct): Schema =
+    new Schema(t.fields.map(f => toArrowField(f.name, f.typ)).asJava)
+
+  def toArrowField(name: String, t: PType): Field = {
+    def field(arrowType: ArrowType, children: Field*): Field =
+      mkField(name, !t.required, arrowType, children: _*)
+
+    def child(name: String, t: PType): Field =
+      toArrowField(name, t)
+
+    t match {
+      case _: PBoolean => field(ArrowType.Bool.INSTANCE)
+      case _: PInt32 => field(Int32)
+      case _: PInt64 => field(Int64)
+      case _: PFloat32 => field(Float32)
+      case _: PFloat64 => field(Float64)
+      case _: PString => field(ArrowType.Utf8.INSTANCE)
+      case _: PBinary => field(ArrowType.Binary.INSTANCE)
+      case _: PCall => field(CallExtensionType.Current)
+      case t: PLocus => field(new LocusExtensionType(t.rg.name, t.rg.contigs))
+      case t: PInterval =>
+        field(
+          IntervalExtensionType,
+          child(IntervalFieldNames(0), t.pointType),
+          child(IntervalFieldNames(1), t.pointType),
+          mkField(IntervalFieldNames(2), nullable = false, ArrowType.Bool.INSTANCE),
+          mkField(IntervalFieldNames(3), nullable = false, ArrowType.Bool.INSTANCE),
+        )
+      case t: PNDArray =>
+        field(
+          VariableShapeTensorExtensionType,
+          mkField(
+            TensorFieldNames(0),
+            nullable = false,
+            ArrowType.List.INSTANCE,
+            child(ElementFieldName, t.elementType),
+          ),
+          mkField(
+            TensorFieldNames(1),
+            nullable = false,
+            new ArrowType.FixedSizeList(t.nDims),
+            mkField(ElementFieldName, nullable = false, Int32),
+          ),
+        )
+      case t: PSet => field(SetExtensionType, child(ElementFieldName, t.elementType))
+      case t: PDict =>
+        field(
+          new ArrowType.Map( /*keysSorted=*/ true),
+          mkField(
+            EntriesFieldName,
+            nullable = false,
+            ArrowType.Struct.INSTANCE,
+            child(KeyFieldName, t.keyType.setRequired(true)),
+            child(ValueFieldName, t.valueType),
+          ),
+        )
+      case t: PArray => field(ArrowType.List.INSTANCE, child(ElementFieldName, t.elementType))
+      case t: PTuple => field(TupleExtensionType, t.fields.map(f => child(f.name, f.typ)): _*)
+      case t: PStruct =>
+        field(ArrowType.Struct.INSTANCE, t.fields.map(f => child(f.name, f.typ)): _*)
+      case _ =>
+        fatal(s"field '$name': Hail type ${t.virtualType} cannot be stored in the columnar format")
+    }
+  }
+
+  def fromArrowSchema(schema: Schema, rgs: ReferenceGenomeEnv): PCanonicalStruct =
+    PCanonicalStruct(structFields(schema.getFields.asScala.toIndexedSeq, rgs), required = true)
+
+  def fromArrowField(field: Field, rgs: ReferenceGenomeEnv): PType = {
+    val name = field.getName
+    val required = !field.isNullable
+    val children = field.getChildren.asScala.toIndexedSeq
+
+    def unsupported(reason: String): Nothing =
+      fatal(
+        s"field '$name' of Arrow type ${field.getType} is not part of the columnar format: $reason"
+      )
+
+    def checkChildren(expected: IndexedSeq[String]): Unit =
+      if (children.map(_.getName) != expected)
+        unsupported(
+          s"expected child fields ${expected.mkString("[", ", ", "]")}, found ${children.map(_.getName).mkString("[", ", ", "]")}"
+        )
+
+    def onlyChild(): Field =
+      if (children.length == 1) children.head
+      else unsupported(s"expected 1 child field, found ${children.length}")
+
+    def convert(f: Field): PType = fromArrowField(f, rgs)
+
+    def requiredChild(f: Field): Field =
+      if (f.isNullable) unsupported(s"child field '${f.getName}' must be non-nullable")
+      else f
+
+    def leaf(t: PType): PType =
+      if (children.isEmpty) t
+      else unsupported(s"expected no child fields, found ${children.length}")
+
+    // Rejects what the Arrow type of `f` does not show: dictionary encoding, and an extension type
+    // that is not registered, which reads as its storage type. Child fields that are not
+    // converted by `fromArrowField` need this check too.
+    def checkEncoding(f: Field): Unit = {
+      if (f.getDictionary != null)
+        unsupported(s"field '${f.getName}' is dictionary encoded")
+      if (!f.getType.isInstanceOf[ArrowType.ExtensionType])
+        Option(f.getMetadata.get(ArrowType.ExtensionType.EXTENSION_METADATA_KEY_NAME)).foreach {
+          n => unsupported(s"field '${f.getName}' has unregistered extension type '$n'")
+        }
+    }
+
+    checkEncoding(field)
+
+    field.getType match {
+      case _: CallExtensionType if field.getType == CallExtensionType.Current =>
+        leaf(PCanonicalCall(required))
+      case t: CallExtensionType =>
+        unsupported(s"unknown call encoding ${t.encoding}")
+      case t: LocusExtensionType =>
+        val rg = rgs.get(t.referenceGenome).getOrElse(
+          fatal(s"field '$name': reference genome '${t.referenceGenome}' is not defined")
+        )
+        if (rg.contigs != t.contigs)
+          fatal(
+            s"field '$name': the stored contigs of reference genome '${t.referenceGenome}' " +
+              "differ from the contigs of the defined reference genome"
+          )
+        leaf(PPackedLocus(rg, required))
+      case SetExtensionType =>
+        PCanonicalSet(convert(onlyChild()), required)
+      case IntervalExtensionType =>
+        checkChildren(IntervalFieldNames)
+        val IndexedSeq(start, end, includesStart, includesEnd) = children
+        val pointType = convert(start)
+        if (convert(end) != pointType)
+          unsupported("interval 'start' and 'end' have different types")
+        for (f <- Seq(includesStart, includesEnd)) {
+          checkEncoding(f)
+          if (f.isNullable || f.getType != ArrowType.Bool.INSTANCE || !f.getChildren.isEmpty)
+            unsupported(s"interval '${f.getName}' must be a non-nullable bool")
+        }
+        PCanonicalInterval(pointType, required)
+      case TupleExtensionType =>
+        val fields = children.map { c =>
+          val idx = c.getName.toIntOption.filter(i => i >= 0 && i.toString == c.getName).getOrElse(
+            unsupported(s"tuple field name '${c.getName}' is not a non-negative integer")
+          )
+          PTupleField(idx, convert(c))
+        }
+        if (fields.map(_.index).distinct.length != fields.length)
+          unsupported("tuple field names are not distinct")
+        PCanonicalTuple(fields, required)
+      case VariableShapeTensorExtensionType =>
+        checkChildren(TensorFieldNames)
+        val IndexedSeq(data, shape) = children
+        checkEncoding(data)
+        checkEncoding(shape)
+        val nDims = shape.getType match {
+          case t: ArrowType.FixedSizeList => t.getListSize
+          case t => unsupported(s"tensor 'shape' must be a fixed_size_list, found $t")
+        }
+        requiredChild(shape): Unit
+        if (shape.getChildren.size != 1 || convert(shape.getChildren.get(0)) != PInt32(true))
+          unsupported("tensor 'shape' must contain non-nullable int32")
+        if (requiredChild(data).getType != ArrowType.List.INSTANCE || data.getChildren.size != 1)
+          unsupported("tensor 'data' must be a list")
+        val elementType = convert(requiredChild(data.getChildren.get(0)))
+        if (elementType.containsPointers)
+          unsupported(s"tensor element type ${elementType.virtualType} is not supported")
+        PCanonicalNDArray(elementType, nDims, required)
+      case t: ArrowType.ExtensionType =>
+        unsupported(s"unknown extension type '${t.extensionName}'")
+      case ArrowType.Bool.INSTANCE => leaf(PBoolean(required))
+      case Int32 => leaf(PInt32(required))
+      case Int64 => leaf(PInt64(required))
+      case Float32 => leaf(PFloat32(required))
+      case Float64 => leaf(PFloat64(required))
+      case ArrowType.Utf8.INSTANCE => leaf(PCanonicalString(required))
+      case ArrowType.Binary.INSTANCE => leaf(PCanonicalBinary(required))
+      case ArrowType.List.INSTANCE => PCanonicalArray(convert(onlyChild()), required)
+      case t: ArrowType.Map =>
+        if (!t.getKeysSorted) unsupported("map keys must be sorted")
+        val entries = requiredChild(onlyChild())
+        checkEncoding(entries)
+        if (entries.getType != ArrowType.Struct.INSTANCE || entries.getChildren.size != 2)
+          unsupported("map entries must be a struct with two fields")
+        val key = requiredChild(entries.getChildren.get(0))
+        PCanonicalDict(convert(key), convert(entries.getChildren.get(1)), required)
+      case ArrowType.Struct.INSTANCE =>
+        PCanonicalStruct(structFields(children, rgs), required)
+      case _ => unsupported("unsupported type")
+    }
+  }
+
+  private def structFields(fields: IndexedSeq[Field], rgs: ReferenceGenomeEnv)
+    : IndexedSeq[PField] = {
+    val duplicates = fields.map(_.getName).groupBy(identity).collect {
+      case (name, names) if names.length > 1 => name
+    }
+    if (duplicates.nonEmpty)
+      fatal(
+        s"duplicate field names are not part of the columnar format: ${duplicates.mkString(", ")}"
+      )
+    fields.zipWithIndex.map { case (f, i) =>
+      PField(f.getName, fromArrowField(f, rgs), i)
+    }
+  }
+
+  private def mkField(name: String, nullable: Boolean, arrowType: ArrowType, children: Field*)
+    : Field =
+    new Field(name, new FieldType(nullable, arrowType, /*dictionary=*/ null), children.asJava)
+}