/**
* 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.hadoop.hive.ql.io.orc;
import java.io.IOException;
import java.sql.Timestamp;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.ql.exec.vector.BytesColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.ColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.DecimalColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.DoubleColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.ListColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.LongColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.MapColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.StructColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.TimestampColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.UnionColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.VectorizedRowBatch;
import org.apache.hadoop.hive.serde2.objectinspector.ListObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.MapObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.StructField;
import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.UnionObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.BinaryObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.BooleanObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.ByteObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.DateObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.DoubleObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.FloatObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.HiveCharObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.HiveDecimalObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.HiveVarcharObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.IntObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.LongObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.ShortObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.StringObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.TimestampObjectInspector;
import org.apache.hadoop.io.BytesWritable;
import org.apache.hadoop.io.Text;
import com.google.common.annotations.VisibleForTesting;
import org.apache.orc.PhysicalWriter;
/**
* An ORC file writer. The file is divided into stripes, which is the natural
* unit of work when reading. Each stripe is buffered in memory until the
* memory reaches the stripe size and then it is written out broken down by
* columns. Each column is written by a TreeWriter that is specific to that
* type of column. TreeWriters may have children TreeWriters that handle the
* sub-types. Each of the TreeWriters writes the column's data as a set of
* streams.
*
* This class is unsynchronized like most Stream objects, so from the creation of an OrcFile and all
* access to a single instance has to be from a single thread.
*
* There are no known cases where these happen between different threads today.
*
* Caveat: the MemoryManager is created during WriterOptions create, that has to be confined to a single
* thread as well.
*
*/
public class WriterImpl extends org.apache.orc.impl.WriterImpl implements Writer {
private final ObjectInspector inspector;
private final VectorizedRowBatch internalBatch;
private final StructField[] fields;
WriterImpl(FileSystem fs,
Path path,
OrcFile.WriterOptions opts) throws IOException {
super(fs, path, opts);
this.inspector = opts.getInspector();
this.internalBatch = opts.getSchema().createRowBatch(opts.getBatchSize());
this.fields = initializeFieldsFromOi(inspector);
}
private static StructField[] initializeFieldsFromOi(ObjectInspector inspector) {
if (inspector instanceof StructObjectInspector) {
List<? extends StructField> fieldList =
((StructObjectInspector) inspector).getAllStructFieldRefs();
StructField[] fields = new StructField[fieldList.size()];
fieldList.toArray(fields);
return fields;
} else {
return null;
}
}
/**
* Set the value for a given column value within a batch.
* @param rowId the row to set
* @param column the column to set
* @param inspector the object inspector to interpret the obj
* @param obj the value to use
*/
static void setColumn(int rowId, ColumnVector column,
ObjectInspector inspector, Object obj) {
if (obj == null) {
column.noNulls = false;
column.isNull[rowId] = true;
} else {
switch (inspector.getCategory()) {
case PRIMITIVE:
switch (((PrimitiveObjectInspector) inspector)
.getPrimitiveCategory()) {
case BOOLEAN: {
LongColumnVector vector = (LongColumnVector) column;
vector.vector[rowId] =
((BooleanObjectInspector) inspector).get(obj) ? 1 : 0;
break;
}
case BYTE: {
LongColumnVector vector = (LongColumnVector) column;
vector.vector[rowId] = ((ByteObjectInspector) inspector).get(obj);
break;
}
case SHORT: {
LongColumnVector vector = (LongColumnVector) column;
vector.vector[rowId] =
((ShortObjectInspector) inspector).get(obj);
break;
}
case INT: {
LongColumnVector vector = (LongColumnVector) column;
vector.vector[rowId] = ((IntObjectInspector) inspector).get(obj);
break;
}
case LONG: {
LongColumnVector vector = (LongColumnVector) column;
vector.vector[rowId] = ((LongObjectInspector) inspector).get(obj);
break;
}
case FLOAT: {
DoubleColumnVector vector = (DoubleColumnVector) column;
vector.vector[rowId] =
((FloatObjectInspector) inspector).get(obj);
break;
}
case DOUBLE: {
DoubleColumnVector vector = (DoubleColumnVector) column;
vector.vector[rowId] =
((DoubleObjectInspector) inspector).get(obj);
break;
}
case BINARY: {
BytesColumnVector vector = (BytesColumnVector) column;
BytesWritable blob = ((BinaryObjectInspector) inspector)
.getPrimitiveWritableObject(obj);
vector.setVal(rowId, blob.getBytes(), 0, blob.getLength());
break;
}
case STRING: {
BytesColumnVector vector = (BytesColumnVector) column;
Text blob = ((StringObjectInspector) inspector)
.getPrimitiveWritableObject(obj);
vector.setVal(rowId, blob.getBytes(), 0, blob.getLength());
break;
}
case VARCHAR: {
BytesColumnVector vector = (BytesColumnVector) column;
Text blob = ((HiveVarcharObjectInspector) inspector)
.getPrimitiveWritableObject(obj).getTextValue();
vector.setVal(rowId, blob.getBytes(), 0, blob.getLength());
break;
}
case CHAR: {
BytesColumnVector vector = (BytesColumnVector) column;
Text blob = ((HiveCharObjectInspector) inspector)
.getPrimitiveWritableObject(obj).getTextValue();
vector.setVal(rowId, blob.getBytes(), 0, blob.getLength());
break;
}
case TIMESTAMP: {
TimestampColumnVector vector = (TimestampColumnVector) column;
Timestamp ts = ((TimestampObjectInspector) inspector)
.getPrimitiveJavaObject(obj);
vector.set(rowId, ts);
break;
}
case DATE: {
LongColumnVector vector = (LongColumnVector) column;
vector.vector[rowId] = ((DateObjectInspector) inspector)
.getPrimitiveWritableObject(obj).getDays();
break;
}
case DECIMAL: {
DecimalColumnVector vector = (DecimalColumnVector) column;
vector.set(rowId, ((HiveDecimalObjectInspector) inspector)
.getPrimitiveWritableObject(obj));
break;
}
}
break;
case STRUCT: {
StructColumnVector vector = (StructColumnVector) column;
StructObjectInspector oi = (StructObjectInspector) inspector;
List<? extends StructField> fields = oi.getAllStructFieldRefs();
for (int c = 0; c < vector.fields.length; ++c) {
StructField field = fields.get(c);
setColumn(rowId, vector.fields[c], field.getFieldObjectInspector(),
oi.getStructFieldData(obj, field));
}
break;
}
case UNION: {
UnionColumnVector vector = (UnionColumnVector) column;
UnionObjectInspector oi = (UnionObjectInspector) inspector;
int tag = oi.getTag(obj);
vector.tags[rowId] = tag;
setColumn(rowId, vector.fields[tag],
oi.getObjectInspectors().get(tag), oi.getField(obj));
break;
}
case LIST: {
ListColumnVector vector = (ListColumnVector) column;
ListObjectInspector oi = (ListObjectInspector) inspector;
int offset = vector.childCount;
int length = oi.getListLength(obj);
vector.offsets[rowId] = offset;
vector.lengths[rowId] = length;
vector.child.ensureSize(offset + length, true);
vector.childCount += length;
for (int c = 0; c < length; ++c) {
setColumn(offset + c, vector.child,
oi.getListElementObjectInspector(),
oi.getListElement(obj, c));
}
break;
}
case MAP: {
MapColumnVector vector = (MapColumnVector) column;
MapObjectInspector oi = (MapObjectInspector) inspector;
int offset = vector.childCount;
Set map = oi.getMap(obj).entrySet();
int length = map.size();
vector.offsets[rowId] = offset;
vector.lengths[rowId] = length;
vector.keys.ensureSize(offset + length, true);
vector.values.ensureSize(offset + length, true);
vector.childCount += length;
for (Object item: map) {
Map.Entry pair = (Map.Entry) item;
setColumn(offset, vector.keys, oi.getMapKeyObjectInspector(),
pair.getKey());
setColumn(offset, vector.values, oi.getMapValueObjectInspector(),
pair.getValue());
offset += 1;
}
break;
}
default:
throw new IllegalArgumentException("Unknown ObjectInspector kind " +
inspector.getCategory());
}
}
}
void flushInternalBatch() throws IOException {
if (internalBatch.size != 0) {
super.addRowBatch(internalBatch);
internalBatch.reset();
}
}
@Override
public void addRow(Object row) throws IOException {
int rowId = internalBatch.size++;
if (fields != null) {
StructObjectInspector soi = (StructObjectInspector) inspector;
for(int i=0; i < fields.length; ++i) {
setColumn(rowId, internalBatch.cols[i],
fields[i].getFieldObjectInspector(),
soi.getStructFieldData(row, fields[i]));
}
} else {
setColumn(rowId, internalBatch.cols[0], inspector, row);
}
if (internalBatch.size == internalBatch.getMaxSize()) {
flushInternalBatch();
}
}
@Override
public long writeIntermediateFooter() throws IOException {
flushInternalBatch();
return super.writeIntermediateFooter();
}
@Override
public void addRowBatch(VectorizedRowBatch batch) throws IOException {
flushInternalBatch();
super.addRowBatch(batch);
}
@Override
public void close() throws IOException {
flushInternalBatch();
super.close();
}
}