implements by option

* implements "by" read/write feature
* add documentation for "by" read/write feature

Signed-off-by: Jorge Aguilera <jorge@edn.es>

Author: jagedn

docs/asciidoc/index.adoc

@@ -139,4 +139,28 @@ workflow{
         .toParquet( params.output, [record:AugmentedRecord])
         | view
 }
-----
+----
+
+== Options
+
+== splitParquet
+
+=== record
+
+"record" option allows specifying a *Record* java class instead of a raw Map. This option can improve parsing time as the
+reader can use only a subset of the fields to read
+
+=== by
+
+"by" option allows specifying a chunk size to read records.
+
+== toParquet
+
+=== record
+
+"record" option allows specifying a *Record* java class instead of a raw Map. This option is required as the writer
+needs to know the structure of the record
+
+=== by
+
+"by" option allows specifying a chunk size to write records.

src/main/groovy/nextflow/parquet/ParquetExtension.groovy

@@ -76,6 +76,8 @@ class ParquetExtension extends PluginExtensionPoint {
     class ParquetSplitter {
         private DataflowWriteChannel target
         private Class<Record> clazz
+        private List batchList = []
+        private int sizeBatch = 0
 
         ParquetSplitter(DataflowWriteChannel target, Map params) {
             this.target = target
@@ -85,6 +87,8 @@ class ParquetExtension extends PluginExtensionPoint {
                 }
                 this.clazz = params.record as Class<Record>
             }
+            batchList = []
+            sizeBatch = params.containsKey('by') && "$params.by".isNumber() ? params.by as int : 0
         }
 
         void apply(Object source) {
@@ -93,7 +97,15 @@ class ParquetExtension extends PluginExtensionPoint {
                 // create parquet reader
                 final reader = new CarpetReader(toFile(source), clazz ?: Map)
                 for (def record : reader) {
-                    target << record
+                    batchList << record
+                    if( batchList.size() >= sizeBatch ){
+                        target << (sizeBatch == 0 ? batchList.first() : batchList.toArray())
+                        batchList.clear()
+                    }
+                }
+                if( batchList.size() ){
+                    target << (sizeBatch == 0 ? batchList.first() : batchList.toArray())
+                    batchList.clear()
                 }
             }
             catch( IOException e ) {
@@ -114,20 +126,33 @@ class ParquetExtension extends PluginExtensionPoint {
     class ParquetWriter implements Closeable{
         private Class<Record> clazz
         CarpetWriter writer
+        private List batchList = []
+        private int sizeBatch = 0
         ParquetWriter(String output, Map params){
             if( !params.record ||!(params.record instanceof Class<Record>)) {
                 throw new IllegalArgumentException("A Record.class is required. Class provided $params.record")
             }
             this.clazz = params.record as Class<Record>
             var outputStream = new FileOutputStream(output)
             writer = new CarpetWriter<>(outputStream, this.clazz)
+            batchList = []
+            sizeBatch = params.containsKey('by') && "$params.by".isNumber() ? params.by as int : 0
         }
 
         void write(Record record){
-            writer.write(record)
+            batchList << record
+            if( batchList.size() >= sizeBatch ) {
+                batchList.each { r ->
+                    writer.write(r)
+                }
+                batchList.clear()
+            }
         }
 
         void close(){
+            batchList.each { r ->
+                writer.write(r)
+            }
             writer.close()
         }
     }

src/test/groovy/nextflow/parquet/PluginTest.groovy

@@ -1,5 +1,6 @@
 package nextflow.parquet
 
+import com.jerolba.carpet.CarpetReader
 import nextflow.Channel
 import nextflow.plugin.Plugins
 import nextflow.plugin.TestPluginDescriptorFinder
@@ -119,4 +120,75 @@ class PluginTest extends Dsl2Spec{
         pathOutput.toFile().length()
     }
 
+    def 'should parse a parquet file in raw mode using by 1'() {
+        given:
+        def path = getClass().getResource('/multiple.parquet').toURI().path
+
+        when:
+        def SCRIPT = """
+        include {splitParquet} from 'plugin/nf-parquet'
+        channel.fromPath("$path").splitParquet(by:1) 
+        """.toString()
+        and:
+        def result = new MockScriptRunner([:]).setScript(SCRIPT).execute()
+        then:
+        result.val == [[id: 1, name: "The 1 record"]]
+        result.val == [[id: 2, name: "The 2 record"]]
+        result.val == [[id: 3, name: "The 3 record"]]
+        result.val == Channel.STOP
+    }
+
+    def 'should parse a parquet file in raw mode using by 2'(){
+        given:
+        def path = getClass().getResource('/multiple.parquet').toURI().path
+        when:
+        def SCRIPT = """
+        include {splitParquet} from 'plugin/nf-parquet'
+        channel.fromPath("$path").splitParquet(by:2) 
+        """.toString()
+        and:
+        def result = new MockScriptRunner([:]).setScript(SCRIPT).execute()
+        then:
+        result.val == [ [id:1, name:"The 1 record"], [id:2, name:"The 2 record"] ]
+        result.val == [ [id:3, name:"The 3 record"] ]
+        result.val == Channel.STOP
+
+    }
+
+    def 'should write #total records to a file using by #by '(){
+        when:
+        def pathOutput = Files.createTempFile("", ".parquet")
+        def SCRIPT = """
+        include {splitParquet; toParquet} from 'plugin/nf-parquet'
+
+        import nextflow.parquet.DemoRecord
+
+        channel.of(1..$total)
+                .map( { new DemoRecord(it, "The \$it record") } )
+                .toParquet("$pathOutput", [record:DemoRecord, by:$by])
+                .view()
+        """.toString()
+        and:
+        def result = new MockScriptRunner([:]).setScript(SCRIPT).execute()
+        def values = (1..total).collect{ result.val }
+        then:
+        values.size() == total
+        result.val == Channel.STOP
+
+        when:
+        final reader = new CarpetReader(pathOutput.toFile(), Map)
+        int count = 0
+        for (def record : reader) {
+            println record
+            count++
+        }
+        then:
+        count == total
+
+        where:
+        by       | total
+        1       | 100
+        10      | 100
+        100     | 100
+    }
 }

src/test/resources/multiple.parquet

@@ -0,0 +1,23 @@
+include { splitParquet } from 'plugin/nf-parquet'
+
+import myrecords.*
+
+
+process processParquetChunk {
+    input:
+    val(input_chunk_parquet)
+    output:
+    stdout 
+    script:
+    """
+        echo ${input_chunk_parquet}
+    """
+}
+
+workflow{
+
+channel.fromPath("*.parquet").splitParquet(by:100)
+        | processParquetChunk
+        | view
+
+}

validation/read_process.nf

@@ -2,5 +2,5 @@ include { splitParquet } from 'plugin/nf-parquet'
 
 import myrecords.*
 
-channel.fromPath("*.parquet").splitParquet()
-        | view
+channel.fromPath("*.parquet").splitParquet(by:1)
+        | view