lucene/facet/src/test/org/apache/lucene/facet/FacetTestCase.java

/*
 * 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.facet;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.apache.lucene.facet.FacetsCollector.MatchingDocs;
import org.apache.lucene.facet.taxonomy.CachedOrdinalsReader;
import org.apache.lucene.facet.taxonomy.DocValuesOrdinalsReader;
import org.apache.lucene.facet.taxonomy.FacetLabel;
import org.apache.lucene.facet.taxonomy.FastTaxonomyFacetCounts;
import org.apache.lucene.facet.taxonomy.OrdinalsReader;
import org.apache.lucene.facet.taxonomy.TaxonomyFacetCounts;
import org.apache.lucene.facet.taxonomy.TaxonomyFacetLabels;
import org.apache.lucene.facet.taxonomy.TaxonomyFacetLabels.FacetLabelReader;
import org.apache.lucene.facet.taxonomy.TaxonomyReader;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.LuceneTestCase;
import org.apache.lucene.util.TestUtil;

public abstract class FacetTestCase extends LuceneTestCase {

  public Facets getTaxonomyFacetCounts(
      TaxonomyReader taxoReader, FacetsConfig config, FacetsCollector c) throws IOException {
    return getTaxonomyFacetCounts(taxoReader, config, c, FacetsConfig.DEFAULT_INDEX_FIELD_NAME);
  }

  public Facets getTaxonomyFacetCounts(
      TaxonomyReader taxoReader, FacetsConfig config, FacetsCollector c, String indexFieldName)
      throws IOException {
    Facets facets;
    if (random().nextBoolean()) {
      facets = new FastTaxonomyFacetCounts(indexFieldName, taxoReader, config, c);
    } else {
      OrdinalsReader ordsReader = new DocValuesOrdinalsReader(indexFieldName);
      if (random().nextBoolean()) {
        ordsReader = new CachedOrdinalsReader(ordsReader);
      }
      facets = new TaxonomyFacetCounts(ordsReader, taxoReader, config, c);
    }

    return facets;
  }

  /**
   * Utility method that uses {@link FacetLabelReader} to get facet labels for each hit in {@link
   * MatchingDocs}. The method returns {@code List<List<FacetLabel>>} where outer list has one entry
   * per document and inner list has all {@link FacetLabel} entries that belong to a document. The
   * inner list may be empty if no {@link FacetLabel} are found for a hit.
   *
   * @param taxoReader {@link TaxonomyReader} used to read taxonomy during search. This instance is
   *     expected to be open for reading.
   * @param fc {@link FacetsCollector} A collector with matching hits.
   * @param dimension facet dimension for which labels are requested. A null value fetches labels
   *     for all dimensions.
   * @return {@code List<List<FacetLabel>} where outer list has one non-null entry per document. and
   *     inner list contain all {@link FacetLabel} entries that belong to a document.
   * @throws IOException when a low-level IO issue occurs.
   */
  public List<List<FacetLabel>> getAllTaxonomyFacetLabels(
      String dimension, TaxonomyReader taxoReader, FacetsCollector fc) throws IOException {
    List<List<FacetLabel>> actualLabels = new ArrayList<>();
    TaxonomyFacetLabels taxoLabels =
        new TaxonomyFacetLabels(taxoReader, FacetsConfig.DEFAULT_INDEX_FIELD_NAME);
    for (MatchingDocs m : fc.getMatchingDocs()) {
      FacetLabelReader facetLabelReader = taxoLabels.getFacetLabelReader(m.context);
      DocIdSetIterator disi = m.bits.iterator();
      while (disi.nextDoc() != DocIdSetIterator.NO_MORE_DOCS) {
        actualLabels.add(allFacetLabels(disi.docID(), dimension, facetLabelReader));
      }
    }
    return actualLabels;
  }

  /**
   * Utility method to get all facet labels for an input docId and dimension using the supplied
   * {@link FacetLabelReader}.
   *
   * @param docId docId for which facet labels are needed.
   * @param dimension Retain facet labels for supplied dimension only. A null value fetches all
   *     facet labels.
   * @param facetLabelReader {@FacetLabelReader} instance use to get facet labels for input docId.
   * @return {@code List<FacetLabel>} containing matching facet labels.
   * @throws IOException when a low-level IO issue occurs while reading facet labels.
   */
  List<FacetLabel> allFacetLabels(int docId, String dimension, FacetLabelReader facetLabelReader)
      throws IOException {
    List<FacetLabel> facetLabels = new ArrayList<>();
    FacetLabel facetLabel;
    if (dimension != null) {
      for (facetLabel = facetLabelReader.nextFacetLabel(docId, dimension); facetLabel != null; ) {
        facetLabels.add(facetLabel);
        facetLabel = facetLabelReader.nextFacetLabel(docId, dimension);
      }
    } else {
      for (facetLabel = facetLabelReader.nextFacetLabel(docId); facetLabel != null; ) {
        facetLabels.add(facetLabel);
        facetLabel = facetLabelReader.nextFacetLabel(docId);
      }
    }
    return facetLabels;
  }

  protected String[] getRandomTokens(int count) {
    String[] tokens = new String[count];
    for (int i = 0; i < tokens.length; i++) {
      tokens[i] = TestUtil.randomRealisticUnicodeString(random(), 1, 10);
      // tokens[i] = _TestUtil.randomSimpleString(random(), 1, 10);
    }
    return tokens;
  }

  protected String pickToken(String[] tokens) {
    for (int i = 0; i < tokens.length; i++) {
      if (random().nextBoolean()) {
        return tokens[i];
      }
    }

    // Move long tail onto first token:
    return tokens[0];
  }

  protected static class TestDoc {
    public String content;
    public String[] dims;
    public float value;
  }

  protected List<TestDoc> getRandomDocs(String[] tokens, int count, int numDims) {
    List<TestDoc> docs = new ArrayList<>();
    for (int i = 0; i < count; i++) {
      TestDoc doc = new TestDoc();
      docs.add(doc);
      doc.content = pickToken(tokens);
      doc.dims = new String[numDims];
      for (int j = 0; j < numDims; j++) {
        doc.dims[j] = pickToken(tokens);
        if (random().nextInt(10) < 3) {
          break;
        }
      }
      if (VERBOSE) {
        System.out.println("  doc " + i + ": content=" + doc.content);
        for (int j = 0; j < numDims; j++) {
          if (doc.dims[j] != null) {
            System.out.println("    dim[" + j + "]=" + doc.dims[j]);
          }
        }
      }
    }

    return docs;
  }

  protected void sortTies(List<FacetResult> results) {
    for (FacetResult result : results) {
      sortTies(result.labelValues);
    }
  }

  protected void sortTies(LabelAndValue[] labelValues) {
    double lastValue = -1;
    int numInRow = 0;
    int i = 0;
    while (i <= labelValues.length) {
      if (i < labelValues.length && labelValues[i].value.doubleValue() == lastValue) {
        numInRow++;
      } else {
        if (numInRow > 1) {
          Arrays.sort(
              labelValues,
              i - numInRow,
              i,
              new Comparator<LabelAndValue>() {
                @Override
                public int compare(LabelAndValue a, LabelAndValue b) {
                  assert a.value.doubleValue() == b.value.doubleValue();
                  return new BytesRef(a.label).compareTo(new BytesRef(b.label));
                }
              });
        }
        numInRow = 1;
        if (i < labelValues.length) {
          lastValue = labelValues[i].value.doubleValue();
        }
      }
      i++;
    }
  }

  protected void sortLabelValues(List<LabelAndValue> labelValues) {
    Collections.sort(
        labelValues,
        new Comparator<LabelAndValue>() {
          @Override
          public int compare(LabelAndValue a, LabelAndValue b) {
            if (a.value.doubleValue() > b.value.doubleValue()) {
              return -1;
            } else if (a.value.doubleValue() < b.value.doubleValue()) {
              return 1;
            } else {
              return new BytesRef(a.label).compareTo(new BytesRef(b.label));
            }
          }
        });
  }

  protected void sortFacetResults(List<FacetResult> results) {
    Collections.sort(
        results,
        new Comparator<FacetResult>() {
          @Override
          public int compare(FacetResult a, FacetResult b) {
            if (a.value.doubleValue() > b.value.doubleValue()) {
              return -1;
            } else if (b.value.doubleValue() > a.value.doubleValue()) {
              return 1;
            } else {
              return 0;
            }
          }
        });
  }

  protected void assertFloatValuesEquals(List<FacetResult> a, List<FacetResult> b) {
    assertEquals(a.size(), b.size());
    float lastValue = Float.POSITIVE_INFINITY;
    Map<String, FacetResult> aByDim = new HashMap<>();
    for (int i = 0; i < a.size(); i++) {
      assertTrue(a.get(i).value.floatValue() <= lastValue);
      lastValue = a.get(i).value.floatValue();
      aByDim.put(a.get(i).dim, a.get(i));
    }
    lastValue = Float.POSITIVE_INFINITY;
    Map<String, FacetResult> bByDim = new HashMap<>();
    for (int i = 0; i < b.size(); i++) {
      bByDim.put(b.get(i).dim, b.get(i));
      assertTrue(b.get(i).value.floatValue() <= lastValue);
      lastValue = b.get(i).value.floatValue();
    }
    for (String dim : aByDim.keySet()) {
      assertFloatValuesEquals(aByDim.get(dim), bByDim.get(dim));
    }
  }

  protected void assertFloatValuesEquals(FacetResult a, FacetResult b) {
    assertEquals(a.dim, b.dim);
    assertTrue(Arrays.equals(a.path, b.path));
    assertEquals(a.childCount, b.childCount);
    assertEquals(a.value.floatValue(), b.value.floatValue(), a.value.floatValue() / 1e5);
    assertEquals(a.labelValues.length, b.labelValues.length);
    for (int i = 0; i < a.labelValues.length; i++) {
      assertEquals(a.labelValues[i].label, b.labelValues[i].label);
      assertEquals(
          a.labelValues[i].value.floatValue(),
          b.labelValues[i].value.floatValue(),
          a.labelValues[i].value.floatValue() / 1e5);
    }
  }
}