lucene/benchmark/src/java/org/apache/lucene/benchmark/byTask/feeds/SpatialDocMaker.java - lucene-solr - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.
  * The ASF licenses this file to You under the Apache License, Version 2.0
  * (the "License"); you may not use this file except in compliance with
  * the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
 package org.apache.lucene.benchmark.byTask.feeds;


 import java.util.AbstractMap;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Random;
 import java.util.Set;

 import org.locationtech.spatial4j.context.SpatialContext;
 import org.locationtech.spatial4j.context.SpatialContextFactory;
 import org.locationtech.spatial4j.shape.Point;
 import org.locationtech.spatial4j.shape.Shape;
 import org.apache.lucene.benchmark.byTask.utils.Config;
 import org.apache.lucene.document.Document;
 import org.apache.lucene.document.Field;
 import org.apache.lucene.spatial.SpatialStrategy;
 import org.apache.lucene.spatial.composite.CompositeSpatialStrategy;
 import org.apache.lucene.spatial.prefix.RecursivePrefixTreeStrategy;
 import org.apache.lucene.spatial.prefix.tree.PackedQuadPrefixTree;
 import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTree;
 import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTreeFactory;
 import org.apache.lucene.spatial.serialized.SerializedDVStrategy;

 /**
  * Indexes spatial data according to a configured {@link SpatialStrategy} with optional
  * shape transformation via a configured {@link ShapeConverter}. The converter can turn points into
  * circles and bounding boxes, in order to vary the type of indexing performance tests.
  * Unless it's subclass-ed to do otherwise, this class configures a {@link SpatialContext},
  * {@link SpatialPrefixTree}, and {@link RecursivePrefixTreeStrategy}. The Strategy is made
  * available to a query maker via the static method {@link #getSpatialStrategy(int)}.
  * See spatial.alg for a listing of spatial parameters, in particular those starting with "spatial."
  * and "doc.spatial".
  */
 public class SpatialDocMaker extends DocMaker {

   public static final String SPATIAL_FIELD = "spatial";

   //cache spatialStrategy by round number
   private static Map<Integer,SpatialStrategy> spatialStrategyCache = new HashMap<>();

   private SpatialStrategy strategy;
   private ShapeConverter shapeConverter;

   /**
    * Looks up the SpatialStrategy from the given round --
    * {@link org.apache.lucene.benchmark.byTask.utils.Config#getRoundNumber()}. It's an error
    * if it wasn't created already for this round -- when SpatialDocMaker is initialized.
    */
   public static SpatialStrategy getSpatialStrategy(int roundNumber) {
     SpatialStrategy result = spatialStrategyCache.get(roundNumber);
     if (result == null) {
       throw new IllegalStateException("Strategy should have been init'ed by SpatialDocMaker by now");
     }
     return result;
   }

   /**
    * Builds a SpatialStrategy from configuration options.
    */
   protected SpatialStrategy makeSpatialStrategy(final Config config) {
     //A Map view of Config that prefixes keys with "spatial."
     Map<String, String> configMap = new AbstractMap<String, String>() {
       @Override
       public Set<Entry<String, String>> entrySet() {
         throw new UnsupportedOperationException();
       }

       @Override
       public String get(Object key) {
         return config.get("spatial." + key, null);
       }
     };

     SpatialContext ctx = SpatialContextFactory.makeSpatialContext(configMap, null);

     return makeSpatialStrategy(config, configMap, ctx);
   }

   protected SpatialStrategy makeSpatialStrategy(final Config config, Map<String, String> configMap,
                                                 SpatialContext ctx) {
     //TODO once strategies have factories, we could use them here.
     final String strategyName = config.get("spatial.strategy", "rpt");
     switch (strategyName) {
       case "rpt": return makeRPTStrategy(SPATIAL_FIELD, config, configMap, ctx);
       case "composite": return makeCompositeStrategy(config, configMap, ctx);
       //TODO add more as-needed
       default: throw new IllegalStateException("Unknown spatial.strategy: " + strategyName);
     }
   }

   protected RecursivePrefixTreeStrategy makeRPTStrategy(String spatialField, Config config,
                                                         Map<String, String> configMap, SpatialContext ctx) {
     //A factory for the prefix tree grid
     SpatialPrefixTree grid = SpatialPrefixTreeFactory.makeSPT(configMap, null, ctx);

     RecursivePrefixTreeStrategy strategy = new RecursivePrefixTreeStrategy(grid, spatialField);
     strategy.setPointsOnly(config.get("spatial.docPointsOnly", false));
     final boolean pruneLeafyBranches = config.get("spatial.pruneLeafyBranches", true);
     if (grid instanceof PackedQuadPrefixTree) {
       ((PackedQuadPrefixTree) grid).setPruneLeafyBranches(pruneLeafyBranches);
       strategy.setPruneLeafyBranches(false);//always leave it to packed grid, even though it isn't the same
     } else {
       strategy.setPruneLeafyBranches(pruneLeafyBranches);
     }

     int prefixGridScanLevel = config.get("query.spatial.prefixGridScanLevel", -4);
     if (prefixGridScanLevel < 0)
       prefixGridScanLevel = grid.getMaxLevels() + prefixGridScanLevel;
     strategy.setPrefixGridScanLevel(prefixGridScanLevel);

     double distErrPct = config.get("spatial.distErrPct", .025);//doc & query; a default
     strategy.setDistErrPct(distErrPct);
     return strategy;
   }

   protected SerializedDVStrategy makeSerializedDVStrategy(String spatialField, Config config,
                                                           Map<String, String> configMap, SpatialContext ctx) {
     return new SerializedDVStrategy(ctx, spatialField);
   }

   protected SpatialStrategy makeCompositeStrategy(Config config, Map<String, String> configMap, SpatialContext ctx) {
     final CompositeSpatialStrategy strategy = new CompositeSpatialStrategy(
         SPATIAL_FIELD, makeRPTStrategy(SPATIAL_FIELD + "_rpt", config, configMap, ctx),
         makeSerializedDVStrategy(SPATIAL_FIELD + "_sdv", config, configMap, ctx)
     );
     strategy.setOptimizePredicates(config.get("query.spatial.composite.optimizePredicates", true));
     return strategy;
   }

   @Override
   public void setConfig(Config config, ContentSource source) {
     super.setConfig(config, source);
     SpatialStrategy existing = spatialStrategyCache.get(config.getRoundNumber());
     if (existing == null) {
       //new round; we need to re-initialize
       strategy = makeSpatialStrategy(config);
       spatialStrategyCache.put(config.getRoundNumber(), strategy);
       //TODO remove previous round config?
       shapeConverter = makeShapeConverter(strategy, config, "doc.spatial.");
       System.out.println("Spatial Strategy: " + strategy);
     }
   }

   /**
    * Optionally converts points to circles, and optionally bbox'es result.
    */
   public static ShapeConverter makeShapeConverter(final SpatialStrategy spatialStrategy,
                                                   Config config, String configKeyPrefix) {
     //by default does no conversion
     final double radiusDegrees = config.get(configKeyPrefix+"radiusDegrees", 0.0);
     final double plusMinus = config.get(configKeyPrefix+"radiusDegreesRandPlusMinus", 0.0);
     final boolean bbox = config.get(configKeyPrefix + "bbox", false);

     return new ShapeConverter() {
       @Override
       public Shape convert(Shape shape) {
         if (shape instanceof Point && (radiusDegrees != 0.0 || plusMinus != 0.0)) {
           Point point = (Point)shape;
           double radius = radiusDegrees;
           if (plusMinus > 0.0) {
             Random random = new Random(point.hashCode());//use hashCode so it's reproducibly random
             radius += random.nextDouble() * 2 * plusMinus - plusMinus;
             radius = Math.abs(radius);//can happen if configured plusMinus > radiusDegrees
           }
           shape = spatialStrategy.getSpatialContext().makeCircle(point, radius);
         }
         if (bbox)
           shape = shape.getBoundingBox();
         return shape;
       }
     };
   }

   /** Converts one shape to another. Created by
    * {@link #makeShapeConverter(org.apache.lucene.spatial.SpatialStrategy, org.apache.lucene.benchmark.byTask.utils.Config, String)} */
   public interface ShapeConverter {
     Shape convert(Shape shape);
   }

   @Override
   public Document makeDocument() throws Exception {

     DocState docState = getDocState();

     Document doc = super.makeDocument();

     // Set SPATIAL_FIELD from body
     DocData docData = docState.docData;
     //   makeDocument() resets docState.getBody() so we can't look there; look in Document
     String shapeStr = doc.getField(DocMaker.BODY_FIELD).stringValue();
     Shape shape = makeShapeFromString(strategy, docData.getName(), shapeStr);
     if (shape != null) {
       shape = shapeConverter.convert(shape);
       //index
       for (Field f : strategy.createIndexableFields(shape)) {
         doc.add(f);
       }
     }

     return doc;
   }

   public static Shape makeShapeFromString(SpatialStrategy strategy, String name, String shapeStr) {
     if (shapeStr != null && shapeStr.length() > 0) {
       try {
         return strategy.getSpatialContext().readShapeFromWkt(shapeStr);
       } catch (Exception e) {//InvalidShapeException TODO
         System.err.println("Shape "+name+" wasn't parseable: "+e+"  (skipping it)");
         return null;
       }
     }
     return null;
   }

   @Override
   public Document makeDocument(int size) throws Exception {
     //TODO consider abusing the 'size' notion to number of shapes per document
     throw new UnsupportedOperationException();
   }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one or more
	* contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership.
	* The ASF licenses this file to You under the Apache License, Version 2.0
	* (the "License"); you may not use this file except in compliance with
	* the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/
	package org.apache.lucene.benchmark.byTask.feeds;


	import java.util.AbstractMap;
	import java.util.HashMap;
	import java.util.Map;
	import java.util.Random;
	import java.util.Set;

	import org.locationtech.spatial4j.context.SpatialContext;
	import org.locationtech.spatial4j.context.SpatialContextFactory;
	import org.locationtech.spatial4j.shape.Point;
	import org.locationtech.spatial4j.shape.Shape;
	import org.apache.lucene.benchmark.byTask.utils.Config;
	import org.apache.lucene.document.Document;
	import org.apache.lucene.document.Field;
	import org.apache.lucene.spatial.SpatialStrategy;
	import org.apache.lucene.spatial.composite.CompositeSpatialStrategy;
	import org.apache.lucene.spatial.prefix.RecursivePrefixTreeStrategy;
	import org.apache.lucene.spatial.prefix.tree.PackedQuadPrefixTree;
	import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTree;
	import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTreeFactory;
	import org.apache.lucene.spatial.serialized.SerializedDVStrategy;

	/**
	* Indexes spatial data according to a configured {@link SpatialStrategy} with optional
	* shape transformation via a configured {@link ShapeConverter}. The converter can turn points into
	* circles and bounding boxes, in order to vary the type of indexing performance tests.
	* Unless it's subclass-ed to do otherwise, this class configures a {@link SpatialContext},
	* {@link SpatialPrefixTree}, and {@link RecursivePrefixTreeStrategy}. The Strategy is made
	* available to a query maker via the static method {@link #getSpatialStrategy(int)}.
	* See spatial.alg for a listing of spatial parameters, in particular those starting with "spatial."
	* and "doc.spatial".
	*/
	public class SpatialDocMaker extends DocMaker {

	public static final String SPATIAL_FIELD = "spatial";

	//cache spatialStrategy by round number
	private static Map<Integer,SpatialStrategy> spatialStrategyCache = new HashMap<>();

	private SpatialStrategy strategy;
	private ShapeConverter shapeConverter;

	/**
	* Looks up the SpatialStrategy from the given round --
	* {@link org.apache.lucene.benchmark.byTask.utils.Config#getRoundNumber()}. It's an error
	* if it wasn't created already for this round -- when SpatialDocMaker is initialized.
	*/
	public static SpatialStrategy getSpatialStrategy(int roundNumber) {
	SpatialStrategy result = spatialStrategyCache.get(roundNumber);
	if (result == null) {
	throw new IllegalStateException("Strategy should have been init'ed by SpatialDocMaker by now");
	}
	return result;
	}

	/**
	* Builds a SpatialStrategy from configuration options.
	*/
	protected SpatialStrategy makeSpatialStrategy(final Config config) {
	//A Map view of Config that prefixes keys with "spatial."
	Map<String, String> configMap = new AbstractMap<String, String>() {
	@Override
	public Set<Entry<String, String>> entrySet() {
	throw new UnsupportedOperationException();
	}

	@Override
	public String get(Object key) {
	return config.get("spatial." + key, null);
	}
	};

	SpatialContext ctx = SpatialContextFactory.makeSpatialContext(configMap, null);

	return makeSpatialStrategy(config, configMap, ctx);
	}

	protected SpatialStrategy makeSpatialStrategy(final Config config, Map<String, String> configMap,
	SpatialContext ctx) {
	//TODO once strategies have factories, we could use them here.
	final String strategyName = config.get("spatial.strategy", "rpt");
	switch (strategyName) {
	case "rpt": return makeRPTStrategy(SPATIAL_FIELD, config, configMap, ctx);
	case "composite": return makeCompositeStrategy(config, configMap, ctx);
	//TODO add more as-needed
	default: throw new IllegalStateException("Unknown spatial.strategy: " + strategyName);
	}
	}

	protected RecursivePrefixTreeStrategy makeRPTStrategy(String spatialField, Config config,
	Map<String, String> configMap, SpatialContext ctx) {
	//A factory for the prefix tree grid
	SpatialPrefixTree grid = SpatialPrefixTreeFactory.makeSPT(configMap, null, ctx);

	RecursivePrefixTreeStrategy strategy = new RecursivePrefixTreeStrategy(grid, spatialField);
	strategy.setPointsOnly(config.get("spatial.docPointsOnly", false));
	final boolean pruneLeafyBranches = config.get("spatial.pruneLeafyBranches", true);
	if (grid instanceof PackedQuadPrefixTree) {
	((PackedQuadPrefixTree) grid).setPruneLeafyBranches(pruneLeafyBranches);
	strategy.setPruneLeafyBranches(false);//always leave it to packed grid, even though it isn't the same
	} else {
	strategy.setPruneLeafyBranches(pruneLeafyBranches);
	}

	int prefixGridScanLevel = config.get("query.spatial.prefixGridScanLevel", -4);
	if (prefixGridScanLevel < 0)
	prefixGridScanLevel = grid.getMaxLevels() + prefixGridScanLevel;
	strategy.setPrefixGridScanLevel(prefixGridScanLevel);

	double distErrPct = config.get("spatial.distErrPct", .025);//doc & query; a default
	strategy.setDistErrPct(distErrPct);
	return strategy;
	}

	protected SerializedDVStrategy makeSerializedDVStrategy(String spatialField, Config config,
	Map<String, String> configMap, SpatialContext ctx) {
	return new SerializedDVStrategy(ctx, spatialField);
	}

	protected SpatialStrategy makeCompositeStrategy(Config config, Map<String, String> configMap, SpatialContext ctx) {
	final CompositeSpatialStrategy strategy = new CompositeSpatialStrategy(
	SPATIAL_FIELD, makeRPTStrategy(SPATIAL_FIELD + "_rpt", config, configMap, ctx),
	makeSerializedDVStrategy(SPATIAL_FIELD + "_sdv", config, configMap, ctx)
	);
	strategy.setOptimizePredicates(config.get("query.spatial.composite.optimizePredicates", true));
	return strategy;
	}

	@Override
	public void setConfig(Config config, ContentSource source) {
	super.setConfig(config, source);
	SpatialStrategy existing = spatialStrategyCache.get(config.getRoundNumber());
	if (existing == null) {
	//new round; we need to re-initialize
	strategy = makeSpatialStrategy(config);
	spatialStrategyCache.put(config.getRoundNumber(), strategy);
	//TODO remove previous round config?
	shapeConverter = makeShapeConverter(strategy, config, "doc.spatial.");
	System.out.println("Spatial Strategy: " + strategy);
	}
	}

	/**
	* Optionally converts points to circles, and optionally bbox'es result.
	*/
	public static ShapeConverter makeShapeConverter(final SpatialStrategy spatialStrategy,
	Config config, String configKeyPrefix) {
	//by default does no conversion
	final double radiusDegrees = config.get(configKeyPrefix+"radiusDegrees", 0.0);
	final double plusMinus = config.get(configKeyPrefix+"radiusDegreesRandPlusMinus", 0.0);
	final boolean bbox = config.get(configKeyPrefix + "bbox", false);

	return new ShapeConverter() {
	@Override
	public Shape convert(Shape shape) {
	if (shape instanceof Point && (radiusDegrees != 0.0 \|\| plusMinus != 0.0)) {
	Point point = (Point)shape;
	double radius = radiusDegrees;
	if (plusMinus > 0.0) {
	Random random = new Random(point.hashCode());//use hashCode so it's reproducibly random
	radius += random.nextDouble() * 2 * plusMinus - plusMinus;
	radius = Math.abs(radius);//can happen if configured plusMinus > radiusDegrees
	}
	shape = spatialStrategy.getSpatialContext().makeCircle(point, radius);
	}
	if (bbox)
	shape = shape.getBoundingBox();
	return shape;
	}
	};
	}

	/** Converts one shape to another. Created by
	* {@link #makeShapeConverter(org.apache.lucene.spatial.SpatialStrategy, org.apache.lucene.benchmark.byTask.utils.Config, String)} */
	public interface ShapeConverter {
	Shape convert(Shape shape);
	}

	@Override
	public Document makeDocument() throws Exception {

	DocState docState = getDocState();

	Document doc = super.makeDocument();

	// Set SPATIAL_FIELD from body
	DocData docData = docState.docData;
	// makeDocument() resets docState.getBody() so we can't look there; look in Document
	String shapeStr = doc.getField(DocMaker.BODY_FIELD).stringValue();
	Shape shape = makeShapeFromString(strategy, docData.getName(), shapeStr);
	if (shape != null) {
	shape = shapeConverter.convert(shape);
	//index
	for (Field f : strategy.createIndexableFields(shape)) {
	doc.add(f);
	}
	}

	return doc;
	}

	public static Shape makeShapeFromString(SpatialStrategy strategy, String name, String shapeStr) {
	if (shapeStr != null && shapeStr.length() > 0) {
	try {
	return strategy.getSpatialContext().readShapeFromWkt(shapeStr);
	} catch (Exception e) {//InvalidShapeException TODO
	System.err.println("Shape "+name+" wasn't parseable: "+e+" (skipping it)");
	return null;
	}
	}
	return null;
	}

	@Override
	public Document makeDocument(int size) throws Exception {
	//TODO consider abusing the 'size' notion to number of shapes per document
	throw new UnsupportedOperationException();
	}
	}