org.projectnessie.gc.iceberg.mocks.MockManifestFile Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of nessie-gc-iceberg-mock Show documentation
There is a newer version: 0.98.0
Show newest version
/*
 * Copyright (C) 2022 Dremio
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.projectnessie.gc.iceberg.mocks;

import static org.apache.iceberg.types.Types.NestedField.required;
import static org.projectnessie.gc.iceberg.mocks.IcebergFileIOMocking.dataFilePath;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.Nonnull;
import jakarta.annotation.Nullable;
import java.nio.ByteBuffer;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import org.apache.avro.Schema;
import org.apache.avro.generic.IndexedRecord;
import org.apache.iceberg.DataFile;
import org.apache.iceberg.ManifestContent;
import org.apache.iceberg.ManifestFile;
import org.apache.iceberg.PartitionSpec;
import org.apache.iceberg.PartitionSpecParser;
import org.apache.iceberg.SchemaParser;
import org.apache.iceberg.avro.Avro;
import org.apache.iceberg.avro.AvroSchemaUtil;
import org.apache.iceberg.io.FileAppender;
import org.apache.iceberg.io.OutputFile;
import org.apache.iceberg.types.Types.StructType;
import org.immutables.value.Value;

@Value.Immutable
public abstract class MockManifestFile implements IndexedRecord {

  @Nonnull
  public abstract String path();

  @Value.Default
  public int partitionSpecId() {
    return 0;
  }

  @Value.Default
  public ManifestContent content() {
    return ManifestContent.DATA;
  }

  @Value.Default
  @Nullable
  public ByteBuffer keyMetadata() {
    return null;
  }

  @Value.Default
  public long length() {
    return 42L;
  }

  @Value.Default
  public long sequenceNumber() {
    return 0;
  }

  @Value.Default
  public long minSequenceNumber() {
    return 0;
  }

  @Value.Default
  @Nullable
  public Long snapshotId() {
    return 0L;
  }

  @Value.Default
  @Nullable
  public Integer addedFilesCount() {
    return 0;
  }

  @Value.Default
  @Nullable
  public Long addedRowsCount() {
    return 0L;
  }

  @Value.Default
  @Nullable
  public Integer existingFilesCount() {
    return 0;
  }

  @Value.Default
  @Nullable
  public Long existingRowsCount() {
    return 0L;
  }

  @Value.Default
  @Nullable
  public Integer deletedFilesCount() {
    return 0;
  }

  @Value.Default
  @Nullable
  public Long deletedRowsCount() {
    return 0L;
  }

  @Value.Default
  public int numEntries() {
    return 1;
  }

  public abstract String baseDataFilePath();

  @Value.Auxiliary
  public List manifestEntries() {
    MockTableMetadata mockTableMeta = MockTableMetadata.empty();
    MockPartitionSpec mockPartitionSpec = mockTableMeta.partitionSpec(0);
    org.apache.iceberg.Schema icebergSchema = mockTableMeta.schema(0).toSchema();
    PartitionSpec partitionSpec = mockPartitionSpec.toPartitionSpec(icebergSchema);
    StructType partitionType = partitionSpec.partitionType();

    return IntStream.range(0, numEntries())
        .mapToObj(
            i ->
                ImmutableMockManifestEntry.builder()
                    .sequenceNumber(0L)
                    .snapshotId(0L)
                    .filePath(dataFilePath(baseDataFilePath(), i))
                    .partitionType(partitionType)
                    .build())
        .collect(Collectors.toList());
  }

  private static final ObjectMapper MAPPER = new ObjectMapper();

  public void write(OutputFile output) {
    MockTableMetadata mockTableMeta = MockTableMetadata.empty();
    MockPartitionSpec mockPartitionSpec = mockTableMeta.partitionSpec(0);
    org.apache.iceberg.Schema icebergSchema = mockTableMeta.schema(0).toSchema();
    PartitionSpec partitionSpec = mockPartitionSpec.toPartitionSpec(icebergSchema);
    StructType partitionType = partitionSpec.partitionType();

    try {
      String partitionSpecJson = PartitionSpecParser.toJson(partitionSpec);
      JsonNode partitionSpecFields =
          MAPPER.readValue(partitionSpecJson, JsonNode.class).get("fields");

      org.apache.iceberg.Schema manifestSchema = entrySchema(partitionType);
      try (FileAppender