com.clickzetta.platform.arrow.writer.MapWriter 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 com.clickzetta.platform.catalyst.data.InternalArray;
import com.clickzetta.platform.catalyst.data.InternalMap;
import org.apache.arrow.vector.complex.MapVector;
import org.apache.arrow.vector.complex.StructVector;

public class MapWriter extends ArrowFieldWriter {

  private MapVector mapVector;
  private StructVector structVector;
  private ArrowFieldWriter keyWriter;
  private ArrowFieldWriter valueWriter;

  public MapWriter(MapVector mapVector, StructVector structVector, ArrowFieldWriter keyWriter, ArrowFieldWriter valueWriter) {
    super(mapVector);
    this.mapVector = mapVector;
    this.structVector = structVector;
    this.keyWriter = keyWriter;
    this.valueWriter = valueWriter;
  }

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

  @Override
  void setDefaultValue() {
    mapVector.startNewValue(count);
    mapVector.endValue(count, 0);
  }

  @Override
  void setValue(DataGetters row, int ordinal) {
    InternalMap internalMap = row.getMap(ordinal);
    mapVector.startNewValue(count);
    InternalArray keyArray = internalMap.keyArray();
    InternalArray valueArray = internalMap.valueArray();
    int i = 0;
    while (i < internalMap.size()) {
      structVector.setIndexDefined(keyWriter.count);
      keyWriter.write(keyArray, i, true);
      valueWriter.write(valueArray, i, true);
      i += 1;
    }
    mapVector.endValue(count, internalMap.size());
  }

  @Override
  public void finish() {
    super.finish();
    keyWriter.finish();
    valueWriter.finish();
  }

  @Override
  public void reset() {
    super.reset();
    keyWriter.reset();
    valueWriter.reset();
  }
}