diff --git a/parquet-column/src/main/java/org/apache/parquet/column/ParquetProperties.java b/parquet-column/src/main/java/org/apache/parquet/column/ParquetProperties.java index 8fe45e01ef..0df47bcde7 100644 --- a/parquet-column/src/main/java/org/apache/parquet/column/ParquetProperties.java +++ b/parquet-column/src/main/java/org/apache/parquet/column/ParquetProperties.java @@ -68,6 +68,13 @@ public class ParquetProperties { public static final boolean DEFAULT_STATISTICS_ENABLED = true; public static final boolean DEFAULT_SIZE_STATISTICS_ENABLED = true; + /** + * Payload size at or below which a {@code FILE} value is stored inline rather than as a + * self-reference. Defaults to the page size: a payload that would fill a page on its own is + * better kept out of the column chunk. + */ + public static final int DEFAULT_FILE_SELF_REFERENCE_THRESHOLD = DEFAULT_PAGE_SIZE; + public static final boolean DEFAULT_PAGE_WRITE_CHECKSUM_ENABLED = true; /** @@ -138,6 +145,7 @@ public static WriterVersion fromString(String name) { private final ColumnProperty sizeStatistics; private final ColumnProperty columnCodecs; private final ColumnProperty columnCompressionLevels; + private final int fileSelfReferenceThreshold; private ParquetProperties(Builder builder) { this.pageSizeThreshold = builder.pageSize; @@ -172,6 +180,7 @@ private ParquetProperties(Builder builder) { this.sizeStatistics = builder.sizeStatistics.build(); this.columnCodecs = builder.columnCodecs.build(); this.columnCompressionLevels = builder.columnCompressionLevels.build(); + this.fileSelfReferenceThreshold = builder.fileSelfReferenceThreshold; } public static Builder builder() { @@ -345,6 +354,14 @@ public int getMaxBloomFilterBytes() { return maxBloomFilterBytes; } + /** + * @return the payload size at or below which a {@code FILE} value is stored inline rather than as + * a self-reference + */ + public int getFileSelfReferenceThreshold() { + return fileSelfReferenceThreshold; + } + public boolean getAdaptiveBloomFilterEnabled(ColumnDescriptor column) { return adaptiveBloomFilterEnabled.getValue(column); } @@ -415,7 +432,8 @@ public String toString() { + "Page row count limit to " + getPageRowCountLimit() + '\n' + "Writing page checksums is: " + (getPageWriteChecksumEnabled() ? "on" : "off") + '\n' + "Statistics enabled: " + statisticsEnabled + '\n' - + "Size statistics enabled: " + sizeStatisticsEnabled; + + "Size statistics enabled: " + sizeStatisticsEnabled + '\n' + + "FILE self-reference threshold is: " + getFileSelfReferenceThreshold(); String perColumn = ""; if (!columnCodecs.toString().equals(Objects.toString(columnCodecs.getDefaultValue()))) { perColumn = "Per-column codecs: " + columnCodecs; @@ -460,6 +478,7 @@ public static class Builder { private final ColumnProperty.Builder sizeStatistics; private final ColumnProperty.Builder columnCodecs; private final ColumnProperty.Builder columnCompressionLevels; + private int fileSelfReferenceThreshold = DEFAULT_FILE_SELF_REFERENCE_THRESHOLD; private Builder() { enableDict = ColumnProperty.builder().withDefaultValue(DEFAULT_IS_DICTIONARY_ENABLED); @@ -511,6 +530,7 @@ private Builder(ParquetProperties toCopy) { this.sizeStatisticsEnabled = toCopy.sizeStatisticsEnabled; this.columnCodecs = ColumnProperty.builder(toCopy.columnCodecs); this.columnCompressionLevels = ColumnProperty.builder(toCopy.columnCompressionLevels); + this.fileSelfReferenceThreshold = toCopy.fileSelfReferenceThreshold; } /** @@ -657,6 +677,27 @@ public Builder withStatisticsTruncateLength(int length) { return this; } + /** + * Set the payload size at or below which a {@code FILE} value is stored inline rather than as a + * self-reference. + * + *

Small payloads are cheaper to keep in the column chunk, where they are read as part of the + * ordinary page stream. Large ones are better stored out of line as self-references, so that + * reading the surrounding columns does not pull the payload bytes along with them. Set to 0 to + * store every payload as a self-reference, or to {@link Integer#MAX_VALUE} to always inline. + * + * @param fileSelfReferenceThreshold the inline size limit in bytes; must not be negative + * @return this builder for method chaining + */ + public Builder withFileSelfReferenceThreshold(int fileSelfReferenceThreshold) { + Preconditions.checkArgument( + fileSelfReferenceThreshold >= 0, + "Invalid FILE self-reference threshold (negative): %s", + fileSelfReferenceThreshold); + this.fileSelfReferenceThreshold = fileSelfReferenceThreshold; + return this; + } + /** * Set max Bloom filter bytes for related columns. * diff --git a/parquet-column/src/main/java/org/apache/parquet/schema/LogicalTypeAnnotation.java b/parquet-column/src/main/java/org/apache/parquet/schema/LogicalTypeAnnotation.java index 625e9fd9d3..9601e3eb84 100644 --- a/parquet-column/src/main/java/org/apache/parquet/schema/LogicalTypeAnnotation.java +++ b/parquet-column/src/main/java/org/apache/parquet/schema/LogicalTypeAnnotation.java @@ -188,6 +188,12 @@ protected LogicalTypeAnnotation fromString(List params) { protected LogicalTypeAnnotation fromString(List params) { return unknownType(); } + }, + FILE { + @Override + protected LogicalTypeAnnotation fromString(List params) { + return fileType(); + } }; protected abstract LogicalTypeAnnotation fromString(List params); @@ -378,6 +384,10 @@ public static UnknownLogicalTypeAnnotation unknownType() { return UnknownLogicalTypeAnnotation.INSTANCE; } + public static FileLogicalTypeAnnotation fileType() { + return FileLogicalTypeAnnotation.INSTANCE; + } + public static class StringLogicalTypeAnnotation extends LogicalTypeAnnotation { private static final StringLogicalTypeAnnotation INSTANCE = new StringLogicalTypeAnnotation(); @@ -1229,6 +1239,100 @@ public boolean equals(Object obj) { } } + /** + * File logical type annotation. Annotates a group (struct) that represents a reference to a + * range of bytes, which may be stored inline in the value, elsewhere within the current file, + * or in an external file. Every field is optional, both in the schema (a writer may omit any + * field from the group definition) and in the data (any field that is present has a field + * repetition type of {@code OPTIONAL}). Fields are identified by name (case sensitively), not by + * field order. A group need only define the fields it uses. The group may contain the following + * fields: + *

    + *
  • {@code uri} (STRING): a URI-reference (RFC 3986) that identifies an external file, for + * example {@code s3://bucket/file.jpg}. If not set, the value refers to the current file + * (a self-reference).
    • + *
  • {@code offset} (INT64): start of the byte range within the referenced data; if not set, + * treated as 0. Must not be negative.
    • + *
  • {@code size} (INT64): byte length of the referenced data. Must be set whenever + * {@code offset} is set or {@code uri} is not set; may be omitted only for a whole-file + * external reference, in which case the range runs to the end of the referenced file. Must + * not be negative.
    • + *
  • {@code content_type} (STRING): the media (MIME) type (RFC 2046) of the resolved bytes; + * when not set, {@code application/octet-stream} is assumed.
    • + *
  • {@code checksum} (STRING): a self-describing integrity token for the resolved bytes, of + * the form {@code :}.
    • + *
  • {@code inline} (BYTE_ARRAY): the referenced bytes stored inline in the value.
    • + *
+ * No fields with names other than the above are permitted. The schema builder additionally + * rejects group definitions that could never produce a valid value: a group that declares + * {@code offset} must also declare {@code size}, and a group must declare at least one of + * {@code inline}, {@code uri}, or {@code offset} (a value resolves to bytes only via one of + * these; a group declaring none of them — even if it declares {@code size} — can never produce a + * resolvable value). A group that declares {@code offset} but not {@code uri} permits only + * self-references (a value with {@code uri} unset that locates bytes within the current file) and + * must therefore also declare {@code inline}: the {@code inline} column chunk of the same row + * group is the reference point whose compression and encryption a self-reference inherits. A + * group that declares {@code uri} is treated as an external-reference schema and is not required + * to declare {@code inline}. Each declared field must also match its required physical type. + * Per-value + * rules that depend on the data in each row — {@code offset} being set for a self-reference + * (unset {@code uri}), {@code size} being set whenever {@code offset} is set, and + * {@code offset}/{@code size} being non-negative — cannot be enforced here and are the + * responsibility of writers and consumers. + */ + public static class FileLogicalTypeAnnotation extends LogicalTypeAnnotation { + private static final FileLogicalTypeAnnotation INSTANCE = new FileLogicalTypeAnnotation(); + + /** Field name holding the URI-reference of an external file. */ + public static final String URI_FIELD = "uri"; + + /** Field name holding the start of the byte range. */ + public static final String OFFSET_FIELD = "offset"; + + /** Field name holding the byte length of the referenced data. */ + public static final String SIZE_FIELD = "size"; + + /** Field name holding the media (MIME) type of the resolved bytes. */ + public static final String CONTENT_TYPE_FIELD = "content_type"; + + /** Field name holding the integrity token for the resolved bytes. */ + public static final String CHECKSUM_FIELD = "checksum"; + + /** Field name holding the referenced bytes stored inline. */ + public static final String INLINE_FIELD = "inline"; + + /** All recognized field names in a FILE-annotated group. All fields are optional. */ + public static final Set FIELD_NAMES = + Set.of(URI_FIELD, OFFSET_FIELD, SIZE_FIELD, CONTENT_TYPE_FIELD, CHECKSUM_FIELD, INLINE_FIELD); + + private FileLogicalTypeAnnotation() {} + + @Override + public OriginalType toOriginalType() { + return null; + } + + @Override + public Optional accept(LogicalTypeAnnotationVisitor logicalTypeAnnotationVisitor) { + return logicalTypeAnnotationVisitor.visit(this); + } + + @Override + LogicalTypeToken getType() { + return LogicalTypeToken.FILE; + } + + @Override + public boolean equals(Object obj) { + return obj instanceof FileLogicalTypeAnnotation; + } + + @Override + public int hashCode() { + return getClass().hashCode(); + } + } + public static class GeometryLogicalTypeAnnotation extends LogicalTypeAnnotation { private final String crs; @@ -1434,5 +1538,9 @@ default Optional visit(GeographyLogicalTypeAnnotation geographyLogicalType) { default Optional visit(UnknownLogicalTypeAnnotation unknownLogicalTypeAnnotation) { return empty(); } + + default Optional visit(FileLogicalTypeAnnotation fileLogicalType) { + return empty(); + } } } diff --git a/parquet-column/src/main/java/org/apache/parquet/schema/Types.java b/parquet-column/src/main/java/org/apache/parquet/schema/Types.java index 2f12991ab0..5c40556d16 100644 --- a/parquet-column/src/main/java/org/apache/parquet/schema/Types.java +++ b/parquet-column/src/main/java/org/apache/parquet/schema/Types.java @@ -821,12 +821,125 @@ public THIS addFields(Type... types) { @Override protected GroupType build(String name) { if (newLogicalTypeSet) { + if (logicalTypeAnnotation instanceof LogicalTypeAnnotation.FileLogicalTypeAnnotation) { + validateFileTypeFields(name, fields); + } return new GroupType(repetition, name, logicalTypeAnnotation, fields, id); } else { return new GroupType(repetition, name, getOriginalType(), fields, id); } } + private static void validateFileTypeFields(String name, List fields) { + boolean hasUri = false; + boolean hasOffset = false; + boolean hasSize = false; + boolean hasInline = false; + for (Type field : fields) { + String fieldName = field.getName(); + if (!LogicalTypeAnnotation.FileLogicalTypeAnnotation.FIELD_NAMES.contains(fieldName)) { + throw new IllegalArgumentException("FILE type group '" + name + "' contains unrecognized field '" + + fieldName + "'. Valid fields are: " + + String.join(", ", LogicalTypeAnnotation.FileLogicalTypeAnnotation.FIELD_NAMES)); + } + Preconditions.checkArgument( + field.isPrimitive() && field.getRepetition() == Type.Repetition.OPTIONAL, + "FILE type field '%s' must be an optional primitive in group '%s'", + fieldName, + name); + validateFileTypeFieldPhysicalType(name, field.asPrimitiveType()); + if (LogicalTypeAnnotation.FileLogicalTypeAnnotation.URI_FIELD.equals(fieldName)) { + hasUri = true; + } else if (LogicalTypeAnnotation.FileLogicalTypeAnnotation.OFFSET_FIELD.equals(fieldName)) { + hasOffset = true; + } else if (LogicalTypeAnnotation.FileLogicalTypeAnnotation.SIZE_FIELD.equals(fieldName)) { + hasSize = true; + } else if (LogicalTypeAnnotation.FileLogicalTypeAnnotation.INLINE_FIELD.equals(fieldName)) { + hasInline = true; + } + } + // The spec requires `size` to be set whenever `offset` is set. A group that declares + // `offset` but not `size` can never produce a valid value, so reject it at schema-build + // time. + Preconditions.checkArgument( + !hasOffset || hasSize, + "FILE type group '%s' declares field 'offset' but not 'size'; 'size' is required whenever 'offset' is set", + name); + // Per the spec resolution table, a value resolves to bytes only if `inline`, `uri`, or + // `offset` is set; `size` on its own never resolves. A group that declares none of `inline`, + // `uri`, or `offset` can therefore never produce a resolvable value, so reject it at + // schema-build time. + Preconditions.checkArgument( + hasInline || hasUri || hasOffset, + "FILE type group '%s' must declare at least one of 'inline', 'uri', or 'offset'; a value " + + "resolves to bytes only via one of these, so a group declaring none of them can " + + "never produce a valid value", + name); + // A schema that permits self-references must declare `inline`. A self-reference (`uri` not + // set) always sets `offset`, and the `inline` column chunk of the same row group is the + // reference point whose compression and encryption a self-reference inherits. A group that + // declares `offset` but not `uri` can only produce self-references (an offset-based read with + // no `uri` is a self-reference), so it must also declare `inline`. A group that declares + // `uri` is not required to declare `inline`: `offset`/`size` there describe an external + // ranged reference, and although the per-value `uri` could be left unset in some rows, the + // schema is treated as an external-reference schema and the `inline` requirement is not + // imposed. A writer must therefore not emit a self-reference under such a schema, since there + // would be no `inline` column chunk to inherit compression and encryption from; that is + // enforced on the write path rather than here. + Preconditions.checkArgument( + !(hasOffset && !hasUri) || hasInline, + "FILE type group '%s' declares field 'offset' but neither 'uri' nor 'inline'; a schema " + + "that permits self-references (offset without uri) must declare 'inline' as the " + + "reference point for storage inheritance", + name); + // The remaining spec rules are per-value constraints that the schema builder cannot verify + // because it sees only which fields are declared, not their values in each row: a + // self-reference (unset `uri`) must set `offset`, `size` must be set whenever `offset` is + // set, and `offset`/`size` must be non-negative. Those are the responsibility of writers and + // consumers of FILE values. + } + + /** + * Validates that a declared FILE field uses the physical type required by the spec: + * {@code uri}, {@code content_type}, and {@code checksum} are STRING (BINARY), {@code offset} + * and {@code size} are INT64, and {@code inline} is BYTE_ARRAY (BINARY). + */ + private static void validateFileTypeFieldPhysicalType(String name, PrimitiveType field) { + String fieldName = field.getName(); + PrimitiveType.PrimitiveTypeName physicalType = field.getPrimitiveTypeName(); + switch (fieldName) { + case LogicalTypeAnnotation.FileLogicalTypeAnnotation.URI_FIELD: + case LogicalTypeAnnotation.FileLogicalTypeAnnotation.CONTENT_TYPE_FIELD: + case LogicalTypeAnnotation.FileLogicalTypeAnnotation.CHECKSUM_FIELD: + Preconditions.checkArgument( + physicalType == PrimitiveType.PrimitiveTypeName.BINARY + && field.getLogicalTypeAnnotation() + instanceof LogicalTypeAnnotation.StringLogicalTypeAnnotation, + "FILE type field '%s' must be a STRING (BINARY annotated as STRING) in group '%s'", + fieldName, + name); + break; + case LogicalTypeAnnotation.FileLogicalTypeAnnotation.OFFSET_FIELD: + case LogicalTypeAnnotation.FileLogicalTypeAnnotation.SIZE_FIELD: + Preconditions.checkArgument( + physicalType == PrimitiveType.PrimitiveTypeName.INT64, + "FILE type field '%s' must be an INT64 in group '%s'", + fieldName, + name); + break; + case LogicalTypeAnnotation.FileLogicalTypeAnnotation.INLINE_FIELD: + Preconditions.checkArgument( + physicalType == PrimitiveType.PrimitiveTypeName.BINARY, + "FILE type field '%s' must be a BYTE_ARRAY (BINARY) in group '%s'", + fieldName, + name); + break; + default: + // Unreachable: field names are validated against FIELD_NAMES before this call. + break; + } + } + public MapBuilder map(Type.Repetition repetition) { return new MapBuilder<>(self()).repetition(repetition); } diff --git a/parquet-column/src/test/java/org/apache/parquet/schema/TestTypeBuildersWithLogicalTypes.java b/parquet-column/src/test/java/org/apache/parquet/schema/TestTypeBuildersWithLogicalTypes.java index 0d7791a19b..1a619ce9b8 100644 --- a/parquet-column/src/test/java/org/apache/parquet/schema/TestTypeBuildersWithLogicalTypes.java +++ b/parquet-column/src/test/java/org/apache/parquet/schema/TestTypeBuildersWithLogicalTypes.java @@ -549,4 +549,285 @@ public void testVariantLogicalTypeWithShredded() { assertThat(((LogicalTypeAnnotation.VariantLogicalTypeAnnotation) annotation).getSpecVersion()) .isEqualTo(specVersion); } + + @Test + public void testFileLogicalTypeUriOnly() { + String name = "file_field"; + GroupType file = new GroupType( + REQUIRED, + name, + LogicalTypeAnnotation.fileType(), + Types.optional(BINARY).as(LogicalTypeAnnotation.stringType()).named("uri")); + + assertThat(file.toString()) + .isEqualTo("required group file_field (FILE) {\n" + " optional binary uri (STRING);\n" + "}"); + + LogicalTypeAnnotation annotation = file.getLogicalTypeAnnotation(); + assertThat(annotation.getType()).isEqualTo(LogicalTypeAnnotation.LogicalTypeToken.FILE); + assertThat(annotation.toOriginalType()).isNull(); + assertThat(annotation).isInstanceOf(LogicalTypeAnnotation.FileLogicalTypeAnnotation.class); + } + + @Test + public void testFileLogicalTypeAllFields() { + String name = "file_field"; + GroupType file = Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("uri") + .optional(INT64) + .named("offset") + .optional(INT64) + .named("size") + .optional(BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("content_type") + .optional(BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("checksum") + .optional(BINARY) + .named("inline") + .named(name); + + LogicalTypeAnnotation annotation = file.getLogicalTypeAnnotation(); + assertThat(annotation).isInstanceOf(LogicalTypeAnnotation.FileLogicalTypeAnnotation.class); + assertThat(file.getFieldCount()).isEqualTo(6); + assertThat(file.getType("uri").getName()).isEqualTo("uri"); + assertThat(file.getType("offset").getName()).isEqualTo("offset"); + assertThat(file.getType("size").getName()).isEqualTo("size"); + assertThat(file.getType("content_type").getName()).isEqualTo("content_type"); + assertThat(file.getType("checksum").getName()).isEqualTo("checksum"); + assertThat(file.getType("inline").getName()).isEqualTo("inline"); + } + + @Test + public void testFileLogicalTypeInlineOnly() { + // Every field is optional, so an inline-only group is valid (spec inline case). + GroupType file = Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(BINARY) + .named("inline") + .named("inline_file"); + + assertThat(file.getLogicalTypeAnnotation()).isInstanceOf(LogicalTypeAnnotation.FileLogicalTypeAnnotation.class); + assertThat(file.getFieldCount()).isEqualTo(1); + assertThat(file.getType("inline").getName()).isEqualTo("inline"); + } + + @Test + public void testFileLogicalTypeSelfReference() { + // A self-reference omits 'uri' and locates bytes within the current file via offset/size. + // A schema that permits self-references must declare 'inline' as the storage-inheritance + // reference point. + GroupType file = Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(INT64) + .named("offset") + .optional(INT64) + .named("size") + .optional(BINARY) + .named("inline") + .named("self_ref_file"); + + assertThat(file.getLogicalTypeAnnotation()).isInstanceOf(LogicalTypeAnnotation.FileLogicalTypeAnnotation.class); + assertThat(file.getFieldCount()).isEqualTo(3); + } + + @Test + public void testFileLogicalTypeOffsetRequiresInline() { + // A schema that permits self-references (declares 'offset' but not 'uri') must declare 'inline' + // as the reference point for storage inheritance. A group declaring 'offset'/'size' with + // neither 'uri' nor 'inline' is rejected at build time. + assertThatThrownBy(() -> Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(INT64) + .named("offset") + .optional(INT64) + .named("size") + .named("self_ref_without_inline")) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + public void testFileLogicalTypeExternalRangedReferenceWithoutInline() { + // An external ranged reference declares 'uri' + 'offset' + 'size' to point at a byte range of + // an external file. Because 'uri' is declared, the schema is treated as an external-reference + // schema and is not required to declare 'inline', even though it declares 'offset'. + GroupType file = Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("uri") + .optional(INT64) + .named("offset") + .optional(INT64) + .named("size") + .named("external_ranged_file"); + + assertThat(file.getLogicalTypeAnnotation()).isInstanceOf(LogicalTypeAnnotation.FileLogicalTypeAnnotation.class); + assertThat(file.getFieldCount()).isEqualTo(3); + } + + @Test + public void testFileLogicalTypeMetadataOnlyRejected() { + // Per the spec resolution table, a value resolves to bytes only via 'inline', 'uri', or + // 'offset'. A group declaring only metadata fields can never produce a resolvable value. + assertThatThrownBy(() -> Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("content_type") + .optional(BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("checksum") + .named("file_metadata_only")) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + public void testFileLogicalTypeSizeOnlyRejected() { + // 'size' alone never resolves to bytes (spec resolution table), so a size-only group is + // rejected: it declares no locator ('inline', 'uri', or 'offset'). + assertThatThrownBy(() -> Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(INT64) + .named("size") + .named("file_size_only")) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + public void testFileLogicalTypeOffsetRequiresSize() { + // The spec requires 'size' whenever 'offset' is set, so a group declaring 'offset' + // without 'size' can never produce a valid value and is rejected at build time. + assertThatThrownBy(() -> Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("uri") + .optional(INT64) + .named("offset") + .named("file_offset_without_size")) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + public void testFileLogicalTypeOffsetWithSize() { + // 'offset' accompanied by 'size' is valid. + GroupType file = Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("uri") + .optional(INT64) + .named("offset") + .optional(INT64) + .named("size") + .named("file_offset_with_size"); + + assertThat(file.getLogicalTypeAnnotation()).isInstanceOf(LogicalTypeAnnotation.FileLogicalTypeAnnotation.class); + assertThat(file.getFieldCount()).isEqualTo(3); + } + + @Test + public void testFileLogicalTypeSizeWithoutOffset() { + // 'uri' + 'size' (without 'offset') is valid: an external reference to '[0, size)'. + GroupType file = Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("uri") + .optional(INT64) + .named("size") + .named("file_size_without_offset"); + + assertThat(file.getLogicalTypeAnnotation()).isInstanceOf(LogicalTypeAnnotation.FileLogicalTypeAnnotation.class); + assertThat(file.getFieldCount()).isEqualTo(2); + } + + @Test + public void testFileLogicalTypeRejectsUnrecognizedField() { + assertThatThrownBy(() -> Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("uri") + .optional(BINARY) + .named("unknown_field") + .named("file_with_bad_field")) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + public void testFileLogicalTypeRejectsRequiredField() { + // All FILE fields must have OPTIONAL repetition under the current spec. + assertThatThrownBy(() -> Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .required(BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("uri") + .named("file_with_required_uri")) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + public void testFileLogicalTypeRejectsGroupField() { + // FILE fields must be primitives, not nested groups. + assertThatThrownBy(() -> Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optionalGroup() + .optional(BINARY) + .named("nested") + .named("uri") + .named("file_with_group_field")) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + public void testFileLogicalTypeRejectsWrongStringPhysicalType() { + // 'uri' must be a STRING (BINARY annotated as STRING); an INT64 is rejected. + assertThatThrownBy(() -> Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(INT64) + .named("uri") + .named("file_uri_wrong_type")) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + public void testFileLogicalTypeRejectsUnannotatedStringField() { + // A STRING field must carry the STRING logical annotation; plain BINARY is rejected. + assertThatThrownBy(() -> Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(BINARY) + .named("uri") + .named("file_uri_unannotated")) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + public void testFileLogicalTypeRejectsWrongInt64PhysicalType() { + // 'offset' and 'size' must be INT64; an INT32 is rejected. + assertThatThrownBy(() -> Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("uri") + .optional(INT32) + .named("size") + .named("file_size_wrong_type")) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + public void testFileLogicalTypeRejectsWrongInlinePhysicalType() { + // 'inline' must be a BYTE_ARRAY (BINARY); an INT64 is rejected. + assertThatThrownBy(() -> Types.requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(INT64) + .named("inline") + .named("file_inline_wrong_type")) + .isInstanceOf(IllegalArgumentException.class); + } } diff --git a/parquet-format-structures/src/main/java/org/apache/parquet/format/LogicalTypes.java b/parquet-format-structures/src/main/java/org/apache/parquet/format/LogicalTypes.java index 8956d3944e..8aa21e0ae3 100644 --- a/parquet-format-structures/src/main/java/org/apache/parquet/format/LogicalTypes.java +++ b/parquet-format-structures/src/main/java/org/apache/parquet/format/LogicalTypes.java @@ -60,4 +60,5 @@ public static LogicalType VARIANT(byte specificationVersion) { public static final LogicalType BSON = LogicalType.BSON(new BsonType()); public static final LogicalType FLOAT16 = LogicalType.FLOAT16(new Float16Type()); public static final LogicalType UUID = LogicalType.UUID(new UUIDType()); + public static final LogicalType FILE = LogicalType.FILE(new FileType()); } diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/AesCipher.java b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/AesCipher.java index 2386412b65..8b4fb77577 100755 --- a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/AesCipher.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/AesCipher.java @@ -120,6 +120,60 @@ public static byte[] createFooterAAD(byte[] aadPrefixBytes) { return createModuleAAD(aadPrefixBytes, ModuleType.Footer, -1, -1, -1); } + /** + * Builds the module AAD for a self-reference (FILE self-reference payload). Unlike pages, which + * are identified by a 2-byte page ordinal, a self-reference is identified by the 8-byte + * little-endian offset of its stored representation within the file, following the row group and + * column ordinals. The column ordinal is that of the {@code inline} column whose encryption the + * self-reference inherits. + * + *

The offset is the value the writer records in the {@code offset} field of the {@code FILE} + * group. Because it is carried in the data, a reader can rebuild this AAD from the value alone, + * without counting the self-references that precede it and therefore without decoding the pages + * it skips. + * + * @param fileAAD the file AAD (AAD prefix concatenated with the AAD file-unique bytes) + * @param rowGroupOrdinal the row group ordinal of the self-reference + * @param columnOrdinal the ordinal of the {@code inline} column the self-reference inherits from + * @param selfReferenceOffset the offset of the stored representation within the file, i.e. the + * value of the self-reference's {@code offset} field + * @return the module AAD bytes + */ + public static byte[] createSelfReferenceAAD( + byte[] fileAAD, int rowGroupOrdinal, int columnOrdinal, long selfReferenceOffset) { + + byte[] typeOrdinalBytes = new byte[1]; + typeOrdinalBytes[0] = ModuleType.SelfReference.getValue(); + + if (rowGroupOrdinal < 0) { + throw new IllegalArgumentException("Wrong row group ordinal: " + rowGroupOrdinal); + } + short shortRGOrdinal = (short) rowGroupOrdinal; + if (shortRGOrdinal != rowGroupOrdinal) { + throw new ParquetCryptoRuntimeException("Encrypted parquet files can't have " + "more than " + + Short.MAX_VALUE + " row groups: " + rowGroupOrdinal); + } + byte[] rowGroupOrdinalBytes = shortToBytesLE(shortRGOrdinal); + + if (columnOrdinal < 0) { + throw new IllegalArgumentException("Wrong column ordinal: " + columnOrdinal); + } + short shortColumnOrdinal = (short) columnOrdinal; + if (shortColumnOrdinal != columnOrdinal) { + throw new ParquetCryptoRuntimeException("Encrypted parquet files can't have " + "more than " + + Short.MAX_VALUE + " columns: " + columnOrdinal); + } + byte[] columnOrdinalBytes = shortToBytesLE(shortColumnOrdinal); + + if (selfReferenceOffset < 0) { + throw new IllegalArgumentException("Wrong self-reference offset: " + selfReferenceOffset); + } + byte[] selfReferenceOffsetBytes = longToBytesLE(selfReferenceOffset); + + return concatByteArrays( + fileAAD, typeOrdinalBytes, rowGroupOrdinalBytes, columnOrdinalBytes, selfReferenceOffsetBytes); + } + // Update last two bytes with new page ordinal (instead of creating new page AAD from scratch) public static void quickUpdatePageAAD(byte[] pageAAD, int newPageOrdinal) { java.util.Objects.requireNonNull(pageAAD); @@ -159,4 +213,13 @@ private static byte[] shortToBytesLE(short input) { return output; } + + private static byte[] longToBytesLE(long input) { + byte[] output = new byte[8]; + for (int i = 0; i < 8; i++) { + output[i] = (byte) (0xff & (input >> (8 * i))); + } + + return output; + } } diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/ModuleCipherFactory.java b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/ModuleCipherFactory.java index 9d258e2825..94c9c68097 100755 --- a/parquet-hadoop/src/main/java/org/apache/parquet/crypto/ModuleCipherFactory.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/crypto/ModuleCipherFactory.java @@ -34,7 +34,8 @@ public enum ModuleType { ColumnIndex((byte) 6), OffsetIndex((byte) 7), BloomFilterHeader((byte) 8), - BloomFilterBitset((byte) 9); + BloomFilterBitset((byte) 9), + SelfReference((byte) 10); private final byte value; diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/format/converter/ParquetMetadataConverter.java b/parquet-hadoop/src/main/java/org/apache/parquet/format/converter/ParquetMetadataConverter.java index 465516e48f..0df057a53d 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/format/converter/ParquetMetadataConverter.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/format/converter/ParquetMetadataConverter.java @@ -591,6 +591,11 @@ public Optional visit(LogicalTypeAnnotation.GeographyLogicalTypeAnn geographyType.setAlgorithm(fromParquetEdgeInterpolationAlgorithm(geographyLogicalType.getAlgorithm())); return of(LogicalType.GEOGRAPHY(geographyType)); } + + @Override + public Optional visit(LogicalTypeAnnotation.FileLogicalTypeAnnotation fileLogicalType) { + return of(LogicalTypes.FILE); + } } private void addRowGroup( @@ -1389,6 +1394,8 @@ LogicalTypeAnnotation getLogicalTypeAnnotation(LogicalType type) { case VARIANT: VariantType variant = type.getVARIANT(); return LogicalTypeAnnotation.variantType(variant.getSpecification_version()); + case FILE: + return LogicalTypeAnnotation.fileType(); default: throw new RuntimeException("Unknown logical type " + type); } diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/CodecFactory.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/CodecFactory.java index c9391201f4..de06bbcfbb 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/CodecFactory.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/CodecFactory.java @@ -281,6 +281,111 @@ public BytesCompressor getCompressor(CompressionCodecName codecName, int level) return comp; } + /** Smallest output buffer tried by {@link #decompressUnknownSize}. */ + private static final int MIN_UNKNOWN_SIZE_BUFFER = 8 * 1024; + + /** + * Largest output buffer tried by {@link #decompressUnknownSize}. Java arrays are indexed by int, + * and some JVMs reserve a few header words, so this stays just below {@link Integer#MAX_VALUE}. + */ + private static final int MAX_UNKNOWN_SIZE_BUFFER = Integer.MAX_VALUE - 8; + + /** + * Decompresses a complete compression block whose decompressed size is not known in advance. This + * is used to resolve FILE self-references, whose stored representation records only the size of + * the (compressed) stored block and not the size of the resolved bytes. Each self-reference is an + * independent compression block, so the entire {@code compressed} range is supplied to the codec + * in one shot. + * + *

All codecs are supported, including those that record no decompressed size of their own. The + * format spec allows a reader to "decompress into a dynamically sized buffer", which is what this + * does: it guesses an output size, and whenever the codec fills the buffer exactly — the signal + * that the output may have been cut off — it doubles the guess and retries. Retries are bounded by + * the 2 GiB ceiling on a Java array. + * + *

The {@link Decompressor} is driven directly rather than through + * {@link BytesDecompressor#decompress(BytesInput, int)} or a stream-drain loop. The former reads + * back exactly the requested number of bytes and so cannot report a short read, and Parquet's + * codec streams are deliberately unframed ({@code NonBlockedDecompressorStream}), signalling a + * fully consumed block by throwing rather than by returning end-of-input. + * + * @param codecName the {@link CompressionCodecName} of the {@code inline} column chunk the + * self-reference inherits from; {@link CompressionCodecName#UNCOMPRESSED} returns the bytes + * unchanged + * @param compressed the complete compressed block + * @return the decompressed (resolved) bytes + * @throws IOException if decompression fails + */ + public BytesInput decompressUnknownSize(CompressionCodecName codecName, BytesInput compressed) throws IOException { + CompressionCodec codec = getCodec(codecName); + if (codec == null) { + // UNCOMPRESSED: the stored bytes are the resolved bytes. + return compressed; + } + + byte[] compressedBytes = compressed.toByteArray(); + if (compressedBytes.length == 0) { + // An empty payload compresses to nothing and resolves back to nothing. + return BytesInput.empty(); + } + + Decompressor decompressor = CodecPool.getDecompressor(codec); + if (decompressor == null) { + // Some codecs (ZSTD) expose no Decompressor and decompress only through their stream, which + // is framed and so reports end-of-input properly. Drain it. + try (InputStream is = codec.createInputStream(compressed.toInputStream(), null); + ByteArrayOutputStream out = new ByteArrayOutputStream(compressedBytes.length * 2)) { + byte[] buffer = new byte[MIN_UNKNOWN_SIZE_BUFFER]; + int read; + while ((read = is.read(buffer)) != -1) { + out.write(buffer, 0, read); + } + return BytesInput.from(out.toByteArray()); + } + } + try { + // Compression rarely achieves better than 2x on the payloads worth storing out of line, so + // the first attempt usually suffices. + long attemptSize = Math.max((long) compressedBytes.length * 2, MIN_UNKNOWN_SIZE_BUFFER); + while (true) { + int candidate = (int) Math.min(attemptSize, MAX_UNKNOWN_SIZE_BUFFER); + boolean lastAttempt = candidate == MAX_UNKNOWN_SIZE_BUFFER; + byte[] output = new byte[candidate]; + int total = 0; + boolean undersized = false; + + decompressor.reset(); + decompressor.setInput(compressedBytes, 0, compressedBytes.length); + try { + while (total < candidate && !decompressor.finished()) { + int written = decompressor.decompress(output, total, candidate - total); + if (written <= 0) { + break; + } + total += written; + } + } catch (IOException | RuntimeException e) { + // Codecs with no length information (e.g. raw LZ4) fail outright when the output buffer is + // too small rather than filling it, so treat a failure as a signal to grow. On the last + // attempt there is nothing left to try, so let it surface. + if (lastAttempt) { + throw e; + } + undersized = true; + } + + // Filling the buffer exactly is also ambiguous: the payload may be complete, or the codec may + // have had more to write. Grow and retry unless the decompressor confirmed it finished. + if (!undersized && (total < candidate || decompressor.finished() || lastAttempt)) { + return BytesInput.from(output, 0, total); + } + attemptSize = (long) candidate * 2; + } + } finally { + CodecPool.returnDecompressor(decompressor); + } + } + @Override public BytesDecompressor getDecompressor(CompressionCodecName codecName) { BytesDecompressor decomp = decompressors.get(codecName); diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/FileValueWriter.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/FileValueWriter.java new file mode 100644 index 0000000000..59d4d75abd --- /dev/null +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/FileValueWriter.java @@ -0,0 +1,197 @@ +/* + * 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.parquet.hadoop; + +import java.io.IOException; +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.format.BlockCipher; +import org.apache.parquet.io.api.Binary; + +/** + * Decides how a {@code FILE} value's payload is stored: inline in the value, or out of line as a + * self-reference within the same Parquet file. + * + *

Object models hand over the resolved (logical) bytes and receive back a {@link Placement} + * describing which fields of the {@code FILE} group to write. Callers do not choose between the two + * forms themselves; the choice follows the configured threshold, so the same writing code produces + * either form: + * + * 

{@code
+ * FileValueWriter.Placement placement = fileValueWriter.write(payload);
+ * if (placement.isInline()) {
+ *   group.add("inline", placement.getInlineBytes());
+ * } else {
+ *   group.add("offset", placement.getOffset());
+ *   group.add("size", placement.getSize());
+ * }
+ * }
+ * + *

Both forms describe the same logical bytes, so {@code content_type} and {@code checksum} are + * written identically either way — they describe the resolved bytes, not the storage. Consumers see + * no difference beyond which fields are set. + * + *

A self-reference payload is written immediately, while the record is being written and before + * the row group's column chunks are flushed. It therefore lands in a contiguous run ahead of those + * chunks, leaving each column chunk contiguous on disk. Writing eagerly is what makes the offset + * knowable in time: {@code offset} and {@code size} are ordinary column values, and once a value has + * been handed to a column writer it is encoded into a buffered page and cannot be revised, so a + * placeholder could never be patched up later. + * + * @see SelfReferenceStorage + */ +public class FileValueWriter { + + /** + * Where a {@code FILE} value's payload was placed, and therefore which fields of the {@code FILE} + * group the caller should write. Either the payload is inline, or it is a self-reference located by + * {@code offset} and {@code size}. + */ + public static final class Placement { + private final Binary inlineBytes; + private final long offset; + private final long size; + + private Placement(Binary inlineBytes, long offset, long size) { + this.inlineBytes = inlineBytes; + this.offset = offset; + this.size = size; + } + + static Placement inline(Binary inlineBytes) { + return new Placement(inlineBytes, -1, -1); + } + + static Placement selfReference(SelfReferenceStorage.StoredRange range) { + return new Placement(null, range.getOffset(), range.getSize()); + } + + /** Whether the payload is stored inline, i.e. whether the {@code inline} field should be set. */ + public boolean isInline() { + return inlineBytes != null; + } + + /** + * The bytes to write to the {@code inline} field. + * + * @throws IllegalStateException if the payload was stored as a self-reference + */ + public Binary getInlineBytes() { + if (!isInline()) { + throw new IllegalStateException("Payload was stored as a self-reference, not inline"); + } + return inlineBytes; + } + + /** + * The value to write to the {@code offset} field. + * + * @throws IllegalStateException if the payload was stored inline + */ + public long getOffset() { + if (isInline()) { + throw new IllegalStateException("Payload was stored inline; it has no offset"); + } + return offset; + } + + /** + * The value to write to the {@code size} field. This is the size of the stored representation + * after compression and encryption, not the size of the resolved bytes. + * + * @throws IllegalStateException if the payload was stored inline + */ + public long getSize() { + if (isInline()) { + throw new IllegalStateException("Payload was stored inline; it has no size"); + } + return size; + } + } + + private final ParquetFileWriter fileWriter; + private final CodecFactory.BytesCompressor inlineColumnCompressor; + private final BlockCipher.Encryptor inlineColumnEncryptor; + private final int inlineColumnOrdinal; + private final int selfReferenceThreshold; + + /** + * @param fileWriter the writer for the file being written; a block must be open when + * {@link #write} is called + * @param inlineColumnCompressor the compressor for the {@code inline} column chunk's codec, whose + * compression a self-reference inherits + * @param inlineColumnEncryptor the data-module encryptor of the {@code inline} column chunk, or + * {@code null} if that column chunk is not encrypted + * @param inlineColumnOrdinal the ordinal of the {@code inline} column within the schema. The + * schema must declare {@code inline}: it is the reference point whose compression and + * encryption a self-reference inherits, so a schema without it can only store payloads inline + * or as external references. Note that the schema builder does not require {@code inline} for + * groups that declare {@code uri}, so an external-reference schema may reach here; pair such a + * schema with a threshold of {@link Integer#MAX_VALUE} so nothing is stored out of line. + * @param selfReferenceThreshold payloads of at most this many bytes are stored inline; larger ones + * become self-references. See + * {@code ParquetProperties.Builder#withFileSelfReferenceThreshold(int)}. + */ + public FileValueWriter( + ParquetFileWriter fileWriter, + CodecFactory.BytesCompressor inlineColumnCompressor, + BlockCipher.Encryptor inlineColumnEncryptor, + int inlineColumnOrdinal, + int selfReferenceThreshold) { + if (selfReferenceThreshold < 0) { + throw new IllegalArgumentException( + "Self-reference threshold must not be negative: " + selfReferenceThreshold); + } + if (inlineColumnOrdinal < 0) { + throw new IllegalArgumentException("Invalid inline column ordinal: " + inlineColumnOrdinal); + } + this.fileWriter = fileWriter; + this.inlineColumnCompressor = inlineColumnCompressor; + this.inlineColumnEncryptor = inlineColumnEncryptor; + this.inlineColumnOrdinal = inlineColumnOrdinal; + this.selfReferenceThreshold = selfReferenceThreshold; + } + + /** + * Stores {@code payload} and returns which {@code FILE} group fields to write for it. Payloads at + * or below the configured threshold are returned for inline storage; larger ones are written to the + * file body immediately as self-references. + * + *

Must be called while a block is open on the underlying writer, and before that block's column + * chunks are flushed. + * + * @param payload the resolved (logical) bytes of the value + * @return the placement describing which fields to write + * @throws IOException if writing the self-reference payload fails + */ + public Placement write(Binary payload) throws IOException { + if (payload == null) { + throw new IllegalArgumentException("FILE payload must not be null"); + } + if (payload.length() <= selfReferenceThreshold) { + return Placement.inline(payload); + } + SelfReferenceStorage.StoredRange range = fileWriter.writeSelfReference( + BytesInput.from(payload.toByteBuffer()), + inlineColumnCompressor, + inlineColumnEncryptor, + inlineColumnOrdinal); + return Placement.selfReference(range); + } +} diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileReader.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileReader.java index 9af4b4ac60..6e70445336 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileReader.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileReader.java @@ -78,6 +78,7 @@ import org.apache.parquet.column.page.PageReadStore; import org.apache.parquet.column.values.bloomfilter.BlockSplitBloomFilter; import org.apache.parquet.column.values.bloomfilter.BloomFilter; +import org.apache.parquet.compression.CompressionCodecFactory; import org.apache.parquet.compression.CompressionCodecFactory.BytesInputDecompressor; import org.apache.parquet.crypto.AesCipher; import org.apache.parquet.crypto.FileDecryptionProperties; @@ -1070,6 +1071,71 @@ public String getFile() { return file.toString(); } + /** + * Resolves a {@code FILE} self-reference to its logical bytes. A self-reference (a {@code FILE} + * value with {@code uri} unset) records the {@code offset} and {@code size} of a stored + * representation within this file; the stored bytes inherit the compression and encryption of the + * {@code inline} column chunk in the same row group. This method reads the stored bytes, decrypts + * them when the {@code inline} column chunk is encrypted, and decompresses them with the column + * chunk's codec, returning the resolved bytes. See {@link SelfReferenceStorage} and the Parquet + * format's "FILE" logical type specification. + * + *

Everything needed to resolve the value comes from the value itself plus the {@code inline} + * column chunk's metadata, so a self-reference can be read without decoding the pages that + * precede it. + * + * @param inlineColumn the {@link ColumnChunkMetaData} of the {@code inline} column chunk whose + * compression and encryption the self-reference inherits + * @param offset the self-reference {@code offset} field (start of the stored representation) + * @param size the self-reference {@code size} field (byte length of the stored representation) + * @return the resolved (logical) bytes of the self-reference + * @throws IOException if reading or resolving fails + */ + public BytesInput resolveSelfReference(ColumnChunkMetaData inlineColumn, long offset, long size) + throws IOException { + if (offset < 0) { + throw new IllegalArgumentException("Self-reference offset must not be negative: " + offset); + } + if (size < 0) { + throw new IllegalArgumentException("Self-reference size must not be negative: " + size); + } + if (size > SelfReferenceStorage.MAX_ENCRYPTED_MODULE_SIZE) { + throw new IllegalArgumentException("Self-reference size exceeds the maximum readable range: " + size); + } + + byte[] stored = new byte[(int) size]; + f.seek(offset); + f.readFully(stored); + + BlockCipher.Decryptor pageDecryptor = null; + byte[] fileAAD = null; + int columnOrdinal = -1; + if (null != fileDecryptor && !fileDecryptor.plaintextFile()) { + InternalColumnDecryptionSetup columnDecryptionSetup = fileDecryptor.getColumnSetup(inlineColumn.getPath()); + if (columnDecryptionSetup.isEncrypted()) { + pageDecryptor = columnDecryptionSetup.getDataDecryptor(); + fileAAD = fileDecryptor.getFileAAD(); + columnOrdinal = columnDecryptionSetup.getOrdinal(); + } + } + + CompressionCodecFactory codecFactory = options.getCodecFactory(); + if (!(codecFactory instanceof CodecFactory)) { + throw new IllegalStateException("Resolving FILE self-references requires a CodecFactory-based " + + "codec factory but found: " + codecFactory.getClass().getName()); + } + + return SelfReferenceStorage.resolve( + BytesInput.from(stored), + inlineColumn.getCodec(), + (CodecFactory) codecFactory, + pageDecryptor, + fileAAD, + inlineColumn.getRowGroupOrdinal(), + columnOrdinal, + offset); + } + private List filterRowGroups(List blocks) throws IOException { FilterCompat.Filter recordFilter = options.getRecordFilter(); if (FilterCompat.isFilteringRequired(recordFilter)) { diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileWriter.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileWriter.java index 82f4577b83..eb4de8d28a 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileWriter.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileWriter.java @@ -609,6 +609,50 @@ public InternalFileEncryptor getEncryptor() { return fileEncryptor; } + /** + * Writes a {@code FILE} self-reference payload into the file body, inheriting the compression and + * encryption of the {@code inline} column chunk, and returns the {@code offset} and {@code size} a + * writer records in the self-reference's {@code offset} and {@code size} fields. See + * {@link SelfReferenceStorage} for the layout and the Parquet format's "FILE" logical type + * specification for the storage-inheritance semantics. + * + *

The payload is compressed as an independent compression block using {@code compressor} (the + * compressor for the {@code inline} column chunk's {@link CompressionCodecName}) and, when + * {@code pageBlockEncryptor} is non-null, encrypted as an independent module with the + * {@code Self-Reference} module type. The row group ordinal is that of the block currently being + * written. + * + *

Payloads are written while a block is open but before its column chunks are flushed, so they + * land in a contiguous run ahead of the row group's chunks. This keeps each column chunk + * contiguous on disk, which the read path relies on when coalescing adjacent chunks into a single + * range read. + * + *

This must be called while a block is open (after {@link #startBlock(long)} and before + * {@link #endBlock()}) so that the returned offset falls within the file body. + * + * @param resolvedBytes the resolved (logical) bytes of the self-reference + * @param compressor the compressor for the {@code inline} column chunk's codec + * @param pageBlockEncryptor the data-module encryptor of the {@code inline} column chunk, or + * {@code null} if the column chunk is not encrypted + * @param columnOrdinal the ordinal of the {@code inline} column the self-reference inherits from + * @return the offset and size of the stored representation + * @throws IOException if writing or compression fails + */ + public SelfReferenceStorage.StoredRange writeSelfReference( + BytesInput resolvedBytes, + CodecFactory.BytesCompressor compressor, + BlockCipher.Encryptor pageBlockEncryptor, + int columnOrdinal) + throws IOException { + return withAbortOnFailure(() -> { + // The block currently being written will be assigned ordinal blocks.size() in endBlock(). + int rowGroupOrdinal = blocks.size(); + byte[] fileAAD = (null == fileEncryptor) ? null : fileEncryptor.getFileAAD(); + return SelfReferenceStorage.write( + resolvedBytes, compressor, pageBlockEncryptor, fileAAD, rowGroupOrdinal, columnOrdinal, out); + }); + } + /** * start a block * diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/SelfReferenceStorage.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/SelfReferenceStorage.java new file mode 100644 index 0000000000..0660b05ab5 --- /dev/null +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/SelfReferenceStorage.java @@ -0,0 +1,209 @@ +/* + * 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.parquet.hadoop; + +import java.io.IOException; +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.crypto.AesCipher; +import org.apache.parquet.format.BlockCipher; +import org.apache.parquet.hadoop.CodecFactory.BytesCompressor; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; + +/** + * Implements the storage-inheritance semantics for {@code FILE} self-references as specified in the + * Parquet format (see {@code LogicalTypes.md}, section "FILE"). A self-reference is a {@code FILE} + * value whose {@code uri} is not set and that locates a byte range within the same Parquet file via + * {@code offset} and {@code size}. + * + *

A self-reference does not point at the resolved (logical) bytes directly. Instead it points at + * a stored representation: the resolved bytes after being compressed and (optionally) + * encrypted, inheriting the {@link CompressionCodecName} and encryption settings of the + * {@code inline} column chunk in the same row group. Each self-reference is an independent + * compression block and (when encrypted) an independent encryption module; state is not shared with + * data pages or with other self-references. + * + *

Layout of a stored self-reference: + * + *

    + *
  • Unencrypted: the compressed block (or the raw bytes when the codec is + * {@link CompressionCodecName#UNCOMPRESSED}). {@code offset}/{@code size} cover exactly these + * bytes. + *
  • Encrypted: the modular-encryption serialization of the compressed block — a 4-byte + * little-endian length, a 12-byte nonce, the ciphertext, and (for AES_GCM_V1) a 16-byte GCM + * tag. {@code offset} points to the beginning of the 4-byte length and {@code size} covers the + * complete encrypted module. The AAD uses the {@code Self-Reference} module type (10) with the + * 8-byte file offset of the stored representation; see + * {@link AesCipher#createSelfReferenceAAD}. + *
+ * + *

Because the AAD is keyed on the file offset — a value the {@code FILE} group already carries in + * its {@code offset} field — a reader can resolve a self-reference directly from the value, without + * decoding the pages preceding it. An encrypted stored representation is therefore bound to one + * column chunk at one offset and must not be shared between column chunks. + * + *

Compression is always applied before encryption on write; decryption is applied before + * decompression on read. + */ +public final class SelfReferenceStorage { + + /** + * The largest encrypted module a writer can serialize: the 4-byte little-endian length field is + * read back as a signed int, so the buffer it describes cannot exceed 2 GiB. + */ + public static final long MAX_ENCRYPTED_MODULE_SIZE = Integer.MAX_VALUE; + + /** + * Bytes an encrypted module adds around the compressed block: the 4-byte length, the 12-byte + * nonce, and the 16-byte GCM tag. AES_GCM_CTR_V1 omits the tag, so this is an upper bound. + */ + private static final long MAX_ENCRYPTION_OVERHEAD = 4 + 12 + 16; + + private SelfReferenceStorage() {} + + /** + * The location of a stored self-reference within the Parquet file. The {@code offset} and + * {@code size} are exactly the values a writer records in the {@code offset} and {@code size} + * fields of the {@code FILE} group. + */ + public static final class StoredRange { + private final long offset; + private final long size; + + public StoredRange(long offset, long size) { + this.offset = offset; + this.size = size; + } + + /** The byte offset of the stored representation within the Parquet file. */ + public long getOffset() { + return offset; + } + + /** The byte length of the stored representation. */ + public long getSize() { + return size; + } + } + + /** + * Compresses (and optionally encrypts) {@code resolvedBytes} as an independent stored block and + * appends it to {@code out}, returning the {@link StoredRange} that a writer records in the + * {@code offset} and {@code size} fields of the self-reference. + * + * @param resolvedBytes the resolved (logical) bytes of the self-reference + * @param compressor the compressor for the {@code inline} column chunk's codec; must not be null + * (use the {@link CompressionCodecName#UNCOMPRESSED} compressor to store bytes uncompressed) + * @param pageBlockEncryptor the data-module encryptor of the {@code inline} column chunk, or + * {@code null} if the column chunk is not encrypted + * @param fileAAD the file AAD, required when {@code pageBlockEncryptor} is non-null + * @param rowGroupOrdinal the row group ordinal of the self-reference + * @param columnOrdinal the ordinal of the {@code inline} column the self-reference inherits from + * @param out the Parquet file output stream, positioned where the stored block should be written + * @return the offset and size of the stored representation + * @throws IOException if writing or compression fails + */ + public static StoredRange write( + BytesInput resolvedBytes, + BytesCompressor compressor, + BlockCipher.Encryptor pageBlockEncryptor, + byte[] fileAAD, + int rowGroupOrdinal, + int columnOrdinal, + org.apache.parquet.io.PositionOutputStream out) + throws IOException { + + // Step 1: compress the resolved bytes as an independent compression block. UNCOMPRESSED leaves + // the bytes unchanged (the NO_OP_COMPRESSOR returns its input). + BytesInput stored = compressor.compress(resolvedBytes); + + // The offset of the stored representation is the current stream position, and it is also the + // AAD's self-reference identity, so it must be read before anything is written. + long offset = out.getPos(); + + // Step 2: when the inline column chunk is encrypted, encrypt the compressed block as an + // independent module keyed on that offset. The encryptor prepends the 4-byte length and the + // nonce and appends the GCM tag (for AES_GCM_V1); the returned byte array is the complete + // stored module. + if (pageBlockEncryptor != null) { + long plaintextSize = stored.size(); + // The 4-byte length field of an encrypted module caps the buffer at 2 GiB. Check before + // encrypting so an oversized value fails with a diagnostic instead of a corrupt length. + long encryptedSize = plaintextSize + MAX_ENCRYPTION_OVERHEAD; + if (encryptedSize > MAX_ENCRYPTED_MODULE_SIZE) { + throw new IllegalArgumentException("Self-reference is too large to encrypt: " + plaintextSize + + " compressed bytes exceed the " + MAX_ENCRYPTED_MODULE_SIZE + + "-byte limit imposed by the 4-byte length field of an encrypted module. " + + "Store this value as an external reference (uri) instead."); + } + byte[] selfReferenceAAD = AesCipher.createSelfReferenceAAD(fileAAD, rowGroupOrdinal, columnOrdinal, offset); + stored = BytesInput.from(pageBlockEncryptor.encrypt(stored.toByteArray(), selfReferenceAAD)); + } + + long size = stored.size(); + stored.writeAllTo(out); + return new StoredRange(offset, size); + } + + /** + * Resolves a stored self-reference back to its logical bytes: decrypts the stored representation + * (when the {@code inline} column chunk is encrypted) and then decompresses it using the column + * chunk's codec. + * + * @param storedBytes the stored representation, i.e. the {@code [offset, offset + size)} range + * read from the Parquet file + * @param codecName the {@link CompressionCodecName} of the {@code inline} column chunk + * @param codecFactory the codec factory used to decompress the block + * @param pageBlockDecryptor the data-module decryptor of the {@code inline} column chunk, or + * {@code null} if the column chunk is not encrypted + * @param fileAAD the file AAD, required when {@code pageBlockDecryptor} is non-null + * @param rowGroupOrdinal the row group ordinal of the self-reference + * @param columnOrdinal the ordinal of the {@code inline} column the self-reference inherits from + * @param selfReferenceOffset the value of the self-reference's {@code offset} field, which is both + * where {@code storedBytes} was read from and the self-reference's AAD identity + * @return the resolved (logical) bytes + * @throws IOException if decompression fails + */ + public static BytesInput resolve( + BytesInput storedBytes, + CompressionCodecName codecName, + CodecFactory codecFactory, + BlockCipher.Decryptor pageBlockDecryptor, + byte[] fileAAD, + int rowGroupOrdinal, + int columnOrdinal, + long selfReferenceOffset) + throws IOException { + + BytesInput compressed = storedBytes; + + // Step 1: decrypt when the inline column chunk is encrypted. The decryptor consumes the 4-byte + // length, nonce, ciphertext, and GCM tag and returns the compressed block. The AAD is rebuilt + // from the offset alone, so no state from preceding values is needed. + if (pageBlockDecryptor != null) { + byte[] selfReferenceAAD = + AesCipher.createSelfReferenceAAD(fileAAD, rowGroupOrdinal, columnOrdinal, selfReferenceOffset); + compressed = BytesInput.from(pageBlockDecryptor.decrypt(storedBytes.toByteArray(), selfReferenceAAD)); + } + + // Step 2: decompress. The resolved size is not stored, so the codec decompresses into a + // dynamically sized buffer. + return codecFactory.decompressUnknownSize(codecName, compressed); + } +} diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestSelfReferenceAAD.java b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestSelfReferenceAAD.java new file mode 100644 index 0000000000..718798c157 --- /dev/null +++ b/parquet-hadoop/src/test/java/org/apache/parquet/crypto/TestSelfReferenceAAD.java @@ -0,0 +1,103 @@ +/* + * 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.parquet.crypto; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import org.apache.parquet.crypto.ModuleCipherFactory.ModuleType; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for the {@code Self-Reference} module type and the {@link AesCipher#createSelfReferenceAAD} + * AAD construction defined for FILE self-references (parquet-format PR #603). + */ +public class TestSelfReferenceAAD { + + private static final byte[] FILE_AAD = new byte[] {1, 2, 3, 4, 5, 6, 7, 8}; + + @Test + public void testSelfReferenceModuleTypeValue() { + // The spec assigns module type 10 to Self-Reference. + assertThat(ModuleType.SelfReference.getValue()).isEqualTo((byte) 10); + } + + @Test + public void testSelfReferenceAADLayout() { + int rowGroupOrdinal = 3; + int columnOrdinal = 7; + long selfReferenceOffset = 0x0102030405060708L; + + byte[] aad = AesCipher.createSelfReferenceAAD(FILE_AAD, rowGroupOrdinal, columnOrdinal, selfReferenceOffset); + + // Layout: fileAAD | moduleType(1) | rowGroupOrdinal(2 LE) | columnOrdinal(2 LE) | + // selfReferenceOffset(8 LE) + assertThat(aad.length).isEqualTo(FILE_AAD.length + 1 + 2 + 2 + 8); + + ByteBuffer buf = ByteBuffer.wrap(aad).order(ByteOrder.LITTLE_ENDIAN); + byte[] filePart = new byte[FILE_AAD.length]; + buf.get(filePart); + assertThat(filePart).isEqualTo(FILE_AAD); + assertThat(buf.get()).isEqualTo((byte) 10); // module type + assertThat(buf.getShort()).isEqualTo((short) rowGroupOrdinal); + assertThat(buf.getShort()).isEqualTo((short) columnOrdinal); + // The self-reference is identified by the 8-byte file offset of its stored representation, + // unlike the 2-byte page ordinal. + assertThat(buf.getLong()).isEqualTo(selfReferenceOffset); + } + + @Test + public void testSelfReferenceAADSupportsLargeOffset() { + // File offsets routinely exceed the 2-byte page-ordinal range, so the field must be 8 bytes. + long largeOffset = 5L * 1024 * 1024 * 1024; + byte[] aad = AesCipher.createSelfReferenceAAD(FILE_AAD, 0, 0, largeOffset); + ByteBuffer buf = ByteBuffer.wrap(aad, FILE_AAD.length + 1 + 2 + 2, 8).order(ByteOrder.LITTLE_ENDIAN); + assertThat(buf.getLong()).isEqualTo(largeOffset); + } + + @Test + public void testDistinctOffsetsProduceDistinctAADs() { + // Two self-references in the same column chunk are distinguished solely by their offsets. + byte[] first = AesCipher.createSelfReferenceAAD(FILE_AAD, 1, 2, 1000L); + byte[] second = AesCipher.createSelfReferenceAAD(FILE_AAD, 1, 2, 1064L); + assertThat(first).isNotEqualTo(second); + } + + @Test + public void testSelfReferenceAADRejectsNegativeValues() { + assertThatThrownBy(() -> AesCipher.createSelfReferenceAAD(FILE_AAD, -1, 0, 0)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> AesCipher.createSelfReferenceAAD(FILE_AAD, 0, -1, 0)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> AesCipher.createSelfReferenceAAD(FILE_AAD, 0, 0, -1)) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + public void testSelfReferenceAADDiffersFromPageAAD() { + // A self-reference and a data page in the same column must not share an AAD, because the module + // type byte differs (and the trailing field differs in width and meaning). + byte[] selfRefAAD = AesCipher.createSelfReferenceAAD(FILE_AAD, 1, 2, 0); + byte[] dataPageAAD = AesCipher.createModuleAAD(FILE_AAD, ModuleType.DataPage, 1, 2, 0); + assertThat(selfRefAAD).isNotEqualTo(dataPageAAD); + } +} diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/format/converter/TestParquetMetadataConverter.java b/parquet-hadoop/src/test/java/org/apache/parquet/format/converter/TestParquetMetadataConverter.java index 4d361d6aa0..31d6a7380e 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/format/converter/TestParquetMetadataConverter.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/format/converter/TestParquetMetadataConverter.java @@ -2279,4 +2279,57 @@ public void testColumnIndexNanCountsRoundTrip() { assertThat(roundTrip).isNotNull(); assertThat(roundTrip.getNanCounts()).containsExactly(1L, 0L, 0L); } + + @Test + public void testFileLogicalType() { + ParquetMetadataConverter parquetMetadataConverter = new ParquetMetadataConverter(); + + MessageType expected = Types.buildMessage() + .requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(PrimitiveTypeName.BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("uri") + .optional(PrimitiveTypeName.INT64) + .named("offset") + .optional(PrimitiveTypeName.INT64) + .named("size") + .optional(PrimitiveTypeName.BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("content_type") + .optional(PrimitiveTypeName.BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("checksum") + .optional(PrimitiveTypeName.BINARY) + .named("inline") + .named("f") + .named("example"); + + List parquetSchema = parquetMetadataConverter.toParquetSchema(expected); + MessageType schema = parquetMetadataConverter.fromParquetSchema(parquetSchema, null); + assertThat(schema).isEqualTo(expected); + LogicalTypeAnnotation logicalType = schema.getType("f").getLogicalTypeAnnotation(); + assertThat(logicalType).isInstanceOf(LogicalTypeAnnotation.FileLogicalTypeAnnotation.class); + assertThat(logicalType).isEqualTo(LogicalTypeAnnotation.fileType()); + } + + @Test + public void testFileLogicalTypeRoundTripUriOnly() { + ParquetMetadataConverter parquetMetadataConverter = new ParquetMetadataConverter(); + + MessageType expected = Types.buildMessage() + .requiredGroup() + .as(LogicalTypeAnnotation.fileType()) + .optional(PrimitiveTypeName.BINARY) + .as(LogicalTypeAnnotation.stringType()) + .named("uri") + .named("f") + .named("example"); + + List parquetSchema = parquetMetadataConverter.toParquetSchema(expected); + MessageType schema = parquetMetadataConverter.fromParquetSchema(parquetSchema, null); + assertThat(schema).isEqualTo(expected); + LogicalTypeAnnotation logicalType = schema.getType("f").getLogicalTypeAnnotation(); + assertThat(logicalType).isInstanceOf(LogicalTypeAnnotation.FileLogicalTypeAnnotation.class); + } } diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestFileValueWriter.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestFileValueWriter.java new file mode 100644 index 0000000000..5110b1b5bb --- /dev/null +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestFileValueWriter.java @@ -0,0 +1,274 @@ +/* + * 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.parquet.hadoop; + +import static org.apache.parquet.hadoop.ParquetFileWriter.Mode.CREATE; +import static org.apache.parquet.hadoop.ParquetWriter.DEFAULT_BLOCK_SIZE; +import static org.apache.parquet.hadoop.ParquetWriter.MAX_PADDING_SIZE_DEFAULT; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.Path; +import org.apache.parquet.ParquetReadOptions; +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.column.Encoding; +import org.apache.parquet.column.ParquetProperties; +import org.apache.parquet.column.statistics.Statistics; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; +import org.apache.parquet.hadoop.metadata.ColumnPath; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; +import org.apache.parquet.hadoop.util.HadoopInputFile; +import org.apache.parquet.hadoop.util.HadoopOutputFile; +import org.apache.parquet.io.InputFile; +import org.apache.parquet.io.api.Binary; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.MessageTypeParser; +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; +import org.apache.parquet.schema.Types; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +/** + * Tests that {@link FileValueWriter} routes a {@code FILE} payload to inline storage or to a + * self-reference according to the configured threshold, and that both forms describe the same logical + * bytes. + */ +public class TestFileValueWriter { + + private static final MessageType SCHEMA = MessageTypeParser.parseMessageType("message m {" + + " optional group file (FILE) {" + + " optional int64 offset;" + + " optional int64 size;" + + " optional binary inline;" + + " }" + + "}"); + + private static final ColumnDescriptor INLINE_COLUMN = SCHEMA.getColumnDescription(new String[] {"file", "inline"}); + + private static final CompressionCodecName CODEC = CompressionCodecName.SNAPPY; + + private static final Statistics EMPTY_STATS = Statistics.getBuilderForReading( + Types.required(PrimitiveTypeName.BINARY).named("inline")) + .build(); + + @TempDir + java.nio.file.Path tempDir; + + @Test + public void testDefaultThresholdIsPageSize() { + assertThat(ParquetProperties.builder().build().getFileSelfReferenceThreshold()) + .isEqualTo(ParquetProperties.DEFAULT_PAGE_SIZE); + } + + @Test + public void testThresholdMustNotBeNegative() { + assertThatThrownBy(() -> ParquetProperties.builder().withFileSelfReferenceThreshold(-1)) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + public void testPayloadAtThresholdIsInlinedAndAboveIsSelfReference() throws IOException { + Configuration conf = new Configuration(); + Path path = new Path(tempDir.resolve("routing.parquet").toUri()); + CodecFactory codecFactory = new CodecFactory(conf, DEFAULT_BLOCK_SIZE); + + int threshold = 64; + Binary atThreshold = Binary.fromConstantByteArray(payload(threshold)); + Binary aboveThreshold = Binary.fromConstantByteArray(payload(threshold + 1)); + + ParquetFileWriter writer = new ParquetFileWriter( + HadoopOutputFile.fromPath(path, conf), SCHEMA, CREATE, DEFAULT_BLOCK_SIZE, MAX_PADDING_SIZE_DEFAULT); + writer.start(); + writer.startBlock(2); + + FileValueWriter valueWriter = new FileValueWriter( + writer, codecFactory.getCompressor(CODEC), null, columnOrdinalOf(INLINE_COLUMN), threshold); + + FileValueWriter.Placement inlined = valueWriter.write(atThreshold); + FileValueWriter.Placement outOfLine = valueWriter.write(aboveThreshold); + + // A payload exactly at the threshold stays inline; one byte more goes out of line. + assertThat(inlined.isInline()).isTrue(); + assertThat(inlined.getInlineBytes()).isEqualTo(atThreshold); + assertThatThrownBy(inlined::getOffset).isInstanceOf(IllegalStateException.class); + + assertThat(outOfLine.isInline()).isFalse(); + assertThat(outOfLine.getSize()).isGreaterThan(0L); + assertThatThrownBy(outOfLine::getInlineBytes).isInstanceOf(IllegalStateException.class); + + // Write the inline column chunk so the reader has metadata carrying the inherited codec. + writer.startColumn(INLINE_COLUMN, 1, CODEC); + writer.writeDataPage( + 1, + (int) inlined.getInlineBytes().length(), + codecFactory + .getCompressor(CODEC) + .compress(BytesInput.from(inlined.getInlineBytes().toByteBuffer())), + EMPTY_STATS, + Encoding.BIT_PACKED, + Encoding.BIT_PACKED, + Encoding.PLAIN); + writer.endColumn(); + writer.endBlock(); + writer.end(new java.util.HashMap<>()); + + // The out-of-line payload resolves back to the original bytes. + InputFile inputFile = HadoopInputFile.fromPath(path, conf); + try (ParquetFileReader reader = + ParquetFileReader.open(inputFile, ParquetReadOptions.builder().build())) { + BlockMetaData block = reader.getFooter().getBlocks().get(0); + ColumnChunkMetaData inlineMeta = findColumn(block, INLINE_COLUMN); + BytesInput resolved = reader.resolveSelfReference(inlineMeta, outOfLine.getOffset(), outOfLine.getSize()); + assertThat(resolved.toByteArray()).isEqualTo(aboveThreshold.getBytes()); + } + codecFactory.release(); + } + + @Test + public void testZeroThresholdAlwaysUsesSelfReferences() throws IOException { + Configuration conf = new Configuration(); + Path path = new Path(tempDir.resolve("always_out_of_line.parquet").toUri()); + CodecFactory codecFactory = new CodecFactory(conf, DEFAULT_BLOCK_SIZE); + + byte[][] payloads = {payload(1), payload(1000)}; + + ParquetFileWriter writer = new ParquetFileWriter( + HadoopOutputFile.fromPath(path, conf), SCHEMA, CREATE, DEFAULT_BLOCK_SIZE, MAX_PADDING_SIZE_DEFAULT); + writer.start(); + writer.startBlock(payloads.length); + + FileValueWriter valueWriter = + new FileValueWriter(writer, codecFactory.getCompressor(CODEC), null, columnOrdinalOf(INLINE_COLUMN), 0); + + List placements = new ArrayList<>(); + for (byte[] p : payloads) { + placements.add(valueWriter.write(Binary.fromConstantByteArray(p))); + } + // Every payload went out of line, including the single-byte one. An empty payload would still be + // inlined, since its length is not greater than the threshold. + assertThat(placements).allMatch(p -> !p.isInline()); + + writer.startColumn(INLINE_COLUMN, 0, CODEC); + writer.writeDataPage( + 0, + 0, + codecFactory.getCompressor(CODEC).compress(BytesInput.empty()), + EMPTY_STATS, + Encoding.BIT_PACKED, + Encoding.BIT_PACKED, + Encoding.PLAIN); + writer.endColumn(); + writer.endBlock(); + writer.end(new java.util.HashMap<>()); + + InputFile inputFile = HadoopInputFile.fromPath(path, conf); + try (ParquetFileReader reader = + ParquetFileReader.open(inputFile, ParquetReadOptions.builder().build())) { + BlockMetaData block = reader.getFooter().getBlocks().get(0); + ColumnChunkMetaData inlineMeta = findColumn(block, INLINE_COLUMN); + for (int i = 0; i < payloads.length; i++) { + FileValueWriter.Placement placement = placements.get(i); + BytesInput resolved = + reader.resolveSelfReference(inlineMeta, placement.getOffset(), placement.getSize()); + assertThat(resolved.toByteArray()).isEqualTo(payloads[i]); + } + } + codecFactory.release(); + } + + @Test + public void testMaxThresholdAlwaysInlines() throws IOException { + Configuration conf = new Configuration(); + Path path = new Path(tempDir.resolve("always_inline.parquet").toUri()); + CodecFactory codecFactory = new CodecFactory(conf, DEFAULT_BLOCK_SIZE); + + ParquetFileWriter writer = new ParquetFileWriter( + HadoopOutputFile.fromPath(path, conf), SCHEMA, CREATE, DEFAULT_BLOCK_SIZE, MAX_PADDING_SIZE_DEFAULT); + writer.start(); + writer.startBlock(1); + + FileValueWriter valueWriter = new FileValueWriter( + writer, codecFactory.getCompressor(CODEC), null, columnOrdinalOf(INLINE_COLUMN), Integer.MAX_VALUE); + + long posBefore = writer.getPos(); + FileValueWriter.Placement placement = valueWriter.write(Binary.fromConstantByteArray(payload(1 << 20))); + + assertThat(placement.isInline()).isTrue(); + // Nothing was written to the file body, because the payload is carried by the value itself. + assertThat(writer.getPos()).isEqualTo(posBefore); + + writer.abort(); + codecFactory.release(); + } + + @Test + public void testNullPayloadIsRejected() throws IOException { + Configuration conf = new Configuration(); + Path path = new Path(tempDir.resolve("null_payload.parquet").toUri()); + CodecFactory codecFactory = new CodecFactory(conf, DEFAULT_BLOCK_SIZE); + + ParquetFileWriter writer = new ParquetFileWriter( + HadoopOutputFile.fromPath(path, conf), SCHEMA, CREATE, DEFAULT_BLOCK_SIZE, MAX_PADDING_SIZE_DEFAULT); + writer.start(); + writer.startBlock(1); + + FileValueWriter valueWriter = new FileValueWriter( + writer, codecFactory.getCompressor(CODEC), null, columnOrdinalOf(INLINE_COLUMN), 64); + assertThatThrownBy(() -> valueWriter.write(null)).isInstanceOf(IllegalArgumentException.class); + + writer.abort(); + codecFactory.release(); + } + + private static int columnOrdinalOf(ColumnDescriptor column) { + List columns = SCHEMA.getColumns(); + for (int i = 0; i < columns.size(); i++) { + if (columns.get(i).equals(column)) { + return i; + } + } + throw new IllegalStateException("Column not found in schema: " + column); + } + + private static ColumnChunkMetaData findColumn(BlockMetaData block, ColumnDescriptor column) { + ColumnPath target = ColumnPath.get(column.getPath()); + for (ColumnChunkMetaData meta : block.getColumns()) { + if (meta.getPath().equals(target)) { + return meta; + } + } + throw new IllegalStateException("Column chunk not found: " + target); + } + + private static byte[] payload(int length) { + StringBuilder sb = new StringBuilder(); + while (sb.length() < length) { + sb.append("file-payload-"); + } + return sb.substring(0, length).getBytes(StandardCharsets.UTF_8); + } +} diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestSelfReferenceFileWrite.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestSelfReferenceFileWrite.java new file mode 100644 index 0000000000..c1ebcdb581 --- /dev/null +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestSelfReferenceFileWrite.java @@ -0,0 +1,216 @@ +/* + * 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.parquet.hadoop; + +import static org.apache.parquet.column.Encoding.BIT_PACKED; +import static org.apache.parquet.column.Encoding.PLAIN; +import static org.apache.parquet.hadoop.ParquetFileWriter.Mode.CREATE; +import static org.apache.parquet.hadoop.ParquetWriter.DEFAULT_BLOCK_SIZE; +import static org.apache.parquet.hadoop.ParquetWriter.MAX_PADDING_SIZE_DEFAULT; +import static org.assertj.core.api.Assertions.assertThat; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.Path; +import org.apache.parquet.ParquetReadOptions; +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.column.page.DataPage; +import org.apache.parquet.column.page.DataPageV1; +import org.apache.parquet.column.page.PageReadStore; +import org.apache.parquet.column.page.PageReader; +import org.apache.parquet.column.statistics.Statistics; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; +import org.apache.parquet.hadoop.metadata.ParquetMetadata; +import org.apache.parquet.hadoop.util.HadoopInputFile; +import org.apache.parquet.hadoop.util.HadoopOutputFile; +import org.apache.parquet.io.InputFile; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.MessageTypeParser; +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; +import org.apache.parquet.schema.Types; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +/** + * End-to-end test that writes normal column data and FILE self-reference payloads into the + * same Parquet file with {@link ParquetFileWriter}, then reopens it and both reads the data pages + * back and resolves the self-references via {@link ParquetFileReader#resolveSelfReference}. This + * exercises the interaction between the storage-inheritance APIs and the ordinary file-write path: + * self-reference payloads are written into the file body while a block is open, and the + * {@code offset}/{@code size} they return would be recorded in the {@code offset} and {@code size} + * columns of a FILE group. + */ +public class TestSelfReferenceFileWrite { + + // A FILE group whose values are self-references: the inline column supplies the codec/encryption + // reference point, and offset/size locate the stored payload within this file. + private static final MessageType SCHEMA = MessageTypeParser.parseMessageType("message m {" + + " required int64 id;" + + " optional group file (FILE) {" + + " optional int64 offset;" + + " optional int64 size;" + + " optional binary inline;" + + " }" + + "}"); + + private static final ColumnDescriptor ID_COLUMN = SCHEMA.getColumnDescription(new String[] {"id"}); + // The inline column is the storage-inheritance reference point for the FILE group. + private static final ColumnDescriptor INLINE_COLUMN = SCHEMA.getColumnDescription(new String[] {"file", "inline"}); + + private static final CompressionCodecName CODEC = CompressionCodecName.SNAPPY; + + private static final Statistics EMPTY_STATS = Statistics.getBuilderForReading( + Types.required(PrimitiveTypeName.INT64).named("id")) + .build(); + + @TempDir + java.nio.file.Path tempDir; + + @Test + public void testWriteDataAlongsideSelfReferences() throws IOException { + Configuration conf = new Configuration(); + Path path = new Path(tempDir.resolve("self_ref.parquet").toUri()); + + // Payloads that will be stored as self-references, inheriting the SNAPPY codec of the inline + // column. Made highly compressible so the stored size differs from the resolved size. + byte[][] payloads = { + repeat("hello self-reference ", 200), + repeat("second blob ", 400), + new byte[0], // empty payload is a valid self-reference + }; + + byte[] idPageBytes = {0, 1, 2, 3, 4, 5, 6, 7}; + + CodecFactory codecFactory = new CodecFactory(conf, DEFAULT_BLOCK_SIZE); + List ranges = new ArrayList<>(); + + ParquetFileWriter writer = new ParquetFileWriter( + HadoopOutputFile.fromPath(path, conf), SCHEMA, CREATE, DEFAULT_BLOCK_SIZE, MAX_PADDING_SIZE_DEFAULT); + + writer.start(); + writer.startBlock(payloads.length); + + // Self-references are written into the file body while the block is open. In a real writer + // these would be interleaved with the column data; the returned offset/size feed the FILE + // group's offset/size columns. + int inlineColumnOrdinal = columnOrdinalOf(INLINE_COLUMN); + for (int i = 0; i < payloads.length; i++) { + ranges.add(writer.writeSelfReference( + BytesInput.from(payloads[i]), + codecFactory.getCompressor(CODEC), + null, // unencrypted file + inlineColumnOrdinal)); + } + + // Write a normal data page for the id column in the same block. + writer.startColumn(ID_COLUMN, 4, CompressionCodecName.UNCOMPRESSED); + writer.writeDataPage(4, idPageBytes.length, BytesInput.from(idPageBytes), EMPTY_STATS, PLAIN, PLAIN, PLAIN); + writer.endColumn(); + + // Write the inline column chunk so the reader has a ColumnChunkMetaData carrying the SNAPPY + // codec that the self-references inherit. (The inline values themselves are empty here because + // the payload lives in the self-reference blocks.) + writer.startColumn(INLINE_COLUMN, 0, CODEC); + BytesInput emptyInline = codecFactory.getCompressor(CODEC).compress(BytesInput.empty()); + writer.writeDataPage(0, 0, emptyInline, EMPTY_STATS, BIT_PACKED, BIT_PACKED, PLAIN); + writer.endColumn(); + + writer.endBlock(); + writer.end(new java.util.HashMap<>()); + + // The stored ranges are non-overlapping and ordered as written. + assertThat(ranges.get(0).getOffset()).isLessThan(ranges.get(1).getOffset()); + assertThat(ranges.get(0).getOffset() + ranges.get(0).getSize()) + .isLessThanOrEqualTo(ranges.get(1).getOffset()); + + // Reopen and verify both the data page and the self-references coexist and resolve correctly. + InputFile inputFile = HadoopInputFile.fromPath(path, conf); + ParquetReadOptions options = ParquetReadOptions.builder().build(); + try (ParquetFileReader reader = ParquetFileReader.open(inputFile, options)) { + ParquetMetadata footer = reader.getFooter(); + assertThat(footer.getBlocks()).hasSize(1); + BlockMetaData block = footer.getBlocks().get(0); + + // The normal id column reads back exactly as written. + ColumnChunkMetaData inlineMeta = findColumn(block, INLINE_COLUMN); + assertThat(inlineMeta.getCodec()).isEqualTo(CODEC); + + try (ParquetFileReader dataReader = ParquetFileReader.open(inputFile, options)) { + PageReadStore pages = dataReader.readNextRowGroup(); + PageReader idPages = pages.getPageReader(ID_COLUMN); + DataPage idPage = idPages.readPage(); + assertThat(((DataPageV1) idPage).getBytes().toByteArray()).isEqualTo(idPageBytes); + } + + // Each self-reference resolves back to its original payload, inheriting the inline column's + // codec. Only the offset and size recorded in the value are needed -- no per-value counter, so + // resolution does not depend on having read the preceding values. + for (int i = 0; i < payloads.length; i++) { + SelfReferenceStorage.StoredRange range = ranges.get(i); + BytesInput resolved = reader.resolveSelfReference(inlineMeta, range.getOffset(), range.getSize()); + assertThat(resolved.toByteArray()).isEqualTo(payloads[i]); + } + + // Resolution order is irrelevant, which is the point of keying on the offset. + for (int i = payloads.length - 1; i >= 0; i--) { + SelfReferenceStorage.StoredRange range = ranges.get(i); + BytesInput resolved = reader.resolveSelfReference(inlineMeta, range.getOffset(), range.getSize()); + assertThat(resolved.toByteArray()).isEqualTo(payloads[i]); + } + } + + codecFactory.release(); + } + + private static int columnOrdinalOf(ColumnDescriptor column) { + List columns = SCHEMA.getColumns(); + for (int i = 0; i < columns.size(); i++) { + if (columns.get(i).equals(column)) { + return i; + } + } + throw new IllegalStateException("Column not found in schema: " + column); + } + + private static ColumnChunkMetaData findColumn(BlockMetaData block, ColumnDescriptor column) { + org.apache.parquet.hadoop.metadata.ColumnPath target = + org.apache.parquet.hadoop.metadata.ColumnPath.get(column.getPath()); + for (ColumnChunkMetaData meta : block.getColumns()) { + if (meta.getPath().equals(target)) { + return meta; + } + } + throw new IllegalStateException("Column chunk not found: " + target); + } + + private static byte[] repeat(String token, int times) { + StringBuilder sb = new StringBuilder(token.length() * times); + for (int i = 0; i < times; i++) { + sb.append(token); + } + return sb.toString().getBytes(StandardCharsets.UTF_8); + } +} diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestSelfReferenceStorage.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestSelfReferenceStorage.java new file mode 100644 index 0000000000..a17daffee0 --- /dev/null +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestSelfReferenceStorage.java @@ -0,0 +1,308 @@ +/* + * 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.parquet.hadoop; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.util.Arrays; +import org.apache.hadoop.conf.Configuration; +import org.apache.parquet.bytes.BytesInput; +import org.apache.parquet.crypto.AesCipher; +import org.apache.parquet.crypto.AesMode; +import org.apache.parquet.crypto.ModuleCipherFactory; +import org.apache.parquet.crypto.ParquetCryptoRuntimeException; +import org.apache.parquet.format.BlockCipher; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; +import org.apache.parquet.io.PositionOutputStream; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; + +/** + * Round-trip tests for {@link SelfReferenceStorage}, the storage-inheritance engine for FILE + * self-references (parquet-format PR #603). Each test writes a stored representation and resolves it + * back, asserting the resolved bytes equal the original and that the recorded {@code offset}/ + * {@code size} cover exactly the stored bytes. + */ +public class TestSelfReferenceStorage { + + private static final int PAGE_SIZE = 64 * 1024; + // A 32-byte AES key. + private static final byte[] COLUMN_KEY = "0123456789012345".getBytes(); + private static final byte[] FILE_AAD = "unique-file-aad!".getBytes(); + + /** A simple in-memory {@link PositionOutputStream} for capturing written bytes. */ + private static final class InMemoryPositionOutputStream extends PositionOutputStream { + private final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + + @Override + public long getPos() { + return baos.size(); + } + + @Override + public void write(int b) { + baos.write(b); + } + + @Override + public void write(byte[] b, int off, int len) { + baos.write(b, off, len); + } + + byte[] toByteArray() { + return baos.toByteArray(); + } + } + + @ParameterizedTest + @EnumSource( + value = CompressionCodecName.class, + names = {"UNCOMPRESSED", "SNAPPY", "GZIP", "ZSTD", "LZ4_RAW"}) + public void testRoundTripUnencrypted(CompressionCodecName codec) throws IOException { + CodecFactory codecFactory = new CodecFactory(new Configuration(), PAGE_SIZE); + byte[] resolved = highlyCompressiblePayload(4096); + + InMemoryPositionOutputStream out = new InMemoryPositionOutputStream(); + // Simulate a leading byte already in the file, so offset is non-zero. + out.write(new byte[] {(byte) 0xAB}, 0, 1); + + SelfReferenceStorage.StoredRange range = SelfReferenceStorage.write( + BytesInput.from(resolved), codecFactory.getCompressor(codec), null, null, 0, 0, out); + + byte[] fileBytes = out.toByteArray(); + assertThat(range.getOffset()).isEqualTo(1L); + assertThat(range.getSize()).isEqualTo(fileBytes.length - 1L); + if (codec == CompressionCodecName.UNCOMPRESSED) { + // Uncompressed: the stored bytes are exactly the resolved bytes. + assertThat(range.getSize()).isEqualTo((long) resolved.length); + } + + byte[] stored = + Arrays.copyOfRange(fileBytes, (int) range.getOffset(), (int) (range.getOffset() + range.getSize())); + BytesInput resolvedBack = + SelfReferenceStorage.resolve(BytesInput.from(stored), codec, codecFactory, null, null, 0, 0, 0L); + + assertThat(resolvedBack.toByteArray()).isEqualTo(resolved); + codecFactory.release(); + } + + /** + * The decompressed size of a self-reference is not stored, so the reader grows its output buffer + * until the payload fits. This exercises payload sizes spanning several doublings, including sizes + * that are exact powers of two, where a full output buffer is ambiguous between "complete" and + * "truncated". + */ + @ParameterizedTest + @EnumSource( + value = CompressionCodecName.class, + names = {"SNAPPY", "GZIP", "ZSTD", "LZ4_RAW"}) + public void testRoundTripAcrossBufferGrowth(CompressionCodecName codec) throws IOException { + CodecFactory codecFactory = new CodecFactory(new Configuration(), PAGE_SIZE); + int[] sizes = {1, 8192, 8193, 16384, 100_000, 1 << 20}; + + for (int size : sizes) { + byte[] resolved = highlyCompressiblePayload(size); + InMemoryPositionOutputStream out = new InMemoryPositionOutputStream(); + SelfReferenceStorage.StoredRange range = SelfReferenceStorage.write( + BytesInput.from(resolved), codecFactory.getCompressor(codec), null, null, 0, 0, out); + + byte[] stored = out.toByteArray(); + assertThat(range.getSize()).isEqualTo((long) stored.length); + BytesInput resolvedBack = + SelfReferenceStorage.resolve(BytesInput.from(stored), codec, codecFactory, null, null, 0, 0, 0L); + assertThat(resolvedBack.toByteArray()) + .as("payload of %s bytes", size) + .isEqualTo(resolved); + } + codecFactory.release(); + } + + /** + * Incompressible data expands slightly under most codecs, so the initial guess of twice the + * compressed size is generous; this simply confirms such payloads round-trip too. + */ + @ParameterizedTest + @EnumSource( + value = CompressionCodecName.class, + names = {"SNAPPY", "GZIP", "ZSTD", "LZ4_RAW"}) + public void testRoundTripIncompressiblePayload(CompressionCodecName codec) throws IOException { + CodecFactory codecFactory = new CodecFactory(new Configuration(), PAGE_SIZE); + byte[] resolved = new byte[64 * 1024]; + new java.util.Random(42).nextBytes(resolved); + + InMemoryPositionOutputStream out = new InMemoryPositionOutputStream(); + SelfReferenceStorage.StoredRange range = SelfReferenceStorage.write( + BytesInput.from(resolved), codecFactory.getCompressor(codec), null, null, 0, 0, out); + + byte[] stored = out.toByteArray(); + assertThat(range.getSize()).isEqualTo((long) stored.length); + BytesInput resolvedBack = + SelfReferenceStorage.resolve(BytesInput.from(stored), codec, codecFactory, null, null, 0, 0, 0L); + assertThat(resolvedBack.toByteArray()).isEqualTo(resolved); + codecFactory.release(); + } + + @ParameterizedTest + @EnumSource(value = AesMode.class) + public void testRoundTripEncrypted(AesMode mode) throws IOException { + CodecFactory codecFactory = new CodecFactory(new Configuration(), PAGE_SIZE); + byte[] resolved = highlyCompressiblePayload(4096); + CompressionCodecName codec = CompressionCodecName.SNAPPY; + + BlockCipher.Encryptor encryptor = ModuleCipherFactory.getEncryptor(mode, COLUMN_KEY); + + InMemoryPositionOutputStream out = new InMemoryPositionOutputStream(); + SelfReferenceStorage.StoredRange range = SelfReferenceStorage.write( + BytesInput.from(resolved), codecFactory.getCompressor(codec), encryptor, FILE_AAD, 1, 2, out); + + byte[] fileBytes = out.toByteArray(); + assertThat(range.getOffset()).isEqualTo(0L); + assertThat(range.getSize()).isEqualTo((long) fileBytes.length); + // The stored module carries the 4-byte length prefix and 12-byte nonce (and a 16-byte GCM tag + // for GCM), so it is larger than the raw compressed payload. + int expectedOverhead = AesCipher.NONCE_LENGTH + 4 + (mode == AesMode.GCM ? AesCipher.GCM_TAG_LENGTH : 0); + assertThat(range.getSize()).isGreaterThan((long) expectedOverhead); + + BlockCipher.Decryptor decryptor = ModuleCipherFactory.getDecryptor(mode, COLUMN_KEY); + byte[] stored = + Arrays.copyOfRange(fileBytes, (int) range.getOffset(), (int) (range.getOffset() + range.getSize())); + BytesInput resolvedBack = SelfReferenceStorage.resolve( + BytesInput.from(stored), codec, codecFactory, decryptor, FILE_AAD, 1, 2, range.getOffset()); + + assertThat(resolvedBack.toByteArray()).isEqualTo(resolved); + codecFactory.release(); + } + + /** + * The AAD binds a stored representation to its offset, so resolving the same bytes as if they lived + * at a different offset must fail rather than silently return data. For GCM the tag check catches + * it; CTR has no tag, so it yields garbage instead -- either way the bytes must not come back + * intact. + */ + @Test + public void testResolveWithWrongOffsetDoesNotReturnPayload() throws IOException { + CodecFactory codecFactory = new CodecFactory(new Configuration(), PAGE_SIZE); + byte[] resolved = highlyCompressiblePayload(4096); + CompressionCodecName codec = CompressionCodecName.SNAPPY; + + BlockCipher.Encryptor encryptor = ModuleCipherFactory.getEncryptor(AesMode.GCM, COLUMN_KEY); + InMemoryPositionOutputStream out = new InMemoryPositionOutputStream(); + SelfReferenceStorage.StoredRange range = SelfReferenceStorage.write( + BytesInput.from(resolved), codecFactory.getCompressor(codec), encryptor, FILE_AAD, 1, 2, out); + + byte[] stored = out.toByteArray(); + BlockCipher.Decryptor decryptor = ModuleCipherFactory.getDecryptor(AesMode.GCM, COLUMN_KEY); + assertThatThrownBy(() -> SelfReferenceStorage.resolve( + BytesInput.from(stored), codec, codecFactory, decryptor, FILE_AAD, 1, 2, range.getOffset() + 1)) + .isInstanceOf(ParquetCryptoRuntimeException.class); + codecFactory.release(); + } + + /** + * Two self-references with identical payloads in the same column chunk sit at different offsets, so + * their AADs differ and their ciphertexts must not be interchangeable. + */ + @Test + public void testIdenticalPayloadsAtDifferentOffsetsAreNotInterchangeable() throws IOException { + CodecFactory codecFactory = new CodecFactory(new Configuration(), PAGE_SIZE); + byte[] resolved = highlyCompressiblePayload(1024); + CompressionCodecName codec = CompressionCodecName.SNAPPY; + + BlockCipher.Encryptor encryptor = ModuleCipherFactory.getEncryptor(AesMode.GCM, COLUMN_KEY); + InMemoryPositionOutputStream out = new InMemoryPositionOutputStream(); + SelfReferenceStorage.StoredRange first = SelfReferenceStorage.write( + BytesInput.from(resolved), codecFactory.getCompressor(codec), encryptor, FILE_AAD, 1, 2, out); + SelfReferenceStorage.StoredRange second = SelfReferenceStorage.write( + BytesInput.from(resolved), codecFactory.getCompressor(codec), encryptor, FILE_AAD, 1, 2, out); + + assertThat(second.getOffset()).isGreaterThan(first.getOffset()); + + byte[] fileBytes = out.toByteArray(); + byte[] firstStored = + Arrays.copyOfRange(fileBytes, (int) first.getOffset(), (int) (first.getOffset() + first.getSize())); + + // The first block's bytes cannot be resolved at the second block's offset. + BlockCipher.Decryptor decryptor = ModuleCipherFactory.getDecryptor(AesMode.GCM, COLUMN_KEY); + assertThatThrownBy(() -> SelfReferenceStorage.resolve( + BytesInput.from(firstStored), + codec, + codecFactory, + decryptor, + FILE_AAD, + 1, + 2, + second.getOffset())) + .isInstanceOf(ParquetCryptoRuntimeException.class); + codecFactory.release(); + } + + @Test + public void testEmptyPayloadRoundTrip() throws IOException { + CodecFactory codecFactory = new CodecFactory(new Configuration(), PAGE_SIZE); + byte[] resolved = new byte[0]; + + InMemoryPositionOutputStream out = new InMemoryPositionOutputStream(); + SelfReferenceStorage.StoredRange range = SelfReferenceStorage.write( + BytesInput.from(resolved), + codecFactory.getCompressor(CompressionCodecName.UNCOMPRESSED), + null, + null, + 0, + 0, + out); + + assertThat(range.getSize()).isEqualTo(0L); + BytesInput resolvedBack = SelfReferenceStorage.resolve( + BytesInput.from(new byte[0]), CompressionCodecName.UNCOMPRESSED, codecFactory, null, null, 0, 0, 0L); + assertThat(resolvedBack.toByteArray()).isEqualTo(resolved); + codecFactory.release(); + } + + /** An empty payload round-trips through a real codec too, not only UNCOMPRESSED. */ + @ParameterizedTest + @EnumSource( + value = CompressionCodecName.class, + names = {"SNAPPY", "GZIP", "ZSTD", "LZ4_RAW"}) + public void testEmptyPayloadRoundTripCompressed(CompressionCodecName codec) throws IOException { + CodecFactory codecFactory = new CodecFactory(new Configuration(), PAGE_SIZE); + InMemoryPositionOutputStream out = new InMemoryPositionOutputStream(); + SelfReferenceStorage.StoredRange range = SelfReferenceStorage.write( + BytesInput.from(new byte[0]), codecFactory.getCompressor(codec), null, null, 0, 0, out); + + byte[] stored = out.toByteArray(); + assertThat(range.getSize()).isEqualTo((long) stored.length); + BytesInput resolvedBack = + SelfReferenceStorage.resolve(BytesInput.from(stored), codec, codecFactory, null, null, 0, 0, 0L); + assertThat(resolvedBack.toByteArray()).isEmpty(); + codecFactory.release(); + } + + private static byte[] highlyCompressiblePayload(int length) { + byte[] payload = new byte[length]; + for (int i = 0; i < length; i++) { + payload[i] = (byte) (i % 16); + } + return payload; + } +}