Merge pull request #1 from amanas/writer-wo-pandas

fix: fixed bug related to save variant keys as pandas dataframe

Author: amanas

src/dnarecords/reader.py

@@ -84,6 +84,16 @@ def _pandas_safe_read_parquet(path):
             return pd.concat(pd.read_parquet(f) for f in files)
         return None
 
+    @staticmethod
+    def _pandas_safe_read_json(path):
+        import pandas as pd
+        import tensorflow as tf
+
+        files = tf.io.gfile.glob(f'{path}/*.json')
+        if files:
+            return pd.concat(pd.read_json(f) for f in files)
+        return None
+
     def metadata(self) -> Dict[str, DataFrame]:
         """Gets the metadata associated to the DNARecords dataset as a dictionary of names to pandas DataFrames.
 
@@ -99,8 +109,10 @@ def metadata(self) -> Dict[str, DataFrame]:
         result = {}
         tree = dr.helper.DNARecordsUtils.dnarecords_tree(self._dnarecords_path)
         for k, v in tree.items():
-            if k in ['skeys', 'vkeys', 'swpfs', 'vwpfs', 'swrfs', 'vwrfs', 'swpsc', 'vwpsc', 'swrsc', 'vwrsc']:
+            if k in ['skeys', 'vkeys', 'swpfs', 'vwpfs', 'swrfs', 'vwrfs']:
                 result.update({k: self._pandas_safe_read_parquet(v)})
+            if k in ['swpsc', 'vwpsc', 'swrsc', 'vwrsc']:
+                result.update({k: self._pandas_safe_read_json(v)})
         return result
 
     def datafiles(self) -> Dict[str, List[str]]:
@@ -132,9 +144,8 @@ def datafiles(self) -> Dict[str, List[str]]:
     def _sw_decoder(dnarecords, schema, gzip):
         import json
         import tensorflow as tf
-
         one_proto = next(iter(tf.data.TFRecordDataset(dnarecords, 'GZIP' if gzip else None)))
-        swrsc_dict = {f['name']: f for f in schema}
+        swrsc_dict = {f['fields']['name']: f['fields'] for _, f in schema.iterrows()}
         features = {'key': tf.io.FixedLenFeature([], tf.int64)}
         for indices_field in [field for field in swrsc_dict.keys() if field.endswith('indices')]:
             feature_name = indices_field.replace('_indices', '')
@@ -152,7 +163,7 @@ def _vw_decoder(dnarecords, schema, gzip):
         import tensorflow as tf
 
         one_proto = next(iter(tf.data.TFRecordDataset(dnarecords, 'GZIP' if gzip else None)))
-        vwrsc_dict = {f['name']: f for f in schema}
+        vwrsc_dict = {f['fields']['name']: f['fields'] for _, f in schema.iterrows()}
         values_type = DNARecordsReader._types_dict()[json.loads(vwrsc_dict['values']['type'])['elementType']]
         dense_shape = tf.io.parse_example(one_proto, {'dense_shape': tf.io.FixedLenFeature([], tf.int64)})[
             'dense_shape']
@@ -193,7 +204,7 @@ def sample_wise_dataset(self, num_parallel_reads: int = -1, num_parallel_calls:
         schema = self.metadata()['swrsc']
         if schema is None or not dnarecords:
             raise Exception(f"No DNARecords found at {self._dnarecords_path}/...")
-        decoder = self._sw_decoder(dnarecords, schema.fields[0], self._gzip)
+        decoder = self._sw_decoder(dnarecords, schema, self._gzip)
         return self._dataset(dnarecords, decoder, num_parallel_reads, num_parallel_calls, deterministic, drop_remainder,
                              batch_size, buffer_size)
 
@@ -219,7 +230,7 @@ def variant_wise_dataset(self, num_parallel_reads: int = -1, num_parallel_calls:
         schema = self.metadata()['vwrsc']
         if schema is None or not dnarecords:
             raise Exception(f"No DNARecords found at {self._dnarecords_path}/...")
-        decoder = self._vw_decoder(dnarecords, schema.fields[0], self._gzip)
+        decoder = self._vw_decoder(dnarecords, schema, self._gzip)
         return self._dataset(dnarecords, decoder, num_parallel_reads, num_parallel_calls, deterministic, drop_remainder,
                              batch_size, buffer_size)
 
@@ -311,8 +322,10 @@ def metadata(self) -> Dict[str, 'DataFrame']:
         tree = dr.helper.DNARecordsUtils.dnarecords_tree(self._dnarecords_path)
         spark = dr.helper.DNARecordsUtils.spark_session()
         for k, v in tree.items():
-            if k in ['skeys', 'vkeys', 'swpfs', 'vwpfs', 'swrfs', 'vwrfs', 'swpsc', 'vwpsc', 'swrsc', 'vwrsc']:
+            if k in ['skeys', 'vkeys', 'swpfs', 'vwpfs', 'swrfs', 'vwrfs']:
                 result.update({k: self._spark_safe_load(spark.read.format("parquet"), v)})
+            if k in ['swpsc', 'vwpsc', 'swrsc', 'vwrsc']:
+                result.update({k: self._spark_safe_load(spark.read.format("json"), v)})
         return result
 
     def sample_wise_dnarecords(self) -> 'DataFrame':

src/dnarecords/writer.py

@@ -111,9 +111,10 @@ def _set_vkeys_skeys(self):
         self._skeys = self._mt.key_cols_by().cols().to_spark().withColumnRenamed('j', 'key').cache()
 
     def _set_chrom_ranges(self):
-        gdf = self._vkeys.toPandas()[['locus.contig', 'key']].groupby('locus.contig', as_index=False)
-        gdf = gdf.agg(start=('key', 'min'), end=('key', 'max'))
-        self._chrom_ranges = {r['locus.contig']: [r['start'], r['end']] for i, r in gdf.iterrows()}
+        from pyspark.sql import functions as F
+        gdf = self._vkeys.select('`locus.contig`', 'key').groupby('`locus.contig`')
+        gdf = gdf.agg(F.min('key').alias('start'), F.max('key').alias('end'))
+        self._chrom_ranges = {r['locus.contig']: [r['start'], r['end']] for i, r in gdf.toPandas().iterrows()}
 
     def _update_vkeys_by_chrom_ranges(self):
         from dnarecords.helper import DNARecordsUtils
@@ -129,11 +130,13 @@ def _select_ijv(self):
         self._mt = self._mt.select_globals().select_rows().select_cols().select_entries('v')
 
     def _filter_out_undefined_entries(self):
-        import hail as hl
+        from dnarecords.helper import DNARecordsUtils
+        hl = DNARecordsUtils.init_hail()
         self._mt = self._mt.filter_entries(hl.is_defined(self._mt.v))
 
     def _filter_out_zeroes(self):
-        import hail as hl
+        from dnarecords.helper import DNARecordsUtils
+        hl = DNARecordsUtils.init_hail()
         self._mt = self._mt.filter_entries(0 != hl.coalesce(self._mt.v, 0))
 
     def _set_max_nrows_ncols(self):
@@ -163,7 +166,8 @@ def _build_ij_blocks(self):
 
     def _set_ij_blocks(self):
         import re
-        import hail as hl
+        from dnarecords.helper import DNARecordsUtils
+        hl = DNARecordsUtils.init_hail()
         all_blocks = [p for p in hl.hadoop_ls(f'{self._kv_blocks_path}/*') if p['is_dir']]
         self._i_blocks = {re.search(r'ib=(\d+)', p['path']).group(1) for p in all_blocks}
         self._j_blocks = {re.search(r'jb=(\d+)', p['path']).group(1) for p in all_blocks}
@@ -280,8 +284,8 @@ def _write_dnarecords(output, output_schema, dna_blocks, write_mode, gzip, tfrec
             if gzip:
                 df_writer = df_writer.option("compression", "gzip")
         df_writer.save(output)
-        sc_writer = spark.read.json(spark.sparkContext.parallelize([df.schema.json()])).repartition(1).write
-        sc_writer.mode(write_mode).parquet(output_schema)
+        sc_writer = spark.read.json(spark.sparkContext.parallelize([df.schema.json()])).coalesce(1).write
+        sc_writer.mode(write_mode).format('json').save(output_schema)
 
     @staticmethod
     def _write_key_files(source, output, tfrecord_format, write_mode):
@@ -294,7 +298,7 @@ def _write_key_files(source, output, tfrecord_format, write_mode):
         else:
             reader = spark.read.format("parquet")
         df = reader.load(source).withColumn("path", F.regexp_extract(F.input_file_name(), f"(.*){source}/(.*)", 2))
-        df.select('key', 'path').repartition(1).write.mode(write_mode).parquet(output)
+        df.select('key', 'path').write.mode(write_mode).parquet(output)
 
     # pylint: disable=too-many-arguments
     # It is reasonable in this case.
@@ -365,5 +369,5 @@ def write(self, output: str, sparse: bool = True, sample_wise: bool = True, vari
                                        gzip, False)
                 self._write_key_files(otree['swpar'], otree['swpfs'], False, write_mode)
 
-        self._vkeys.repartition(1).write.mode(write_mode).parquet(otree['vkeys'])
-        self._skeys.repartition(1).write.mode(write_mode).parquet(otree['skeys'])
+        self._vkeys.write.mode(write_mode).parquet(otree['vkeys'])
+        self._skeys.write.mode(write_mode).parquet(otree['skeys'])