fix spatial indexing with multi spatial dims

processing/src/main/java/io/druid/segment/incremental/SpatialDimensionRowFormatter.java

@@ -49,26 +49,15 @@ public class SpatialDimensionRowFormatter
   private static final Joiner JOINER = Joiner.on(",");
   private static final Splitter SPLITTER = Splitter.on(",");
 
-  private final List<SpatialDimensionSchema> spatialDimensions;
+  private final Map<String, SpatialDimensionSchema> spatialDimensionMap;
-  private final Set<String> spatialDimNames;
   private final Set<String> spatialPartialDimNames;
 
   public SpatialDimensionRowFormatter(List<SpatialDimensionSchema> spatialDimensions)
   {
-    this.spatialDimensions = spatialDimensions;
+    this.spatialDimensionMap = Maps.newHashMap();
-    this.spatialDimNames = Sets.newHashSet(
+    for (SpatialDimensionSchema spatialDimension : spatialDimensions) {
-        Lists.transform(
+      this.spatialDimensionMap.put(spatialDimension.getDimName(), spatialDimension);
-            spatialDimensions,
-            new Function<SpatialDimensionSchema, String>()
-            {
-              @Override
-              public String apply(SpatialDimensionSchema input)
-              {
-                return input.getDimName();
     }
-            }
-        )
-    );
     this.spatialPartialDimNames = Sets.newHashSet(
         Iterables.concat(
             Lists.transform(
@@ -110,7 +99,7 @@ public class SpatialDimensionRowFormatter
               @Override
               public boolean apply(String input)
               {
-                return !spatialDimNames.contains(input) && !spatialPartialDimNames.contains(input);
+                return !spatialDimensionMap.containsKey(input) && !spatialPartialDimNames.contains(input);
               }
             }
         )
@@ -173,32 +162,32 @@ public class SpatialDimensionRowFormatter
       }
     };
 
-    if (!spatialPartialDimNames.isEmpty()) {
+    for (Map.Entry<String, SpatialDimensionSchema> entry : spatialDimensionMap.entrySet()) {
-      for (SpatialDimensionSchema spatialDimension : spatialDimensions) {
+      final String spatialDimName = entry.getKey();
-        List<String> spatialDimVals = Lists.newArrayList();
+      final SpatialDimensionSchema spatialDim = entry.getValue();
 
-        for (String partialSpatialDim : spatialDimension.getDims()) {
-          List<String> dimVals = row.getDimension(partialSpatialDim);
-          if (isSpatialDimValsValid(dimVals)) {
-            spatialDimVals.addAll(dimVals);
-          }
-        }
-
-        if (spatialDimVals.size() == spatialPartialDimNames.size()) {
-          spatialLookup.put(spatialDimension.getDimName(), Arrays.asList(JOINER.join(spatialDimVals)));
-          finalDims.add(spatialDimension.getDimName());
-        }
-      }
-    } else {
-      for (String spatialDimName : spatialDimNames) {
       List<String> dimVals = row.getDimension(spatialDimName);
+      if (dimVals != null && !dimVals.isEmpty()) {
         if (dimVals.size() != 1) {
-          throw new ISE("Cannot have a spatial dimension value with size[%d]", dimVals.size());
+          throw new ISE("Spatial dimension value must be in an array!");
         }
         if (isJoinedSpatialDimValValid(dimVals.get(0))) {
           spatialLookup.put(spatialDimName, dimVals);
           finalDims.add(spatialDimName);
         }
+      } else {
+        List<String> spatialDimVals = Lists.newArrayList();
+        for (String dim : spatialDim.getDims()) {
+          List<String> partialDimVals = row.getDimension(dim);
+          if (isSpatialDimValsValid(partialDimVals)) {
+            spatialDimVals.addAll(partialDimVals);
+          }
+        }
+
+        if (spatialDimVals.size() == spatialDim.getDims().size()) {
+          spatialLookup.put(spatialDimName, Arrays.asList(JOINER.join(spatialDimVals)));
+          finalDims.add(spatialDimName);
+        }
       }
     }
 

processing/src/test/java/io/druid/segment/filter/SpatialFilterTest.java

@@ -76,7 +76,7 @@ public class SpatialFilterTest
       new LongSumAggregatorFactory("val", "val")
   };
 
-  private static List<String> DIMS = Lists.newArrayList("dim", "lat", "long");
+  private static List<String> DIMS = Lists.newArrayList("dim", "lat", "long", "lat2", "long2");
 
   @Parameterized.Parameters
   public static Collection<?> constructorFeeder() throws IOException
@@ -110,6 +110,10 @@ public class SpatialFilterTest
                                                     new SpatialDimensionSchema(
                                                         "dim.geo",
                                                         Arrays.asList("lat", "long")
+                                                    ),
+                                                    new SpatialDimensionSchema(
+                                                        "spatialIsRad",
+                                                        Arrays.asList("lat2", "long2")
                                                     )
                                                 )
                                             ).build()
@@ -204,6 +208,18 @@ public class SpatialFilterTest
             )
         )
     );
+    theIndex.add(
+        new MapBasedInputRow(
+            new DateTime("2013-01-05").getMillis(),
+            DIMS,
+            ImmutableMap.<String, Object>of(
+                "timestamp", new DateTime("2013-01-05").toString(),
+                "lat2", 0.0f,
+                "long2", 0.0f,
+                "val", 13l
+            )
+        )
+    );
 
     // Add a bunch of random points
     Random rand = new Random();
@@ -250,6 +266,10 @@ public class SpatialFilterTest
                                                       new SpatialDimensionSchema(
                                                           "dim.geo",
                                                           Arrays.asList("lat", "long")
+                                                      ),
+                                                      new SpatialDimensionSchema(
+                                                          "spatialIsRad",
+                                                          Arrays.asList("lat2", "long2")
                                                       )
                                                   )
                                               ).build()
@@ -263,6 +283,10 @@ public class SpatialFilterTest
                                                       new SpatialDimensionSchema(
                                                           "dim.geo",
                                                           Arrays.asList("lat", "long")
+                                                      ),
+                                                      new SpatialDimensionSchema(
+                                                          "spatialIsRad",
+                                                          Arrays.asList("lat2", "long2")
                                                       )
                                                   )
                                               ).build()
@@ -276,6 +300,10 @@ public class SpatialFilterTest
                                                       new SpatialDimensionSchema(
                                                           "dim.geo",
                                                           Arrays.asList("lat", "long")
+                                                      ),
+                                                      new SpatialDimensionSchema(
+                                                          "spatialIsRad",
+                                                          Arrays.asList("lat2", "long2")
                                                       )
                                                   )
                                               ).build()
@@ -372,6 +400,18 @@ public class SpatialFilterTest
               )
           )
       );
+      second.add(
+          new MapBasedInputRow(
+              new DateTime("2013-01-05").getMillis(),
+              DIMS,
+              ImmutableMap.<String, Object>of(
+                  "timestamp", new DateTime("2013-01-05").toString(),
+                  "lat2", 0.0f,
+                  "long2", 0.0f,
+                  "val", 13l
+              )
+          )
+      );
 
       // Add a bunch of random points
       Random rand = new Random();
@@ -486,6 +526,58 @@ public class SpatialFilterTest
     }
   }
 
+
+  @Test
+  public void testSpatialQueryWithOtherSpatialDim()
+  {
+    TimeseriesQuery query = Druids.newTimeseriesQueryBuilder()
+                                  .dataSource("test")
+                                  .granularity(QueryGranularity.ALL)
+                                  .intervals(Arrays.asList(new Interval("2013-01-01/2013-01-07")))
+                                  .filters(
+                                      new SpatialDimFilter(
+                                          "spatialIsRad",
+                                          new RadiusBound(new float[]{0.0f, 0.0f}, 5)
+                                      )
+                                  )
+                                  .aggregators(
+                                      Arrays.<AggregatorFactory>asList(
+                                          new CountAggregatorFactory("rows"),
+                                          new LongSumAggregatorFactory("val", "val")
+                                      )
+                                  )
+                                  .build();
+
+    List<Result<TimeseriesResultValue>> expectedResults = Arrays.asList(
+        new Result<TimeseriesResultValue>(
+            new DateTime("2013-01-01T00:00:00.000Z"),
+            new TimeseriesResultValue(
+                ImmutableMap.<String, Object>builder()
+                            .put("rows", 1L)
+                            .put("val", 13l)
+                            .build()
+            )
+        )
+    );
+    try {
+      TimeseriesQueryRunnerFactory factory = new TimeseriesQueryRunnerFactory(
+          new TimeseriesQueryQueryToolChest(new QueryConfig()),
+          new TimeseriesQueryEngine(),
+          QueryRunnerTestHelper.NOOP_QUERYWATCHER
+      );
+
+      QueryRunner runner = new FinalizeResultsQueryRunner(
+          factory.createRunner(segment),
+          factory.getToolchest()
+      );
+
+      TestHelper.assertExpectedResults(expectedResults, runner.run(query));
+    }
+    catch (Exception e) {
+      throw Throwables.propagate(e);
+    }
+  }
+
   @Test
   public void testSpatialQueryMorePoints()
   {