Skip to content

Commit 71c451e

Browse files
committed
[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. Both directions take the reference genome map: loci need the contig list to write, and reading checks that the stored contigs match the defined genome. 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 - 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 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 drives 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.
1 parent 1b30333 commit 71c451e

4 files changed

Lines changed: 857 additions & 0 deletions

File tree

‎hail/hail/package.mill‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ object `package` extends HailScalaModule:
7373
`breeze` :: "2.1.0",
7474
`blas` :: "3.0.4",
7575
`arpack` :: "3.0.4",
76+
`arrow-vector` :: "18.3.0",
7677
`avro` :: "1.12.1",
7778
`commons-codec` :: "1.19.0",
7879
`commons-io` :: "2.21.0",
Lines changed: 240 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,240 @@
1+
package is.hail.types.arrow
2+
3+
import is.hail.types.physical._
4+
import is.hail.utils.fatal
5+
import is.hail.variant.ReferenceGenomeEnv
6+
7+
import scala.jdk.CollectionConverters._
8+
9+
import org.apache.arrow.vector.types.FloatingPointPrecision
10+
import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema}
11+
12+
/** Conversion between Hail physical types and the Arrow schema of the columnar format.
13+
*
14+
* Arrow nullability carries Hail requiredness at every level, with one exception: dict keys are
15+
* always non-nullable, so a dict whose key type is optional reads back with a required key type.
16+
*
17+
* Both locus representations write the same Arrow type, whose storage is a packed locus, and a
18+
* locus always reads back as a [[PPackedLocus]].
19+
*
20+
* Arrow -> Hail accepts exactly the Arrow types that Hail -> Arrow produces.
21+
*/
22+
object ArrowSchemaConversion {
23+
ExtensionTypes.register()
24+
25+
val ElementFieldName: String = "element"
26+
val EntriesFieldName: String = "entries"
27+
val KeyFieldName: String = "key"
28+
val ValueFieldName: String = "value"
29+
30+
private val Int32 = new ArrowType.Int(32, true)
31+
private val Int64 = new ArrowType.Int(64, true)
32+
private val Float32 = new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)
33+
private val Float64 = new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)
34+
35+
private val IntervalFieldNames =
36+
IndexedSeq("start", "end", "includes_start", "includes_end")
37+
38+
private val TensorFieldNames = IndexedSeq("data", "shape")
39+
40+
def toArrowSchema(t: PStruct): Schema =
41+
new Schema(t.fields.map(f => toArrowField(f.name, f.typ)).asJava)
42+
43+
def toArrowField(name: String, t: PType): Field = {
44+
def field(arrowType: ArrowType, children: Field*): Field =
45+
mkField(name, !t.required, arrowType, children: _*)
46+
47+
def child(name: String, t: PType): Field =
48+
toArrowField(name, t)
49+
50+
t match {
51+
case _: PBoolean => field(ArrowType.Bool.INSTANCE)
52+
case _: PInt32 => field(Int32)
53+
case _: PInt64 => field(Int64)
54+
case _: PFloat32 => field(Float32)
55+
case _: PFloat64 => field(Float64)
56+
case _: PString => field(ArrowType.Utf8.INSTANCE)
57+
case _: PBinary => field(ArrowType.Binary.INSTANCE)
58+
case _: PCall => field(CallExtensionType.Current)
59+
case t: PLocus => field(new LocusExtensionType(t.rg.name, t.rg.contigs))
60+
case t: PInterval =>
61+
field(
62+
IntervalExtensionType,
63+
child(IntervalFieldNames(0), t.pointType),
64+
child(IntervalFieldNames(1), t.pointType),
65+
mkField(IntervalFieldNames(2), nullable = false, ArrowType.Bool.INSTANCE),
66+
mkField(IntervalFieldNames(3), nullable = false, ArrowType.Bool.INSTANCE),
67+
)
68+
case t: PNDArray if !t.elementType.isInstanceOf[PLocus] =>
69+
field(
70+
VariableShapeTensorExtensionType,
71+
mkField(
72+
TensorFieldNames(0),
73+
nullable = false,
74+
ArrowType.List.INSTANCE,
75+
child(ElementFieldName, t.elementType),
76+
),
77+
mkField(
78+
TensorFieldNames(1),
79+
nullable = false,
80+
new ArrowType.FixedSizeList(t.nDims),
81+
mkField(ElementFieldName, nullable = false, Int32),
82+
),
83+
)
84+
case t: PSet => field(SetExtensionType, child(ElementFieldName, t.elementType))
85+
case t: PDict =>
86+
field(
87+
new ArrowType.Map( /*keysSorted=*/ true),
88+
mkField(
89+
EntriesFieldName,
90+
nullable = false,
91+
ArrowType.Struct.INSTANCE,
92+
child(KeyFieldName, t.keyType.setRequired(true)),
93+
child(ValueFieldName, t.valueType),
94+
),
95+
)
96+
case t: PArray => field(ArrowType.List.INSTANCE, child(ElementFieldName, t.elementType))
97+
case t: PTuple => field(TupleExtensionType, t.fields.map(f => child(f.name, f.typ)): _*)
98+
case t: PStruct =>
99+
field(ArrowType.Struct.INSTANCE, t.fields.map(f => child(f.name, f.typ)): _*)
100+
case _ =>
101+
fatal(s"field '$name': Hail type ${t.virtualType} cannot be stored in the columnar format")
102+
}
103+
}
104+
105+
def fromArrowSchema(schema: Schema, references: ReferenceGenomeEnv): PCanonicalStruct =
106+
PCanonicalStruct(
107+
structFields(schema.getFields.asScala.toIndexedSeq, references),
108+
required = true,
109+
)
110+
111+
def fromArrowField(field: Field, references: ReferenceGenomeEnv): PType = {
112+
val name = field.getName
113+
val required = !field.isNullable
114+
val children = field.getChildren.asScala.toIndexedSeq
115+
116+
def unsupported(reason: String): Nothing =
117+
fatal(
118+
s"field '$name' of Arrow type ${field.getType} is not part of the columnar format: $reason"
119+
)
120+
121+
def checkChildren(expected: IndexedSeq[String]): Unit =
122+
if (children.map(_.getName) != expected)
123+
unsupported(
124+
s"expected child fields ${expected.mkString("[", ", ", "]")}, found ${children.map(_.getName).mkString("[", ", ", "]")}"
125+
)
126+
127+
def onlyChild(): Field =
128+
if (children.length == 1) children.head
129+
else unsupported(s"expected 1 child field, found ${children.length}")
130+
131+
def convert(f: Field): PType = fromArrowField(f, references)
132+
133+
def requiredChild(f: Field): Field =
134+
if (f.isNullable) unsupported(s"child field '${f.getName}' must be non-nullable")
135+
else f
136+
137+
if (field.getDictionary != null)
138+
unsupported("dictionary encoding is not supported")
139+
140+
field.getType match {
141+
case _: CallExtensionType if field.getType == CallExtensionType.Current =>
142+
PCanonicalCall(required)
143+
case t: CallExtensionType =>
144+
unsupported(s"unknown call encoding ${t.encoding}")
145+
case t: LocusExtensionType =>
146+
val rg = references.get(t.referenceGenome).getOrElse(
147+
fatal(s"field '$name': reference genome '${t.referenceGenome}' is not defined")
148+
)
149+
if (rg.contigs != t.contigs)
150+
fatal(
151+
s"field '$name': the stored contigs of reference genome '${t.referenceGenome}' " +
152+
"differ from the contigs of the defined reference genome"
153+
)
154+
PPackedLocus(rg, required)
155+
case SetExtensionType =>
156+
PCanonicalSet(convert(onlyChild()), required)
157+
case IntervalExtensionType =>
158+
checkChildren(IntervalFieldNames)
159+
val IndexedSeq(start, end, includesStart, includesEnd) = children
160+
val pointType = convert(start)
161+
if (convert(end) != pointType)
162+
unsupported("interval 'start' and 'end' have different types")
163+
for (f <- Seq(includesStart, includesEnd))
164+
if (f.isNullable || f.getType != ArrowType.Bool.INSTANCE || !f.getChildren.isEmpty)
165+
unsupported(s"interval '${f.getName}' must be a non-nullable bool")
166+
PCanonicalInterval(pointType, required)
167+
case TupleExtensionType =>
168+
val fields = children.map { c =>
169+
val idx = c.getName.toIntOption.filter(i => i >= 0 && i.toString == c.getName).getOrElse(
170+
unsupported(s"tuple field name '${c.getName}' is not a non-negative integer")
171+
)
172+
PTupleField(idx, convert(c))
173+
}
174+
if (fields.map(_.index).distinct.length != fields.length)
175+
unsupported("tuple field names are not distinct")
176+
PCanonicalTuple(fields, required)
177+
case VariableShapeTensorExtensionType =>
178+
checkChildren(TensorFieldNames)
179+
val IndexedSeq(data, shape) = children
180+
val nDims = shape.getType match {
181+
case t: ArrowType.FixedSizeList => t.getListSize
182+
case t => unsupported(s"tensor 'shape' must be a fixed_size_list, found $t")
183+
}
184+
requiredChild(shape)
185+
if (shape.getChildren.size != 1 || convert(shape.getChildren.get(0)) != PInt32(true))
186+
unsupported("tensor 'shape' must contain non-nullable int32")
187+
if (requiredChild(data).getType != ArrowType.List.INSTANCE || data.getChildren.size != 1)
188+
unsupported("tensor 'data' must be a list")
189+
val elementType = convert(requiredChild(data.getChildren.get(0)))
190+
if (elementType.containsPointers || elementType.isInstanceOf[PLocus])
191+
unsupported(s"tensor element type ${elementType.virtualType} is not supported")
192+
PCanonicalNDArray(elementType, nDims, required)
193+
case t: ArrowType.ExtensionType =>
194+
unsupported(s"unknown extension type '${t.extensionName}'")
195+
case _
196+
if field.getMetadata.containsKey(ArrowType.ExtensionType.EXTENSION_METADATA_KEY_NAME) =>
197+
unsupported(
198+
s"unregistered extension type '${field.getMetadata.get(ArrowType.ExtensionType.EXTENSION_METADATA_KEY_NAME)}'"
199+
)
200+
case ArrowType.Bool.INSTANCE => PBoolean(required)
201+
case Int32 => PInt32(required)
202+
case Int64 => PInt64(required)
203+
case Float32 => PFloat32(required)
204+
case Float64 => PFloat64(required)
205+
case ArrowType.Utf8.INSTANCE => PCanonicalString(required)
206+
case ArrowType.Binary.INSTANCE => PCanonicalBinary(required)
207+
case ArrowType.List.INSTANCE => PCanonicalArray(convert(onlyChild()), required)
208+
case t: ArrowType.Map =>
209+
if (!t.getKeysSorted) unsupported("map keys must be sorted")
210+
val entries = requiredChild(onlyChild())
211+
if (entries.getType != ArrowType.Struct.INSTANCE || entries.getChildren.size != 2)
212+
unsupported("map entries must be a struct with two fields")
213+
val key = requiredChild(entries.getChildren.get(0))
214+
PCanonicalDict(convert(key), convert(entries.getChildren.get(1)), required)
215+
case ArrowType.Struct.INSTANCE =>
216+
PCanonicalStruct(structFields(children, references), required)
217+
case _ => unsupported("unsupported type")
218+
}
219+
}
220+
221+
private def structFields(
222+
fields: IndexedSeq[Field],
223+
references: ReferenceGenomeEnv,
224+
): IndexedSeq[PField] = {
225+
val duplicates = fields.map(_.getName).groupBy(identity).collect {
226+
case (name, names) if names.length > 1 => name
227+
}
228+
if (duplicates.nonEmpty)
229+
fatal(
230+
s"duplicate field names are not part of the columnar format: ${duplicates.mkString(", ")}"
231+
)
232+
fields.zipWithIndex.map { case (f, i) =>
233+
PField(f.getName, fromArrowField(f, references), i)
234+
}
235+
}
236+
237+
private def mkField(name: String, nullable: Boolean, arrowType: ArrowType, children: Field*)
238+
: Field =
239+
new Field(name, new FieldType(nullable, arrowType, /*dictionary=*/ null), children.asJava)
240+
}

0 commit comments

Comments
 (0)