/*

* Licensed to the Apache Software Foundation (ASF) under one

* or more contributor license agreements. See the NOTICE file

* distributed with this work for additional information

* regarding copyright ownership. The ASF licenses this file

* to you 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.apache.iceberg;

import java.io.Serializable;

import java.net.URLEncoder;

import java.nio.charset.StandardCharsets;

import java.util.AbstractMap;

import java.util.Arrays;

import java.util.List;

import java.util.Map;

import java.util.Set;

import java.util.concurrent.atomic.AtomicInteger;

import java.util.stream.Collectors;

import org.apache.iceberg.exceptions.ValidationException;

import org.apache.iceberg.relocated.com.google.common.base.Preconditions;

import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;

import org.apache.iceberg.relocated.com.google.common.collect.ListMultimap;

import org.apache.iceberg.relocated.com.google.common.collect.Lists;

import org.apache.iceberg.relocated.com.google.common.collect.Maps;

import org.apache.iceberg.relocated.com.google.common.collect.Multimaps;

import org.apache.iceberg.relocated.com.google.common.collect.Sets;

import org.apache.iceberg.transforms.Transform;

import org.apache.iceberg.transforms.Transforms;

import org.apache.iceberg.transforms.UnknownTransform;

import org.apache.iceberg.types.Type;

import org.apache.iceberg.types.TypeUtil;

import org.apache.iceberg.types.Types;

import org.apache.iceberg.types.Types.StructType;

/**

* Represents how to produce partition data for a table.

*

* <p>Partition data is produced by transforming columns in a table. Each column transform is

* represented by a named {@link PartitionField}.

*

* <p>Partition specs are created using a builder obtained from {@link #builderFor(Schema)}:

*

* <pre>{@code

* PartitionSpec spec = PartitionSpec.builderFor(schema)

* .hour("ts")

* .bucket("id", 10)

* .build();

* }</pre>

*/

public class PartitionSpec implements Serializable {

// IDs for partition fields start at 1000

private static final int PARTITION_DATA_ID_START = 1000;

private final Schema schema;

// this is ordered so that DataFile has a consistent schema

private final int specId;

private final PartitionField[] fields;

private transient volatile ListMultimap<Integer, PartitionField> fieldsBySourceId = null;

private transient volatile Class<?>[] lazyJavaClasses = null;

private transient volatile StructType lazyPartitionType = null;

private transient volatile StructType lazyRawPartitionType = null;

private transient volatile List<PartitionField> fieldList = null;

private final int lastAssignedFieldId;

private PartitionSpec(

Schema schema, int specId, List<PartitionField> fields, int lastAssignedFieldId) {

this.schema = schema;

this.specId = specId;

this.fields = fields.toArray(new PartitionField[0]);

this.lastAssignedFieldId = lastAssignedFieldId;

}

/** Returns the {@link Schema} for this spec. */

public Schema schema() {

return schema;

}

/** Returns the ID of this spec. */

public int specId() {

return specId;

}

/** Returns the list of {@link PartitionField partition fields} for this spec. */

public List<PartitionField> fields() {

return lazyFieldList();

}

public boolean isPartitioned() {

return fields.length > 0 && fields().stream().anyMatch(f -> !f.transform().isVoid());

}

public boolean isUnpartitioned() {

return !isPartitioned();

}

int lastAssignedFieldId() {

return lastAssignedFieldId;

}

public UnboundPartitionSpec toUnbound() {

UnboundPartitionSpec.Builder builder = UnboundPartitionSpec.builder().withSpecId(specId);

for (PartitionField field : fields) {

builder.addField(

field.transform().toString(), field.sourceId(), field.fieldId(), field.name());

}

return builder.build();

}

/**

* Returns the {@link PartitionField field} that partitions the given source field

*

* @param fieldId a field id from the source schema

* @return the {@link PartitionField field} that partitions the given source field

*/

public List<PartitionField> getFieldsBySourceId(int fieldId) {

return lazyFieldsBySourceId().get(fieldId);

}

/** Returns a {@link StructType} for partition data defined by this spec. */

public StructType partitionType() {

if (lazyPartitionType == null) {

synchronized (this) {

if (lazyPartitionType == null) {

List<Types.NestedField> structFields = Lists.newArrayListWithExpectedSize(fields.length);

for (PartitionField field : fields) {

structFields.add(

Types.NestedField.optional(field.fieldId(), field.name(), resultType(field)));

}

this.lazyPartitionType = Types.StructType.of(structFields);

}

}

}

return lazyPartitionType;

}

/**

* Returns a struct matching partition information as written into manifest files. See {@link

* #partitionType()} for a struct with field ID's potentially re-assigned to avoid conflict.

*/

public StructType rawPartitionType() {

if (schema.idsToOriginal().isEmpty()) {

// not re-assigned.

return partitionType();

}

if (lazyRawPartitionType == null) {

synchronized (this) {

if (lazyRawPartitionType == null) {

this.lazyRawPartitionType =

StructType.of(

partitionType().fields().stream()

.map(f -> f.withFieldId(schema.idsToOriginal().get(f.fieldId())))

.collect(Collectors.toList()));

}

}

}

return lazyRawPartitionType;

}

public Class<?>[] javaClasses() {

if (lazyJavaClasses == null) {

synchronized (this) {

if (lazyJavaClasses == null) {

Class<?>[] classes = new Class<?>[fields.length];

for (int i = 0; i < fields.length; i += 1) {

PartitionField field = fields[i];

if (field.transform() instanceof UnknownTransform) {

classes[i] = Object.class;

} else {

classes[i] = resultType(field).typeId().javaClass();

}

}

this.lazyJavaClasses = classes;

}

}

}

return lazyJavaClasses;

}

private Type resultType(PartitionField field) {

Type sourceType = schema.findType(field.sourceId());

if (sourceType == null) {

// When the source field has been dropped, the source type has been lost

// Transforms with a fixed result type still work, so use unknown for source

sourceType = Types.UnknownType.get();

}

return field.transform().getResultType(sourceType);

}

@SuppressWarnings("unchecked")

private <T> T get(StructLike data, int pos, Class<?> javaClass) {

return data.get(pos, (Class<T>) javaClass);

}

private String escape(String string) {

return URLEncoder.encode(string, StandardCharsets.UTF_8);

}

public String partitionToPath(StructLike data) {

StringBuilder sb = new StringBuilder();

Class<?>[] javaClasses = javaClasses();

List<Types.NestedField> outputFields = partitionType().fields();

for (int i = 0; i < javaClasses.length; i += 1) {

PartitionField field = fields[i];

Type type = outputFields.get(i).type();

String valueString = field.transform().toHumanString(type, get(data, i, javaClasses[i]));

if (i > 0) {

sb.append("/");

}

sb.append(escape(field.name())).append("=").append(escape(valueString));

}

return sb.toString();

}

/**

* Returns true if this spec is equivalent to the other, with partition field ids ignored. That

* is, if both specs have the same number of fields, field order, field name, source columns, and

* transforms.

*

* @param other another PartitionSpec

* @return true if the specs have the same fields, source columns, and transforms.

*/

public boolean compatibleWith(PartitionSpec other) {

if (equals(other)) {

return true;

}

if (fields.length != other.fields.length) {

return false;

}

for (int i = 0; i < fields.length; i += 1) {

PartitionField thisField = fields[i];

PartitionField thatField = other.fields[i];

if (thisField.sourceId() != thatField.sourceId()

|| !thisField.transform().toString().equals(thatField.transform().toString())

|| !thisField.name().equals(thatField.name())) {

return false;

}

}

return true;

}

@Override

public boolean equals(Object other) {

if (this == other) {

return true;

} else if (!(other instanceof PartitionSpec)) {

return false;

}

PartitionSpec that = (PartitionSpec) other;

if (this.specId != that.specId) {

return false;

}

return Arrays.equals(fields, that.fields);

}

@Override

public int hashCode() {

return 31 * Integer.hashCode(specId) + Arrays.hashCode(fields);

}

private List<PartitionField> lazyFieldList() {

if (fieldList == null) {

synchronized (this) {

if (fieldList == null) {

this.fieldList = ImmutableList.copyOf(fields);

}

}

}

return fieldList;

}

private ListMultimap<Integer, PartitionField> lazyFieldsBySourceId() {

if (fieldsBySourceId == null) {

synchronized (this) {

if (fieldsBySourceId == null) {

ListMultimap<Integer, PartitionField> multiMap =

Multimaps.newListMultimap(

Maps.newHashMap(), () -> Lists.newArrayListWithCapacity(fields.length));

for (PartitionField field : fields) {

multiMap.put(field.sourceId(), field);

}

this.fieldsBySourceId = multiMap;

}

}

}

return fieldsBySourceId;

}

/**

* Returns the source field ids for identity partitions.

*

* @return a set of source ids for the identity partitions.

*/

public Set<Integer> identitySourceIds() {

Set<Integer> sourceIds = Sets.newHashSet();

for (PartitionField field : fields()) {

if (field.transform().isIdentity()) {

sourceIds.add(field.sourceId());

}

}

return sourceIds;

}

@Override

public String toString() {

StringBuilder sb = new StringBuilder();

sb.append("[");

for (PartitionField field : fields) {

sb.append("\n");

sb.append(" ").append(field);

}

if (fields.length > 0) {

sb.append("\n");

}

sb.append("]");

return sb.toString();

}

private static final PartitionSpec UNPARTITIONED_SPEC =

new PartitionSpec(new Schema(), 0, ImmutableList.of(), unpartitionedLastAssignedId());

/**

* Returns a spec for unpartitioned tables.

*

* @return a partition spec with no partitions

*/

public static PartitionSpec unpartitioned() {

return UNPARTITIONED_SPEC;

}

private static int unpartitionedLastAssignedId() {

return PARTITION_DATA_ID_START - 1;

}

/**

* Creates a new {@link Builder partition spec builder} for the given {@link Schema}.

*

* @param schema a schema

* @return a partition spec builder for the given schema

*/

public static Builder builderFor(Schema schema) {

return new Builder(schema);

}

/**

* Used to create valid {@link PartitionSpec partition specs}.

*

* <p>Call {@link #builderFor(Schema)} to create a new builder.

*/

public static class Builder {

private final Schema schema;

private final List<PartitionField> fields = Lists.newArrayList();

private final Set<String> partitionNames = Sets.newHashSet();

private final Map<Map.Entry<Integer, String>, PartitionField> dedupFields = Maps.newHashMap();

private int specId = 0;

private final AtomicInteger lastAssignedFieldId =

new AtomicInteger(unpartitionedLastAssignedId());

// check if there are conflicts between partition and schema field name

private boolean checkConflicts = true;

private boolean caseSensitive = true;

private Builder(Schema schema) {

this.schema = schema;

}

private int nextFieldId() {

return lastAssignedFieldId.incrementAndGet();

}

private void checkAndAddPartitionName(String name) {

checkAndAddPartitionName(name, null);

}

Builder checkConflicts(boolean check) {

checkConflicts = check;

return this;

}

private void checkAndAddPartitionName(String name, Integer sourceColumnId) {

Types.NestedField schemaField =

this.caseSensitive ? schema.findField(name) : schema.caseInsensitiveFindField(name);

if (checkConflicts) {

if (sourceColumnId == null) {

Preconditions.checkArgument(

schemaField == null,

"Cannot create partition from name that exists in schema: %s",

name);

} else {

boolean sourceFieldExists = schema.findField(sourceColumnId) != null;

// For identity transforms, require the partition name to match the source column when it

// still exists in the schema. When the source was dropped, the spec may be historical;

// skip the identity name check in that case.

if (sourceFieldExists) {

Preconditions.checkArgument(

schemaField == null || schemaField.fieldId() == sourceColumnId,

"Cannot create identity partition sourced from different field in schema: %s",

name);

}

}

}

Preconditions.checkArgument(!name.isEmpty(), "Cannot use empty partition name: %s", name);

Preconditions.checkArgument(

!partitionNames.contains(name), "Cannot use partition name more than once: %s", name);

partitionNames.add(name);

}

private void checkForRedundantPartitions(PartitionField field) {

Map.Entry<Integer, String> dedupKey =

new AbstractMap.SimpleEntry<>(field.sourceId(), field.transform().dedupName());

PartitionField partitionField = dedupFields.get(dedupKey);

Preconditions.checkArgument(

partitionField == null,

"Cannot add redundant partition: %s conflicts with %s",

partitionField,

field);

dedupFields.put(dedupKey, field);

}

public Builder caseSensitive(boolean sensitive) {

this.caseSensitive = sensitive;

return this;

}

public Builder withSpecId(int newSpecId) {

this.specId = newSpecId;

return this;

}

private Types.NestedField findSourceColumn(String sourceName) {

Types.NestedField sourceColumn =

this.caseSensitive

? schema.findField(sourceName)

: schema.caseInsensitiveFindField(sourceName);

Preconditions.checkArgument(

sourceColumn != null, "Cannot find source column: %s", sourceName);

return sourceColumn;

}

public Builder identity(String sourceName, String targetName) {

return identity(findSourceColumn(sourceName), targetName);

}

private Builder identity(Types.NestedField sourceColumn, String targetName) {

checkAndAddPartitionName(targetName, sourceColumn.fieldId());

PartitionField field =

new PartitionField(

sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.identity());

checkForRedundantPartitions(field);

fields.add(field);

return this;

}

public Builder identity(String sourceName) {

Types.NestedField sourceColumn = findSourceColumn(sourceName);

return identity(sourceColumn, schema.findColumnName(sourceColumn.fieldId()));

}

public Builder year(String sourceName, String targetName) {

return year(findSourceColumn(sourceName), targetName);

}

private Builder year(Types.NestedField sourceColumn, String targetName) {

checkAndAddPartitionName(targetName);

PartitionField field =

new PartitionField(sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.year());

checkForRedundantPartitions(field);

fields.add(field);

return this;

}

public Builder year(String sourceName) {

Types.NestedField sourceColumn = findSourceColumn(sourceName);

String columnName = schema.findColumnName(sourceColumn.fieldId());

return year(sourceColumn, columnName + "_year");

}

public Builder month(String sourceName, String targetName) {

return month(findSourceColumn(sourceName), targetName);

}

private Builder month(Types.NestedField sourceColumn, String targetName) {

checkAndAddPartitionName(targetName);

PartitionField field =

new PartitionField(sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.month());

checkForRedundantPartitions(field);

fields.add(field);

return this;

}

public Builder month(String sourceName) {

Types.NestedField sourceColumn = findSourceColumn(sourceName);

String columnName = schema.findColumnName(sourceColumn.fieldId());

return month(sourceColumn, columnName + "_month");

}

public Builder day(String sourceName, String targetName) {

return day(findSourceColumn(sourceName), targetName);

}

private Builder day(Types.NestedField sourceColumn, String targetName) {

checkAndAddPartitionName(targetName);

PartitionField field =

new PartitionField(sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.day());

checkForRedundantPartitions(field);

fields.add(field);

return this;

}

public Builder day(String sourceName) {

Types.NestedField sourceColumn = findSourceColumn(sourceName);

String columnName = schema.findColumnName(sourceColumn.fieldId());

return day(sourceColumn, columnName + "_day");

}

public Builder hour(String sourceName, String targetName) {

return hour(findSourceColumn(sourceName), targetName);

}

private Builder hour(Types.NestedField sourceColumn, String targetName) {

checkAndAddPartitionName(targetName);

PartitionField field =

new PartitionField(sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.hour());

checkForRedundantPartitions(field);

fields.add(field);

return this;

}

public Builder hour(String sourceName) {

Types.NestedField sourceColumn = findSourceColumn(sourceName);

String columnName = schema.findColumnName(sourceColumn.fieldId());

return hour(sourceColumn, columnName + "_hour");

}

public Builder bucket(String sourceName, int numBuckets, String targetName) {

return bucket(findSourceColumn(sourceName), numBuckets, targetName);

}

private Builder bucket(Types.NestedField sourceColumn, int numBuckets, String targetName) {

checkAndAddPartitionName(targetName);

fields.add(

new PartitionField(

sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.bucket(numBuckets)));

return this;

}

public Builder bucket(String sourceName, int numBuckets) {

Types.NestedField sourceColumn = findSourceColumn(sourceName);

String columnName = schema.findColumnName(sourceColumn.fieldId());

return bucket(sourceColumn, numBuckets, columnName + "_bucket");

}

public Builder truncate(String sourceName, int width, String targetName) {

return truncate(findSourceColumn(sourceName), width, targetName);

}

private Builder truncate(Types.NestedField sourceColumn, int width, String targetName) {

checkAndAddPartitionName(targetName);

fields.add(

new PartitionField(

sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.truncate(width)));

return this;

}

public Builder truncate(String sourceName, int width) {

Types.NestedField sourceColumn = findSourceColumn(sourceName);

String columnName = schema.findColumnName(sourceColumn.fieldId());

return truncate(sourceColumn, width, columnName + "_trunc");

}

public Builder alwaysNull(String sourceName, String targetName) {

return alwaysNull(findSourceColumn(sourceName), targetName);

}

private Builder alwaysNull(Types.NestedField sourceColumn, String targetName) {

checkAndAddPartitionName(

targetName, sourceColumn.fieldId()); // can duplicate a source column name

fields.add(

new PartitionField(

sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.alwaysNull()));

return this;

}

public Builder alwaysNull(String sourceName) {

Types.NestedField sourceColumn = findSourceColumn(sourceName);

String columnName = schema.findColumnName(sourceColumn.fieldId());

return alwaysNull(sourceColumn, columnName + "_null");

}

// add a partition field with an auto-increment partition field id starting from

// PARTITION_DATA_ID_START

Builder add(int sourceId, String name, Transform<?, ?> transform) {

return add(sourceId, nextFieldId(), name, transform);

}

Builder add(int sourceId, int fieldId, String name, Transform<?, ?> transform) {

checkAndAddPartitionName(name, sourceId);

fields.add(new PartitionField(sourceId, fieldId, name, transform));

lastAssignedFieldId.getAndAccumulate(fieldId, Math::max);

return this;

}

public PartitionSpec build() {

return build(false);

}

public PartitionSpec build(boolean allowMissingFields) {

PartitionSpec spec = buildUnchecked();

checkCompatibility(spec, schema, allowMissingFields);

return spec;

}

PartitionSpec buildUnchecked() {

return new PartitionSpec(schema, specId, fields, lastAssignedFieldId.get());

}

}

static void checkCompatibility(PartitionSpec spec, Schema schema) {

checkCompatibility(spec, schema, false);

}

static void checkCompatibility(PartitionSpec spec, Schema schema, boolean allowMissingFields) {

final Map<Integer, Integer> parents = TypeUtil.indexParents(schema.asStruct());

for (PartitionField field : spec.fields) {

Type sourceType = schema.findType(field.sourceId());

Transform<?, ?> transform = field.transform();

// In the case the underlying field is dropped, we cannot check if they are compatible

if (allowMissingFields && sourceType == null) {

continue;

}

// In the case of a Version 1 partition-spec field gets deleted,

// it is replaced with a void transform, see:

// https://iceberg.apache.org/spec/#partition-transforms

// We don't care about the source type since a VoidTransform is always compatible and skip the

// checks

if (!transform.equals(Transforms.alwaysNull())) {

ValidationException.check(

sourceType != null, "Cannot find source column for partition field: %s", field);

ValidationException.check(

sourceType.isPrimitiveType(),

"Cannot partition by non-primitive source field: %s",

sourceType);

ValidationException.check(

transform.canTransform(sourceType),

"Invalid source type %s for transform: %s",

sourceType,

transform);

// The only valid parent types for a PartitionField are StructTypes. This must be checked

// recursively.

Integer parentId = parents.get(field.sourceId());

while (parentId != null) {

Type parentType = schema.findType(parentId);

ValidationException.check(

parentType.isStructType(), "Invalid partition field parent: %s", parentType);

parentId = parents.get(parentId);

}

}

}

}

static boolean hasSequentialIds(PartitionSpec spec) {

for (int i = 0; i < spec.fields.length; i += 1) {

if (spec.fields[i].fieldId() != PARTITION_DATA_ID_START + i) {

return false;

}

}

return true;

}

}