From 44f6adada57717ff85a6281fa88dfb4e493b5cdb Mon Sep 17 00:00:00 2001 From: Chris Vittal Date: Thu, 29 Oct 2026 14:53:49 -0400 Subject: [PATCH] [query] Remove EStructOfArrays This was an important experiment in columnar encodings but it no longer aligns with the direction I would like to take file formats and hail data serialization in. Aditionally, with the creation of packed loci in pr #15880, I would like people to be able to opt in to that with unstable encodings without also selecting struct of arrays encoding which could have much more severe performance and correctness implications compared to the relatively small change that packed loci bring. --- .../hail/types/encoded/EStructOfArrays.scala | 265 ------------------ .../src/is/hail/types/encoded/EType.scala | 27 -- .../is/hail/types/encoded/ETypeSuite.scala | 51 ---- .../hail/matrixtable/test_matrix_table.py | 27 -- 4 files changed, 370 deletions(-) delete mode 100644 hail/hail/src/is/hail/types/encoded/EStructOfArrays.scala diff --git a/hail/hail/src/is/hail/types/encoded/EStructOfArrays.scala b/hail/hail/src/is/hail/types/encoded/EStructOfArrays.scala deleted file mode 100644 index cbd74fec05d..00000000000 --- a/hail/hail/src/is/hail/types/encoded/EStructOfArrays.scala +++ /dev/null @@ -1,265 +0,0 @@ -package is.hail.types.encoded - -import is.hail.annotations.{Region, UnsafeUtils} -import is.hail.asm4s._ -import is.hail.asm4s.implicits.{ - valueToRichCodeInputBuffer, valueToRichCodeOutputBuffer, valueToRichCodeRegion, -} -import is.hail.collection.implicits.toRichIterable -import is.hail.expr.ir.{EmitCodeBuilder, IEmitCode} -import is.hail.io.{InputBuffer, OutputBuffer} -import is.hail.types._ -import is.hail.types.physical._ -import is.hail.types.physical.stypes.{SType, SValue} -import is.hail.types.physical.stypes.concrete._ -import is.hail.types.physical.stypes.interfaces._ -import is.hail.types.virtual.{Field => TField, _} - -object EStructOfArrays { - // expand this as more types are supported - def supportsFieldType(typ: EType): Boolean = typ match { - case _: EBoolean => true - case _: EFloat32 => true - case _: EFloat64 => true - case _: EInt32 => true - case _: EInt64 => true - case _ => false - } - - def supportsFieldType(typ: Type): Boolean = typ match { - case TBoolean | TFloat32 | TFloat64 | TInt32 | TInt64 => true - case _ => false - } - - def fromTypeAndRequiredness(t: TIterable, r: RIterable): EStructOfArrays = { - val et = tcoerce[TBaseStruct](t.elementType) - val ret = tcoerce[RBaseStruct](r.elementType) - val fields = et.fields.zip(ret.fields).map { case (TField(name, typ, index), r) => - val encodedType = typ match { - case TBoolean => EArray(EBoolean(r.typ.required), required = true) - case TFloat32 => EArray(EFloat32(r.typ.required), required = true) - case TFloat64 => EArray(EFloat64(r.typ.required), required = true) - case TInt32 => EArray(EInt32(r.typ.required), required = true) - case TInt64 => EArray(EInt64(r.typ.required), required = true) - } - EField(name, encodedType, index) - } - - EStructOfArrays(fields, required = r.required, structRequired = ret.required) - } -} - -final case class EStructOfArrays( - fields: IndexedSeq[EField], - val required: Boolean = false, - val structRequired: Boolean = false, -) extends EContainer { - assert(fields.zipWithIndex.forall { case (f, i) => f.index == i }) - require(fields.forall(f => f.typ.isInstanceOf[EContainer])) - - // We don't care about the requiredness of these EContainers, but at the top level, - // they should all be present, so we require them to be required - require(fields.forall(f => f.typ.required)) - - val elementType = EBaseStruct( - fields.map { field => - EField(field.name, field.typ.asInstanceOf[EContainer].elementType, field.index) - }, - required = structRequired, - ) - - elementType.fields.foreach { fld => - if (!EStructOfArrays.supportsFieldType(fld.typ)) - throw new UnsupportedOperationException( - s"${fld.typ} not supported as field type for struct of arrays" - ) - } - - override def _decodedSType(requestedType: Type): SType = { - require(requestedType.isInstanceOf[TArray]) - val elementPType = elementType.decodedPType(requestedType.asInstanceOf[TContainer].elementType) - SIndexablePointer(PCanonicalArray(elementPType, required = false)) - } - - override def _buildDecoder( - cb: EmitCodeBuilder, - t: Type, - region: Value[Region], - in: Value[InputBuffer], - ): SValue = { - val st = tcoerce[SIndexablePointer](decodedSType(t)) - val pt = tcoerce[PCanonicalArray](st.pType) - val ept = tcoerce[PBaseStruct](pt.elementType) - assert( - pt.elementType.required == elementType.required, - s"${pt.elementType.required} | ${elementType.required}", - ) - - val scratchRegion: Value[Region] = cb.memoize(region.getPool().invoke[Region]("getRegion")) - val length = cb.memoize(in.readVarint()) - - val arrayPtr = cb.memoize(pt.allocate(region, length)) - cb += Region.setMemory( - arrayPtr, - pt.pastLastElementOffset(arrayPtr, length) - arrayPtr, - 0.toByte, - ) - pt.storeLength(cb, arrayPtr, length) - val nMissingBytes = cb.memoize(pt.nMissingBytes(length)) - if (!elementType.required) - cb += in.readBytes(region, arrayPtr + pt.missingBytesOffset, nMissingBytes) - - // scratch variables - val i = cb.newLocal[Int]("i") - val structPtr = cb.newLocal[Long]("element_out") // pointer to struct we're constructing - fields.foreach { field => - cb += scratchRegion.clearRegion() - if (!ept.hasField(field.name)) { - field.typ.buildSkip(cb.emb.ecb)(cb, scratchRegion, in) - } else { - val pFieldIdx = ept.fieldIdx(field.name) - val pFieldType = tcoerce[PPrimitive](ept.types(pFieldIdx)) - - val array = field.typ.buildDecoder(t, cb.emb.ecb)(cb, scratchRegion, in).asIndexable - cb.if_( - array.loadLength.cne(length), - cb._fatal( - "Mismatch in length for decoded array of field `", - field.name, - "` expected ", - length.toS, - ", was ", - array.loadLength.toS, - ), - ) - cb.assign(i, 0) - cb.assign(structPtr, pt.firstElementOffset(arrayPtr)) - cb.while_( - i < length, { - cb.if_( - pt.isElementMissing(arrayPtr, i), - if (!pFieldType.required) ept.setFieldMissing(cb, structPtr, pFieldIdx), - array.loadElement(cb, i).consume( - cb, - if (!pFieldType.required) ept.setFieldMissing(cb, structPtr, pFieldIdx), - { sv => - val fieldPtr = ept.fieldOffset(structPtr, pFieldIdx) - pFieldType.storeAtAddress(cb, fieldPtr, region, sv, true) - }, - ), - ) - cb.assign(i, i + 1) - cb.assign(structPtr, pt.nextElementAddress(structPtr)) - }, - ) - } - } - - cb += scratchRegion.invalidate() - new SIndexablePointerValue( - st, - arrayPtr, - length, - cb.memoize(pt.firstElementOffset(arrayPtr, length)), - ) - } - - override def _buildEncoder(cb: EmitCodeBuilder, v: SValue, out: Value[OutputBuffer]): Unit = - v match { - case sv: SIndexablePointerValue => - val pArray = sv.st.pType.asInstanceOf[PCanonicalArrayBackedContainer].arrayRep - val r: Value[Region] = // scratch region - cb.memoize(cb.emb.ecb.pool().invoke[Region]("getRegion")) - cb += out.writeVarint(sv.length) - if (!elementType.required) { - val nMissingBytes = cb.memoize(pArray.nMissingBytes(sv.length)) - cb += out.writeBytes(sv.a + pArray.missingBytesOffset, nMissingBytes) - } - - fields.foreach { field => - val pFieldType = tcoerce[PBaseStruct](pArray.elementType).fieldByName(field.name).typ - require(EStructOfArrays.supportsFieldType(pFieldType.virtualType)) - val arrayType = PCanonicalArray(pFieldType, required = false) - transposeAndWriteField(cb, field, arrayType, sv, r, out) - } - - cb += r.invalidate() - } - - private[this] def transposeAndWriteField( - cb: EmitCodeBuilder, - field: EField, - arrayType: PCanonicalArray, - sv: SIndexablePointerValue, - r: Value[Region], - out: Value[OutputBuffer], - ): Unit = { - cb += r.clearRegion() - val arrayValue = - arrayType.constructFromElements(cb, r, sv.length, deepCopy = false) { (cb, i) => - if (arrayType.elementType.required) { - sv.loadElement(cb, i).consumeI( - cb, { - val elementSType = tcoerce[SBaseStruct](sv.st.elementType) - val fieldSType = elementSType.fieldTypes(elementSType.fieldIdx(field.name)) - IEmitCode.present( - cb, - cb.newSLocal( - fieldSType, - s"${fieldSType.asIdent}_present_for_required_element_in_missing_struct", - ), - ) - }, - { case sbsv: SBaseStructValue => - sbsv.loadField(cb, field.name) - }, - ) - } else { - sv.loadElement(cb, i).flatMap(cb) { case sbsv: SBaseStructValue => - sbsv.loadField(cb, field.name) - } - } - } - - val arrayEncoder = field.typ.buildEncoder(arrayValue.st, cb.emb.ecb) - arrayEncoder(cb, arrayValue, out) - } - - override def _buildSkip(cb: EmitCodeBuilder, r: Value[Region], in: Value[InputBuffer]): Unit = { - val length = cb.memoize(in.readVarint()) - val nMissingBytes = - cb.memoize(UnsafeUtils.packBitsToBytes(length)) // valid for all top level arrays - if (!elementType.required) - cb += in.skipBytes(nMissingBytes) - - fields.foreach(field => field.typ.buildSkip(cb.emb.ecb)(cb, r, in)) - } - - override def setRequired(newRequired: Boolean): EStructOfArrays = - EStructOfArrays(fields, required = newRequired, structRequired = structRequired) - - override def _asIdent: String = s"struct_of_arrays_from_${elementType.asIdent}" - - override def _toPretty: String = { - val sb = new StringBuilder - _pretty(sb, 0, compact = true) - sb.result() - } - - override def _pretty(sb: StringBuilder, indent: Int, compact: Boolean): Unit = { - sb ++= "EStructOfArrays[" - if (structRequired) sb += '+' - sb += '{' - if (compact) { - fields.foreachBetween(_.pretty(sb, indent, compact))(sb += ',') - } else if (fields.isEmpty) { - sb += ' ' - } else { - sb += '\n' - fields.foreachBetween(_.pretty(sb, indent + 4, compact))(sb ++= ",\n") - sb += '\n' - sb ++= (" " * indent) - } - sb ++= "}]" - } -} diff --git a/hail/hail/src/is/hail/types/encoded/EType.scala b/hail/hail/src/is/hail/types/encoded/EType.scala index 5d20f3ffd7d..29cb319e9bc 100644 --- a/hail/hail/src/is/hail/types/encoded/EType.scala +++ b/hail/hail/src/is/hail/types/encoded/EType.scala @@ -389,13 +389,6 @@ object EType extends Logging { ), required = rinterval.required, ) - case t: TArray - if (ctx.flags.isDefined( - EType.Flags.UseUnstableEncodings - ) && t.elementType.isInstanceOf[TBaseStruct] && t.elementType.asInstanceOf[ - TBaseStruct - ].fields.forall(fld => EStructOfArrays.supportsFieldType(fld.typ))) => - EStructOfArrays.fromTypeAndRequiredness(t, tcoerce[RIterable](r)) case t: TIterable => EArray( fromTypeAndAnalysis(ctx, t.elementType, tcoerce[RIterable](r).elementType, useVarint), @@ -512,26 +505,6 @@ object EType extends Logging { val nDims = IRParser.int32_literal(it) IRParser.punctuation(it, "]") ENDArrayColumnMajor(elementType, nDims, req) - case "EStructOfArrays" => - IRParser.punctuation(it, "[") - val structReq = it.head match { - case x: PunctuationToken if x.value == "+" => - IRParser.punctuation(it, "+") - true - case _ => false - } - IRParser.punctuation(it, "{") - val args = IRParser.repsepUntil( - it, - IRParser.struct_field(eTypeParser), - PunctuationToken(","), - PunctuationToken("}"), - ) - IRParser.punctuation(it, "}") - IRParser.punctuation(it, "]") - val fields = args.zipWithIndex.map { case ((name, t), i) => EField(name, t, i) } - EStructOfArrays(fields, required = req, structRequired = structReq) - case x => throw new UnsupportedOperationException(s"Couldn't parse $x ${it.toIndexedSeq}") } diff --git a/hail/hail/test/src/is/hail/types/encoded/ETypeSuite.scala b/hail/hail/test/src/is/hail/types/encoded/ETypeSuite.scala index 492991bc597..b54f025f8fb 100644 --- a/hail/hail/test/src/is/hail/types/encoded/ETypeSuite.scala +++ b/hail/hail/test/src/is/hail/types/encoded/ETypeSuite.scala @@ -65,14 +65,6 @@ class ETypeSuite { required = true, ), ENDArrayColumnMajor(EFloat64Required, 3), - EStructOfArrays( - FastSeq( - EField("a", EArray(EInt32Required, true), 0), - EField("b", EArray(EInt32Optional, true), 1), - ), - required = true, - structRequired = false, - ), ) @ParameterizedTest @@ -326,49 +318,6 @@ class ETypeSuite { assertNoVarint(EType.fromPythonTypeEncoding(t), t) } - @Test def testStructOfArrays(implicit ctx: ExecuteContext): Unit = { - val etype = - EStructOfArrays( - FastSeq( - EField("a", EArray(EInt32Required, true), 0), - EField("b", EArray(EFloat32Optional, true), 1), - EField("c", EArray(EBooleanRequired, true), 2), - ) - ) - val toEncode = - PCanonicalArray(PCanonicalStruct( - false, - "a" -> PInt32Required, - "b" -> PFloat32Optional, - "c" -> PBooleanRequired, - )) - val toDecode = toEncode - val data = FastSeq( - RowSeq(1, 2f, true), - null, - RowSeq(3, null, false), - RowSeq(4, 5f, true), - null, - RowSeq(6, 7f, false), - RowSeq(8, null, true), - RowSeq(9, 10f, false), - ) - - assertEqualEncodeDecode(toEncode, etype, toDecode, data) - } - - @Test def testUnstableEncodingsOfContainersOfStructs(implicit ctx: ExecuteContext): Unit = - withUnstableEncodings { implicit ctx => - val elt = PCanonicalStruct(false, "a" -> PInt32Optional, "b" -> PFloat64Optional) - val cases = ArraySeq[(PType, Annotation)]( - PCanonicalArray(elt) -> FastSeq(RowSeq(1, 2.0), null, RowSeq(3, null)), - PCanonicalSet(elt) -> Set(RowSeq(1, 2.0), null, RowSeq(3, null)), - PCanonicalDict(PInt32Optional, PFloat64Optional) -> Map(1 -> 2.0, 3 -> null), - ) - for ((pt, a) <- cases) - assertEqualEncodeDecode(pt, EType.defaultFromPType(ctx, pt), pt, a) - } - // Loci of either representation, nested anywhere, decode and store into either representation, // on the same reference genome or an equal copy of it. The session does not hold the custom // reference genome: locus types carry their own. diff --git a/hail/python/test/hail/matrixtable/test_matrix_table.py b/hail/python/test/hail/matrixtable/test_matrix_table.py index e8e309c0ef9..05e634b4235 100644 --- a/hail/python/test/hail/matrixtable/test_matrix_table.py +++ b/hail/python/test/hail/matrixtable/test_matrix_table.py @@ -1,8 +1,5 @@ -import gzip -import json import math import operator -import os import random import unittest @@ -2361,30 +2358,6 @@ def test_sample_entries(): ht._force_count() -def test_struct_of_arrays_encoding(): - mt = hl.utils.range_matrix_table(10, 10) - mt = mt.annotate_entries( - p=hl.or_missing(hl.rand_bool(0.6), hl.rand_bool(0.5)), - d=mt.row_idx * mt.col_idx / math.pi, - x=mt.row_idx + mt.col_idx, - y=mt.row_idx * mt.col_idx, - z=hl.or_missing(hl.rand_bool(0.8), mt.row_idx - mt.col_idx), - ) - std_mt = mt.checkpoint(new_temp_file(extension='mt')) - with hl._with_flags(use_unstable_encodings='1'): - out_path = new_temp_file(extension='mt') - mt = std_mt.checkpoint(out_path) - md_path = os.path.join(out_path, 'entries/rows/metadata.json.gz') - with ( - Env.fs().open(md_path, 'rb') as f, - gzip.open(f) as md_file, - ): - md = json.load(md_file) - etype = md['_codecSpec']['_eType'] - assert 'EStructOfArrays' in etype - assert mt._same(std_mt) - - @pytest.fixture(scope='module') def query_mt_mt(): path = new_temp_file(extension='mt')