[query] Add semhash support for Table-to-Table Aggregations (hail-is#…

…13922) Namely, TableKeyByAndAggregate and TableAggregateByKey
jigold · Oct 31, 2023 · a950ede · a950ede
1 parent 5d42835
commit a950ede
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/hail/src/main/scala/is/hail/expr/ir/analyses/SemanticHash.scala b/hail/src/main/scala/is/hail/expr/ir/analyses/SemanticHash.scala
@@ -233,6 +233,13 @@ case object SemanticHash extends Logging {
         val getFieldIndex = table.typ.rowType.fieldIdx
         keys.map(getFieldIndex).foreach(buffer ++= Bytes.fromInt(_))
 
+
+      case TableKeyByAndAggregate(_, _, _, nPartitions, bufferSize) =>
+        nPartitions.foreach {
+          buffer ++= Bytes.fromInt(_)
+        }
+        buffer ++= Bytes.fromInt(bufferSize)
+
       case TableJoin(_, _, joinop, key) =>
         buffer ++= joinop.getBytes ++= Bytes.fromInt(key)
 
@@ -339,6 +346,7 @@ case object SemanticHash extends Logging {
            _: StreamTakeWhile |
            _: TableGetGlobals |
            _: TableAggregate |
+           _: TableAggregateByKey |
            _: TableCollect |
            _: TableCount |
            _: TableDistinct |

diff --git a/hail/src/test/scala/is/hail/expr/ir/analyses/SemanticHashSuite.scala b/hail/src/test/scala/is/hail/expr/ir/analyses/SemanticHashSuite.scala
@@ -169,7 +169,7 @@ class SemanticHashSuite extends HailSuite {
     )
 
   def isTableIRSemanticallyEquivalent: Array[Array[Any]] = {
-    val ttype = TableType(TStruct("a" -> TInt32, "b" -> TStruct()), IndexedSeq(), TStruct())
+    val ttype = TableType(TStruct("a" -> TInt32, "b" -> TStruct()), IndexedSeq("a"), TStruct())
     val ttypeb = TableType(TStruct("c" -> TInt32, "d" -> TStruct()), IndexedSeq(), TStruct())
 
     def mkTableRead(reader: TableReader): TableIR =
@@ -205,12 +205,14 @@ class SemanticHashSuite extends HailSuite {
       Array(
         TableGetGlobals,
         TableAggregate(_, Void()),
+        TableAggregateByKey(_, MakeStruct(FastSeq())),
+        TableKeyByAndAggregate(_, MakeStruct(FastSeq()), MakeStruct(FastSeq("idx" -> I32(0))), None, 256),
         TableCollect,
         TableCount,
         TableDistinct,
         TableFilter(_, Void()),
         TableMapGlobals(_, MakeStruct(IndexedSeq.empty)),
-        TableMapRows(_, MakeStruct(IndexedSeq.empty)),
+        TableMapRows(_, MakeStruct(FastSeq("a" -> I32(0)))),
         TableRename(_, Map.empty, Map.empty),
       ).map { wrap =>
         Array(wrap(tir), wrap(tir), true, "")