/**
* Copyright 2012 Twitter, Inc.
*
* Licensed 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 parquet.io;
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 parquet.Log;
import parquet.column.ColumnReader;
import parquet.column.impl.ColumnReadStoreImpl;
import parquet.io.api.Converter;
import parquet.io.api.GroupConverter;
import parquet.io.api.PrimitiveConverter;
import parquet.io.api.RecordConsumer;
import parquet.io.api.RecordMaterializer;
import parquet.schema.MessageType;
import parquet.schema.PrimitiveType.PrimitiveTypeName;
/**
* used to read reassembled records
* @author Julien Le Dem
*
* @param <T> the type of the materialized record
*/
class RecordReaderImplementation<T> extends RecordReader<T> {
private static final Log LOG = Log.getLog(RecordReaderImplementation.class);
public static class Case {
private int id;
private final int startLevel;
private final int depth;
private final int nextLevel;
private final boolean goingUp;
private final boolean goingDown;
private final int nextState;
private final boolean defined;
public Case(int startLevel, int depth, int nextLevel, int nextState, boolean defined) {
this.startLevel = startLevel;
this.depth = depth;
this.nextLevel = nextLevel;
this.nextState = nextState;
this.defined = defined;
// means going up the tree (towards the leaves) of the record
// true if we need to open up groups in this case
goingUp = startLevel <= depth;
// means going down the tree (towards the root) of the record
// true if we need to close groups in this case
goingDown = depth + 1 > nextLevel;
}
public void setID(int id) {
this.id = id;
}
@Override
// this implementation is buggy but the simpler one bellow has duplicates.
// it still works but generates more code than necessary
// a middle ground is necessary
// public int hashCode() {
// int hashCode = 0;
// if (goingUp) {
// hashCode += 1 * (1 + startLevel) + 2 * (1 + depth);
// }
// if (goingDown) {
// hashCode += 3 * (1 + depth) + 5 * (1 + nextLevel);
// }
// return hashCode;
// }
public int hashCode() {
int hashCode = 17;
hashCode += 31 * startLevel;
hashCode += 31 * depth;
hashCode += 31 * nextLevel;
hashCode += 31 * nextState;
hashCode += 31 * (defined ? 0 : 1);
return hashCode;
}
@Override
public boolean equals(Object obj) {
if (obj instanceof Case) {
return equals((Case)obj);
}
return false;
};
// see comment for hashCode above
// public boolean equals(Case other) {
// if (goingUp && !other.goingUp || !goingUp && other.goingUp) {
// return false;
// }
// if (goingUp && other.goingUp && (startLevel != other.startLevel || depth != other.depth)) {
// return false;
// }
// if (goingDown && !other.goingDown || !goingDown && other.goingDown) {
// return false;
// }
// if (goingDown && other.goingDown && (depth != other.depth || nextLevel != other.nextLevel)) {
// return false;
// }
// return true;
// }
public boolean equals(Case other) {
return startLevel == other.startLevel
&& depth == other.depth
&& nextLevel == other.nextLevel
&& nextState == other.nextState
&& ((defined && other.defined) || (!defined && !other.defined));
}
public int getID() {
return id;
}
public int getStartLevel() {
return startLevel;
}
public int getDepth() {
return depth;
}
public int getNextLevel() {
return nextLevel;
}
public int getNextState() {
return nextState;
}
public boolean isGoingUp() {
return goingUp;
}
public boolean isGoingDown() {
return goingDown;
}
public boolean isDefined() {
return defined;
}
@Override
public String toString() {
return "Case " + startLevel + " -> " + depth + " -> " + nextLevel + "; goto sate_"+getNextState();
}
}
public static class State {
public final int id;
public final PrimitiveColumnIO primitiveColumnIO;
public final int maxDefinitionLevel;
public final int maxRepetitionLevel;
public final PrimitiveTypeName primitive;
public final ColumnReader column;
public final String[] fieldPath; // indexed by currentLevel
public final int[] indexFieldPath; // indexed by currentLevel
public final GroupConverter[] groupConverterPath;
public final PrimitiveConverter primitiveConverter;
public final String primitiveField;
public final int primitiveFieldIndex;
public final int[] nextLevel; //indexed by next r
private int[] definitionLevelToDepth; // indexed by current d
private State[] nextState; // indexed by next r
private Case[][][] caseLookup;
private List<Case> definedCases;
private List<Case> undefinedCases;
private State(int id, PrimitiveColumnIO primitiveColumnIO, ColumnReader column, int[] nextLevel, GroupConverter[] groupConverterPath, PrimitiveConverter primitiveConverter) {
this.id = id;
this.primitiveColumnIO = primitiveColumnIO;
this.maxDefinitionLevel = primitiveColumnIO.getDefinitionLevel();
this.maxRepetitionLevel = primitiveColumnIO.getRepetitionLevel();
this.column = column;
this.nextLevel = nextLevel;
this.groupConverterPath = groupConverterPath;
this.primitiveConverter = primitiveConverter;
this.primitive = primitiveColumnIO.getType().asPrimitiveType().getPrimitiveTypeName();
this.fieldPath = primitiveColumnIO.getFieldPath();
this.primitiveField = fieldPath[fieldPath.length - 1];
this.indexFieldPath = primitiveColumnIO.getIndexFieldPath();
this.primitiveFieldIndex = indexFieldPath[indexFieldPath.length - 1];
}
public int getDepth(int definitionLevel) {
return definitionLevelToDepth[definitionLevel];
}
public List<Case> getDefinedCases() {
return definedCases;
}
public List<Case> getUndefinedCases() {
return undefinedCases;
}
public Case getCase(int currentLevel, int d, int nextR) {
return caseLookup[currentLevel][d][nextR];
}
public State getNextState(int nextR) {
return nextState[nextR];
}
}
private final GroupConverter recordRootConverter;
private final RecordMaterializer<T> recordMaterializer;
private State[] states;
private ColumnReader[] columnReaders;
/**
* @param root the root of the schema
* @param recordMaterializer responsible of materializing the records
* @param validating whether we should validate against the schema
* @param columnStore where to read the column data from
*/
public RecordReaderImplementation(MessageColumnIO root, RecordMaterializer<T> recordMaterializer, boolean validating, ColumnReadStoreImpl columnStore) {
this.recordMaterializer = recordMaterializer;
this.recordRootConverter = recordMaterializer.getRootConverter(); // TODO: validator(wrap(recordMaterializer), validating, root.getType());
PrimitiveColumnIO[] leaves = root.getLeaves().toArray(new PrimitiveColumnIO[root.getLeaves().size()]);
columnReaders = new ColumnReader[leaves.length];
int[][] nextColumnIdxForRepLevel = new int[leaves.length][];
int[][] levelToClose = new int[leaves.length][];
GroupConverter[][] groupConverterPaths = new GroupConverter[leaves.length][];
PrimitiveConverter[] leafConverters = new PrimitiveConverter[leaves.length];
int[] firstIndexForLevel = new int[256]; // "256 levels of nesting ought to be enough for anybody"
// build the automaton
for (int i = 0; i < leaves.length; i++) {
PrimitiveColumnIO leafColumnIO = leaves[i];
//generate converters along the path from root to leaf
final int[] indexFieldPath = leafColumnIO.getIndexFieldPath();
groupConverterPaths[i] = new GroupConverter[indexFieldPath.length - 1];
GroupConverter current = this.recordRootConverter;
for (int j = 0; j < indexFieldPath.length - 1; j++) {
current = current.getConverter(indexFieldPath[j]).asGroupConverter();
groupConverterPaths[i][j] = current;
}
leafConverters[i] = current.getConverter(indexFieldPath[indexFieldPath.length - 1]).asPrimitiveConverter();
columnReaders[i] = columnStore.getColumnReader(leafColumnIO.getColumnDescriptor());
int maxRepetitionLevel = leafColumnIO.getRepetitionLevel();
nextColumnIdxForRepLevel[i] = new int[maxRepetitionLevel+1];
levelToClose[i] = new int[maxRepetitionLevel+1]; //next level
for (int nextRepLevel = 0; nextRepLevel <= maxRepetitionLevel; ++nextRepLevel) {
// remember which is the first for this level
if (leafColumnIO.isFirst(nextRepLevel)) {
firstIndexForLevel[nextRepLevel] = i;
}
int nextColIdx;
//TODO: when we use nextColumnIdxForRepLevel, should we provide current rep level or the rep level for next item
// figure out automaton transition
if (nextRepLevel == 0) { // 0 always means jump to the next (the last one being a special case)
nextColIdx = i + 1;
} else if (leafColumnIO.isLast(nextRepLevel)) { // when we are at the last of the next repetition level we jump back to the first
nextColIdx = firstIndexForLevel[nextRepLevel];
} else { // otherwise we just go back to the next.
nextColIdx = i + 1;
}
// figure out which level down the tree we need to go back
if (nextColIdx == leaves.length) { // reached the end of the record => close all levels
levelToClose[i][nextRepLevel] = 0;
} else if (leafColumnIO.isLast(nextRepLevel)) { // reached the end of this level => close the repetition level
ColumnIO parent = leafColumnIO.getParent(nextRepLevel);
levelToClose[i][nextRepLevel] = parent.getFieldPath().length - 1;
} else { // otherwise close until the next common parent
levelToClose[i][nextRepLevel] = getCommonParentLevel(
leafColumnIO.getFieldPath(),
leaves[nextColIdx].getFieldPath());
}
// sanity check: that would be a bug
if (levelToClose[i][nextRepLevel] > leaves[i].getFieldPath().length-1) {
throw new ParquetEncodingException(Arrays.toString(leaves[i].getFieldPath())+" -("+nextRepLevel+")-> "+levelToClose[i][nextRepLevel]);
}
nextColumnIdxForRepLevel[i][nextRepLevel] = nextColIdx;
}
}
states = new State[leaves.length];
for (int i = 0; i < leaves.length; i++) {
states[i] = new State(i, leaves[i], columnReaders[i], levelToClose[i], groupConverterPaths[i], leafConverters[i]);
int[] definitionLevelToDepth = new int[states[i].primitiveColumnIO.getDefinitionLevel() + 1];
// for each possible definition level, determine the depth at which to create groups
final ColumnIO[] path = states[i].primitiveColumnIO.getPath();
int depth = 0;
for (int d = 0; d < definitionLevelToDepth.length; ++d) {
while (depth < (states[i].fieldPath.length - 1)
&& d >= path[depth + 1].getDefinitionLevel()
) {
++ depth;
}
definitionLevelToDepth[d] = depth - 1;
}
states[i].definitionLevelToDepth = definitionLevelToDepth;
}
for (int i = 0; i < leaves.length; i++) {
State state = states[i];
int[] nextStateIds = nextColumnIdxForRepLevel[i];
state.nextState = new State[nextStateIds.length];
for (int j = 0; j < nextStateIds.length; j++) {
state.nextState[j] = nextStateIds[j] == states.length ? null : states[nextStateIds[j]];
}
}
for (int i = 0; i < states.length; i++) {
State state = states[i];
final Map<Case, Case> definedCases = new HashMap<Case, Case>();
final Map<Case, Case> undefinedCases = new HashMap<Case, Case>();
Case[][][] caseLookup = new Case[state.fieldPath.length][][];
for (int currentLevel = 0; currentLevel < state.fieldPath.length; ++ currentLevel) {
caseLookup[currentLevel] = new Case[state.maxDefinitionLevel+1][];
for (int d = 0; d <= state.maxDefinitionLevel; ++ d) {
caseLookup[currentLevel][d] = new Case[state.maxRepetitionLevel+1];
for (int nextR = 0; nextR <= state.maxRepetitionLevel; ++ nextR) {
int caseStartLevel = currentLevel;
int caseDepth = Math.max(state.getDepth(d), caseStartLevel - 1);
int caseNextLevel = Math.min(state.nextLevel[nextR], caseDepth + 1);
Case currentCase = new Case(caseStartLevel, caseDepth, caseNextLevel, getNextReader(state.id, nextR), d == state.maxDefinitionLevel);
Map<Case, Case> cases = currentCase.isDefined() ? definedCases : undefinedCases;
if (!cases.containsKey(currentCase)) {
currentCase.setID(cases.size());
cases.put(currentCase, currentCase);
} else {
currentCase = cases.get(currentCase);
}
caseLookup[currentLevel][d][nextR] = currentCase;
}
}
}
state.caseLookup = caseLookup;
state.definedCases = new ArrayList<Case>(definedCases.values());
state.undefinedCases = new ArrayList<Case>(undefinedCases.values());
Comparator<Case> caseComparator = new Comparator<Case>() {
@Override
public int compare(Case o1, Case o2) {
return o1.id - o2.id;
}
};
Collections.sort(state.definedCases, caseComparator);
Collections.sort(state.undefinedCases, caseComparator);
}
}
//TODO: have those wrappers for a converter
private RecordConsumer validator(RecordConsumer recordConsumer, boolean validating, MessageType schema) {
return validating ? new ValidatingRecordConsumer(recordConsumer, schema) : recordConsumer;
}
private RecordConsumer wrap(RecordConsumer recordConsumer) {
if (Log.DEBUG) {
return new RecordConsumerLoggingWrapper(recordConsumer);
}
return recordConsumer;
}
/**
* @see parquet.io.RecordReader#read()
*/
@Override
public T read() {
int currentLevel = 0;
recordRootConverter.start();
State currentState = states[0];
do {
ColumnReader columnReader = currentState.column;
int d = columnReader.getCurrentDefinitionLevel();
// creating needed nested groups until the current field (opening tags)
int depth = currentState.definitionLevelToDepth[d];
for (; currentLevel <= depth; ++currentLevel) {
currentState.groupConverterPath[currentLevel].start();
}
// currentLevel = depth + 1 at this point
// set the current value
if (d >= currentState.maxDefinitionLevel) {
// not null
columnReader.writeCurrentValueToConverter();
}
columnReader.consume();
int nextR = currentState.maxRepetitionLevel == 0 ? 0 : columnReader.getCurrentRepetitionLevel();
// level to go to close current groups
int next = currentState.nextLevel[nextR];
for (; currentLevel > next; currentLevel--) {
currentState.groupConverterPath[currentLevel - 1].end();
}
currentState = currentState.nextState[nextR];
} while (currentState != null);
recordRootConverter.end();
return recordMaterializer.getCurrentRecord();
}
private static void log(String string) {
LOG.debug(string);
}
int getNextReader(int current, int nextRepetitionLevel) {
State nextState = states[current].nextState[nextRepetitionLevel];
return nextState == null ? states.length : nextState.id;
}
int getNextLevel(int current, int nextRepetitionLevel) {
return states[current].nextLevel[nextRepetitionLevel];
}
private int getCommonParentLevel(String[] previous, String[] next) {
int i = 0;
while (i < Math.min(previous.length, next.length) && previous[i].equals(next[i])) {
++i;
}
return i;
}
protected int getStateCount() {
return states.length;
}
protected State getState(int i) {
return states[i];
}
protected RecordMaterializer<T> getMaterializer() {
return recordMaterializer;
}
protected Converter getRecordConsumer() {
return recordRootConverter;
}
protected Iterable<ColumnReader> getColumnReaders() {
// Converting the array to an iterable ensures that the array cannot be altered
return Arrays.asList(columnReaders);
}
}