/**
* 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.trevni.avro;
import java.io.IOException;
import java.io.Closeable;
import java.io.File;
import java.nio.ByteBuffer;
import java.util.Iterator;
import java.util.Map;
import java.util.HashMap;
import java.util.List;
import org.apache.trevni.ColumnMetaData;
import org.apache.trevni.ColumnFileReader;
import org.apache.trevni.ColumnValues;
import org.apache.trevni.Input;
import org.apache.trevni.InputFile;
import org.apache.trevni.TrevniRuntimeException;
import org.apache.avro.Schema;
import org.apache.avro.Schema.Field;
import org.apache.avro.generic.GenericData;
import static org.apache.trevni.avro.AvroColumnator.isSimple;
/** Read files written with {@link AvroColumnWriter}. A subset of the schema
* used for writing may be specified when reading. In this case only columns
* of the subset schema are read. */
public class AvroColumnReader<D>
implements Iterator<D>, Iterable<D>, Closeable {
private ColumnFileReader reader;
private GenericData model;
private Schema fileSchema;
private Schema readSchema;
private ColumnValues[] values;
private int[] arrayWidths;
private int column; // current index in values
private Map<String,Map<String,Object>> defaults =
new HashMap<String,Map<String,Object>>();
/** Parameters for reading an Avro column file. */
public static class Params {
Input input;
Schema schema;
GenericData model = GenericData.get();
/** Construct reading from a file. */
public Params(File file) throws IOException {
this(new InputFile(file));
}
/** Construct reading from input. */
public Params(Input input) { this.input = input; }
/** Set subset schema to project data down to. */
public Params setSchema(Schema schema) {
this.schema = schema;
return this;
}
/** Set data representation. */
public Params setModel(GenericData model) {
this.model = model;
return this;
}
}
/** Construct a reader for a file. */
public AvroColumnReader(Params params)
throws IOException {
this.reader = new ColumnFileReader(params.input);
this.model = params.model;
this.fileSchema =
Schema.parse(reader.getMetaData().getString(AvroColumnWriter.SCHEMA_KEY));
this.readSchema = params.schema == null ? fileSchema : params.schema;
initialize();
}
/** Return the schema for data in this file. */
public Schema getFileSchema() { return fileSchema; }
void initialize() throws IOException {
// compute a mapping from column name to number for file
Map<String,Integer> fileColumnNumbers = new HashMap<String,Integer>();
int i = 0;
for (ColumnMetaData c : new AvroColumnator(fileSchema).getColumns())
fileColumnNumbers.put(c.getName(), i++);
// create iterator for each column in readSchema
AvroColumnator readColumnator = new AvroColumnator(readSchema);
this.arrayWidths = readColumnator.getArrayWidths();
ColumnMetaData[] readColumns = readColumnator.getColumns();
this.values = new ColumnValues[readColumns.length];
int j = 0;
for (ColumnMetaData c : readColumns) {
Integer n = fileColumnNumbers.get(c.getName());
if (n != null)
values[j++] = reader.getValues(n);
}
findDefaults(readSchema, fileSchema);
}
// get defaults for fields in read that are not in write
private void findDefaults(Schema read, Schema write) {
switch (read.getType()) {
case NULL: case BOOLEAN:
case INT: case LONG:
case FLOAT: case DOUBLE:
case BYTES: case STRING:
case ENUM: case FIXED:
if (read.getType() != write.getType())
throw new TrevniRuntimeException("Type mismatch: "+read+" & "+write);
break;
case MAP:
findDefaults(read.getValueType(), write.getValueType());
break;
case ARRAY:
findDefaults(read.getElementType(), write.getElementType());
break;
case UNION:
for (Schema s : read.getTypes()) {
Integer index = write.getIndexNamed(s.getFullName());
if (index == null)
throw new TrevniRuntimeException("No matching branch: "+s);
findDefaults(s, write.getTypes().get(index));
}
break;
case RECORD:
for (Field f : read.getFields()) {
Field g = write.getField(f.name());
if (g == null)
setDefault(read, f);
else
findDefaults(f.schema(), g.schema());
}
break;
default:
throw new TrevniRuntimeException("Unknown schema: "+read);
}
}
private void setDefault(Schema record, Field f) {
String recordName = record.getFullName();
Map<String,Object> recordDefaults = defaults.get(recordName);
if (recordDefaults == null) {
recordDefaults = new HashMap<String,Object>();
defaults.put(recordName, recordDefaults);
}
recordDefaults.put(f.name(), model.getDefaultValue(f));
}
@Override
public Iterator<D> iterator() { return this; }
@Override
public boolean hasNext() {
return values[0].hasNext();
}
/** Return the number of rows in this file. */
public long getRowCount() { return reader.getRowCount(); }
@Override
public D next() {
try {
for (int i = 0; i < values.length; i++)
if (values[i] != null)
values[i].startRow();
this.column = 0;
return (D)read(readSchema);
} catch (IOException e) {
throw new TrevniRuntimeException(e);
}
}
private Object read(Schema s) throws IOException {
if (isSimple(s))
return nextValue(s, column++);
final int startColumn = column;
switch (s.getType()) {
case MAP:
int size = values[column].nextLength();
Map map = (Map)new HashMap(size);
for (int i = 0; i < size; i++) {
this.column = startColumn;
values[column++].nextValue(); // null in parent
String key = (String)values[column++].nextValue(); // key
map.put(key, read(s.getValueType())); // value
}
column = startColumn + arrayWidths[startColumn];
return map;
case RECORD:
Object record = model.newRecord(null, s);
Map<String,Object> rDefaults = defaults.get(s.getFullName());
for (Field f : s.getFields()) {
Object value = ((rDefaults != null) && rDefaults.containsKey(f.name()))
? model.deepCopy(f.schema(), rDefaults.get(f.name()))
: read(f.schema());
model.setField(record, f.name(), f.pos(), value);
}
return record;
case ARRAY:
int length = values[column].nextLength();
List elements = (List)new GenericData.Array(length, s);
for (int i = 0; i < length; i++) {
this.column = startColumn;
Object value = nextValue(s, column++);
if (!isSimple(s.getElementType()))
value = read(s.getElementType());
elements.add(value);
}
column = startColumn + arrayWidths[startColumn];
return elements;
case UNION:
Object value = null;
for (Schema branch : s.getTypes()) {
if (branch.getType() == Schema.Type.NULL) continue;
if (values[column].nextLength() == 1) {
value = nextValue(branch, column);
column++;
if (!isSimple(branch))
value = read(branch);
} else {
column += arrayWidths[column];
}
}
return value;
default:
throw new TrevniRuntimeException("Unknown schema: "+s);
}
}
private Object nextValue(Schema s, int column) throws IOException {
Object v = values[column].nextValue();
switch (s.getType()) {
case ENUM:
return model.createEnum(s.getEnumSymbols().get((Integer)v), s);
case FIXED:
return model.createFixed(null, ((ByteBuffer)v).array(), s);
}
return v;
}
@Override
public void remove() { throw new UnsupportedOperationException(); }
@Override
public void close() throws IOException {
reader.close();
}
}