AvroColumnWriter.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.trevni.avro;

import java.io.IOException;
import java.io.File;
import java.io.OutputStream;
import java.util.Collection;
import java.util.Map;

import org.apache.trevni.ColumnFileMetaData;
import org.apache.trevni.ColumnFileWriter;
import org.apache.trevni.TrevniRuntimeException;

import org.apache.avro.Schema;
import org.apache.avro.Schema.Field;
import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericFixed;
import org.apache.avro.util.Utf8;

import static org.apache.trevni.avro.AvroColumnator.isSimple;

/** Write Avro records to a Trevni column file.
 *
 * <p>Each primitive type is written to a separate column.
 *
 * <p>Output is buffered until {@link #writeTo(OutputStream)} is called.  The
 * {@link #sizeEstimate()} indicates both the amount of data buffered and the
 * size of the file that will be written.
 */
public class AvroColumnWriter<D> {
  private Schema schema;
  private GenericData model;
  private ColumnFileWriter writer;
  private int[] arrayWidths;

  public static final String SCHEMA_KEY = "avro.schema";

  public AvroColumnWriter(Schema s, ColumnFileMetaData meta)
    throws IOException {
    this(s, meta, GenericData.get());
  }

  public AvroColumnWriter(Schema s, ColumnFileMetaData meta, GenericData model)
    throws IOException {
    this.schema = s;
    AvroColumnator columnator = new AvroColumnator(s);
    meta.set(SCHEMA_KEY, s.toString());           // save schema in file
    this.writer = new ColumnFileWriter(meta, columnator.getColumns());
    this.arrayWidths = columnator.getArrayWidths();
    this.model = model;
  }

  /** Return the approximate size of the file that will be written.  Tries to
   * slightly over-estimate.  Indicates both the size in memory of the buffered
   * data as well as the size of the file that will be written by {@link
   * #writeTo(OutputStream)}. */
  public long sizeEstimate() { return writer.sizeEstimate(); }

  /** Write all rows added to the named output stream. */
  public void writeTo(OutputStream out) throws IOException {
    writer.writeTo(out);
  }

  /** Write all rows added to the named file. */
  public void writeTo(File file) throws IOException {
    writer.writeTo(file);
  }

  /** Add a row to the file. */
  public void write(D value) throws IOException {
    writer.startRow();
    int count = write(value, schema, 0);
    assert(count == writer.getColumnCount());
    writer.endRow();
  }

  private int write(Object o, Schema s, int column) throws IOException {
    if (isSimple(s)) {
      writeValue(o, s, column);
      return column+1;
    }
    switch (s.getType()) {
    case MAP:
      Map<?,?> map = (Map)o;
      writer.writeLength(map.size(), column);
      for (Map.Entry e : map.entrySet()) {
        writer.writeValue(null, column);
        writer.writeValue(e.getKey(), column+1);
        int c = write(e.getValue(), s.getValueType(), column+2);
        assert(c == column+arrayWidths[column]);
      }
      return column+arrayWidths[column];
    case RECORD:
      for (Field f : s.getFields())
        column = write(model.getField(o,f.name(),f.pos()), f.schema(), column);
      return column;
    case ARRAY:
      Collection elements = (Collection)o;
      writer.writeLength(elements.size(), column);
      if (isSimple(s.getElementType())) {         // optimize simple arrays
        for (Object element : elements)
          writeValue(element, s.getElementType(), column);
        return column+1;
      }
      for (Object element : elements) {
        writer.writeValue(null, column);
        int c = write(element, s.getElementType(), column+1);
        assert(c == column+arrayWidths[column]);
      }
      return column+arrayWidths[column];
    case UNION:
      int b = model.resolveUnion(s, o);
      int i = 0;
      for (Schema branch : s.getTypes()) {
        boolean selected = i++ == b;
        if (branch.getType() == Schema.Type.NULL) continue;
        if (!selected) {
          writer.writeLength(0, column);
          column+=arrayWidths[column];
        } else {
          writer.writeLength(1, column);
          if (isSimple(branch)) {
            writeValue(o, branch, column++);
          } else {
            writer.writeValue(null, column);
            column = write(o, branch, column+1);
          }
        }
      }
      return column;
    default:
      throw new TrevniRuntimeException("Unknown schema: "+s);
    }
  }

  private void writeValue(Object value, Schema s, int column)
    throws IOException {

    switch (s.getType()) {
    case STRING:
      if (value instanceof Utf8)                    // convert Utf8 to String
        value = value.toString();
      break;
    case ENUM:
      if (value instanceof Enum)
        value = ((Enum)value).ordinal();
      else
        value = s.getEnumOrdinal(value.toString());
      break;
    case FIXED:
      value = ((GenericFixed)value).bytes();
      break;
    }
    writer.writeValue(value, column);
  }

}