<commit_before>package se.l4.silo.engine.internal.index;

import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.Comparator;
import java.util.Iterator;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.atomic.LongAdder;
import java.util.function.Predicate;

import org.h2.mvstore.MVMap;
import org.h2.mvstore.type.DataType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Multimap;

import se.l4.silo.StorageException;
import se.l4.silo.engine.DataEncounter;
import se.l4.silo.engine.FieldDef;
import se.l4.silo.engine.Fields;
import se.l4.silo.engine.MVStoreManager;
import se.l4.silo.engine.QueryEncounter;
import se.l4.silo.engine.QueryEngine;
import se.l4.silo.engine.config.IndexConfig;
import se.l4.silo.engine.types.ArrayFieldType;
import se.l4.silo.engine.types.FieldType;
import se.l4.silo.engine.types.LongFieldType;
import se.l4.silo.engine.types.MaxMin;
import se.l4.silo.engine.types.MergedFieldType;
import se.l4.silo.index.IndexQueryRequest;
import se.l4.silo.index.IndexQueryRequest.Criterion;
import se.l4.silo.index.IndexQueryRequest.SortOnField;

/**
 * {@link QueryEngine} that creates a queryable index similiar in functionality
 * to traditional databases.
 * 
 * @author <NAME>
 *
 */
public class IndexQueryEngine
	implements QueryEngine<IndexQueryRequest>
{
	private static final Logger logger = LoggerFactory.getLogger(IndexQueryEngine.class);
	
	private final String name;
	
	private final MVStoreManager store;
	private final MVMap<Long, Object[]> indexedData;
	private final MVMap<Object[], Object[]> index;
	
	private final String[] fields;
	private final Set<String> fieldsSet;
	private final FieldType[] fieldTypes;
	private final String[] sortFields;
	private final FieldType[] sortFieldTypes;

	public IndexQueryEngine(String name, Fields fields, MVStoreManager store, IndexConfig config)
	{
		this.name = name;
		this.store = store;
		
		this.fields = config.getFields();
		this.sortFields = config.getSortFields();
		
		this.fieldsSet = ImmutableSet.copyOf(this.fields);
		
		MergedFieldType dataFieldType = createMultiFieldType(name, fields, config.getFields(), false);
		indexedData = store.openMap("data:" + name, LongFieldType.INSTANCE, dataFieldType);
		
		MergedFieldType indexKey = createFieldType(name, fields, config.getFields(), true);
		this.fieldTypes = indexKey.getTypes();
		
		MergedFieldType indexData = createFieldType(name, fields, config.getSortFields(), false);
		this.sortFieldTypes = indexData.getTypes();
		index = store.openMap("index:" + name, indexKey, indexData);
	}
	
	/**
	 * Create a {@link MergedFieldType} for the given fields.
	 * 
	 * @param name
	 * @param fields
	 * @param fieldNames
	 * @param appendId
	 * @return
	 */
	private MergedFieldType createFieldType(String name, Fields fields, String[] fieldNames, boolean appendId)
	{
		FieldType[] result = new FieldType[fieldNames.length + (appendId ? 1 : 0)];
		for(int i=0, n=fieldNames.length; i<n; i++)
		{
			String field = fieldNames[i];
			Optional<FieldDef> def = fields.get(field);
			if(! def.isPresent())
			{
				throw new StorageException("Problem creating index query engine `" + name + "`, trying to use unknown field `" + field + "`");
			}
			result[i] = def.get().getType();
		}
		
		if(appendId)
		{
			result[fieldNames.length] = LongFieldType.INSTANCE;
		}
		
		return new MergedFieldType(result);
	}
	
	/**
	 * Create a {@link MergedFieldType} that allows multiple values for
	 * every field.
	 * 
	 * @param name
	 * @param fields
	 * @param fieldNames
	 * @param appendId
	 * @return
	 */
	private MergedFieldType createMultiFieldType(String name, Fields fields, String[] fieldNames, boolean appendId)
	{
		FieldType[] result = new FieldType[fieldNames.length + (appendId ? 1 : 0)];
		for(int i=0, n=fieldNames.length; i<n; i++)
		{
			String field = fieldNames[i];
			Optional<FieldDef> def = fields.get(field);
			if(! def.isPresent())
			{
				throw new StorageException("Problem creating index query engine `" + name + "`, trying to use unknown field `" + field + "`");
			}
			result[i] = new ArrayFieldType(def.get().getType());
		}
		
		if(appendId)
		{
			result[fieldNames.length] = LongFieldType.INSTANCE;
		}
		
		return new MergedFieldType(result);
	}
	
	
	@Override
	public void close()
		throws IOException
	{
		store.close();
	}

	@Override
	public void update(long id, DataEncounter encounter)
	{
		logger.debug("{}: Update entry for id {}", name, id);

		// Fetch all the values that need to be indexed
		Multimap<String, Object> values = HashMultimap.create();
		encounter.findStructuredKeys(fieldsSet, values::put);
		
		// Generate the sort data
		Object[] data = encounter.getStructuredArray(sortFields);
		
		// Look up our previously indexed key to see if we need to delete it
		Object[] previousKey = indexedData.get(id);
		remove(previousKey, id);
		
		// Store the new key
		store(values, data, id);
	}
	
	private void store(Multimap<String, Object> values, Object[] sortData, long id)
	{
		Object[] generatedKey = new Object[fields.length + 1];
		generatedKey[fields.length] = id;
		
		recursiveStore(values, sortData, generatedKey, 0);
		
		// Create a combined key for indexedData
		Object[][] data = new Object[fields.length][];
		for(int i=0, n=fields.length; i<n; i++)
		{
			data[i] = values.get(fields[i]).toArray();
		}
		indexedData.put(id, data);
	}
	
	private void recursiveStore(Multimap<String, Object> values, Object[] sortData, Object[] generatedKey, int i)
	{
		String fieldName = fields[i];
		if(i == fields.length - 1)
		{
			for(Object o : values.get(fieldName))
			{
				generatedKey[i] = o;
				if(logger.isTraceEnabled())
				{
					logger.trace("  storing key=" + Arrays.toString(generatedKey) + ", sort=" + Arrays.toString(sortData));<commit_msg>Handling case where only sort fields are used<commit_after>package se.l4.silo.engine.internal.index;

import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.Comparator;
import java.util.Iterator;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.atomic.LongAdder;
import java.util.function.Predicate;

import org.h2.mvstore.MVMap;
import org.h2.mvstore.type.DataType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Multimap;

import se.l4.silo.StorageException;
import se.l4.silo.engine.DataEncounter;
import se.l4.silo.engine.FieldDef;
import se.l4.silo.engine.Fields;
import se.l4.silo.engine.MVStoreManager;
import se.l4.silo.engine.QueryEncounter;
import se.l4.silo.engine.QueryEngine;
import se.l4.silo.engine.config.IndexConfig;
import se.l4.silo.engine.types.ArrayFieldType;
import se.l4.silo.engine.types.FieldType;
import se.l4.silo.engine.types.LongFieldType;
import se.l4.silo.engine.types.MaxMin;
import se.l4.silo.engine.types.MergedFieldType;
import se.l4.silo.index.IndexQueryRequest;
import se.l4.silo.index.IndexQueryRequest.Criterion;
import se.l4.silo.index.IndexQueryRequest.SortOnField;

/**
 * {@link QueryEngine} that creates a queryable index similiar in functionality
 * to traditional databases.
 * 
 * @author <NAME>
 *
 */
public class IndexQueryEngine
	implements QueryEngine<IndexQueryRequest>
{
	private static final Logger logger = LoggerFactory.getLogger(IndexQueryEngine.class);
	
	private final String name;
	
	private final MVStoreManager store;
	private final MVMap<Long, Object[]> indexedData;
	private final MVMap<Object[], Object[]> index;
	
	private final String[] fields;
	private final Set<String> fieldsSet;
	private final FieldType[] fieldTypes;
	private final String[] sortFields;
	private final FieldType[] sortFieldTypes;

	public IndexQueryEngine(String name, Fields fields, MVStoreManager store, IndexConfig config)
	{
		this.name = name;
		this.store = store;
		
		this.fields = config.getFields();
		this.sortFields = config.getSortFields();
		
		this.fieldsSet = ImmutableSet.copyOf(this.fields);
		
		MergedFieldType dataFieldType = createMultiFieldType(name, fields, config.getFields(), false);
		indexedData = store.openMap("data:" + name, LongFieldType.INSTANCE, dataFieldType);
		
		MergedFieldType indexKey = createFieldType(name, fields, config.getFields(), true);
		this.fieldTypes = indexKey.getTypes();
		
		MergedFieldType indexData = createFieldType(name, fields, config.getSortFields(), false);
		this.sortFieldTypes = indexData.getTypes();
		index = store.openMap("index:" + name, indexKey, indexData);
		
		if(logger.isDebugEnabled())
		{
			logger.debug(name + ": fields=" + Arrays.toString(this.fields) + ", sortFields=" + Arrays.toString(this.sortFields));
		}
	}
	
	/**
	 * Create a {@link MergedFieldType} for the given fields.
	 * 
	 * @param name
	 * @param fields
	 * @param fieldNames
	 * @param appendId
	 * @return
	 */
	private MergedFieldType createFieldType(String name, Fields fields, String[] fieldNames, boolean appendId)
	{
		FieldType[] result = new FieldType[fieldNames.length + (appendId ? 1 : 0)];
		for(int i=0, n=fieldNames.length; i<n; i++)
		{
			String field = fieldNames[i];
			Optional<FieldDef> def = fields.get(field);
			if(! def.isPresent())
			{
				throw new StorageException("Problem creating index query engine `" + name + "`, trying to use unknown field `" + field + "`");
			}
			result[i] = def.get().getType();
		}
		
		if(appendId)
		{
			result[fieldNames.length] = LongFieldType.INSTANCE;
		}
		
		return new MergedFieldType(result);
	}
	
	/**
	 * Create a {@link MergedFieldType} that allows multiple values for
	 * every field.
	 * 
	 * @param name
	 * @param fields
	 * @param fieldNames
	 * @param appendId
	 * @return
	 */
	private MergedFieldType createMultiFieldType(String name, Fields fields, String[] fieldNames, boolean appendId)
	{
		FieldType[] result = new FieldType[fieldNames.length + (appendId ? 1 : 0)];
		for(int i=0, n=fieldNames.length; i<n; i++)
		{
			String field = fieldNames[i];
			Optional<FieldDef> def = fields.get(field);
			if(! def.isPresent())
			{
				throw new StorageException("Problem creating index query engine `" + name + "`, trying to use unknown field `" + field + "`");
			}
			result[i] = new ArrayFieldType(def.get().getType());
		}
		
		if(appendId)
		{
			result[fieldNames.length] = LongFieldType.INSTANCE;
		}
		
		return new MergedFieldType(result);
	}
	
	
	@Override
	public void close()
		throws IOException
	{
		store.close();
	}

	@Override
	public void update(long id, DataEncounter encounter)
	{
		logger.debug("{}: Update entry for id {}", name, id);

		// Fetch all the values that need to be indexed
		Multimap<String, Object> values = HashMultimap.create();
		encounter.findStructuredKeys(fieldsSet, values::put);
		
		// Generate the sort data
		Object[] data = encounter.getStructuredArray(sortFields);
		
		// Look up our previously indexed key to see if we need to delete it
		Object[] previousKey = indexedData.get(id);
		remove(previousKey, id);
		
		// Store the new key
		store(values, data, id);
	}
	
	private void store(Multimap<String, Object> values, Object[] sortData, long id)
	{
		Object[] generatedKey = new Object[fields.length + 1];
		generatedKey[fields.length] = id;
		
		if(fields.length == 0)
		{
			// Special case for only sort data
			if(logger.isTraceEnabled())
			{
				logger.trace("  storing key=" + Arrays.toString(generatedKey) + ", sort=" + Arrays.toString(sortData));
			}
			index.put(generatedKey, sortData);
		}
		else
		{
			recursiveStore(values, sortData, generatedKey, 0);
		}
		
		// Create a combined key for indexedData
		Object[][] data = new Object[fields.length][];
		for(int i=0, n=fields.length; i<n; i++)
		{
			data[i] = values.get(fields[i]).toArray();
		}
		indexedData.put(id, data);
	}
	
	private void recursiveStore(Multimap<String, Object> values, Object[] sortData, Object[] generatedKey, int i)
	{
		String fieldName = fields[i];
		if(i == fields.length - 1)
		{
			for(Object o : values.get(fieldName))
			{
				generatedKey[i] = o;
				if(logger.isTraceEnabled())
				{
					logger.trace("  storing key=" + Arrays.toString(generatedKey) + ", sort=" + Arrays.toString(sortData));
{'fraction_non_alphanumeric': 0.12730479397146063, 'fraction_numerical': 0.005050505050505051, 'mean_word_length': 4.310770540655598, 'pattern_counts': {'":': 1, '<': 29, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4673318', 'n_tokens_mistral': 4418, 'n_tokens_neox': 3994, 'n_words': 958}