code: def import_bgen(
      path,
      entry_fields,
      sample_file=None,
      n_partitions=None,
      block_size=None,
      index_file_map=None,
      variants=None,
      _row_fields=["varid", "rsid"],
  ) -> MatrixTable:
      """Import BGEN file(s) as a :class:`.MatrixTable`.
  
      Examples
      --------
  
      Import a BGEN file as a matrix table with GT and GP entry fields:
  
      >>> ds_result = hl.import_bgen("data/example.8bits.bgen",
      ...                            entry_fields=['GT', 'GP'],
      ...                            sample_file="data/example.8bits.sample")
  
      Import a BGEN file as a matrix table with genotype dosage entry field:
  
      >>> ds_result = hl.import_bgen("data/example.8bits.bgen",
      ...                             entry_fields=['dosage'],
      ...                             sample_file="data/example.8bits.sample")
  
      Load a single variant from a BGEN file:
  
      >>> ds_result = hl.import_bgen("data/example.8bits.bgen",
      ...                            entry_fields=['dosage'],
      ...                            sample_file="data/example.8bits.sample",
      ...                            variants=[hl.eval(hl.parse_variant('1:2000:A:G'))])
  
      Load a set of variants specified by a table expression from a BGEN file:
  
      >>> variants = hl.import_table("data/bgen-variants.txt")
      >>> variants = variants.annotate(v=hl.parse_variant(variants.v)).key_by('v')
      >>> ds_result = hl.import_bgen("data/example.8bits.bgen",
      ...                            entry_fields=['dosage'],
      ...                            sample_file="data/example.8bits.sample",
      ...                            variants=variants.v)
  
      Load a set of variants specified by a table keyed by 'locus' and 'alleles' from a BGEN file:
  
      >>> ds_result = hl.import_bgen("data/example.8bits.bgen",
      ...                            entry_fields=['dosage'],
      ...                            sample_file="data/example.8bits.sample",
      ...                            variants=variants_table)
  
      Notes
      -----
  
      Hail supports importing data from v1.2 of the `BGEN file format
      <http://www.well.ox.ac.uk/~gav/bgen_format/bgen_format.html>`__.
      Genotypes must be **unphased** and **diploid**, genotype
      probabilities must be stored with 8 bits, and genotype probability
      blocks must be compressed with zlib or uncompressed. All variants
      must be bi-allelic.
  
      Each BGEN file must have a corresponding index file, which can be generated
      with :func:`.index_bgen`. All files must have been indexed with the same
      reference genome. To load multiple files at the same time,
      use :ref:`Hadoop Glob Patterns <sec-hadoop-glob>`.
  
      If n_partitions and block_size are both specified, block_size is
      used. If neither are specified, the default is a 128MB block
      size.
  
      **Column Fields**
  
      - `s` (:py:data:`.tstr`) -- Column key. This is the sample ID imported
        from the first column of the sample file if given. Otherwise, the sample
        ID is taken from the sample identifying block in the first BGEN file if it
        exists; else IDs are assigned from `_0`, `_1`, to `_N`.
  
      **Row Fields**
  
      Between two and four row fields are created. The `locus` and `alleles` are
      always included. `_row_fields` determines if `varid` and `rsid` are also
      included. For best performance, only include fields necessary for your
      analysis. NOTE: the `_row_fields` parameter is considered an experimental
      feature and may be removed without warning.
  
      - `locus` (:class:`.tlocus` or :class:`.tstruct`) -- Row key. The chromosome
        and position. If `reference_genome` is defined, the type will be
        :class:`.tlocus` parameterized by `reference_genome`. Otherwise, the type
        will be a :class:`.tstruct` with two fields: `contig` with type
        :py:data:`.tstr` and `position` with type :py:data:`.tint32`.
      - `alleles` (:class:`.tarray` of :py:data:`.tstr`) -- Row key. An
        array containing the alleles of the variant. The reference
        allele is the first element in the array.
      - `varid` (:py:data:`.tstr`) -- The variant identifier. The third field in
        each variant identifying block.
      - `rsid` (:py:data:`.tstr`) -- The rsID for the variant. The fifth field in
        each variant identifying block.
  
      **Entry Fields**
  
      Up to three entry fields are created, as determined by
      `entry_fields`.  For best performance, include precisely those
      fields required for your analysis. It is also possible to pass an
      empty tuple or list for `entry_fields`, which can greatly
      accelerate processing speed if your workflow does not use the
      genotype data.
  
      - `GT` (:py:data:`.tcall`) -- The hard call corresponding to the genotype with
        the greatest probability. If there is not a unique maximum probability, the
        hard call is set to missing.
      - `GP` (:class:`.tarray` of :py:data:`.tfloat64`) -- Genotype probabilities
        as defined by the BGEN file spec. For bi-allelic variants, the array has
        three elements giving the probabilities of homozygous reference,
        heterozygous, and homozygous alternate genotype, in that order.
      - `dosage` (:py:data:`.tfloat64`) -- The expected value of the number of
        alternate alleles, given by the probability of heterozygous genotype plus
        twice the probability of homozygous alternate genotype. All variants must
        be bi-allelic.
  
      See Also
      --------
      :func:`.index_bgen`
  
      Parameters
      ----------
      path : :class:`str` or :obj:`list` of :obj:`str`
          BGEN file(s) to read.
      entry_fields : :obj:`list` of :class:`str`
          List of entry fields to create.
          Options: ``'GT'``, ``'GP'``, ``'dosage'``.
      sample_file : :class:`str`, optional
          Sample file to read the sample ids from. If specified, the number of
          samples in the file must match the number in the BGEN file(s).
      n_partitions : :obj:`int`, optional
          Number of partitions.
      block_size : :obj:`int`, optional
          Block size, in MB.
      index_file_map : :obj:`dict` of :class:`str` to :obj:`str`, optional
          Dict of BGEN file to index file location. Cannot use Hadoop glob
          patterns in file names.
      variants : :class:`.StructExpression` or :class:`.LocusExpression` or :obj:`list` of :class:`.Struct` or :obj:`list` of :class:`.Locus` or :class:`.Table`
          Variants to filter to. The underlying type of the input (row key in the case of a :class:`.Table`)
          must be either a :class:`.tlocus`, a struct with one field `locus`, or a struct with two fields:
          `locus` and `alleles`. The type of `locus` can either be a :class:`.tlocus` or a :class:`.tstruct`
          with two fields: `contig` of type :obj:`.tstr` and `position` of type :obj:`.tint`. If the
          type of `locus` is :class:`.tlocus`, the reference genome must match that used to index the BGEN
          file(s). The type of `alleles` is a :class:`.tarray` of :obj:`.tstr`.
      _row_fields : :obj:`list` of :class:`str`
          List of non-key row fields to create.
          Options: ``'varid'``, ``'rsid'``
  
      Returns
      -------
      :class:`.MatrixTable`
  
      """
  
      if n_partitions is None and block_size is None:
          block_size = 128
  
      if index_file_map is None:
          index_file_map = {}
  
      entry_set = set(entry_fields)
      row_set = set(_row_fields)
  
      if variants is not None:
          mt_type = Env.backend().matrix_type(
              ir.MatrixRead(
                  ir.MatrixBGENReader(
                      path,
                      sample_file,
                      index_file_map,
                      n_partitions,
                      block_size,
                      None,
                  )
              )
          )
          lt = mt_type.row_type["locus"]
  
          expected_vtype = tstruct(locus=lt, alleles=tarray(tstr))
  
          if isinstance(variants, StructExpression) or isinstance(
              variants, LocusExpression
          ):
              if isinstance(variants, LocusExpression):
                  variants = hl.struct(locus=variants)
  
              if len(
                  variants.dtype
              ) == 0 or not variants.dtype._is_prefix_of(
                  expected_vtype
              ):
                  raise TypeError(
                      "'import_bgen' requires the expression type for 'variants' is a non-empty prefix of the BGEN key type: \n"
                      + f"\tFound: {repr(variants.dtype)}\n"
                      + f"\tExpected: {repr(expected_vtype)}\n"
                  )
  
              uid = Env.get_uid()
              fnames = list(variants.dtype)
              name, variants = variants._to_table(
                  uid
              )  # This will add back the other key fields of the source, which we don't want
              variants = variants.key_by(
                  **{fname: variants[name][fname] for fname in fnames}
              )
              variants = variants.select()
          elif isinstance(variants, Table):
              if len(
                  variants.key
              ) == 0 or not variants.key.dtype._is_prefix_of(
                  expected_vtype
              ):
                  raise TypeError(
                      "'import_bgen' requires the row key type for 'variants' is a non-empty prefix of the BGEN key type: \n"
                      + f"\tFound: {repr(variants.key.dtype)}\n"
                      + f"\tExpected: {repr(expected_vtype)}\n"
                  )
              variants = variants.select()
          else:
              assert isinstance(variants, list)
              try:
                  if len(variants) == 0:
                      variants = hl.Table.parallelize(
                          variants,
                          schema=expected_vtype,
                          key=["locus", "alleles"],
                      )
                  else:
                      first_v = variants[0]
                      if isinstance(first_v, hl.Locus):
                          variants = hl.Table.parallelize(
                              [hl.Struct(locus=v) for v in variants],
                              schema=hl.tstruct(locus=lt),
                              key="locus",
                          )
                      else:
                          assert isinstance(first_v, hl.utils.Struct)
                          if len(first_v) == 1:
                              variants = hl.Table.parallelize(
                                  variants,
                                  schema=hl.tstruct(locus=lt),
                                  key="locus",
                              )
                          else:
                              variants = hl.Table.parallelize(
                                  variants,
                                  schema=expected_vtype,
                                  key=["locus", "alleles"],
                              )
              except Exception:
                  raise TypeError(
                      f"'import_bgen' requires all elements in 'variants' are a non-empty prefix of the BGEN key type: {repr(expected_vtype)}"
                  )
  
          vir = variants._tir
          if (
              isinstance(vir, ir.TableRead)
              and isinstance(vir.reader, ir.TableNativeReader)
              and vir.reader.intervals is None
              and variants.count() == variants.distinct().count()
          ):
              variants_path = vir.reader.path
          else:
              variants_path = new_temp_file(
                  prefix="bgen_included_vars", extension="ht"
              )
              variants.distinct().write(variants_path)
      else:
          variants_path = None
  
      reader = ir.MatrixBGENReader(
          path,
          sample_file,
          index_file_map,
          n_partitions,
          block_size,
          variants_path,
      )
  
      mt = MatrixTable(ir.MatrixRead(reader)).drop(
          *[fd for fd in ["GT", "GP", "dosage"] if fd not in entry_set],
          *[
              fd
              for fd in ["rsid", "varid", "offset", "file_idx"]
              if fd not in row_set
          ],
      )
  
      return mt
  
imports: import os
  import re
  from collections import defaultdict
  import avro.schema
  from avro.datafile import DataFileReader
  from avro.io import DatumReader
  import hail
  from hail import ir
  from hail.expr import StructExpression, LocusExpression, expr_array, expr_float64, expr_str, expr_numeric, expr_call, expr_bool, expr_int32, to_expr, analyze
  from hail.expr.types import hail_type, tarray, tfloat64, tstr, tint32, tstruct, tcall, tbool, tint64, tfloat32
  from hail.genetics.reference_genome import reference_genome_type
  from hail.ir.utils import parse_type
  from hail.matrixtable import MatrixTable
  from hail.methods.misc import require_biallelic, require_row_key_variant, require_col_key_str
  from hail.table import Table
  from hail.typecheck import typecheck, nullable, oneof, dictof, anytype, sequenceof, enumeration, sized_tupleof, numeric, table_key_type, char
  from hail.utils import new_temp_file
  from hail.utils.deduplicate import deduplicate
  from hail.utils.java import Env, FatalError, jindexed_seq_args, warning
  from hail.utils.java import info
  from hail.utils.misc import wrap_to_list, plural
  from .import_lines_helpers import split_lines, should_remove_line
label_desc: Use image node_id 11 for calling a global function with example usage: import_bgen(path, entry_fields, sample_file, n_partitions, block_size, index_file_map, variants, _row_fields) and returns: mt