saths008 commented on code in PR #18553:
URL: https://github.com/apache/iceberg/pull/18553#discussion_r4359762743
##########
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:
added a tests
--
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]
