diff --git a/src/main/java/datawave/query/composite/CompositeMetadata.java b/src/main/java/datawave/query/composite/CompositeMetadata.java index aed85ec..22b5c3b 100644 --- a/src/main/java/datawave/query/composite/CompositeMetadata.java +++ b/src/main/java/datawave/query/composite/CompositeMetadata.java @@ -124,12 +124,23 @@ public CompositeMetadata filter(Set componentFields) { return filter(ingestTypes, componentFields); } + /** + * Returns a new instance of {@link CompositeMetadata} filtered by the provided ingest types and fields. The fields can be component fields or composite + * fields. + * + * @param ingestTypes + * the ingest types + * @param componentFields + * the component or composite fields + * @return an instance of filtered composite metadata + */ public CompositeMetadata filter(Set ingestTypes, Set componentFields) { if (!isEmpty()) { CompositeMetadata compositeMetadata = new CompositeMetadata(); for (String ingestType : ingestTypes) { for (String componentField : componentFields) { if (this.compositeFieldMapByType.containsKey(ingestType)) { + // add composites via the requisite component fields for (String compositeField : this.compositeFieldMapByType.get(ingestType).keySet()) { if (this.compositeFieldMapByType.get(ingestType).get(compositeField).contains(componentField)) { compositeMetadata.setCompositeFieldMappingByType(ingestType, compositeField, @@ -146,6 +157,22 @@ public CompositeMetadata filter(Set ingestTypes, Set componentFi this.compositeFieldSeparatorsByType.get(ingestType).get(compositeField)); } } + + // also check for the case when the filter field is also a composite field + if (this.compositeFieldMapByType.get(ingestType).containsKey(componentField)) { + Collection parts = this.compositeFieldMapByType.get(ingestType).get(componentField); + compositeMetadata.setCompositeFieldMappingByType(ingestType, componentField, parts); + + if (this.compositeTransitionDatesByType.containsKey(ingestType) + && this.compositeTransitionDatesByType.get(ingestType).containsKey(componentField)) + compositeMetadata.addCompositeTransitionDateByType(ingestType, componentField, + this.compositeTransitionDatesByType.get(ingestType).get(componentField)); + + if (this.compositeFieldSeparatorsByType.containsKey(ingestType) + && this.compositeFieldSeparatorsByType.get(ingestType).containsKey(componentField)) + compositeMetadata.addCompositeFieldSeparatorByType(ingestType, componentField, + this.compositeFieldSeparatorsByType.get(ingestType).get(componentField)); + } } } } diff --git a/src/test/java/datawave/query/composite/CompositeMetadataTest.java b/src/test/java/datawave/query/composite/CompositeMetadataTest.java index 6574ea6..78ee643 100644 --- a/src/test/java/datawave/query/composite/CompositeMetadataTest.java +++ b/src/test/java/datawave/query/composite/CompositeMetadataTest.java @@ -8,6 +8,7 @@ import java.util.Arrays; import java.util.Date; import java.util.HashSet; +import java.util.List; import java.util.Set; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -22,10 +23,14 @@ public class CompositeMetadataTest { private CompositeMetadata compositeMetadata; - private String[] ingestTypes = new String[] {"EARTH", "FIRE", "WIND", "WATER", "HEART"}; - private String[][] compositeFields = new String[][] {new String[] {"CAPTAIN_PLANET", "HES", "A", "HERO"}, + private final String[] ingestTypes = new String[] {"EARTH", "FIRE", "WIND", "WATER", "HEART"}; + // @formatter:off + private final String[][] compositeFields = new String[][] { + new String[] {"CAPTAIN_PLANET", "HES", "A", "HERO"}, new String[] {"GONNA_TAKE", "POLLUTION", "DOWN", "TO", "ZERO"}, - new String[] {"CAPTAIN_POLLUTION", "RADIATION", "DEFORESTATION", "SMOG", "TOXICS", "HATE"}}; + new String[] {"CAPTAIN_POLLUTION", "RADIATION", "DEFORESTATION", "SMOG", "TOXICS", "HATE"} + }; + // @formatter:on @BeforeEach public void setup() { @@ -75,6 +80,21 @@ public void filterCompositeMetadataTest() { assertTrue(filteredCompMetadata.compositeTransitionDatesByType.isEmpty()); } + @Test + public void filterCompositeMetadataViaCompositeField() { + CompositeMetadata metadata = new CompositeMetadata(); + metadata.setCompositeFieldMappingByType("country", "CITY_STATE", List.of("CITY", "STATE")); + metadata.addCompositeTransitionDateByType("country", "CITY_STATE", new Date(0)); + + CompositeMetadata filtered = metadata.filter(Set.of("CITY", "STATE")); + assertEquals(metadata.getCompositeFieldMapByType(), filtered.getCompositeFieldMapByType()); + assertEquals(metadata.getCompositeTransitionDatesByType(), filtered.getCompositeTransitionDatesByType()); + + filtered = metadata.filter(Set.of("CITY_STATE")); + assertEquals(metadata.getCompositeFieldMapByType(), filtered.getCompositeFieldMapByType()); + assertEquals(metadata.getCompositeTransitionDatesByType(), filtered.getCompositeTransitionDatesByType()); + } + @Test public void readWriteCompositeMetadataTest() { byte[] compMetadataBytes = CompositeMetadata.toBytes(compositeMetadata);