com.clickzetta.platform.test.MockStructStreamMeta 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.test;

import cz.proto.*;
import cz.proto.ingestion.Ingestion;
import cz.proto.ingestion.v2.IngestionV2;
import scala.Tuple2;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class MockStructStreamMeta {

  private static final Map, Ingestion.IGSTableType> tableNameMaps = new HashMap<>();
  private static final Map structStreamSchemaMaps = new HashMap<>();

  static {
    // tableNameMaps
    tableNameMaps.put(new Tuple2<>("mock-db-struct", "mock-tbl-common-struct"), Ingestion.IGSTableType.NORMAL);
    tableNameMaps.put(new Tuple2<>("mock-db-struct", "mock-tbl-cluster-struct"), Ingestion.IGSTableType.CLUSTER);
    tableNameMaps.put(new Tuple2<>("mock-db-struct", "mock-tbl-pk-struct"), Ingestion.IGSTableType.ACID);

    // build stream schema base.
    IngestionV2.StreamSchema streamSchema = IngestionV2.StreamSchema.newBuilder()
        .addDataFields(IngestionV2.DataField.newBuilder()
            .setName("col1")
            .setType(DataType.newBuilder().setCategory(DataTypeCategory.INT32).setNullable(false).setFieldId(1).build())
            .build())
        .addDataFields(IngestionV2.DataField.newBuilder()
            .setName("col2")
            .setType(DataType.newBuilder().setCategory(DataTypeCategory.ARRAY).setNullable(false).setFieldId(2)
                .setArrayTypeInfo(ArrayTypeInfo.newBuilder().setElementType(
                    DataType.newBuilder().setCategory(DataTypeCategory.INT64).setNullable(false).build()).build()).build())
            .build())
        .addDataFields(IngestionV2.DataField.newBuilder()
            .setName("col3")
            .setType(DataType.newBuilder().setCategory(DataTypeCategory.MAP).setNullable(false).setFieldId(3)
                .setMapTypeInfo(MapTypeInfo.newBuilder()
                    .setKeyType(DataType.newBuilder().setCategory(DataTypeCategory.STRING).setNullable(false).build())
                    .setValueType(DataType.newBuilder().setCategory(DataTypeCategory.STRING).setNullable(true).build()).build()).build())
            .build())
        .addDataFields(IngestionV2.DataField.newBuilder()
            .setName("col4")
            .setType(DataType.newBuilder().setCategory(DataTypeCategory.ARRAY).setNullable(true).setFieldId(4)
                .setArrayTypeInfo(ArrayTypeInfo.newBuilder().setElementType(
                    DataType.newBuilder().setCategory(DataTypeCategory.ARRAY).setArrayTypeInfo(
                        ArrayTypeInfo.newBuilder().setElementType(
                            DataType.newBuilder().setCategory(DataTypeCategory.STRING).build()).build()).build()).build()).build())
            .build())
        .addDataFields(IngestionV2.DataField.newBuilder()
            .setName("col5")
            .setType(DataType.newBuilder().setCategory(DataTypeCategory.MAP).setNullable(true).setFieldId(5)
                .setMapTypeInfo(MapTypeInfo.newBuilder()
                    .setKeyType(DataType.newBuilder().setCategory(DataTypeCategory.ARRAY)
                        .setArrayTypeInfo(ArrayTypeInfo.newBuilder().setElementType(
                            DataType.newBuilder().setCategory(DataTypeCategory.STRING).build()).build()).setNullable(false).build())
                    .setValueType(DataType.newBuilder().setCategory(DataTypeCategory.STRING).setNullable(true).build()).build()).build()))
        .addDataFields(IngestionV2.DataField.newBuilder()
            .setName("col6")
            .setType(DataType.newBuilder().setCategory(DataTypeCategory.STRUCT).setNullable(true).setFieldId(6)
                .setStructTypeInfo(StructTypeInfo.newBuilder()
                    .addFields(StructTypeInfo.Field.newBuilder()
                        .setName("col6StructInternal1")
                        .setType(DataType.newBuilder().setCategory(DataTypeCategory.STRING).setNullable(true).build()).build())
                    .addFields(StructTypeInfo.Field.newBuilder()
                        .setName("col6StructInternal2")
                        .setType(DataType.newBuilder().setCategory(DataTypeCategory.ARRAY)
                            .setArrayTypeInfo(ArrayTypeInfo.newBuilder().setElementType(
                                DataType.newBuilder().setCategory(DataTypeCategory.INT64).build()).build())
                            .setNullable(true).build()).build())
                    .addFields(StructTypeInfo.Field.newBuilder()
                        .setName("col6StructInternal3")
                        .setType(DataType.newBuilder().setCategory(DataTypeCategory.STRUCT).setNullable(false)
                            .setStructTypeInfo(StructTypeInfo.newBuilder()
                                .addFields(StructTypeInfo.Field.newBuilder()
                                    .setName("col6StructInternal3inInternal1")
                                    .setType(DataType.newBuilder().setCategory(DataTypeCategory.STRING).setNullable(true).build()))
                                .build()).build()))
                    .build())))
        .addDataFields(IngestionV2.DataField.newBuilder()
            .setName("col7")
            .setType(DataType.newBuilder().setCategory(DataTypeCategory.STRUCT).setNullable(false).setFieldId(7)
                .setStructTypeInfo(StructTypeInfo.newBuilder()
                    .addFields(StructTypeInfo.Field.newBuilder()
                        .setName("col7StructInternal1")
                        .setType(DataType.newBuilder().setCategory(DataTypeCategory.STRING).setNullable(false).build()).build())
                    .build())))
        .addDataFields(IngestionV2.DataField.newBuilder()
            .setName("col8")
            .setType(DataType.newBuilder().setCategory(DataTypeCategory.ARRAY).setNullable(false).setFieldId(8)
                .setArrayTypeInfo(ArrayTypeInfo.newBuilder().setElementType(
                    DataType.newBuilder().setCategory(DataTypeCategory.DATE).setNullable(false).build()).build()).build())
            .build())
        .build();

    // build DistributionSpec.
    IngestionV2.DistributionSpec distributionSpec = IngestionV2.DistributionSpec.newBuilder()
        .addFieldIds(1)
        .setNumBuckets(1)
        .build();

    // build PrimaryKeySpec.
    IngestionV2.PrimaryKeySpec primaryKeySpec = IngestionV2.PrimaryKeySpec.newBuilder()
        .addFieldIds(1)
        .build();

    // build common table schema & info.
    structStreamSchemaMaps.put(Ingestion.IGSTableType.NORMAL, streamSchema);

    // build cluster table schema & info.
    IngestionV2.StreamSchema clusterStreamSchema = streamSchema.toBuilder()
        .setDistSpec(distributionSpec).build();
    structStreamSchemaMaps.put(Ingestion.IGSTableType.CLUSTER, clusterStreamSchema);

    // build pk table schema & info.
    IngestionV2.StreamSchema pkStreamSchema = streamSchema.toBuilder()
        .setDistSpec(distributionSpec)
        .setPrimaryKeySpec(primaryKeySpec).build();
    structStreamSchemaMaps.put(Ingestion.IGSTableType.ACID, pkStreamSchema);
  }

  public static IngestionV2.StreamSchema getTableMeta(Ingestion.IGSTableType tableType) {
    return structStreamSchemaMaps.get(tableType);
  }

  public static Ingestion.IGSTableType getTableType(String dbName, String tblName) {
    return tableNameMaps.get(new Tuple2<>(dbName, tblName));
  }

  public static Tuple2 getTableNameByType(Ingestion.IGSTableType tableType) {
    Tuple2 tuple2 = null;
    for (Map.Entry, Ingestion.IGSTableType> entry : tableNameMaps.entrySet()) {
      if (entry.getValue() == tableType) {
        tuple2 = entry.getKey();
        break;
      }
    }
    return tuple2;
  }

  public static List getAllIGSTableTypes() {
    return new ArrayList() {{
      add(Ingestion.IGSTableType.NORMAL);
      add(Ingestion.IGSTableType.CLUSTER);
      add(Ingestion.IGSTableType.ACID);
    }};
  }

  public static List> getAllTables() {
    return new ArrayList<>(tableNameMaps.keySet());
  }
}