Skip to content

Commit 16b89d1

Browse files
jagednjagedn
authored
refactor of reader and writer
Signed-off-by: Jorge Aguilera <jorge@edn.es>
1 parent 58099d1 commit 16b89d1

9 files changed

Lines changed: 349 additions & 89 deletions

File tree

‎src/main/groovy/nextflow/parquet/EmitterRecordChannel.groovy‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
package nextflow.parquet
2+
3+
import groovy.transform.CompileStatic
4+
import groovy.util.logging.Slf4j
5+
import groovyx.gpars.dataflow.DataflowWriteChannel
6+
import nextflow.Channel
7+
import nextflow.parquet.impl.ReadRecordAware
8+
9+
@CompileStatic
10+
@Slf4j
11+
class EmitterRecordChannel implements ReadRecordAware{
12+
13+
private final DataflowWriteChannel target
14+
15+
EmitterRecordChannel(DataflowWriteChannel target) {
16+
this.target = target
17+
}
18+
19+
@Override
20+
void recordRead(Object row) {
21+
target << row
22+
}
23+
24+
void complete(){
25+
target << Channel.STOP
26+
}
27+
}

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

Lines changed: 14 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
package nextflow.parquet
22

3-
4-
import java.nio.file.Path
3+
import nextflow.parquet.impl.ParquetReader
4+
import nextflow.parquet.impl.ParquetWriter
55

66
import groovy.transform.CompileStatic
77
import groovy.util.logging.Slf4j
@@ -14,8 +14,8 @@ import nextflow.extension.DataflowHelper
1414
import nextflow.plugin.extension.Operator
1515
import nextflow.plugin.extension.PluginExtensionPoint
1616

17-
import com.jerolba.carpet.CarpetReader
18-
import com.jerolba.carpet.CarpetWriter
17+
import java.nio.file.Path
18+
1919

2020
/**
2121
* Implements extensions for reading and writing Parquet files.
@@ -42,10 +42,12 @@ class ParquetExtension extends PluginExtensionPoint {
4242
@Operator
4343
DataflowWriteChannel splitParquet(DataflowReadChannel source, Map params=[:]) {
4444
final target = CH.create()
45-
final splitter = new ParquetSplitter(target, params)
45+
final emitter = new EmitterRecordChannel(target)
46+
final reader = new ParquetReader(emitter, params)
47+
48+
final onNext = { path -> reader.readFile(toFile(path)) }
49+
final onComplete = { emitter.complete() }
4650

47-
final onNext = { it -> splitter.apply(it) }
48-
final onComplete = { target << Channel.STOP }
4951
DataflowHelper.subscribeImpl(source, [onNext: onNext, onComplete: onComplete])
5052

5153
return target
@@ -73,87 +75,11 @@ class ParquetExtension extends PluginExtensionPoint {
7375
return target
7476
}
7577

76-
class ParquetSplitter {
77-
private DataflowWriteChannel target
78-
private Class<Record> clazz
79-
private List batchList = []
80-
private int sizeBatch = 0
81-
82-
ParquetSplitter(DataflowWriteChannel target, Map params) {
83-
this.target = target
84-
if( params.record ) {
85-
if (!(params.record instanceof Class<Record>)) {
86-
throw new IllegalArgumentException("A Record.class is required. Class provided $params.record")
87-
}
88-
this.clazz = params.record as Class<Record>
89-
}
90-
batchList = []
91-
sizeBatch = params.containsKey('by') && "$params.by".isNumber() ? params.by as int : 0
92-
}
93-
94-
void apply(Object source) {
95-
try {
96-
log.debug "Start reading $source, with projection ${clazz ?: 'raw'}"
97-
// create parquet reader
98-
final reader = new CarpetReader(toFile(source), clazz ?: Map)
99-
for (def record : reader) {
100-
batchList << record
101-
if( batchList.size() >= sizeBatch ){
102-
target << (sizeBatch == 0 ? batchList.first() : batchList.toArray())
103-
batchList.clear()
104-
}
105-
}
106-
if( batchList.size() ){
107-
target << (sizeBatch == 0 ? batchList.first() : batchList.toArray())
108-
batchList.clear()
109-
}
110-
}
111-
catch( IOException e ) {
112-
throw new IllegalStateException("Error while reading Parquet file - cause: ${e.message ?: e}", e)
113-
}
114-
}
115-
116-
private File toFile(Object source){
117-
return switch( source ){
118-
case {it instanceof String}->Path.of(source as String).toFile()
119-
case {it instanceof Path} -> (source as Path).toFile()
120-
default->throw new IllegalArgumentException("Invalid input for splitParquet operator: ${source}")
121-
}
122-
}
123-
124-
}
125-
126-
class ParquetWriter implements Closeable{
127-
private Class<Record> clazz
128-
CarpetWriter writer
129-
private List batchList = []
130-
private int sizeBatch = 0
131-
ParquetWriter(String output, Map params){
132-
if( !params.record ||!(params.record instanceof Class<Record>)) {
133-
throw new IllegalArgumentException("A Record.class is required. Class provided $params.record")
134-
}
135-
this.clazz = params.record as Class<Record>
136-
var outputStream = new FileOutputStream(output)
137-
writer = new CarpetWriter<>(outputStream, this.clazz)
138-
batchList = []
139-
sizeBatch = params.containsKey('by') && "$params.by".isNumber() ? params.by as int : 0
140-
}
141-
142-
void write(Record record){
143-
batchList << record
144-
if( batchList.size() >= sizeBatch ) {
145-
batchList.each { r ->
146-
writer.write(r)
147-
}
148-
batchList.clear()
149-
}
150-
}
151-
152-
void close(){
153-
batchList.each { r ->
154-
writer.write(r)
155-
}
156-
writer.close()
78+
private File toFile(Object source) {
79+
return switch (source) {
80+
case { it instanceof String } -> Path.of(source as String).toFile()
81+
case { it instanceof Path } -> (source as Path).toFile()
82+
default -> throw new IllegalArgumentException("Invalid input for splitParquet operator: ${source}")
15783
}
15884
}
15985
}

‎src/main/groovy/nextflow/parquet/impl/ParquetReader.groovy‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
package nextflow.parquet.impl
2+
3+
import com.jerolba.carpet.CarpetReader
4+
import groovy.transform.CompileStatic
5+
import groovy.util.logging.Slf4j
6+
7+
8+
/**
9+
* Implements a parquet files reader
10+
*
11+
* @author Jorge Aguilera <jorge@edn.es>
12+
*/
13+
@Slf4j
14+
@CompileStatic
15+
class ParquetReader {
16+
private ReadRecordAware recordAware
17+
private Class<Record> clazz
18+
private List batchList = []
19+
private int sizeBatch = 0
20+
21+
ParquetReader(ReadRecordAware recordAware, Map params) {
22+
if( recordAware == null){
23+
throw new IllegalArgumentException("RecordAware is required")
24+
}
25+
this.recordAware = recordAware
26+
parseArgs(params)
27+
}
28+
29+
private List<String> ARGUMENTS = [
30+
'record',
31+
'by'
32+
]
33+
34+
void parseArgs(Map params){
35+
for(def key : params.keySet()){
36+
if( !ARGUMENTS.contains(key) ){
37+
throw new IllegalArgumentException("Unknow ${key} arguments. Posible values ${ARGUMENTS}")
38+
}
39+
}
40+
if (params.record) {
41+
if (!(params.record instanceof Class<Record>)) {
42+
throw new IllegalArgumentException("A Record.class is required. Class provided $params.record")
43+
}
44+
this.clazz = params.record as Class<Record>
45+
}
46+
sizeBatch = params.containsKey('by') && "$params.by".isNumber() ? params.by as int : 0
47+
}
48+
49+
void readFile(File source) {
50+
try {
51+
log.debug "Start reading $source, with projection ${clazz ?: 'raw'}"
52+
53+
final reader = new CarpetReader(source, clazz ?: Map)
54+
for (def record : reader) {
55+
batchList.add(record)
56+
if (batchList.size() >= sizeBatch) {
57+
recordAware.recordRead( sizeBatch == 0 ? batchList.first() : batchList.toArray() )
58+
batchList.clear()
59+
}
60+
}
61+
if (batchList.size()) {
62+
recordAware.recordRead( sizeBatch == 0 ? batchList.first() : batchList.toArray() )
63+
batchList.clear()
64+
}
65+
}
66+
catch (IOException e) {
67+
throw new IllegalStateException("Error while reading Parquet file - cause: ${e.message ?: e}", e)
68+
}
69+
}
70+
71+
72+
}

‎src/main/groovy/nextflow/parquet/impl/ParquetWriter.groovy‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
package nextflow.parquet.impl
2+
3+
import com.jerolba.carpet.CarpetWriter
4+
import groovy.transform.CompileStatic
5+
import groovy.util.logging.Slf4j
6+
7+
/**
8+
* Implements a parquet files writer
9+
*
10+
* @author Jorge Aguilera <jorge@edn.es>
11+
*/
12+
@Slf4j
13+
@CompileStatic
14+
class ParquetWriter implements Closeable {
15+
private Class<Record> clazz
16+
CarpetWriter writer
17+
private List batchList = []
18+
private int sizeBatch = 0
19+
20+
ParquetWriter(String output, Map params) {
21+
if( output == null){
22+
throw new IllegalArgumentException("Parquet file is required")
23+
}
24+
parseArgs(params)
25+
this.clazz = params.record as Class<Record>
26+
var outputStream = new FileOutputStream(output)
27+
writer = new CarpetWriter<>(outputStream, this.clazz)
28+
}
29+
30+
private List<String> ARGUMENTS = [
31+
'record',
32+
'by'
33+
]
34+
35+
void parseArgs(Map params){
36+
for(def key : params.keySet()){
37+
if( !ARGUMENTS.contains(key) ){
38+
throw new IllegalArgumentException("Unknow ${key} arguments. Posible values ${ARGUMENTS}")
39+
}
40+
}
41+
if (!params.record || !(params.record instanceof Class<Record>)) {
42+
throw new IllegalArgumentException("A Record.class is required. Class provided $params.record")
43+
}
44+
sizeBatch = params.containsKey('by') && "$params.by".isNumber() ? params.by as int : 0
45+
}
46+
47+
void write(Record record) {
48+
batchList << record
49+
if (batchList.size() >= sizeBatch) {
50+
batchList.each { r ->
51+
writer.write(r)
52+
}
53+
batchList.clear()
54+
}
55+
}
56+
57+
void close() {
58+
batchList.each { r ->
59+
writer.write(r)
60+
}
61+
writer.close()
62+
}
63+
}

‎src/main/groovy/nextflow/parquet/impl/ReadRecordAware.groovy‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
package nextflow.parquet.impl
2+
3+
interface ReadRecordAware {
4+
5+
void recordRead(Object row )
6+
7+
}

0 commit comments

Comments
 (0)