com.clickzetta.platform.bulkload.PartitionSpecUtils 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.bulkload;

import com.clickzetta.platform.client.api.Row;
import cz.proto.ingestion.v2.IngestionV2;

import java.math.BigDecimal;
import java.math.RoundingMode;
import java.nio.ByteBuffer;
import java.sql.Date;
import java.sql.Timestamp;
import java.text.ParseException;
import java.text.SimpleDateFormat;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

class PartitionSpecUtils {

  private static final String PARTITION_SPEC_DELIMITER = ",";

  public static Map parsePartitionSpec(IngestionV2.StreamSchema schema, Optional partitionSpec) {
    if (partitionSpec.isPresent()) {
      return parsePartitionSpec(schema, partitionSpec.get());
    } else {
      return Collections.EMPTY_MAP;
    }
  }

  public static Map parsePartitionSpec(IngestionV2.StreamSchema schema, String partitionSpec) {
    Map ret = new HashMap<>();
    if (partitionSpec != null && !partitionSpec.isEmpty()) {
      String[] partKeyValues = partitionSpec
        .trim()
        .replace("\'", "")
        .replace("\"", "")
        .split(PARTITION_SPEC_DELIMITER);
      for (String part : partKeyValues) {
        String[] kv = part.split("=");
        if (kv.length != 2) {
          throw new IllegalArgumentException("Invalid partition key-value: " + part);
        }
        String partKey = kv[0].toLowerCase();
        String partValue = kv[1];
        boolean found = false;
        for (int i = 0; i < schema.getDataFieldsCount(); ++i) {
          // TODO(gang.wu): validate partition key via schema.getPartitionSpec()
          IngestionV2.DataField field = schema.getDataFields(i);
          if (field.getName().toLowerCase().equals(partKey)) {
            ret.put(field.getName(), parsePartitionValue(partValue, field));
            found = true;
            break;
          }
        }
        if (!found) {
          throw new IllegalArgumentException("Partition key " + partKey + " not found in schema");
        }
      }
    }
    return ret;
  }

  public static void fillPartitionValues(Map partitionValues, Row... rows) {
    if (partitionValues != null) {
      for (Row row : rows) {
        row.setValues(partitionValues);
      }
    }
  }

  private static Object parsePartitionValue(String value, IngestionV2.DataField field) {
    switch (field.getType().getCategory()) {
      case BOOLEAN:
        return Boolean.valueOf(value);
      case INT8:
        return Byte.valueOf(value);
      case INT16:
        return Short.valueOf(value);
      case INT32:
        return Integer.valueOf(value);
      case INT64:
        return Long.valueOf(value);
      case FLOAT32:
        return Float.valueOf(value);
      case FLOAT64:
        return Double.valueOf(value);
      case STRING:
        return value;
      case VARCHAR:
        long length = field.getType().getVarCharTypeInfo().getLength();
        return value.substring(0, length > value.length() ? value.length() : (int) length);
      case CHAR:
        length = field.getType().getCharTypeInfo().getLength();
        return value.substring(0, length > value.length() ? value.length() : (int) length);
      case BINARY:
        return ByteBuffer.wrap(value.getBytes());
      case DATE: {
        // TODO: should we return LocalDate?
        try {
          return new Date(new SimpleDateFormat("yyyy-MM-dd").parse(value).getTime());
        } catch (ParseException e) {
          throw new IllegalArgumentException("Invalid date value: " + e);
        }
      }
      case TIMESTAMP_LTZ:
        // TODO: should we return OffsetDateTime?
        return Timestamp.valueOf(value);
      case DECIMAL:
        // TODO: should we throw if scale is shorter?
        long scale = field.getType().getDecimalTypeInfo().getScale();
        return new BigDecimal(value).setScale((int) scale, RoundingMode.HALF_UP);
      default:
        throw new IllegalArgumentException("Unsupported partition column type: " + field.getType());
    }
  }

}