content: /*
   * 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.asterix.external.operators;
  
  import static org.msgpack.core.MessagePack.Code.ARRAY16;
  import static org.msgpack.core.MessagePack.Code.ARRAY32;
  import static org.msgpack.core.MessagePack.Code.FIXARRAY_PREFIX;
  import static org.msgpack.core.MessagePack.Code.isFixedArray;
  
  import java.io.IOException;
  import java.nio.ByteBuffer;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.List;
  
  import org.apache.asterix.common.exceptions.AsterixException;
  import org.apache.asterix.common.exceptions.ErrorCode;
  import org.apache.asterix.common.exceptions.RuntimeDataException;
  import org.apache.asterix.external.api.IExternalLangIPCProto;
  import org.apache.asterix.external.api.ILibraryEvaluator;
  import org.apache.asterix.external.library.PythonLibraryEvaluatorFactory;
  import org.apache.asterix.external.library.msgpack.MessageUnpackerToADM;
  import org.apache.asterix.external.library.msgpack.MsgPackPointableVisitor;
  import org.apache.asterix.external.util.ExternalDataConstants;
  import org.apache.asterix.external.util.ExternalDataUtils;
  import org.apache.asterix.om.functions.IExternalFunctionDescriptor;
  import org.apache.asterix.om.pointables.PointableAllocator;
  import org.apache.asterix.om.types.ATypeTag;
  import org.apache.hyracks.algebricks.common.utils.Pair;
  import org.apache.hyracks.algebricks.core.algebra.base.Counter;
  import org.apache.hyracks.algebricks.runtime.operators.base.AbstractOneInputOneOutputOneFramePushRuntime;
  import org.apache.hyracks.algebricks.runtime.operators.base.AbstractOneInputOneOutputPushRuntime;
  import org.apache.hyracks.algebricks.runtime.operators.base.AbstractOneInputOneOutputRuntimeFactory;
  import org.apache.hyracks.api.context.IHyracksTaskContext;
  import org.apache.hyracks.api.exceptions.HyracksDataException;
  import org.apache.hyracks.api.exceptions.Warning;
  import org.apache.hyracks.data.std.primitive.TaggedValuePointable;
  import org.apache.hyracks.data.std.primitive.VoidPointable;
  import org.apache.hyracks.data.std.util.ArrayBackedValueStorage;
  import org.apache.hyracks.dataflow.common.comm.io.ArrayTupleBuilder;
  import org.apache.hyracks.dataflow.common.data.accessors.FrameTupleReference;
  import org.msgpack.core.MessagePack;
  import org.msgpack.core.MessagePackException;
  import org.msgpack.core.MessageUnpacker;
  import org.msgpack.core.buffer.ArrayBufferInput;
  
  public final class ExternalAssignBatchRuntimeFactory extends AbstractOneInputOneOutputRuntimeFactory {
  
      private static final long serialVersionUID = 1L;
      private final int[] outColumns;
      private final IExternalFunctionDescriptor[] fnDescs;
      private final int[][] fnArgColumns;
  
      public ExternalAssignBatchRuntimeFactory(int[] outColumns, IExternalFunctionDescriptor[] fnDescs,
              int[][] fnArgColumns, int[] projectionList) {
          super(projectionList);
          this.outColumns = outColumns;
          this.fnDescs = fnDescs;
          this.fnArgColumns = fnArgColumns;
      }
  
      @Override
      public AbstractOneInputOneOutputPushRuntime createOneOutputPushRuntime(IHyracksTaskContext ctx) {
  
          final int[] projectionToOutColumns = new int[projectionList.length];
          for (int j = 0; j < projectionList.length; j++) {
              projectionToOutColumns[j] = Arrays.binarySearch(outColumns, projectionList[j]);
          }
  
          return new AbstractOneInputOneOutputOneFramePushRuntime() {
  
              private ArrayBackedValueStorage outputWrapper;
              private List<ArrayBackedValueStorage> argHolders;
              ArrayTupleBuilder tupleBuilder;
              private List<Pair<Long, ILibraryEvaluator>> libraryEvaluators;
              private ATypeTag[][] nullCalls;
              private int[] numCalls;
              private VoidPointable ref;
              private MessageUnpacker unpacker;
              private ArrayBufferInput unpackerInput;
              private List<Pair<ByteBuffer, Counter>> batchResults;
              private MessageUnpackerToADM unpackerToADM;
              private PointableAllocator pointableAllocator;
              private MsgPackPointableVisitor pointableVisitor;
              private TaggedValuePointable anyPointer;
  
              @Override
              public void open() throws HyracksDataException {
                  super.open();
                  initAccessAppend(ctx);
                  tupleBuilder = new ArrayTupleBuilder(projectionList.length);
                  tRef = new FrameTupleReference();
                  ref = VoidPointable.FACTORY.createPointable();
                  libraryEvaluators = new ArrayList<>();
                  try {
                      PythonLibraryEvaluatorFactory evalFactory = new PythonLibraryEvaluatorFactory(ctx);
                      for (IExternalFunctionDescriptor fnDesc : fnDescs) {
                          ILibraryEvaluator eval = evalFactory.getEvaluator(fnDesc.getFunctionInfo(), sourceLoc);
                          long id = eval.initialize(fnDesc.getFunctionInfo());
                          libraryEvaluators.add(new Pair<>(id, eval));
                      }
                  } catch (IOException | AsterixException e) {
                      throw RuntimeDataException.create(ErrorCode.EXTERNAL_UDF_EXCEPTION, e, sourceLoc, e.getMessage());
                  }
                  argHolders = new ArrayList<>(fnArgColumns.length);
                  for (int i = 0; i < fnArgColumns.length; i++) {
                      argHolders.add(new ArrayBackedValueStorage());
                  }
                  outputWrapper = new ArrayBackedValueStorage();
                  nullCalls = new ATypeTag[argHolders.size()][0];
                  numCalls = new int[fnArgColumns.length];
                  batchResults = new ArrayList<>(argHolders.size());
                  for (int i = 0; i < argHolders.size(); i++) {
                      batchResults.add(new Pair<>(ByteBuffer.allocate(ExternalDataConstants.DEFAULT_BUFFER_SIZE),
                              new Counter(-1)));
                  }
                  unpackerInput = new ArrayBufferInput(new byte[0]);
                  unpacker = MessagePack.newDefaultUnpacker(unpackerInput);
                  unpackerToADM = new MessageUnpackerToADM();
                  pointableAllocator = new PointableAllocator();
                  pointableVisitor = new MsgPackPointableVisitor();
                  anyPointer = TaggedValuePointable.FACTORY.createPointable();
              }
  
              private void resetBuffers(int numTuples, int[] numCalls) {
                  for (int func = 0; func < fnArgColumns.length; func++) {
                      argHolders.get(func).reset();
                      if (nullCalls[func].length < numTuples) {
                          nullCalls[func] = new ATypeTag[numTuples];
                      }
                      numCalls[func] = numTuples;
                      Arrays.fill(nullCalls[func], ATypeTag.TYPE);
                      for (Pair<ByteBuffer, Counter> batch : batchResults) {
                          batch.getFirst().clear();
                          batch.getFirst().position(0);
                          batch.getSecond().set(-1);
                      }
                  }
              }
  
              private ATypeTag handleNullMatrix(int func, int t, ATypeTag argumentPresence, ATypeTag argumentStatus) {
                  //If any argument is unknown, skip call. If any argument is null, return null, first.
                  //However, if any argument is missing, return missing instead.
                  if (nullCalls[func][t] == ATypeTag.TYPE && argumentPresence != ATypeTag.TYPE) {
                      if (argumentPresence == ATypeTag.NULL && argumentStatus != ATypeTag.MISSING) {
                          nullCalls[func][t] = argumentPresence;
                          return ATypeTag.NULL;
                      } else {
                          nullCalls[func][t] = argumentPresence;
                          return ATypeTag.MISSING;
                      }
                  }
                  return argumentPresence;
              }
  
              private void collectFunctionWarnings(List<Pair<ByteBuffer, Counter>> batchResults) throws IOException {
                  for (Pair<ByteBuffer, Counter> result : batchResults) {
                      if (result.getSecond().get() > -1) {
                          ByteBuffer resBuf = result.getFirst();
                          unpackerInput.reset(resBuf.array(), resBuf.position() + resBuf.arrayOffset(),
                                  resBuf.remaining());
                          unpacker.reset(unpackerInput);
                          try {
                              int numEntries = unpacker.unpackArrayHeader();
                              for (int j = 0; j < numEntries; j++) {
                                  if (ctx.getWarningCollector().shouldWarn()) {
                                      //TODO: in domain socket mode, a NUL can appear at the end of the stacktrace strings.
                                      //      this should probably not happen but warnings with control characters should
                                      //      also be properly escaped
                                      ctx.getWarningCollector()
                                              .warn(Warning.of(sourceLoc, ErrorCode.EXTERNAL_UDF_EXCEPTION,
                                                      unpacker.unpackString().replace('\0', ' ')));
                                  }
                              }
                          } catch (MessagePackException e) {
                              if (ctx.getWarningCollector().shouldWarn()) {
                                  ctx.getWarningCollector().warn(Warning.of(sourceLoc, ErrorCode.EXTERNAL_UDF_EXCEPTION,
                                          "Error retrieving returned warnings from Python UDF"));
                              }
                          }
                      }
                  }
              }
  
              @Override
              public void nextFrame(ByteBuffer buffer) throws HyracksDataException {
                  /*TODO: this whole transposition stuff is not necessary
                          the evaulator should accept a format that is a collection of rows, logically*/
                  tAccess.reset(buffer);
                  tupleBuilder.reset();
                  try {
                      int numTuples = tAccess.getTupleCount();
                      resetBuffers(numTuples, numCalls);
                      //build columns of arguments for each function
                      for (int t = 0; t < numTuples; t++) {
                          for (int func = 0; func < fnArgColumns.length; func++) {
                              tRef.reset(tAccess, t);
                              int[] cols = fnArgColumns[func];
                              //TODO: switch between fixarray/array16/array32 where appropriate
                              ATypeTag argumentStatus = ATypeTag.TYPE;
                              if (!fnDescs[func].getFunctionInfo().getNullCall()) {
                                  for (int colIdx = 0; colIdx < cols.length; colIdx++) {
                                      ref.set(buffer.array(), tRef.getFieldStart(cols[colIdx]),
                                              tRef.getFieldLength(cols[colIdx]));
                                      ATypeTag argumentPresence = ExternalDataUtils
                                              .peekArgument(fnDescs[func].getArgumentTypes()[colIdx], ref, anyPointer);
                                      argumentStatus = handleNullMatrix(func, t, argumentPresence, argumentStatus);
                                  }
                              }
                              if (argumentStatus == ATypeTag.TYPE) {
                                  if (cols.length > 0) {
                                      argHolders.get(func).getDataOutput().writeByte(ARRAY16);
                                      argHolders.get(func).getDataOutput().writeShort((short) cols.length);
                                  }
                                  for (int colIdx = 0; colIdx < cols.length; colIdx++) {
                                      ref.set(buffer.array(), tRef.getFieldStart(cols[colIdx]),
                                              tRef.getFieldLength(cols[colIdx]));
                                      IExternalLangIPCProto.visitValueRef(fnDescs[func].getArgumentTypes()[colIdx],
                                              argHolders.get(func).getDataOutput(), ref, pointableAllocator,
                                              pointableVisitor, fnDescs[func].getFunctionInfo().getNullCall());
                                  }
                              } else {
                                  numCalls[func]--;
                              }
                              if (cols.length == 0) {
                                  ExternalDataUtils.setVoidArgument(argHolders.get(func));
                              }
                          }
                      }
  
                      //TODO: maybe this could be done in parallel for each unique library evaluator?
                      for (int argHolderIdx = 0; argHolderIdx < argHolders.size(); argHolderIdx++) {
                          Pair<Long, ILibraryEvaluator> fnEval = libraryEvaluators.get(argHolderIdx);
                          ByteBuffer columnResult = fnEval.getSecond().callMulti(fnEval.getFirst(),
                                  argHolders.get(argHolderIdx), numCalls[argHolderIdx]);
                          if (columnResult != null) {
                              Pair<ByteBuffer, Counter> resultholder = batchResults.get(argHolderIdx);
                              if (resultholder.getFirst().capacity() < columnResult.remaining()) {
                                  ByteBuffer realloc =
                                          ctx.reallocateFrame(resultholder.getFirst(),
                                                  ctx.getInitialFrameSize()
                                                          * ((columnResult.remaining() / ctx.getInitialFrameSize()) + 1),
                                                  false);
                                  realloc.limit(columnResult.limit());
                                  resultholder.setFirst(realloc);
                              }
                              ByteBuffer resultBuf = resultholder.getFirst();
                              //offset 1 to skip message type
                              System.arraycopy(columnResult.array(), 1, resultBuf.array(), 0,
                                      columnResult.remaining() - 1);
                              //wrapper for results and warnings arrays. always length 2
                              consumeAndGetBatchLength(resultBuf);
                              int numResults = (int) consumeAndGetBatchLength(resultBuf);
                              resultholder.getSecond().set(numResults);
                          } else {
                              if (ctx.getWarningCollector().shouldWarn()) {
                                  ctx.getWarningCollector()
                                          .warn(Warning.of(sourceLoc, ErrorCode.EXTERNAL_UDF_EXCEPTION,
                                                  "Function "
                                                          + fnDescs[argHolderIdx].getFunctionInfo()
                                                                  .getFunctionIdentifier().toString()
                                                          + " failed to execute"));
                              }
                          }
                      }
  
                      //decompose returned function columns into frame tuple format
                      for (int i = 0; i < numTuples; i++) {
                          tupleBuilder.reset();
                          for (int f = 0; f < projectionList.length; f++) {
                              int k = projectionToOutColumns[f];
                              if (k >= 0) {
                                  outputWrapper.reset();
                                  Pair<ByteBuffer, Counter> result = batchResults.get(k);
                                  ATypeTag functionCalled = nullCalls[k][i];
                                  if (functionCalled == ATypeTag.TYPE) {
                                      if (result.getSecond().get() > 0) {
                                          unpackerToADM.unpack(result.getFirst(), outputWrapper.getDataOutput(), true);
                                          result.getSecond().set(result.getSecond().get() - 1);
                                      } else {
                                          //emit NULL for functions which failed with a warning
                                          outputWrapper.getDataOutput().writeByte(ATypeTag.SERIALIZED_NULL_TYPE_TAG);
                                      }
                                  } else if (functionCalled == ATypeTag.NULL) {
                                      outputWrapper.getDataOutput().writeByte(ATypeTag.SERIALIZED_NULL_TYPE_TAG);
                                  } else {
                                      outputWrapper.getDataOutput().writeByte(ATypeTag.SERIALIZED_MISSING_TYPE_TAG);
                                  }
                                  tupleBuilder.addField(outputWrapper.getByteArray(), 0, outputWrapper.getLength());
                              } else {
                                  tupleBuilder.addField(tAccess, i, projectionList[f]);
                              }
                          }
                          appendToFrameFromTupleBuilder(tupleBuilder);
                      }
                      collectFunctionWarnings(batchResults);
                  } catch (IOException e) {
                      throw HyracksDataException.create(e);
                  }
              }
  
              private long consumeAndGetBatchLength(ByteBuffer buf) {
                  byte tag = buf.get();
                  if (isFixedArray(tag)) {
                      return tag ^ FIXARRAY_PREFIX;
                  } else if (tag == ARRAY16) {
                      return Short.toUnsignedInt(buf.getShort());
                  } else if (tag == ARRAY32) {
                      return Integer.toUnsignedLong(buf.getInt());
                  }
                  return -1L;
              }
  
              @Override
              public void flush() throws HyracksDataException {
                  appender.flush(writer);
              }
          };
      }
  }
  