Fix possible leakage when streaming SC vectors

Author: arteymix

gemma-cli/src/main/java/ubic/gemma/apps/SingleCellDataWriterCli.java

@@ -295,16 +295,14 @@ private int aggregate( ExpressionExperiment ee, QuantitationType qt, @Nullable L
         if ( useStreaming ) {
             log.info( "Single-cell data will be streamed by batch of " + fetchSize + " vectors." );
             long numberOfVectors = singleCellExpressionExperimentService.getNumberOfSingleCellDataVectors( ee, qt );
-            Stream<SingleCellExpressionDataVector> scVecs;
-            if ( assays != null ) {
-                scVecs = singleCellExpressionExperimentService.streamSingleCellDataVectors( ee, assays, qt, fetchSize, useCursorFetchIfSupported, true, config );
-            } else {
-                scVecs = singleCellExpressionExperimentService.streamSingleCellDataVectors( ee, qt, fetchSize, useCursorFetchIfSupported, true, config );
+            try ( Stream<SingleCellExpressionDataVector> scVecs = assays != null ?
+                    singleCellExpressionExperimentService.streamSingleCellDataVectors( ee, assays, qt, fetchSize, useCursorFetchIfSupported, true, config ) :
+                    singleCellExpressionExperimentService.streamSingleCellDataVectors( ee, qt, fetchSize, useCursorFetchIfSupported, true, config ) ) {
+                vecs = scVecs
+                        .peek( createStreamMonitor( ee, qt, getClass().getName(), 100, numberOfVectors ) )
+                        .map( createAggregator( aggregationMethod, cellLevelCharacteristics, aggregateUnknownCharacteristics ) )
+                        .collect( Collectors.toList() );
             }
-            vecs = scVecs
-                    .peek( createStreamMonitor( ee, qt, getClass().getName(), 100, numberOfVectors ) )
-                    .map( createAggregator( aggregationMethod, cellLevelCharacteristics, aggregateUnknownCharacteristics ) )
-                    .collect( Collectors.toList() );
         } else {
             log.info( "Single-cell data will be loaded into memory. This process can use a lot of memory, press Ctrl-C at any time to interrupt." );
             Collection<SingleCellExpressionDataVector> scVecs;

gemma-core/src/main/java/ubic/gemma/core/analysis/service/ExpressionDataFileHelperService.java

@@ -118,7 +118,7 @@ public ExpressionDataDoubleMatrix getDataMatrix( ExpressionExperiment ee, @Nulla
         return matrix;
     }
 
-    public Stream<SingleCellExpressionDataVector> getSingleCellVectors( ExpressionExperiment ee, @Nullable List<BioAssay> samples, QuantitationType qt, Map<CompositeSequence, Set<Gene>> cs2gene, AtomicLong numVecs1, int fetchSize, boolean useCursorFetchIfSupported ) {
+    public Stream<SingleCellExpressionDataVector> streamSingleCellVectors( ExpressionExperiment ee, @Nullable List<BioAssay> samples, QuantitationType qt, Map<CompositeSequence, Set<Gene>> cs2gene, AtomicLong numVecs1, int fetchSize, boolean useCursorFetchIfSupported ) {
         long numVecs = singleCellExpressionExperimentService.getNumberOfSingleCellDataVectors( ee, qt );
         if ( numVecs == 0 ) {
             throw new IllegalStateException( "There are no vectors for " + qt + " in " + ee + "." );
@@ -142,7 +142,7 @@ public Collection<SingleCellExpressionDataVector> getSingleCellVectors( Expressi
         return vectors;
     }
 
-    public Stream<SingleCellExpressionDataVector> getSingleCellVectors( ExpressionExperiment ee, @Nullable List<BioAssay> samples, QuantitationType qt, Map<CompositeSequence, Set<Gene>> cs2gene, AtomicLong numVecs1, Map<BioAssay, Long> nnzBySample, int fetchSize, boolean useCursorFetchIfSupported ) {
+    public Stream<SingleCellExpressionDataVector> streamSingleCellVectors( ExpressionExperiment ee, @Nullable List<BioAssay> samples, QuantitationType qt, Map<CompositeSequence, Set<Gene>> cs2gene, AtomicLong numVecs1, Map<BioAssay, Long> nnzBySample, int fetchSize, boolean useCursorFetchIfSupported ) {
         long numVecs = singleCellExpressionExperimentService.getNumberOfSingleCellDataVectors( ee, qt );
         if ( numVecs == 0 ) {
             throw new IllegalStateException( "There are no vectors for " + qt + " in " + ee + "." );

gemma-core/src/main/java/ubic/gemma/core/analysis/service/ExpressionDataFileServiceImpl.java

@@ -547,7 +547,7 @@ private int writeTabularSingleCellExpressionDataInternal( ExpressionExperiment e
         }
         if ( fetchSize > 0 ) {
             AtomicLong numVecs = new AtomicLong();
-            try ( Stream<SingleCellExpressionDataVector> vectors = helperService.getSingleCellVectors( ee, samples, qt, cs2gene, numVecs, fetchSize, useCursorFetchIfSupported ) ) {
+            try ( Stream<SingleCellExpressionDataVector> vectors = helperService.streamSingleCellVectors( ee, samples, qt, cs2gene, numVecs, fetchSize, useCursorFetchIfSupported ) ) {
                 return matrixWriter.write( vectors.peek( createStreamMonitor( ee, qt, ExpressionDataFileServiceImpl.class.getName(), 100, numVecs.get() ) ), cs2gene, writer );
             }
         } else {
@@ -569,7 +569,7 @@ public int writeCellBrowserSingleCellExpressionData( ExpressionExperiment ee, Qu
         }
         if ( fetchSize > 0 ) {
             AtomicLong numVecs = new AtomicLong();
-            try ( Stream<SingleCellExpressionDataVector> vectors = helperService.getSingleCellVectors( ee, null, qt, cs2gene, numVecs, fetchSize, useCursorFetchIfSupported ) ) {
+            try ( Stream<SingleCellExpressionDataVector> vectors = helperService.streamSingleCellVectors( ee, null, qt, cs2gene, numVecs, fetchSize, useCursorFetchIfSupported ) ) {
                 return matrixWriter.write( vectors.peek( createStreamMonitor( ee, qt, ExpressionDataFileServiceImpl.class.getName(), 100, numVecs.get() ) ), cs2gene, writer );
             }
         } else {
@@ -657,7 +657,7 @@ private int writeMexSingleCellExpressionDataInternal( ExpressionExperiment ee, @
             if ( fetchSize > 0 ) {
                 Map<BioAssay, Long> nnzBySample = new HashMap<>();
                 AtomicLong numVecs = new AtomicLong();
-                try ( Stream<SingleCellExpressionDataVector> vectors = helperService.getSingleCellVectors( ee, samples, qt, cs2gene, numVecs, nnzBySample, fetchSize, useCursorFetchIfSupported ) ) {
+                try ( Stream<SingleCellExpressionDataVector> vectors = helperService.streamSingleCellVectors( ee, samples, qt, cs2gene, numVecs, nnzBySample, fetchSize, useCursorFetchIfSupported ) ) {
                     if ( Files.exists( destDir ) ) {
                         log.info( destDir + " already exists, removing..." );
                         PathUtils.deleteDirectory( destDir );