wombatu-kun commented on code in PR #18553:
URL: https://github.com/apache/iceberg/pull/18553#discussion_r4356497933
##########
spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/data/SparkParquetWriters.java:
##########
@@ -371,7 +372,22 @@ private UTF8StringWriter(ColumnDescriptor desc) {
@Override
public void write(int repetitionLevel, UTF8String value) {
- column.writeBinary(repetitionLevel,
Binary.fromReusedByteArray(value.getBytes()));
+ long baseOffset = value.getBaseOffset();
+ if (value.getBaseObject() instanceof byte[] bytes
+ && baseOffset >= Platform.BYTE_ARRAY_OFFSET) {
+ // Write the backing array in place. UTF8String is often a slice of a
larger array, so
+ // getBytes would copy the value out before Parquet reads it. Use the
smaller wrapper for
+ // whole-array strings to avoid extra per-value allocation.
+ column.writeBinary(
+ repetitionLevel,
+ baseOffset == Platform.BYTE_ARRAY_OFFSET && value.numBytes() ==
bytes.length
Review Comment:
`UTF8String.getBytes` already returns the backing array without copying when
the string covers the whole array, so this arm does the same thing as the
`Binary.fromReusedByteArray(value.getBytes())` call in the fallback. Could the
`if` take only `byte[]` slices and leave whole arrays to that call, which keeps
the same allocations with one branch fewer?
##########
spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/data/SparkParquetWriters.java:
##########
@@ -371,7 +372,22 @@ private UTF8StringWriter(ColumnDescriptor desc) {
@Override
public void write(int repetitionLevel, UTF8String value) {
- column.writeBinary(repetitionLevel,
Binary.fromReusedByteArray(value.getBytes()));
+ long baseOffset = value.getBaseOffset();
+ if (value.getBaseObject() instanceof byte[] bytes
+ && baseOffset >= Platform.BYTE_ARRAY_OFFSET) {
+ // Write the backing array in place. UTF8String is often a slice of a
larger array, so
+ // getBytes would copy the value out before Parquet reads it. Use the
smaller wrapper for
+ // whole-array strings to avoid extra per-value allocation.
+ column.writeBinary(
+ repetitionLevel,
+ baseOffset == Platform.BYTE_ARRAY_OFFSET && value.numBytes() ==
bytes.length
+ ? Binary.fromReusedByteArray(bytes)
+ : Binary.fromReusedByteArray(
+ bytes, (int) (baseOffset - Platform.BYTE_ARRAY_OFFSET),
value.numBytes()));
+ } else {
+ // Off-heap values must be copied onto the heap.
Review Comment:
Off-heap memory is not the only thing that reaches this copy: in Spark's
default on-heap mode, rows coming out of an in-memory sort point into `long[]`
pages, so their strings land here too. Could the comment say that anything not
backed by a `byte[]` is copied, or be dropped?
##########
spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/TestSparkParquetWriter.java:
##########
@@ -350,4 +358,123 @@ public void testNdv() throws IOException,
NoSuchFieldException, IllegalAccessExc
assertThat(bloomFilterNDV.getAsLong()).isEqualTo(expectedNdv);
}
}
+
+ /** Where the bytes handed to the writer live. */
+ private enum Storage {
+ /** An exactly sized array, which the writer may pass straight to Parquet.
*/
+ ARRAY,
+ /** A window into a larger array, as produced by {@code
UnsafeRow#getUTF8String}. */
+ SLICE,
+ /** Off-heap memory, which the writer must copy onto the heap. */
+ OFFHEAP
+ }
+
+ @ParameterizedTest(name = "storage = {0}, dictionary = {1}, dictionaryBytes
= {2}")
+ @CsvSource({
+ "ARRAY, false, 4096",
+ "ARRAY, true, 4096",
+ "ARRAY, true, 1048576",
+ "SLICE, false, 4096",
+ "SLICE, true, 4096",
+ "SLICE, true, 1048576",
+ "OFFHEAP, false, 4096",
+ "OFFHEAP, true, 4096",
+ "OFFHEAP, true, 1048576"
+ })
+ void stringsWithReusedBackingMemory(
+ Storage storage, boolean dictionaryEnabled, int dictionaryBytes) throws
IOException {
+ Schema schema = new Schema(optional(1, "text", Types.StringType.get()));
+ File testFile = File.createTempFile("junit", null, temp.toFile());
+ assertThat(testFile.delete()).isTrue();
+
+ int numRows = 128;
+ int padding = 13;
+ String[] expectedValues = new String[numRows];
+ byte[][] expectedBytes = new byte[numRows][];
+ int longestValue = 0;
+ for (int id = 0; id < numRows; id += 1) {
+ expectedValues[id] = stringValue(id);
+ expectedBytes[id] =
+ expectedValues[id] == null
+ ? new byte[0]
+ : expectedValues[id].getBytes(StandardCharsets.UTF_8);
+ longestValue = Math.max(longestValue, expectedBytes[id].length);
+ }
+
+ // Pad on both sides so a wrong offset or length reads poison rather than
running off the end.
+ int sharedSize = padding + longestValue + padding;
+ byte[] shared = storage == Storage.SLICE ? new byte[sharedSize] : null;
+ long address = storage == Storage.OFFHEAP ?
Platform.allocateMemory(sharedSize) : 0;
+
+ try (FileAppender<InternalRow> writer =
+ Parquet.write(Files.localOutput(testFile))
+ .schema(schema)
+ .set(ParquetOutputFormat.ENABLE_DICTIONARY,
Boolean.toString(dictionaryEnabled))
+ .set(PARQUET_DICT_SIZE_BYTES, Integer.toString(dictionaryBytes))
+ .createWriterFunc(
+ type ->
SparkParquetWriters.buildWriter(SparkSchemaUtil.convert(schema), type))
+ .build()) {
+ for (int id = 0; id < numRows; id += 1) {
+ byte[] bytes = expectedBytes[id];
+ UTF8String value = stringWithStorage(storage, bytes, shared, address,
padding);
+ writer.add(
+ new GenericInternalRow(new Object[] {expectedValues[id] == null ?
null : value}));
+
+ // Writers must consume or copy reused storage before the next row is
written.
+ switch (storage) {
+ case ARRAY -> Arrays.fill(bytes, (byte) '?');
+ case SLICE -> Arrays.fill(shared, (byte) '?');
+ case OFFHEAP -> Platform.setMemory(address, (byte) '?', sharedSize);
+ }
+ }
+ } finally {
+ if (storage == Storage.OFFHEAP) {
+ Platform.freeMemory(address);
+ }
+ }
+
+ try (CloseableIterable<InternalRow> reader =
Review Comment:
With dictionary encoding off, the column statistics are the only values
Parquet keeps, and the test never reads them back, so those rows cannot catch a
dropped reuse flag. Could it also compare the file's lower and upper bounds
(e.g. via `ParquetUtil.fileMetrics`) with the expected values?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
