| /* |
| * 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.nio.ByteBuffer; |
| import java.util.Arrays; |
| import java.util.List; |
| import java.util.Objects; |
| import java.util.stream.Stream; |
| import org.apache.avro.Schema; |
| import org.apache.avro.generic.IndexedRecord; |
| import org.apache.avro.specific.SpecificData; |
| import org.apache.avro.util.Utf8; |
| import org.apache.iceberg.avro.AvroSchemaUtil; |
| import org.apache.iceberg.relocated.com.google.common.base.Preconditions; |
| import org.apache.iceberg.relocated.com.google.common.hash.Hasher; |
| import org.apache.iceberg.relocated.com.google.common.hash.Hashing; |
| import org.apache.iceberg.types.Type; |
| import org.apache.iceberg.types.Types; |
| |
| class PartitionData |
| implements IndexedRecord, StructLike, SpecificData.SchemaConstructable, Serializable { |
| |
| static Schema partitionDataSchema(Types.StructType partitionType) { |
| return AvroSchemaUtil.convert(partitionType, PartitionData.class.getName()); |
| } |
| |
| private final Types.StructType partitionType; |
| private final int size; |
| private final Object[] data; |
| private final String stringSchema; |
| private transient Schema schema = null; |
| |
| /** |
| * Used by Avro reflection to instantiate this class when reading manifest files. |
| */ |
| PartitionData(Schema schema) { |
| this.partitionType = AvroSchemaUtil.convert(schema).asNestedType().asStructType(); |
| this.size = partitionType.fields().size(); |
| this.data = new Object[size]; |
| this.stringSchema = schema.toString(); |
| this.schema = schema; |
| } |
| |
| PartitionData(Types.StructType partitionType) { |
| for (Types.NestedField field : partitionType.fields()) { |
| Preconditions.checkArgument(field.type().isPrimitiveType(), |
| "Partitions cannot contain nested types: %s", field.type()); |
| } |
| |
| this.partitionType = partitionType; |
| this.size = partitionType.fields().size(); |
| this.data = new Object[size]; |
| this.schema = partitionDataSchema(partitionType); |
| this.stringSchema = schema.toString(); |
| } |
| |
| /** |
| * Copy constructor |
| */ |
| private PartitionData(PartitionData toCopy) { |
| this.partitionType = toCopy.partitionType; |
| this.size = toCopy.size; |
| this.data = copyData(toCopy.partitionType, toCopy.data); |
| this.stringSchema = toCopy.stringSchema; |
| this.schema = toCopy.schema; |
| } |
| |
| public Types.StructType getPartitionType() { |
| return partitionType; |
| } |
| |
| @Override |
| public Schema getSchema() { |
| if (schema == null) { |
| this.schema = new Schema.Parser().parse(stringSchema); |
| } |
| return schema; |
| } |
| |
| public Type getType(int pos) { |
| return partitionType.fields().get(pos).type(); |
| } |
| |
| public void clear() { |
| Arrays.fill(data, null); |
| } |
| |
| @Override |
| public int size() { |
| return size; |
| } |
| |
| @Override |
| @SuppressWarnings("unchecked") |
| public <T> T get(int pos, Class<T> javaClass) { |
| Object value = get(pos); |
| if (value == null || javaClass.isInstance(value)) { |
| return javaClass.cast(value); |
| } |
| |
| throw new IllegalArgumentException(String.format( |
| "Wrong class, %s, for object: %s", |
| javaClass.getName(), String.valueOf(value))); |
| } |
| |
| @Override |
| public Object get(int pos) { |
| if (pos >= data.length) { |
| return null; |
| } |
| |
| if (data[pos] instanceof byte[]) { |
| return ByteBuffer.wrap((byte[]) data[pos]); |
| } |
| |
| return data[pos]; |
| } |
| |
| @Override |
| public <T> void set(int pos, T value) { |
| if (value instanceof Utf8) { |
| // Utf8 is not Serializable |
| data[pos] = value.toString(); |
| } else if (value instanceof ByteBuffer) { |
| // ByteBuffer is not Serializable |
| ByteBuffer buffer = (ByteBuffer) value; |
| byte[] bytes = new byte[buffer.remaining()]; |
| buffer.duplicate().get(bytes); |
| data[pos] = bytes; |
| } else { |
| data[pos] = value; |
| } |
| } |
| |
| @Override |
| public void put(int i, Object v) { |
| set(i, v); |
| } |
| |
| @Override |
| public String toString() { |
| StringBuilder sb = new StringBuilder(); |
| sb.append("PartitionData{"); |
| for (int i = 0; i < data.length; i += 1) { |
| if (i > 0) { |
| sb.append(", "); |
| } |
| sb.append(partitionType.fields().get(i).name()) |
| .append("=") |
| .append(data[i]); |
| } |
| sb.append("}"); |
| return sb.toString(); |
| } |
| |
| public PartitionData copy() { |
| return new PartitionData(this); |
| } |
| |
| @Override |
| public boolean equals(Object o) { |
| if (this == o) { |
| return true; |
| } else if (!(o instanceof PartitionData)) { |
| return false; |
| } |
| |
| PartitionData that = (PartitionData) o; |
| return partitionType.equals(that.partitionType) && Arrays.equals(data, that.data); |
| } |
| |
| @Override |
| public int hashCode() { |
| Hasher hasher = Hashing.goodFastHash(32).newHasher(); |
| Stream.of(data).map(Objects::hashCode).forEach(hasher::putInt); |
| partitionType.fields().stream().map(Objects::hashCode).forEach(hasher::putInt); |
| return hasher.hash().hashCode(); |
| } |
| |
| public static Object[] copyData(Types.StructType type, Object[] data) { |
| List<Types.NestedField> fields = type.fields(); |
| Object[] copy = new Object[data.length]; |
| for (int i = 0; i < data.length; i += 1) { |
| if (data[i] == null) { |
| copy[i] = null; |
| } else { |
| Types.NestedField field = fields.get(i); |
| switch (field.type().typeId()) { |
| case STRUCT: |
| case LIST: |
| case MAP: |
| throw new IllegalArgumentException("Unsupported type in partition data: " + type); |
| case BINARY: |
| case FIXED: |
| byte[] buffer = (byte[]) data[i]; |
| copy[i] = Arrays.copyOf(buffer, buffer.length); |
| break; |
| case STRING: |
| copy[i] = data[i].toString(); |
| break; |
| default: |
| // no need to copy the object |
| copy[i] = data[i]; |
| } |
| } |
| } |
| |
| return copy; |
| } |
| } |