package io.github.vmzakharov.ecdataframe.dataframe;
import org.eclipse.collections.api.list.ListIterable;
import org.eclipse.collections.api.list.MutableList;
import org.eclipse.collections.api.list.primitive.IntList;
import org.eclipse.collections.api.list.primitive.MutableIntList;
import org.eclipse.collections.api.map.MutableMap;
import org.eclipse.collections.impl.factory.Lists;
import org.eclipse.collections.impl.factory.Maps;
import org.eclipse.collections.impl.factory.primitive.IntLists;
public class DfIndex
{
private final MutableMap<ListIterable<Object>, MutableIntList> rowIndexByKey = Maps.mutable.of();
public DfIndex(DataFrame newIndexedDataFrame, ListIterable<String> indexByColumnNames)
{
this.buildIndex(newIndexedDataFrame, indexByColumnNames);
}
private void buildIndex(DataFrame dataFrame, ListIterable<String> indexByColumnNames)
{
ListIterable<DfColumn> indexByColumns = indexByColumnNames.collect(dataFrame::getColumnNamed);
int rowCount = dataFrame.rowCount();
for (int rowIndex = 0; rowIndex < rowCount; rowIndex++)
{
this.addRowIndex(
this.computeKeyFrom(indexByColumns, rowIndex),
rowIndex
);
}
}
public IntList getRowIndicesAtKey(ListIterable<Object> key)
{
return this.rowIndexByKey.getIfAbsent(key, IntLists.mutable::of);
}
private ListIterable<Object> computeKeyFrom(ListIterable<DfColumn> indexByColumns, int rowIndex)
{
MutableList<Object> key = Lists.fixedSize.of(new Object[indexByColumns.size()]);
indexByColumns.forEachWithIndex((col, index) -> key.set(index, col.getObject(rowIndex)));
return key;
}
public void addRowIndex(ListIterable<Object> key, int rowIndex)
{
MutableIntList rowIndicesAtKey = this.rowIndexByKey.getIfAbsentPut(key, IntLists.mutable::of);
rowIndicesAtKey.add(rowIndex);
}
}