content: /*
   * Copyright (c) 2022. Alibaba Group Holding Limited
   */
  
  package com.alibaba.hologres.client.copy.in;
  
  import com.alibaba.hologres.client.model.Column;
  import com.alibaba.hologres.client.model.Record;
  import com.alibaba.hologres.client.model.TableSchema;
  import org.postgresql.core.BaseConnection;
  import org.postgresql.jdbc.ArrayUtil;
  import org.postgresql.jdbc.TimestampUtil;
  
  import java.io.IOException;
  import java.io.OutputStream;
  import java.math.BigDecimal;
  import java.math.RoundingMode;
  import java.sql.Array;
  import java.sql.Date;
  import java.sql.SQLException;
  import java.sql.Types;
  import java.time.LocalTime;
  import java.util.TimeZone;
  
  /** Record转pg binary流. */
  public class RecordBinaryOutputStream extends RecordOutputStream {
  
      protected boolean fillHeader = false;
  
      public RecordBinaryOutputStream(
              OutputStream os, TableSchema schema, BaseConnection conn, int maxCellBufferSize) {
          super(os, schema, conn, maxCellBufferSize);
      }
  
      protected void fillHeader() throws IOException {
          write("PGCOPY\n".getBytes(UTF8));
          write(0xFF);
          write("\r\n".getBytes(UTF8));
          write(0);
          writeInt(0);
          writeInt(0);
      }
  
      @Override
      public void close() throws IOException {
          if (!fillHeader) {
              fillHeader = true;
              fillHeader();
              writeCellBuffer();
          }
          super.close();
      }
  
      @Override
      protected void fillByteBuffer(Record record) throws IOException {
          if (!fillHeader) {
              fillHeader = true;
              fillHeader();
          }
  
          writeShort((short) record.getBitSet().cardinality());
          int index = 0;
          for (Column column : record.getSchema().getColumnSchema()) {
              try {
                  if (record.isSet(index)) {
                      fillByteBuffer(record.getObject(index), column);
                  }
              } catch (Exception e) {
                  throw new IOException(
                          "fail to convert column "
                                  + column.getName()
                                  + " type "
                                  + column.getTypeName()
                                  + " value "
                                  + record.getObject(index)
                                  + " to binary",
                          e);
              }
              ++index;
          }
      }
  
      private void fillByteBuffer(Object obj, Column column) throws IOException {
          if (obj == null) {
              writeInt(-1);
              return;
          }
          int type = column.getType();
          String typeName = column.getTypeName();
          String columnName = column.getName();
          switch (type) {
              case Types.SMALLINT:
                  if (obj instanceof Number) {
                      writeInt(2);
                      writeShort(((Number) obj).shortValue());
                      break;
                  } else {
                      throw new IOException(
                              "unsupported class for int2 : " + obj.getClass().getName());
                  }
              case Types.INTEGER:
                  if (obj instanceof Number) {
                      writeInt(4);
                      writeInt(((Number) obj).intValue());
                      break;
                  } else {
                      throw new IOException(
                              "unsupported class for int4 : " + obj.getClass().getName());
                  }
              case Types.BIGINT:
                  if (obj instanceof Number) {
                      writeInt(8);
                      writeLong(((Number) obj).longValue());
                      break;
                  } else {
                      throw new IOException(
                              "unsupported class for int8 : " + obj.getClass().getName());
                  }
              case Types.VARCHAR:
              case Types.CHAR:
                  byte[] bytes = obj.toString().getBytes(UTF8);
                  writeInt(bytes.length);
                  write(bytes);
                  break;
              case Types.BIT:
                  if ("bool".equals(typeName)) {
                      if (obj instanceof Boolean) {
                          writeInt(1);
                          write((Boolean) obj ? 1 : 0);
                      } else if (obj instanceof Number) {
                          writeInt(1);
                          write(((Number) obj).intValue() > 0 ? 1 : 0);
                      } else {
                          throw new IOException(
                                  "unsupported class for bool : " + obj.getClass().getName());
                      }
                  } else {
                      throw new IOException("unsupported type:" + typeName);
                  }
                  break;
  
              case Types.REAL:
                  if ("float4".equals(typeName)) {
                      if (obj instanceof Float) {
                          writeInt(4);
                          writeFloat((Float) obj);
                      } else if (obj instanceof Number) {
                          writeInt(4);
                          writeFloat(((Number) obj).floatValue());
                      } else {
                          throw new IOException(
                                  "unsupported class for bool : " + obj.getClass().getName());
                      }
                  } else {
                      throw new IOException("unsupported type:" + typeName);
                  }
                  break;
              case Types.DOUBLE:
                  if ("float8".equals(typeName)) {
                      if (obj instanceof Double) {
                          writeInt(8);
                          writeDouble((Double) obj);
                      } else if (obj instanceof Number) {
                          writeInt(8);
                          writeDouble(((Number) obj).doubleValue());
                      } else {
                          throw new IOException(
                                  "unsupported class for bool : " + obj.getClass().getName());
                      }
                  } else {
                      throw new IOException("unsupported type:" + typeName);
                  }
                  break;
              case Types.DATE:
                  byte[] val = new byte[4];
                  try {
                      if (obj instanceof java.sql.Date) {
                          timestampUtils.toBinDate(null, val, (Date) obj);
                      } else if (obj instanceof java.util.Date) {
                          Date tmpd = new java.sql.Date(((java.util.Date) obj).getTime());
                          timestampUtils.toBinDate(null, val, tmpd);
                      } else if (obj instanceof String) {
                          timestampUtils.toBinDate(
                                  null, val, timestampUtils.toDate(null, (String) obj));
                      } else {
                          throw new IOException(
                                  "unsupported class for date : " + obj.getClass().getName());
                      }
                      writeInt(4);
                      write(val);
                  } catch (SQLException e) {
                      throw new IOException(e);
                  }
                  break;
              case Types.TIMESTAMP:
                  writeInt(8);
                  writeLong(TimestampUtil.timestampToPgEpochMicroSecond(obj, typeName));
                  break;
              case Types.TIME:
              case Types.TIME_WITH_TIMEZONE:
                  {
                      long tVal;
                      int timezone;
                      if (obj instanceof java.sql.Time) {
                          java.sql.Time timeObj = (java.sql.Time) obj;
                          LocalTime localTime = timeObj.toLocalTime();
                          long nanos = (timeObj.getTime() % 1000) * 1000000L;
                          localTime = localTime.plusNanos(nanos);
                          tVal = localTime.toNanoOfDay() / 1000;
                          timezone = timeObj.getTimezoneOffset() * 60;
                      } else if (obj instanceof java.time.LocalTime) {
                          java.time.LocalTime localTime = (java.time.LocalTime) obj;
                          tVal = localTime.toNanoOfDay() / 1000;
                          timezone = TimeZone.getDefault().getRawOffset() / -1000;
                      } else {
                          throw new IOException(
                                  "unsupported class for time : " + obj.getClass().getName());
                      }
                      if (column.getTypeName().equals("time")) {
                          writeInt(8);
                          writeLong(tVal);
                      } else {
                          writeInt(12);
                          writeLong(tVal);
                          writeInt(timezone);
                      }
                  }
                  break;
              case Types.BINARY:
                  if (obj instanceof byte[]) {
                      byte[] binary = (byte[]) obj;
                      writeInt(binary.length);
                      write(binary);
                  }
                  break;
              case Types.NUMERIC:
                  BigDecimal decimal;
                  if (obj instanceof String) {
                      decimal = new BigDecimal((String) obj);
                  } else if (obj instanceof BigDecimal) {
                      decimal = (BigDecimal) obj;
                  } else if (obj instanceof Integer) {
                      decimal = new BigDecimal((Integer) obj);
                  } else if (obj instanceof Long) {
                      decimal = new BigDecimal((Long) obj);
                  } else {
                      throw new RuntimeException(
                              "unsupported type for numeric " + obj.getClass().getName());
                  }
                  decimal = decimal.setScale(column.getScale(), RoundingMode.HALF_UP);
                  String num = decimal.toPlainString();
  
                  short[] info = new short[3];
                  short[] digits = encodeFromString(num, info);
                  writeInt((4 + digits.length) * 2);
                  writeShort((short) digits.length);
  
                  writeShort(info[0]); // weight
                  writeShort(info[1]); // sign
                  writeShort(info[2]); // displayScale
  
                  for (short digit : digits) {
                      writeShort(digit);
                  }
                  break;
  
              case Types.OTHER:
                  if ("jsonb".equals(typeName)) {
                      byte[] jsonBytes = String.valueOf(obj).getBytes(UTF8);
                      writeInt(jsonBytes.length + 1);
                      write(1);
                      write(jsonBytes);
                      break;
                  } else if ("json".equals(typeName)) {
                      byte[] jsonBytes = String.valueOf(obj).getBytes(UTF8);
                      writeInt(jsonBytes.length);
                      write(jsonBytes);
                      break;
                  } else if ("roaringbitmap".equals(typeName)) {
                      if (obj instanceof byte[]) {
                          byte[] rbBytes = (byte[]) obj;
                          writeInt(rbBytes.length);
                          write(rbBytes);
                      } else {
                          throw new RuntimeException(
                                  "unsupported type for roaringbitmap " + obj.getClass().getName());
                      }
                      break;
                  } else {
                      throw new IOException("unsupported type:" + typeName + "(" + type + ")");
                  }
              case Types.ARRAY:
                  if (conn == null) {
                      throw new IOException(
                              "unsupported type:"
                                      + typeName
                                      + "("
                                      + type
                                      + "). Please call RecordBinaryOutputSteam constructor with BaseConnection Param");
                  }
                  try {
                      // obj如果是List<>或Object[]，都尝试转成Array
                      Array array = ArrayUtil.objectToArray(conn, obj, column.getTypeName());
                      byte[] arrayBytes =
                              ArrayUtil.arrayToBinary(
                                      conn,
                                      array != null
                                              ? ArrayUtil.objectToArray(
                                                      conn, obj, column.getTypeName())
                                              : obj,
                                      column.getTypeName());
                      writeInt(arrayBytes.length);
                      write(arrayBytes);
                  } catch (SQLException e) {
                      throw new IOException(e);
                  }
                  break;
              default:
                  throw new IOException("unsupported type:" + typeName + "(" + type + ")");
          }
      }
  
      private static final short NUMERIC_POS = (short) 0x0000;
      private static final short NUMERIC_NEG = (short) 0x4000;
      private static final short DEC_DIGITS = 4;
  
      private static short[] encodeFromString(String num, short[] info) {
  
          char[] numChars = num.toCharArray();
          byte[] numDigs = new byte[numChars.length - 1 + DEC_DIGITS * 2];
          int ch = 0;
          int digs = DEC_DIGITS;
          boolean haveDP = false;
  
          // Swallow leading zeros
          while (ch < numChars.length && numChars[ch] == '0') {
              ch++;
          }
  
          short sign = NUMERIC_POS;
          short displayWeight = -1;
          short displayScale = 0;
  
          if (ch < numChars.length && numChars[ch] == '-') {
              sign = NUMERIC_NEG;
              ++ch;
          }
  
          /*
           * Copy to array of single byte digits
           */
  
          while (ch < numChars.length) {
  
              if (numChars[ch] == '.') {
  
                  haveDP = true;
                  ch++;
              } else {
  
                  numDigs[digs++] = (byte) (numChars[ch++] - '0');
                  if (!haveDP) {
                      displayWeight++;
                  } else {
                      displayScale++;
                  }
              }
          }
  
          digs -= DEC_DIGITS;
  
          /*
           * Pack into NBASE format
           */
  
          short weight;
  
          if (displayWeight >= 0) {
              weight = (short) ((displayWeight + 1 + DEC_DIGITS - 1) / DEC_DIGITS - 1);
          } else {
              weight = (short) -((-displayWeight - 1) / DEC_DIGITS + 1);
          }
  
          int offset = (weight + 1) * DEC_DIGITS - (displayWeight + 1);
          int digitCount = (digs + offset + DEC_DIGITS - 1) / DEC_DIGITS;
  
          int i = DEC_DIGITS - offset;
          short[] digits = new short[digitCount];
          int d = 0;
  
          while (digitCount-- > 0) {
              digits[d++] =
                      (short)
                              (((numDigs[i] * 10 + numDigs[i + 1]) * 10 + numDigs[i + 2]) * 10
                                      + numDigs[i + 3]);
              i += DEC_DIGITS;
          }
  
          info[0] = weight;
          info[1] = sign;
          info[2] = displayScale;
          return digits;
      }
  }
  