com.clickzetta.platform.arrow.writer.StringWriter 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.arrow.ArrowSchemaConvert;
import com.clickzetta.platform.catalyst.data.DataGetters;
import org.apache.arrow.vector.VarCharVector;

import java.util.Map;

public class StringWriter extends ArrowFieldWriter {

  private static final byte[] EMPTY_BINARY = new byte[0];
  private VarCharVector varCharVector;

  // if varchar type.
  private boolean isVarcharType;
  private int varcharLength;

  public StringWriter(VarCharVector varCharVector) {
    super(varCharVector);
    this.varCharVector = varCharVector;

    // check if varchar length;
    Map metadata = varCharVector.getField().getMetadata();
    parserLengthFromMetadata(metadata);
  }

  private void parserLengthFromMetadata(Map metadata) {
    if (metadata == null || metadata.isEmpty()) {
      return;
    }
    if (metadata.containsKey(ArrowSchemaConvert.TYPE_KEY)) {
      String value = metadata.get(ArrowSchemaConvert.TYPE_KEY);
      if (value.startsWith("char(")) {
        isVarcharType = true;
        varcharLength = Integer.parseInt(value.substring("char(".length(), value.length() - 1));
      } else if (value.startsWith("varchar(")) {
        isVarcharType = true;
        varcharLength = Integer.parseInt(value.substring("varchar(".length(), value.length() - 1));
      }
    }
  }

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

  @Override
  void setDefaultValue() {
    varCharVector.setSafe(count, EMPTY_BINARY, 0, EMPTY_BINARY.length);
  }

  @Override
  void setValue(DataGetters row, int ordinal) {
    // varchar type.
    if (isVarcharType) {
      String binaryString = row.getString(ordinal);
      if (varcharLength < binaryString.length()) {
        binaryString = binaryString.substring(0, varcharLength);
      }
      byte[] bytes = binaryString.getBytes();
      varCharVector.setSafe(count, bytes, 0, bytes.length);
    } else {
      // original string type.
      byte[] bytes = row.getStringBytes(ordinal);
      varCharVector.setSafe(count, bytes, 0, bytes.length);
    }
  }
}