Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,8 @@
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.LocalTime;
import java.util.List;
import java.util.Map;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
Expand Down Expand Up @@ -76,6 +78,24 @@ class ParquetIcebergWriterTest {
LocalDate.ofEpochDay(1), LocalTime.ofSecondOfDay(0)
);

private static final String ID_FIELD_NAME = "id";

private static final String TAGS_FIELD_NAME = "tags";

private static final String ADDRESS_FIELD_NAME = "address";

private static final String CITY_FIELD_NAME = "city";

private static final String ATTRIBUTES_FIELD_NAME = "attributes";

private static final String ID_FIELD_VALUE = "row-1";

private static final String CITY_FIELD_VALUE = "Berlin";

private static final List<String> TAGS_FIELD_VALUE = List.of("a", "b");

private static final Map<String, String> ATTRIBUTES_FIELD_VALUE = Map.of("k", "v");

private ParquetIcebergWriter parquetIcebergWriter;

private TestRunner runner;
Expand Down Expand Up @@ -194,6 +214,43 @@ void testWriteDataFilesPartitionedTimestamp() throws IOException {
assertEquals(microsecondsExpected, partitionField);
}

@Test
void testWriteDataFilesComplexTypes() throws IOException {
runner.enableControllerService(parquetIcebergWriter);

final Types.StructType nestedStruct = Types.StructType.of(
Types.NestedField.optional(10, CITY_FIELD_NAME, Types.StringType.get())
);
final Schema schema = new Schema(
Types.NestedField.required(1, ID_FIELD_NAME, Types.StringType.get()),
Types.NestedField.optional(2, TAGS_FIELD_NAME,
Types.ListType.ofOptional(3, Types.StringType.get())),
Types.NestedField.optional(4, ADDRESS_FIELD_NAME, nestedStruct),
Types.NestedField.optional(5, ATTRIBUTES_FIELD_NAME,
Types.MapType.ofOptional(6, 7, Types.StringType.get(), Types.StringType.get()))
);
Comment on lines +224 to +231

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The field names should be declared statically and reused in both the input Schema and output row.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Moved hardcoded constants to static ones: ef3e162

final InMemoryOutputFile outputFile = new InMemoryOutputFile();
final PartitionSpec partitionSpec = PartitionSpec.unpartitioned();
setTable(schema, partitionSpec, outputFile);
when(locationProvider.newDataLocation(anyString())).thenReturn(LOCATION);

final IcebergRowWriter rowWriter = parquetIcebergWriter.getRowWriter(table);

final GenericRecord address = GenericRecord.create(nestedStruct);
address.setField(CITY_FIELD_NAME, CITY_FIELD_VALUE);

final GenericRecord row = GenericRecord.create(schema);
row.setField(ID_FIELD_NAME, ID_FIELD_VALUE);
row.setField(TAGS_FIELD_NAME, TAGS_FIELD_VALUE);
row.setField(ADDRESS_FIELD_NAME, address);
row.setField(ATTRIBUTES_FIELD_NAME, ATTRIBUTES_FIELD_VALUE);
rowWriter.write(row);

final DataFile[] dataFiles = rowWriter.dataFiles();
final byte[] serialized = outputFile.toByteArray();
assertDataFilesFound(dataFiles, serialized);
}

private void writeRow(final Schema schema, final IcebergRowWriter rowWriter) throws IOException {
final GenericRecord row = GenericRecord.create(schema);
row.setField(FIRST_FIELD_NAME, FIRST_FIELD_VALUE);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,8 @@ public DelegatedRecord(
final org.apache.nifi.serialization.record.Record record,
final Types.StructType struct
) {
this.record = RecordConverter.getConvertedRecord(Objects.requireNonNull(record));
this.struct = Objects.requireNonNull(struct);
this.record = RecordConverter.getConvertedRecord(Objects.requireNonNull(record), struct);
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@
*/
package org.apache.nifi.processors.iceberg.record;

import org.apache.nifi.serialization.record.DataType;
import org.apache.iceberg.types.Type;
import org.apache.iceberg.types.Types;
import org.apache.nifi.serialization.record.MapRecord;
import org.apache.nifi.serialization.record.Record;
import org.apache.nifi.serialization.record.RecordField;
Expand All @@ -26,6 +27,9 @@
import java.sql.Date;
import java.sql.Time;
import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
Expand All @@ -39,63 +43,107 @@ class RecordConverter {
private static final Set<RecordFieldType> CONVERSION_REQUIRED_FIELD_TYPES = Set.of(
RecordFieldType.TIMESTAMP,
RecordFieldType.DATE,
RecordFieldType.TIME
RecordFieldType.TIME,
RecordFieldType.ARRAY,
RecordFieldType.RECORD,
RecordFieldType.MAP
);

/**
* Get Converted Record with conditional handling for field values requiring translation
* Get Converted Record with recursive, schema-aware handling for field values requiring translation
*
* @param inputRecord Input Record to be converted
* @param struct Iceberg Struct Type describing the target field types (may be null for scalar-only conversion)
* @return Input Record or new Record with converted field values
*/
static Record getConvertedRecord(final Record inputRecord) {
final Record convertedRecord;

static Record getConvertedRecord(final Record inputRecord, final Types.StructType struct) {
final RecordSchema recordSchema = inputRecord.getSchema();
if (isConversionRequired(recordSchema)) {
final Map<String, Object> values = inputRecord.toMap();
convertedRecord = getConvertedRecord(recordSchema, values);
} else {
convertedRecord = inputRecord;
if (!isConversionRequired(recordSchema)) {
return inputRecord;
}

return convertedRecord;
}

private static Record getConvertedRecord(final RecordSchema recordSchema, final Map<String, Object> values) {
final Map<String, Object> values = inputRecord.toMap();
final Map<String, Object> convertedValues = new LinkedHashMap<>();

for (final Map.Entry<String, Object> entry : values.entrySet()) {
final String field = entry.getKey();
final Object value = entry.getValue();
final Object converted = getConvertedValue(value);
convertedValues.put(field, converted);
final Type fieldType = fieldType(struct, field);
convertedValues.put(field, convertValue(entry.getValue(), fieldType));
}

return new MapRecord(recordSchema, convertedValues);
}

private static Object getConvertedValue(final Object value) {
static Object convertValue(final Object value, final Type icebergType) {
return switch (value) {
// Convert java.sql types to corresponding java.time types for Apache Iceberg
case Timestamp timestamp -> timestamp.toLocalDateTime();
case Date date -> date.toLocalDate();
case Time time -> time.toLocalTime();
case null, default -> value;
// Recursively convert complex types against the matching Iceberg type
case null, default -> convertComplexValue(value, icebergType);
};
}

private static boolean isConversionRequired(final RecordSchema recordSchema) {
final List<RecordField> fields = recordSchema.getFields();
/**
* Recursively convert array, collection, nested record, and map values against the matching Iceberg type.
* The value is returned unchanged when the target Iceberg type is unknown or does not describe a complex type
* matching the value.
*/
private static Object convertComplexValue(final Object value, final Type icebergType) {
if (icebergType == null) {
return value;
}

if (icebergType.isListType()) {
final Type elementType = icebergType.asListType().elementType();
if (value instanceof Object[] array) {
return convertList(Arrays.asList(array), elementType);
}
if (value instanceof Collection<?> collection) {
return convertList(collection, elementType);
}
} else if (icebergType.isStructType() && value instanceof Record nestedRecord) {
return new DelegatedRecord(nestedRecord, icebergType.asStructType());
} else if (icebergType.isMapType() && value instanceof Map<?, ?> map) {
return convertMap(map, icebergType.asMapType());
}

for (final RecordField field : fields) {
final DataType dataType = field.getDataType();
final RecordFieldType recordFieldType = dataType.getFieldType();
return value;
}

private static List<Object> convertList(final Collection<?> collection, final Type elementType) {
final List<Object> converted = new ArrayList<>(collection.size());
for (final Object element : collection) {
converted.add(convertValue(element, elementType));
}
return converted;
}

private static Map<Object, Object> convertMap(final Map<?, ?> map, final Types.MapType mapType) {
final Map<Object, Object> converted = new LinkedHashMap<>();

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LinkedHashMap vs HashMap might be debateable?

for (final Map.Entry<?, ?> entry : map.entrySet()) {
final Object key = convertValue(entry.getKey(), mapType.keyType());
final Object mappedValue = convertValue(entry.getValue(), mapType.valueType());
converted.put(key, mappedValue);
}
return converted;
}

private static Type fieldType(final Types.StructType struct, final String fieldName) {
if (struct == null) {
return null;
}
final Types.NestedField nestedField = struct.field(fieldName);
return nestedField == null ? null : nestedField.type();
}

private static boolean isConversionRequired(final RecordSchema recordSchema) {
for (final RecordField field : recordSchema.getFields()) {
final RecordFieldType recordFieldType = field.getDataType().getFieldType();
if (CONVERSION_REQUIRED_FIELD_TYPES.contains(recordFieldType)) {
return true;
}
}

return false;
}
}
Loading
Loading