shapefileindexer.java

shape file read and write

/*
 *    GeoTools - OpenSource mapping toolkit
 *    http://geotools.org
 *    (C) 2003-2006, GeoTools Project Managment Committee (PMC)
 *    (C) 2002, Centre for Computational Geography
 *
 *    This library is free software; you can redistribute it and/or
 *    modify it under the terms of the GNU Lesser General Public
 *    License as published by the Free Software Foundation;
 *    version 2.1 of the License.
 *
 *    This library is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 *    Lesser General Public License for more details.
 */
package org.geotools.data.shapefile.indexed;

import java.io.File;
import java.io.IOException;
import java.net.MalformedURLException;

import org.geotools.data.shapefile.FileWriter;
import org.geotools.data.shapefile.ShpFiles;
import org.geotools.data.shapefile.StorageFile;
import org.geotools.data.shapefile.shp.IndexFile;
import org.geotools.data.shapefile.shp.ShapefileHeader;
import org.geotools.data.shapefile.shp.ShapefileReader;
import org.geotools.data.shapefile.shp.ShapefileReader.Record;
import org.geotools.index.Data;
import org.geotools.index.DataDefinition;
import org.geotools.index.LockTimeoutException;
import org.geotools.index.TreeException;
import org.geotools.index.quadtree.QuadTree;
import org.geotools.index.quadtree.StoreException;
import org.geotools.index.quadtree.fs.FileSystemIndexStore;
import org.geotools.index.quadtree.fs.IndexHeader;
import org.geotools.index.rtree.PageStore;
import org.geotools.index.rtree.RTree;
import org.geotools.index.rtree.cachefs.FileSystemPageStore;
import org.geotools.util.NullProgressListener;
import org.opengis.util.ProgressListener;

import com.vividsolutions.jts.geom.Envelope;

/**
 * Utility class for Shapefile spatial indexing
 * 
 * @author Tommaso Nolli
 * @source $URL:
 *         http://svn.geotools.org/geotools/trunk/gt/modules/plugin/shapefile/src/main/java/org/geotools/data/shapefile/indexed/ShapeFileIndexer.java $
 */
public class ShapeFileIndexer implements FileWriter {
    private IndexType idxType;
    private int max = 50;
    private int min = 25;
    private short split = PageStore.SPLIT_QUADRATIC;
    private String byteOrder;

    private ShpFiles shpFiles;

    public static void main(String[] args) throws IOException {
        if ((args.length < 1) || (((args.length - 1) % 2) != 0)) {
            usage();
        }

        long start = System.currentTimeMillis();

        ShapeFileIndexer idx = new ShapeFileIndexer();

        for (int i = 0; i < args.length; i++) {
            if (args[i].equals("-t")) {
                idx.setIdxType(IndexType.valueOf(args[++i]));
            } else if (args[i].equals("-M")) {
                idx.setMax(Integer.parseInt(args[++i]));
            } else if (args[i].equals("-m")) {
                idx.setMin(Integer.parseInt(args[++i]));
            } else if (args[i].equals("-s")) {
                idx.setSplit(Short.parseShort(args[++i]));
            } else if (args[i].equals("-b")) {
                idx.setByteOrder(args[++i]);
            } else {
                if (!args[i].toLowerCase().endsWith(".shp")) {
                    System.out.println("File extension must be '.shp'");
                    System.exit(1);
                }

                idx.setShapeFileName(new ShpFiles(args[i]));
            }
        }

        try {
            System.out.print("Indexing ");

            int cnt = idx.index(true, new NullProgressListener());
            System.out.println();
            System.out.print(cnt + " features indexed ");
            System.out.println("in " + (System.currentTimeMillis() - start)
                    + "ms.");
            System.out.println();
        } catch (Exception e) {
            e.printStackTrace();
            usage();
            System.exit(1);
        }
    }

    private static void usage() {
        System.out.println("Usage: ShapeFileIndexer " + "-t <QIX | GRX> "
                + "[-M <max entries per node>] "
                + "[-m <min entries per node>] " + "[-s <split algorithm>] "
                + "[-b <byte order NL | NM>] " + "<shape file>");

        System.out.println();

        System.out.println("Options:");
        System.out.println("\t-t Index type: RTREE or QUADTREE");
        System.out.println();
        System.out.println("Following options apllies only to RTREE:");
        System.out.println("\t-M maximum number of entries per node");
        System.out.println("\t-m minimum number of entries per node");
        System.out.println("\t-s split algorithm to use");
        System.out.println();
        System.out.println("Following options apllies only to QUADTREE:");
        System.out.println("\t-b byte order to use: NL = LSB; "
                + "NM = MSB (default)");

        System.exit(1);
    }

    /**
     * Index the shapefile denoted by setShapeFileName(String fileName) If when
     * a thread starts, another thread is indexing the same file, this thread
     * will wait that the first thread ends indexing; in this case <b>zero</b>
     * is reurned as result of the indexing process.
     * 
     * @param verbose
     *                enable/disable printing of dots every 500 indexed records
     * @param listener
     *                DOCUMENT ME!
     * 
     * @return The number of indexed records (or zero)
     * 
     * @throws MalformedURLException
     * @throws IOException
     * @throws TreeException
     * @throws StoreException
     *                 DOCUMENT ME!
     * @throws LockTimeoutException
     */
    public int index(boolean verbose, ProgressListener listener)
            throws MalformedURLException, IOException, TreeException,
            StoreException, LockTimeoutException {

        if (this.shpFiles == null) {
            throw new IOException("You have to set a shape file name!");
        }

        int cnt = 0;

        ShapefileReader reader = null;

        // Temporary file for building...
        StorageFile storage = shpFiles.getStorageFile(this.idxType.shpFileType);
        File treeFile = storage.getFile();

        try {
            reader = new ShapefileReader(shpFiles, true, false);

            switch (idxType) {
            case EXPERIMENTAL_UNSUPPORTED_GRX:
                cnt = this.buildRTree(reader, treeFile, verbose);
                break;
            case QIX:
                cnt = this.buildQuadTree(reader, treeFile, verbose);
                break;

            default:
                throw new IllegalArgumentException(
                        "NONE is not a legal index choice");
            }
        } finally {
            if (reader != null)
                reader.close();
        }

        // Final index file
        storage.replaceOriginal();

        return cnt;
    }

    private int buildRTree(ShapefileReader reader, File rtreeFile,
            boolean verbose) throws TreeException, LockTimeoutException,
            IOException {
        DataDefinition keyDef = new DataDefinition("US-ASCII");
        keyDef.addField(Integer.class);
        keyDef.addField(Long.class);

        FileSystemPageStore fps = new FileSystemPageStore(rtreeFile, keyDef,
                this.max, this.min, this.split);
        RTree rtree = new RTree(fps);

        Record record = null;
        Data data = null;

        int cnt = 0;

        while (reader.hasNext()) {
            record = reader.nextRecord();
            data = new Data(keyDef);
            data.addValue(new Integer(++cnt));
            data.addValue(new Long(record.offset()));

            rtree.insert(new Envelope(record.minX, record.maxX, record.minY,
                    record.maxY), data);

            if (verbose && ((cnt % 500) == 0)) {
                System.out.print('.');
            }
        }

        rtree.close();

        return cnt;
    }

    private int buildQuadTree(ShapefileReader reader, File file, boolean verbose)
            throws IOException, StoreException {
        byte order = 0;

        if ((this.byteOrder == null) || this.byteOrder.equalsIgnoreCase("NM")) {
            order = IndexHeader.NEW_MSB_ORDER;
        } else if (this.byteOrder.equalsIgnoreCase("NL")) {
            order = IndexHeader.NEW_LSB_ORDER;
        } else {
            throw new StoreException("Asked byte order '" + this.byteOrder
                    + "' must be 'NL' or 'NM'!");
        }

        IndexFile shpIndex = new IndexFile(shpFiles, false);
        QuadTree tree = null;
        int cnt = 0;
        int numRecs = shpIndex.getRecordCount();
        ShapefileHeader header = reader.getHeader();
        Envelope bounds = new Envelope(header.minX(), header.maxX(), header
                .minY(), header.maxY());

        tree = new QuadTree(numRecs, max, bounds, shpIndex);
        try {
            Record rec = null;

            while (reader.hasNext()) {
                rec = reader.nextRecord();
                tree.insert(cnt++, new Envelope(rec.minX, rec.maxX, rec.minY,
                        rec.maxY));

                if (verbose && ((cnt % 1000) == 0)) {
                    System.out.print('.');
                }
                if (cnt % 100000 == 0)
                    System.out.print('\n');
            }
            if (verbose)
                System.out.println("done");
            FileSystemIndexStore store = new FileSystemIndexStore(file, order);
            store.store(tree);
        } finally {
            tree.close();
        }
        return cnt;
    }

    /**
     * For quad tree this is the max depth. I don't know what it is for RTree
     * 
     * @param i
     */
    public void setMax(int i) {
        max = i;
    }

    /**
     * DOCUMENT ME!
     * 
     * @param i
     */
    public void setMin(int i) {
        min = i;
    }

    /**
     * DOCUMENT ME!
     * 
     * @param s
     */
    public void setSplit(short s) {
        split = s;
    }

    /**
     * DOCUMENT ME!
     * 
     * @param shpFiles
     */
    public void setShapeFileName(ShpFiles shpFiles) {
        this.shpFiles = shpFiles;
    }

    /**
     * Sets the type of index to create
     * 
     * @param indexType
     *                The idxType to set.
     */
    public void setIdxType(IndexType indexType) {
        this.idxType = indexType;
    }

    /**
     * DOCUMENT ME!
     * 
     * @param byteOrder
     *                The byteOrder to set.
     */
    public void setByteOrder(String byteOrder) {
        this.byteOrder = byteOrder;
    }

    public String id() {
        return getClass().getName();
    }
}