org.apache.tajo.plan.function.stream.OutputHandler.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.tajo.plan.function.stream.OutputHandler.java

Source

/*
 * Lisensed 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.tajo.plan.function.stream;

import com.google.common.base.Charsets;
import io.netty.buffer.ByteBuf;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.tajo.plan.function.FunctionContext;
import org.apache.tajo.storage.Tuple;
import org.apache.tajo.storage.VTuple;
import org.apache.tajo.util.FileUtil;

import java.io.Closeable;
import java.io.IOException;
import java.io.InputStream;

/**
 * {@link OutputHandler} is responsible for handling the output of the
 * Tajo-Streaming external command.
 */
public class OutputHandler implements Closeable {

    private static final Log LOG = LogFactory.getLog(OutputHandler.class);
    private static int DEFAULT_BUFFER = 64 * 1024;
    private final static byte[] END_OF_RECORD_DELIM = "|_\n".getBytes();

    private final TextLineDeserializer deserializer;

    private ByteBufLineReader in = null;

    private String currValue = null;

    private InputStream istream;

    private ByteBuf buf;

    // Both of these ignore the trailing "\n".  So if the default delimiter is "\n", recordDelimStr is "".
    private String recordDelimStr = null;
    private int recordDelimLength = 0;
    private final Tuple tuple;

    // flag to mark if close() has already been called
    private boolean alreadyClosed = false;
    private final String FIELD_DELIM;

    public OutputHandler(TextLineDeserializer deserializer) {
        this.deserializer = deserializer;
        FIELD_DELIM = new String(CSVLineSerDe.getFieldDelimiter(deserializer.meta));
        tuple = new VTuple(deserializer.schema.size());
    }

    /**
     * Bind the <code>OutputHandler</code> to the <code>InputStream</code>
     * from which to read the output data of the managed process.
     *
     * @param is <code>InputStream</code> from which to read the output data
     *           of the managed process
     * @throws IOException
     */
    public void bindTo(InputStream is) throws IOException {
        this.istream = is;
        this.in = new ByteBufLineReader(new ByteBufInputChannel(istream));
        this.buf = BufferPool.directBuffer(DEFAULT_BUFFER);
    }

    /**
     * Get the next output <code>Tuple</code> of the managed process.
     *
     * @return the next output <code>Tuple</code> of the managed process
     * @throws IOException
     */
    public Tuple getNext() throws IOException {
        if (in == null) {
            return null;
        }

        currValue = null;
        if (!readValue()) {
            return null;
        }
        buf.clear();
        buf.writeBytes(currValue.getBytes());
        try {
            deserializer.deserialize(buf, tuple);
        } catch (TextLineParsingError textLineParsingError) {
            throw new IOException(textLineParsingError);
        }
        return tuple;
    }

    public String getPartialResultString() throws IOException {
        if (in == null) {
            return null;
        }

        currValue = null;
        if (!readValue()) {
            return null;
        }
        return currValue;
    }

    public FunctionContext getNext(FunctionContext context) throws IOException {
        if (in == null) {
            return null;
        }

        currValue = null;
        if (!readValue()) {
            return null;
        }
        buf.clear();
        buf.writeBytes(currValue.getBytes());
        try {
            deserializer.deserialize(buf, context);
        } catch (TextLineParsingError textLineParsingError) {
            throw new IOException(textLineParsingError);
        }

        return context;
    }

    private boolean readValue() throws IOException {
        currValue = in.readLine();
        if (currValue == null) {
            return false;
        }

        while (!isEndOfRow()) {
            // Need to add back the newline character we ate.
            currValue += '\n';

            byte[] lineBytes = readNextLine();
            if (lineBytes == null) {
                // We have no more input, so just break;
                break;
            }
            currValue += new String(lineBytes);
        }

        String line = currValue;

        int candidate = -1;
        StringBuffer sb = new StringBuffer();
        while ((candidate = line.indexOf("|")) >= 0) {
            if (line.substring(candidate, candidate + 2).equals("|_")) {
                // record end
                sb.append(line.substring(0, candidate));
                break;
            } else if (line.substring(candidate, candidate + 3).equals("|,_")) {
                sb.append(line.substring(0, candidate)).append(FIELD_DELIM);
                line = line.substring(candidate + 3, line.length());
            } else if (line.substring(candidate, candidate + 3).equals("|-_")) {
                // null value
                sb.append(FIELD_DELIM);
                line = line.substring(candidate + 3, line.length());
            }
        }
        currValue = sb.toString();

        return true;
    }

    private byte[] readNextLine() throws IOException {
        String line = in.readLine();
        if (line == null) {
            return null;
        }

        return line.getBytes();
    }

    private boolean isEndOfRow() {
        if (recordDelimStr == null) {
            byte[] recordDelimBa = END_OF_RECORD_DELIM;
            recordDelimLength = recordDelimBa.length - 1; //Ignore trailing \n
            recordDelimStr = new String(recordDelimBa, 0, recordDelimLength, Charsets.UTF_8);
        }
        if (recordDelimLength == 0 || currValue.length() < recordDelimLength) {
            return true;
        }
        return currValue.contains(recordDelimStr);
    }

    /**
     * Close the <code>OutputHandler</code>.
     * @throws IOException
     */
    public void close() throws IOException {
        if (!alreadyClosed) {
            FileUtil.cleanup(LOG, istream, in);
            in = null;
            istream = null;
            alreadyClosed = true;
            if (this.buf.refCnt() > 0) {
                this.buf.release();
            }
        }
    }
}