<commit_before>package is.hail.io.index

import is.hail.annotations.{Annotation, Region, RegionValueBuilder}
import is.hail.expr.types._
import is.hail.io.CodecSpec
import is.hail.utils._
import is.hail.utils.richUtils.ByteTrackingOutputStream
import org.apache.hadoop.conf.Configuration
import org.json4s.Formats
import org.json4s.jackson.Serialization

case class IndexMetadata(
  fileVersion: Int,
  branchingFactor: Int,
  height: Int,
  keyType: String,
  annotationType: String,
  nKeys: Long,
  indexPath: String,
  rootOffset: Long,
  attributes: Map[String, Any]
)

case class IndexNodeInfo(
  indexFileOffset: Long,
  firstIndex: Long,
  firstKey: Annotation,
  firstRecordOffset: Long,
  firstAnnotation: Annotation
)

object IndexWriter {
  val version: SemanticVersion = SemanticVersion(1, 0, 0)
}

class IndexWriter(
  hConf: Configuration,
  path: String,
  keyType: Type,
  annotationType: Type,
  branchingFactor: Int = 1024,
  attributes: Map[String, Any] = Map.empty[String, Any]) extends AutoCloseable {
  require(branchingFactor > 1)

  private var elementIdx = 0L
  private val region = new Region()
  private val rvb = new RegionValueBuilder(region)

  private val leafNodeBuilder = new LeafNodeBuilder(keyType, annotationType, 0L)
  private val internalNodeBuilders = new ArrayBuilder[InternalNodeBuilder]()
  internalNodeBuilders += new InternalNodeBuilder(keyType, annotationType)

  private val trackedOS = new ByteTrackingOutputStream(hConf.unsafeWriter(path + "/index"))
  private val codecSpec = CodecSpec.default
  private val leafEncoder = codecSpec.buildEncoder(leafNodeBuilder.typ)(trackedOS)
  private val internalEncoder = codecSpec.buildEncoder(InternalNodeBuilder.typ(keyType, annotationType))(trackedOS)

  private def height: Int = internalNodeBuilders.length + 1 // have one leaf node layer

  private def writeInternalNode(node: InternalNodeBuilder, level: Int, isRoot: Boolean = false): Long = {
    val indexFileOffset = trackedOS.bytesWritten

    val info = if (node.size > 0) {
      val child = node.getChild(0)
      val firstIndex = node.firstIdx
      val firstKey = child.firstKey
      val firstRecordOffset = child.firstRecordOffset
      val firstAnnotation = child.firstAnnotation
      IndexNodeInfo(indexFileOffset, firstIndex, firstKey, firstRecordOffset, firstAnnotation)
    } else {
      assert(isRoot && level == 0)
      null
    }

    internalEncoder.writeByte(1)

    val regionOffset = node.write(rvb)
    internalEncoder.writeRegionValue(region, regionOffset)
    internalEncoder.flush()

    region.clear()
    node.clear()

    if (!isRoot) {
      if (level + 1 == internalNodeBuilders.length)
        internalNodeBuilders += new InternalNodeBuilder(keyType, annotationType) // , info.firstIndex
      val parent = internalNodeBuilders(level + 1)
      if (parent.size == branchingFactor)
        writeInternalNode(parent, level + 1)
      parent += info
    }

    indexFileOffset
  }

  private def writeLeafNode(): Long = {
    val indexFileOffset = trackedOS.bytesWritten

    assert(leafNodeBuilder.size > 0)
    val child = leafNodeBuilder.getChild(0)
    val firstIndex = leafNodeBuilder.firstIdx
    val firstKey = child.key
    val firstRecordOffset = child.recordOffset
    val firstAnnotation = child.annotation

    leafEncoder.writeByte(0)

    val regionOffset = leafNodeBuilder.write(rvb)
    leafEncoder.writeRegionValue(region, regionOffset)
    leafEncoder.flush()

    region.clear()
    leafNodeBuilder.clear(elementIdx)

    val parent = internalNodeBuilders(0)
    if (parent.size == branchingFactor)
      writeInternalNode(parent, 0)
    parent += IndexNodeInfo(indexFileOffset, firstIndex, firstKey, firstRecordOffset, firstAnnotation)

    indexFileOffset
  }

  private def flush(): Long = {
    var offsetLastBlockWritten = 0L

    if (leafNodeBuilder.size > 0)
      writeLeafNode()

    var level = 0
    while (level < internalNodeBuilders.size) {
      val node = internalNodeBuilders(level)
      val isRoot = level == internalNodeBuilders.size - 1
      if (node.size > 0 || isRoot) {
        offsetLastBlockWritten = writeInternalNode(node, level, isRoot)
      }
      level += 1
    }

    offsetLastBlockWritten
  }

  private def writeMetadata(rootOffset: Long) = {
    hConf.writeTextFile(path + "/metadata.json.gz") { out =>
      val metadata = IndexMetadata(IndexWriter.version.rep, branchingFactor, height, keyType._toPretty, annotationType._toPretty, elementIdx, "index", rootOffset, attributes)
      implicit val formats: Formats = defaultJSONFormats
      Serialization.write(metadata, out)
    }
  }
<commit_msg>fix type in index metadata (#4199)<commit_after>package is.hail.io.index

import is.hail.annotations.{Annotation, Region, RegionValueBuilder}
import is.hail.expr.types._
import is.hail.io.CodecSpec
import is.hail.utils._
import is.hail.utils.richUtils.ByteTrackingOutputStream
import org.apache.hadoop.conf.Configuration
import org.json4s.Formats
import org.json4s.jackson.Serialization

case class IndexMetadata(
  fileVersion: Int,
  branchingFactor: Int,
  height: Int,
  keyType: String,
  annotationType: String,
  nKeys: Long,
  indexPath: String,
  rootOffset: Long,
  attributes: Map[String, Any]
)

case class IndexNodeInfo(
  indexFileOffset: Long,
  firstIndex: Long,
  firstKey: Annotation,
  firstRecordOffset: Long,
  firstAnnotation: Annotation
)

object IndexWriter {
  val version: SemanticVersion = SemanticVersion(1, 0, 0)
}

class IndexWriter(
  hConf: Configuration,
  path: String,
  keyType: Type,
  annotationType: Type,
  branchingFactor: Int = 1024,
  attributes: Map[String, Any] = Map.empty[String, Any]) extends AutoCloseable {
  require(branchingFactor > 1)

  private var elementIdx = 0L
  private val region = new Region()
  private val rvb = new RegionValueBuilder(region)

  private val leafNodeBuilder = new LeafNodeBuilder(keyType, annotationType, 0L)
  private val internalNodeBuilders = new ArrayBuilder[InternalNodeBuilder]()
  internalNodeBuilders += new InternalNodeBuilder(keyType, annotationType)

  private val trackedOS = new ByteTrackingOutputStream(hConf.unsafeWriter(path + "/index"))
  private val codecSpec = CodecSpec.default
  private val leafEncoder = codecSpec.buildEncoder(leafNodeBuilder.typ)(trackedOS)
  private val internalEncoder = codecSpec.buildEncoder(InternalNodeBuilder.typ(keyType, annotationType))(trackedOS)

  private def height: Int = internalNodeBuilders.length + 1 // have one leaf node layer

  private def writeInternalNode(node: InternalNodeBuilder, level: Int, isRoot: Boolean = false): Long = {
    val indexFileOffset = trackedOS.bytesWritten

    val info = if (node.size > 0) {
      val child = node.getChild(0)
      val firstIndex = node.firstIdx
      val firstKey = child.firstKey
      val firstRecordOffset = child.firstRecordOffset
      val firstAnnotation = child.firstAnnotation
      IndexNodeInfo(indexFileOffset, firstIndex, firstKey, firstRecordOffset, firstAnnotation)
    } else {
      assert(isRoot && level == 0)
      null
    }

    internalEncoder.writeByte(1)

    val regionOffset = node.write(rvb)
    internalEncoder.writeRegionValue(region, regionOffset)
    internalEncoder.flush()

    region.clear()
    node.clear()

    if (!isRoot) {
      if (level + 1 == internalNodeBuilders.length)
        internalNodeBuilders += new InternalNodeBuilder(keyType, annotationType) // , info.firstIndex
      val parent = internalNodeBuilders(level + 1)
      if (parent.size == branchingFactor)
        writeInternalNode(parent, level + 1)
      parent += info
    }

    indexFileOffset
  }

  private def writeLeafNode(): Long = {
    val indexFileOffset = trackedOS.bytesWritten

    assert(leafNodeBuilder.size > 0)
    val child = leafNodeBuilder.getChild(0)
    val firstIndex = leafNodeBuilder.firstIdx
    val firstKey = child.key
    val firstRecordOffset = child.recordOffset
    val firstAnnotation = child.annotation

    leafEncoder.writeByte(0)

    val regionOffset = leafNodeBuilder.write(rvb)
    leafEncoder.writeRegionValue(region, regionOffset)
    leafEncoder.flush()

    region.clear()
    leafNodeBuilder.clear(elementIdx)

    val parent = internalNodeBuilders(0)
    if (parent.size == branchingFactor)
      writeInternalNode(parent, 0)
    parent += IndexNodeInfo(indexFileOffset, firstIndex, firstKey, firstRecordOffset, firstAnnotation)

    indexFileOffset
  }

  private def flush(): Long = {
    var offsetLastBlockWritten = 0L

    if (leafNodeBuilder.size > 0)
      writeLeafNode()

    var level = 0
    while (level < internalNodeBuilders.size) {
      val node = internalNodeBuilders(level)
      val isRoot = level == internalNodeBuilders.size - 1
      if (node.size > 0 || isRoot) {
        offsetLastBlockWritten = writeInternalNode(node, level, isRoot)
      }
      level += 1
    }

    offsetLastBlockWritten
  }

  private def writeMetadata(rootOffset: Long) = {
    hConf.writeTextFile(path + "/metadata.json.gz") { out =>
      val metadata = IndexMetadata(
        IndexWriter.version.rep,
        branchingFactor,
        height, 
        keyType.parsableString(),
        annotationType.parsableString(),
        elementIdx,
        "index",
        rootOffset,
        attributes)
      implicit val formats: Formats = defaultJSONFormats
      Serialization.write(metadata, out)
    }
  }

{'fraction_non_alphanumeric': 0.08819451850271942, 'fraction_numerical': 0.007038498453663219, 'mean_word_length': 3.834020618556701, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7052838', 'n_tokens_mistral': 2813, 'n_tokens_neox': 2672, 'n_words': 666}