com.clickzetta.platform.arrow.writer.DecimalWriter Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of clickzetta-java Show documentation

The java SDK for clickzetta's Lakehouse

There is a newer version: 2.0.0

package com.clickzetta.platform.arrow.writer;

import com.clickzetta.platform.catalyst.data.DataGetters;
import org.apache.arrow.vector.DecimalVector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.math.BigDecimal;
import java.math.RoundingMode;

public class DecimalWriter extends ArrowFieldWriter {

  private static final Logger LOG = LoggerFactory.getLogger(DecimalWriter.class);

  private static final BigDecimal EMPTY_DECIMAL = new BigDecimal(0);
  private DecimalVector decimalVector;
  private int precision;
  private int scale;

  public DecimalWriter(DecimalVector decimalVector) {
    super(decimalVector);
    this.decimalVector = decimalVector;
    this.precision = decimalVector.getPrecision();
    this.scale = decimalVector.getScale();
  }

  @Override
  public void setNull() {
    decimalVector.setNull(count);
  }

  @Override
  void setDefaultValue() {
    decimalVector.setSafe(count, EMPTY_DECIMAL.setScale(scale, RoundingMode.HALF_UP));
  }

  @Override
  void setValue(DataGetters row, int ordinal) {
    BigDecimal decimal = row.getDecimal(ordinal, precision, scale);
    Tuple2 tuple2 = changePrecision(decimal, precision, scale, RoundingMode.HALF_UP);
    if (tuple2._1) {
      decimalVector.setSafe(count, tuple2._2);
    } else {
      setDefaultValue();
    }
  }

  private Tuple2 changePrecision(BigDecimal decimal, int precision, int scale, RoundingMode roundMode) {
    if (decimal.precision() != precision || decimal.scale() != scale) {
      decimal = decimal.setScale(scale, roundMode);
      if (decimal.precision() > precision) {
        LOG.warn("data precision {} > DecimalType precision {}. Cast Overflow.", decimal.precision(), precision);
        return new Tuple2<>(false, null);
      }
    }
    return new Tuple2<>(true, decimal);
  }
}