diff --git a/hail/hail/package.mill b/hail/hail/package.mill index b9fc516683c..fd8b4fcc27f 100644 --- a/hail/hail/package.mill +++ b/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", diff --git a/hail/hail/src/is/hail/types/arrow/ArrowSchemaConversion.scala b/hail/hail/src/is/hail/types/arrow/ArrowSchemaConversion.scala new file mode 100644 index 00000000000..58db5c2ad5e --- /dev/null +++ b/hail/hail/src/is/hail/types/arrow/ArrowSchemaConversion.scala @@ -0,0 +1,240 @@ +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, with one exception: dict keys are + * always non-nullable, so a dict whose key type is optional reads back with a required key type. + * + * Both locus representations write the same Arrow type, whose storage is a packed locus, and a + * locus always reads back as a [[PPackedLocus]]. + * + * 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 if !t.elementType.isInstanceOf[PLocus] => + 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, references: ReferenceGenomeEnv): PCanonicalStruct = + PCanonicalStruct( + structFields(schema.getFields.asScala.toIndexedSeq, references), + required = true, + ) + + def fromArrowField(field: Field, references: 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, references) + + def requiredChild(f: Field): Field = + if (f.isNullable) unsupported(s"child field '${f.getName}' must be non-nullable") + else f + + if (field.getDictionary != null) + unsupported("dictionary encoding is not supported") + + field.getType match { + case _: CallExtensionType if field.getType == CallExtensionType.Current => + PCanonicalCall(required) + case t: CallExtensionType => + unsupported(s"unknown call encoding ${t.encoding}") + case t: LocusExtensionType => + val rg = references.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" + ) + 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)) + 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 + 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) + 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 || elementType.isInstanceOf[PLocus]) + 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 _ + if field.getMetadata.containsKey(ArrowType.ExtensionType.EXTENSION_METADATA_KEY_NAME) => + unsupported( + s"unregistered extension type '${field.getMetadata.get(ArrowType.ExtensionType.EXTENSION_METADATA_KEY_NAME)}'" + ) + case ArrowType.Bool.INSTANCE => PBoolean(required) + case Int32 => PInt32(required) + case Int64 => PInt64(required) + case Float32 => PFloat32(required) + case Float64 => PFloat64(required) + case ArrowType.Utf8.INSTANCE => PCanonicalString(required) + case ArrowType.Binary.INSTANCE => 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()) + 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, references), required) + case _ => unsupported("unsupported type") + } + } + + private def structFields( + fields: IndexedSeq[Field], + references: 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, references), i) + } + } + + private def mkField(name: String, nullable: Boolean, arrowType: ArrowType, children: Field*) + : Field = + new Field(name, new FieldType(nullable, arrowType, /*dictionary=*/ null), children.asJava) +} diff --git a/hail/hail/src/is/hail/types/arrow/ExtensionTypes.scala b/hail/hail/src/is/hail/types/arrow/ExtensionTypes.scala new file mode 100644 index 00000000000..cea961a73c4 --- /dev/null +++ b/hail/hail/src/is/hail/types/arrow/ExtensionTypes.scala @@ -0,0 +1,186 @@ +package is.hail.types.arrow + +import org.apache.arrow.memory.BufferAllocator +import org.apache.arrow.vector.FieldVector +import org.apache.arrow.vector.extension.InvalidExtensionMetadataException +import org.apache.arrow.vector.types.pojo.{ArrowType, ExtensionTypeRegistry, FieldType} +import org.json4s._ +import org.json4s.jackson.JsonMethods + +/** Arrow extension types used by the columnar format. + * + * Extension metadata is a JSON object. Readers ignore keys they do not recognise, so keys can be + * added without breaking existing readers. + * + * The columnar format is written without Arrow vectors, so none of these types can create one. + */ +object ExtensionTypes { + private lazy val registered: Unit = + Seq[ArrowType.ExtensionType]( + CallExtensionType.Current, + LocusExtensionType.Prototype, + SetExtensionType, + IntervalExtensionType, + TupleExtensionType, + VariableShapeTensorExtensionType, + ).foreach(ExtensionTypeRegistry.register) + + /** Registers every extension type with Arrow's global registry, so that schemas read by Arrow + * contain extension types rather than their storage types. Idempotent. + */ + def register(): Unit = registered +} + +sealed abstract class HailExtensionType extends ArrowType.ExtensionType { + protected def checkStorageType(storageType: ArrowType): Unit = + if (storageType != this.storageType()) + throw new InvalidExtensionMetadataException( + s"${extensionName()}: expected storage type ${this.storageType()}, found $storageType" + ) + + protected def parseMetadata(serializedData: String): JObject = + try + JsonMethods.parse(serializedData) match { + case o: JObject => o + case _ => throw new InvalidExtensionMetadataException( + s"${extensionName()}: metadata is not a JSON object: $serializedData" + ) + } + catch { + case e: com.fasterxml.jackson.core.JsonProcessingException => + throw new InvalidExtensionMetadataException( + s"${extensionName()}: metadata is not valid JSON: $serializedData", + e, + ) + } + + override def getNewVector(name: String, fieldType: FieldType, allocator: BufferAllocator) + : FieldVector = + throw new UnsupportedOperationException( + s"Arrow vectors are not implemented for extension type ${extensionName()}" + ) +} + +/** An extension type with no parameters. Everything needed to reconstruct the Hail type is in the + * storage type and the field's children. + */ +sealed abstract class MarkerExtensionType(name: String, storage: ArrowType) + extends HailExtensionType { + override def extensionName(): String = name + + override def storageType(): ArrowType = storage + + override def extensionEquals(other: ArrowType.ExtensionType): Boolean = other eq this + + override def serialize(): String = "{}" + + override def deserialize(storageType: ArrowType, serializedData: String): ArrowType = { + checkStorageType(storageType) + parseMetadata(serializedData): Unit + this + } +} + +/** `list` whose elements are sorted and distinct. */ +object SetExtensionType extends MarkerExtensionType("hail.set", ArrowType.List.INSTANCE) + +/** `struct`. */ +object IntervalExtensionType extends MarkerExtensionType("hail.interval", ArrowType.Struct.INSTANCE) + +/** A struct whose field names are the tuple indices `"0"`, `"1"`, .... */ +object TupleExtensionType extends MarkerExtensionType("hail.tuple", ArrowType.Struct.INSTANCE) + +/** Arrow's canonical variable shape tensor: `struct, shape: + * fixed_size_list[ndim]>`, row-major. arrow-java does not provide it. + */ +object VariableShapeTensorExtensionType + extends MarkerExtensionType("arrow.variable_shape_tensor", ArrowType.Struct.INSTANCE) + +object CallExtensionType { + val Name: String = "hail.call" + + /** Hail's packed call representation, as produced by `is.hail.variant.Call`. */ + val Current: CallExtensionType = new CallExtensionType(1) +} + +/** A Hail call packed into a signed `int32`. `encoding` identifies the bit layout. */ +final class CallExtensionType(val encoding: Int) extends HailExtensionType { + override def extensionName(): String = CallExtensionType.Name + + override def storageType(): ArrowType = new ArrowType.Int(32, true) + + override def extensionEquals(other: ArrowType.ExtensionType): Boolean = + other match { + case o: CallExtensionType => o.encoding == encoding + case _ => false + } + + override def serialize(): String = + JsonMethods.compact(JObject("encoding" -> JInt(encoding))) + + override def deserialize(storageType: ArrowType, serializedData: String): ArrowType = { + checkStorageType(storageType) + parseMetadata(serializedData) \ "encoding" match { + case JInt(e) if e.isValidInt => new CallExtensionType(e.toInt) + case _ => throw new InvalidExtensionMetadataException( + s"$extensionName: missing or invalid 'encoding': $serializedData" + ) + } + } + + override def toString: String = s"CallExtensionType(encoding=$encoding)" +} + +object LocusExtensionType { + val Name: String = "hail.locus" + + private[arrow] val Prototype: LocusExtensionType = new LocusExtensionType("", IndexedSeq.empty) +} + +/** A locus packed into a signed `int64` as `(contig_index << 32) | position`, where `contig_index` + * indexes `contigs`, the reference genome's contigs in order. + */ +final class LocusExtensionType(val referenceGenome: String, val contigs: IndexedSeq[String]) + extends HailExtensionType { + override def extensionName(): String = LocusExtensionType.Name + + override def storageType(): ArrowType = new ArrowType.Int(64, true) + + override def extensionEquals(other: ArrowType.ExtensionType): Boolean = + other match { + case o: LocusExtensionType => o.referenceGenome == referenceGenome && o.contigs == contigs + case _ => false + } + + override def serialize(): String = + JsonMethods.compact(JObject( + "reference_genome" -> JString(referenceGenome), + "contigs" -> JArray(contigs.map(JString(_)).toList), + )) + + override def deserialize(storageType: ArrowType, serializedData: String): ArrowType = { + checkStorageType(storageType) + val metadata = parseMetadata(serializedData) + val rg = metadata \ "reference_genome" match { + case JString(rg) => rg + case _ => throw new InvalidExtensionMetadataException( + s"$extensionName: missing or invalid 'reference_genome'" + ) + } + val contigs = metadata \ "contigs" match { + case JArray(cs) => cs.toIndexedSeq.map { + case JString(c) => c + case _ => throw new InvalidExtensionMetadataException( + s"$extensionName: 'contigs' must be an array of strings" + ) + } + case _ => throw new InvalidExtensionMetadataException( + s"$extensionName: missing or invalid 'contigs'" + ) + } + new LocusExtensionType(rg, contigs) + } + + // Deliberately omits the contig list, which can run to thousands of entries. + override def toString: String = s"LocusExtensionType($referenceGenome)" +} diff --git a/hail/hail/test/src/is/hail/types/arrow/ArrowSchemaConversionSuite.scala b/hail/hail/test/src/is/hail/types/arrow/ArrowSchemaConversionSuite.scala new file mode 100644 index 00000000000..400f38015b8 --- /dev/null +++ b/hail/hail/test/src/is/hail/types/arrow/ArrowSchemaConversionSuite.scala @@ -0,0 +1,430 @@ +package is.hail.types.arrow + +import is.hail.ParameterizedTest +import is.hail.TestUtils._ +import is.hail.collection.FastSeq +import is.hail.scalacheck._ +import is.hail.types.physical._ +import is.hail.types.physical.stypes.concrete.SRNGState +import is.hail.variant.{ReferenceGenome, ReferenceGenomeEnv} +import is.hail.variant.BuiltinLoci.ppackedLocus + +import scala.collection.immutable.ArraySeq +import scala.jdk.CollectionConverters._ + +import java.nio.ByteBuffer + +import org.apache.arrow.vector.types.FloatingPointPrecision +import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field, FieldType, Schema} +import org.junit.jupiter.api.Test +import org.scalacheck.Prop.forAll + +class ArrowSchemaConversionSuite { + + private val small = + ReferenceGenome("small", FastSeq("1", "2", "X"), Map("1" -> 10, "2" -> 10, "X" -> 10)) + + private val references: ReferenceGenomeEnv = ReferenceGenomeEnv(FastSeq(small)) + + private def toField(t: PType, name: String = "x"): Field = + ArrowSchemaConversion.toArrowField(name, t) + + private def fromField(f: Field): PType = + ArrowSchemaConversion.fromArrowField(f, references) + + private def leaf(name: String, nullable: Boolean, t: ArrowType): Field = + new Field(name, new FieldType(nullable, t, null), FastSeq.empty[Field].asJava) + + private def node(name: String, nullable: Boolean, t: ArrowType, children: Field*): Field = + new Field(name, new FieldType(nullable, t, null), children.asJava) + + private def children(f: Field): IndexedSeq[Field] = f.getChildren.asScala.toIndexedSeq + + def roundTripTypes() = ArraySeq[PType]( + PBoolean(true), + PBoolean(false), + PInt32(true), + PInt32(false), + PInt64(true), + PInt64(false), + PFloat32(true), + PFloat64(false), + PCanonicalString(true), + PCanonicalString(false), + PCanonicalBinary(true), + PCanonicalBinary(false), + PCanonicalCall(true), + PCanonicalCall(false), + ppackedLocus(ReferenceGenome.GRCh37, true), + PPackedLocus(small, false), + PCanonicalArray(PInt32(true), false), + PCanonicalArray(PCanonicalArray(PCanonicalString(false), true), true), + PCanonicalSet(PInt32(false), true), + PCanonicalSet(PPackedLocus(small, true), false), + PCanonicalDict(PCanonicalString(true), PInt32(false), false), + PCanonicalDict(PCanonicalStruct(true, "a" -> PInt32(true)), PCanonicalCall(true), true), + PCanonicalInterval(PInt32(false), false), + PCanonicalInterval(PPackedLocus(small, true), true), + PCanonicalNDArray(PFloat64(true), 2, false), + PCanonicalNDArray(PInt32(true), 0, true), + PCanonicalTuple(true, PInt32(false), PCanonicalString(true)), + PCanonicalTuple( + FastSeq(PTupleField(1, PInt32(true)), PTupleField(3, PCanonicalString(false))), + false, + ), + PCanonicalStruct(false), + PCanonicalStruct( + true, + "locus" -> PPackedLocus(small, true), + "alleles" -> PCanonicalArray(PCanonicalString(true), false), + "info" -> PCanonicalStruct( + false, + "AC" -> PCanonicalArray(PInt32(false), false), + "filters" -> PCanonicalSet(PCanonicalString(true), false), + ), + "entries" -> PCanonicalArray( + PCanonicalStruct( + true, + "GT" -> PCanonicalCall(false), + "PL" -> PCanonicalNDArray(PInt32(true), 1), + ), + true, + ), + ), + ) + + @Test def testRoundTripTypesDataProvider(): Unit = roundTripTypes(): Unit + + @ParameterizedTest("roundTripTypes") + def testRoundTrip(t: PType): Unit = + assertEq(fromField(toField(t)), t) + + @ParameterizedTest("roundTripTypes") + def testRoundTripThroughIpcSchema(t: PType): Unit = { + val schema = new Schema(FastSeq(toField(t)).asJava) + val bytes = schema.serializeAsMessage() + val read = Schema.deserializeMessage(ByteBuffer.wrap(bytes)) + assertEq(read, schema) + assertEq(fromField(read.getFields.get(0)), t) + } + + // Hail -> Arrow -> Hail is not the identity, since loci read back packed and dict keys read back + // required, but the Arrow representation must be stable. + @Test def testArrowFieldRoundTripProperty(): Unit = + check(forAll { (t: PType) => + val f = toField(t) + assertEq(toField(fromField(f)), f) + }) + + @Test def testSchemaRoundTrip(): Unit = { + val t = PCanonicalStruct( + true, + "a" -> PInt32(true), + "b" -> PCanonicalArray(PCanonicalString(false), false), + ) + val schema = ArrowSchemaConversion.toArrowSchema(t) + assertEq(schema.getFields.asScala.map(_.getName).toSeq, Seq("a", "b")) + assertEq(ArrowSchemaConversion.fromArrowSchema(schema, references), t) + } + + @Test def testPrimitives(): Unit = { + assertEq(toField(PBoolean(true)), leaf("x", false, ArrowType.Bool.INSTANCE)) + assertEq(toField(PInt32(false)), leaf("x", true, new ArrowType.Int(32, true))) + assertEq(toField(PInt64(true)), leaf("x", false, new ArrowType.Int(64, true))) + assertEq( + toField(PFloat32(true)), + leaf("x", false, new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)), + ) + assertEq( + toField(PFloat64(true)), + leaf("x", false, new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + ) + assertEq(toField(PCanonicalString(false)), leaf("x", true, ArrowType.Utf8.INSTANCE)) + assertEq(toField(PCanonicalBinary(false)), leaf("x", true, ArrowType.Binary.INSTANCE)) + } + + @Test def testArray(): Unit = + assertEq( + toField(PCanonicalArray(PInt32(true), false)), + node( + "x", + true, + ArrowType.List.INSTANCE, + leaf("element", false, new ArrowType.Int(32, true)), + ), + ) + + @Test def testSet(): Unit = { + val f = toField(PCanonicalSet(PInt32(false), true)) + assertEq( + f, + node( + "x", + false, + SetExtensionType, + leaf("element", true, new ArrowType.Int(32, true)), + ), + ) + assertEq(SetExtensionType.storageType(), ArrowType.List.INSTANCE) + } + + @Test def testDict(): Unit = + assertEq( + toField(PCanonicalDict(PCanonicalString(true), PInt32(false), false)), + node( + "x", + true, + new ArrowType.Map(true), + node( + "entries", + false, + ArrowType.Struct.INSTANCE, + leaf("key", false, ArrowType.Utf8.INSTANCE), + leaf("value", true, new ArrowType.Int(32, true)), + ), + ), + ) + + @Test def testDictKeysAreAlwaysRequired(): Unit = { + val t = PCanonicalDict(PCanonicalString(false), PInt32(false), false) + val f = toField(t) + assert(!children(children(f).head).head.isNullable) + assertEq(fromField(f), PCanonicalDict(PCanonicalString(true), PInt32(false), false)) + } + + @Test def testStruct(): Unit = + assertEq( + toField(PCanonicalStruct(false, "a" -> PInt32(true), "b" -> PCanonicalString(false))), + node( + "x", + true, + ArrowType.Struct.INSTANCE, + leaf("a", false, new ArrowType.Int(32, true)), + leaf("b", true, ArrowType.Utf8.INSTANCE), + ), + ) + + @Test def testTuple(): Unit = + assertEq( + toField( + PCanonicalTuple( + FastSeq(PTupleField(0, PInt32(true)), PTupleField(2, PCanonicalString(false))), + true, + ) + ), + node( + "x", + false, + TupleExtensionType, + leaf("0", false, new ArrowType.Int(32, true)), + leaf("2", true, ArrowType.Utf8.INSTANCE), + ), + ) + + @Test def testCall(): Unit = { + val f = toField(PCanonicalCall(false)) + assertEq(f, leaf("x", true, CallExtensionType.Current)) + assertEq(CallExtensionType.Current.storageType(), new ArrowType.Int(32, true)) + assertEq(CallExtensionType.Current.serialize(), """{"encoding":1}""") + } + + @Test def testLocus(): Unit = { + val f = toField(PPackedLocus(small, true)) + val ext = new LocusExtensionType(small.name, small.contigs) + assertEq(f, leaf("x", false, ext)) + assertEq(ext.storageType(), new ArrowType.Int(64, true)) + assertEq(ext.serialize(), """{"reference_genome":"small","contigs":["1","2","X"]}""") + } + + @Test def testCanonicalLocusReadsBackPacked(): Unit = { + val f = toField(PCanonicalLocus(small, false)) + assertEq(f, toField(PPackedLocus(small, false))) + assertEq(fromField(f), PPackedLocus(small, false)) + } + + @Test def testInterval(): Unit = { + val locus = new LocusExtensionType(small.name, small.contigs) + assertEq( + toField(PCanonicalInterval(PPackedLocus(small, true), false)), + node( + "x", + true, + IntervalExtensionType, + leaf("start", false, locus), + leaf("end", false, locus), + leaf("includes_start", false, ArrowType.Bool.INSTANCE), + leaf("includes_end", false, ArrowType.Bool.INSTANCE), + ), + ) + } + + @Test def testNDArray(): Unit = + assertEq( + toField(PCanonicalNDArray(PFloat64(true), 3, false)), + node( + "x", + true, + VariableShapeTensorExtensionType, + node( + "data", + false, + ArrowType.List.INSTANCE, + leaf("element", false, new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + ), + node( + "shape", + false, + new ArrowType.FixedSizeList(3), + leaf("element", false, new ArrowType.Int(32, true)), + ), + ), + ) + + @Test def testExtensionNames(): Unit = { + assertEq(CallExtensionType.Current.extensionName(), "hail.call") + assertEq(new LocusExtensionType(small.name, small.contigs).extensionName(), "hail.locus") + assertEq(SetExtensionType.extensionName(), "hail.set") + assertEq(IntervalExtensionType.extensionName(), "hail.interval") + assertEq(TupleExtensionType.extensionName(), "hail.tuple") + assertEq(VariableShapeTensorExtensionType.extensionName(), "arrow.variable_shape_tensor") + } + + @Test def testUnstorableTypes(): Unit = { + interceptFatal("cannot be stored")(toField(PVoid)) + interceptFatal("cannot be stored")(toField(StoredSTypePType(SRNGState(None), true))) + interceptFatal("cannot be stored")( + toField(PCanonicalStruct(false, "rng" -> StoredSTypePType(SRNGState(None), true))) + ) + interceptFatal("cannot be stored")(toField(PCanonicalNDArray(PPackedLocus(small, true), 1))) + } + + @Test def testUnknownReferenceGenomeOnRead(): Unit = + interceptFatal("reference genome 'nope'")( + fromField(leaf("x", false, new LocusExtensionType("nope", FastSeq("1")))) + ) + + @Test def testContigMismatchOnRead(): Unit = + interceptFatal("contigs")( + fromField(leaf("x", false, new LocusExtensionType(small.name, FastSeq("1", "X", "2")))) + ) + + @Test def testUnknownCallEncoding(): Unit = + interceptFatal("call encoding 2")(fromField(leaf("x", false, new CallExtensionType(2)))) + + def unsupportedArrowFields() = ArraySeq[Field]( + leaf("x", true, ArrowType.LargeUtf8.INSTANCE), + leaf("x", true, ArrowType.LargeBinary.INSTANCE), + leaf("x", true, new ArrowType.Int(8, true)), + leaf("x", true, new ArrowType.Int(32, false)), + leaf("x", true, new ArrowType.FloatingPoint(FloatingPointPrecision.HALF)), + node("x", true, ArrowType.LargeList.INSTANCE, leaf("element", true, ArrowType.Bool.INSTANCE)), + node( + "x", + true, + new ArrowType.Map(false), + node( + "entries", + false, + ArrowType.Struct.INSTANCE, + leaf("key", false, ArrowType.Utf8.INSTANCE), + leaf("value", true, ArrowType.Utf8.INSTANCE), + ), + ), + node( + "x", + true, + new ArrowType.Map(true), + node( + "entries", + false, + ArrowType.Struct.INSTANCE, + leaf("key", true, ArrowType.Utf8.INSTANCE), + leaf("value", true, ArrowType.Utf8.INSTANCE), + ), + ), + node( + "x", + true, + TupleExtensionType, + leaf("a", true, ArrowType.Utf8.INSTANCE), + ), + node( + "x", + true, + IntervalExtensionType, + leaf("start", true, ArrowType.Utf8.INSTANCE), + leaf("end", true, new ArrowType.Int(32, true)), + leaf("includes_start", false, ArrowType.Bool.INSTANCE), + leaf("includes_end", false, ArrowType.Bool.INSTANCE), + ), + node( + "x", + true, + VariableShapeTensorExtensionType, + node( + "data", + false, + ArrowType.List.INSTANCE, + leaf("element", true, new ArrowType.Int(32, true)), + ), + node( + "shape", + false, + new ArrowType.FixedSizeList(1), + leaf("element", false, new ArrowType.Int(32, true)), + ), + ), + node( + "x", + true, + VariableShapeTensorExtensionType, + node( + "data", + false, + ArrowType.List.INSTANCE, + leaf("element", false, new LocusExtensionType(small.name, small.contigs)), + ), + node( + "shape", + false, + new ArrowType.FixedSizeList(1), + leaf("element", false, new ArrowType.Int(32, true)), + ), + ), + new Field( + "x", + new FieldType( + true, + new ArrowType.Int(64, true), + null, + Map( + ArrowType.ExtensionType.EXTENSION_METADATA_KEY_NAME -> "hail.unknown", + ArrowType.ExtensionType.EXTENSION_METADATA_KEY_METADATA -> "{}", + ).asJava, + ), + FastSeq.empty[Field].asJava, + ), + new Field( + "x", + new FieldType( + true, + ArrowType.Utf8.INSTANCE, + new DictionaryEncoding(0, false, new ArrowType.Int(32, true)), + ), + FastSeq.empty[Field].asJava, + ), + node( + "x", + true, + ArrowType.Struct.INSTANCE, + leaf("a", true, ArrowType.Utf8.INSTANCE), + leaf("a", true, ArrowType.Bool.INSTANCE), + ), + ) + + @Test def testUnsupportedArrowFieldsDataProvider(): Unit = unsupportedArrowFields(): Unit + + @ParameterizedTest("unsupportedArrowFields") + def testUnsupportedArrowField(f: Field): Unit = + interceptFatal("columnar format")(fromField(f)) +}