Write batches of Features using a GeoTools FeatureWriter

FeatureWriter.groovy

import geoscript.layer.*
import geoscript.feature.*
import geoscript.workspace.*

import org.geotools.data.Transaction
import org.geotools.data.DefaultTransaction
import org.geotools.data.FeatureStore
import org.geotools.data.DataUtilities
import org.geotools.data.FeatureWriter as GtFeatureWriter

/**
 * Write batches of Features using a GeoTools FeatureWriter
 * @author Jared Erickson
 */
class FeatureWriter {

    /**
     * The GeoTools FeatureWriter
     */
    private GtFeatureWriter writer 
    
    /**
     * The number of features to write in one go
     */
    private int batch

    /**
     * The FeatureStore to write to
     */
    private FeatureStore store
    
    /**
     * The Transaction used during the writing
     */
    private Transaction transaction
    
    /**
     * The internal counter that keeps track of how many features we have added
     */
    private int counter = 0
    
    /**
     * Create a new FeatureWriter for a Layer
     * @param options The named parameters
     * <ul>
     *   <li>batch: The number of features to write at one time (defaults to 1000)</li>
     *   <li>mode: The mode can either be write or append</li>
     *   <li>autoCommit: Whether to use an auto commit transaction or not (defaults to false unless the Layer comes from a Shapefile)</li>
     * </ul>
     * @param layer The Layer to write to
     */
    FeatureWriter(Map options = [:], Layer layer) {
        this.batch = options.get("batch", 1000)
        String mode = options.get("mode", "write") // or append
        boolean autoCommit = options.get("autoCommit", 
            layer.fs instanceof org.geotools.data.directory.DirectoryFeatureStore ? true : false)
        this.store = layer.fs as FeatureStore
        this.transaction = autoCommit ? Transaction.AUTO_COMMIT : new DefaultTransaction("${layer.name}Transaction")
        if (mode.equalsIgnoreCase("write")) {
            this.writer = layer.fs.dataStore.getFeatureWriter(layer.name,transaction)
        } else {
            this.writer = layer.fs.dataStore.getFeatureWriterAppend(layer.name,transaction)
        }
    }
    
    /**
     * Closing the writing also commit's the transaction
     */
    void close() {
        this.writer.close()
    }
   
    /**
     * Get a new blank Feature
     * @return A new Feature
     */
    Feature feature() {
        new Feature(this.writer.next())
    }
    
    /**
     * Write the Feature created by feature()
     */
    void write() {
        this.writer.write()
        // If the counter exceeds the batch size, commit the transaction
        // and reset the counter
        counter++
        if (counter >= batch) {
            transaction.commit()
            counter = 0
        }
    }
    
}

/**
 * Add Layer.withFeatureWriter(Closure c) method
 */
Layer.metaClass.withFeatureWriter << {closure ->
    def writer = new FeatureWriter(delegate, batch: 10)
    try {
        closure.call(writer)
    } finally {
        writer.close()
    }
}

/**
 * Add Layer.withFeatureAppender(Closure c) method
 */
Layer.metaClass.withFeatureAppender << {closure ->
    def writer = new FeatureWriter(delegate, mode: "append")
    try {
        closure.call(writer)
    } finally {
        writer.close()
    }
}

/**
 * Add Feature.set(Feature f) method
 */
Feature.metaClass.set << {Feature f ->
    f.attributes.each {k,v ->
        delegate.set(k,v)
    }
}

def dir = new Directory(new File("/Users/jericks/Projects/GeoScript/scripts/"))
def inLayer = dir.get("states")
def outLayer = dir.create(inLayer.schema.changeGeometryType("point","state_points"))

outLayer.withFeatureWriter{w ->
    inLayer.eachFeature{f ->
        def newF = w.feature().set(f)
        newF.geom = f.geom.centroid
        w.write() 
    }
}

outLayer = dir.get("state_points")
println outLayer.count