unikmer count: fix setting global taxid

shenwei356 · Jul 3, 2020 · 84091af · 84091af
1 parent 161a0ae
commit 84091af
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -4,7 +4,7 @@
     - `unikmer dump`: fix a nil pointer bug.
     - `unikmer count`:
         - fix checking taxid in sequence header.
-        - disable global taxid when using parsing taxid from sequence header.
+        - fix setting global taxid.
     - `unikmer count/diff/union`: slightly reduce memory and speedup when sorting k-mers.
 - v0.10.0
     - `unikmer`: fix loading custom taxonomy files.

diff --git a/unikmer/cmd/count.go b/unikmer/cmd/count.go
@@ -59,6 +59,7 @@ var countCmd = &cobra.Command{
 		sortKmers := getFlagBool(cmd, "sort")
 
 		taxid := getFlagUint32(cmd, "taxid")
+		var setGlobalTaxid bool
 
 		parseTaxid := getFlagBool(cmd, "parse-taxid")
 		parseTaxidRegexp := getFlagString(cmd, "parse-taxid-regexp")
@@ -86,6 +87,8 @@ var countCmd = &cobra.Command{
 			if err != nil {
 				checkError(fmt.Errorf("invalid regular express: %s", parseTaxidRegexp))
 			}
+		} else if taxid > 0 {
+			setGlobalTaxid = true
 		}
 
 		if opt.Verbose {
@@ -116,6 +119,10 @@ var countCmd = &cobra.Command{
 		var mode uint32
 		var writer *unikmer.Writer
 
+		if setGlobalTaxid && opt.Verbose {
+			log.Infof("set global taxid: %d", taxid)
+		}
+
 		if !parseTaxid && !sortKmers {
 			if sortKmers {
 				mode |= unikmer.UNIK_SORTED
@@ -131,7 +138,7 @@ var countCmd = &cobra.Command{
 			writer, err = unikmer.NewWriter(outfh, k, mode)
 			checkError(err)
 			writer.SetMaxTaxid(opt.MaxTaxid)
-			if taxid > 0 {
+			if setGlobalTaxid {
 				checkError(writer.SetGlobalTaxid(taxid))
 			}
 		}
@@ -309,25 +316,31 @@ var countCmd = &cobra.Command{
 
 		if sortKmers || parseTaxid {
 			var mode uint32
+			if sortKmers {
+				mode |= unikmer.UNIK_SORTED
+			} else if opt.Compact {
+				mode |= unikmer.UNIK_COMPACT
+			}
 			if canonical {
 				mode |= unikmer.UNIK_CANONICAL
 			}
 			if parseTaxid {
 				mode |= unikmer.UNIK_INCLUDETAXID
 			}
-			if sortKmers {
-				mode |= unikmer.UNIK_SORTED
-			}
 			writer, err = unikmer.NewWriter(outfh, k, mode)
 			checkError(err)
 			writer.SetMaxTaxid(opt.MaxTaxid)
+			if setGlobalTaxid {
+				checkError(writer.SetGlobalTaxid(taxid))
+			}
 
 			if parseTaxid {
 				n = int64(len(mt))
 			} else {
 				n = int64(len(m))
 			}
 			writer.Number = int64(n)
+
 		}
 
 		var code uint64