Skip to content

[common] Extract a variant into a MAP whose STRING key is NOT NULL #10876

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
thswlsqls wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[common] Extract a variant into a MAP whose STRING key is NOT NULL #10876

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
8 changes: 2 additions & 6 deletions paimon-common/src/main/java/org/apache/paimon/data/variant/BaseVariantReader.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,12 +154,8 @@ public static BaseVariantReader create(
} else if (targetType instanceof ArrayType) {
return new ArrayReader(schema, (ArrayType) targetType, castArgs);
} else if (targetType instanceof MapType
&& ((MapType) targetType).getKeyType().equals(DataTypes.STRING())) {
if (((MapType) targetType).getKeyType().equals(DataTypes.STRING())) {
return new MapReader(schema, (MapType) targetType, castArgs);
} else {
throw new UnsupportedOperationException();
}
&& ((MapType) targetType).getKeyType().equalsIgnoreNullable(DataTypes.STRING())) {
return new MapReader(schema, (MapType) targetType, castArgs);
} else if (targetType instanceof VariantType) {
return new VariantReader(
schema, (VariantType) targetType, castArgs, isTopLevelUnshredded);
Expand Down
2 changes: 1 addition & 1 deletion paimon-common/src/main/java/org/apache/paimon/data/variant/VariantGet.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ public static Object cast(GenericVariant v, DataType dataType, VariantCastArgs c
} else if (dataType instanceof MapType) {
MapType mapType = (MapType) dataType;
DataType valueType = mapType.getValueType();
if (mapType.getKeyType().equals(DataTypes.STRING())) {
if (mapType.getKeyType().equalsIgnoreNullable(DataTypes.STRING())) {
if (variantType == Type.OBJECT) {
int size = v.objectSize();
HashMap<BinaryString, Object> map = new HashMap<>();
Expand Down
27 changes: 27 additions & 0 deletions paimon-common/src/test/java/org/apache/paimon/data/variant/BaseVariantReaderTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,15 +19,19 @@
package org.apache.paimon.data.variant;

import org.apache.paimon.data.BinaryString;
import org.apache.paimon.data.GenericMap;
import org.apache.paimon.data.GenericRow;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.types.DataType;
import org.apache.paimon.types.DataTypes;
import org.apache.paimon.types.MapType;
import org.apache.paimon.types.RowType;

import org.junit.jupiter.api.Test;

import java.time.ZoneOffset;
import java.util.HashMap;
import java.util.Map;

import static org.apache.paimon.data.variant.PaimonShreddingUtils.buildVariantSchema;
import static org.apache.paimon.data.variant.PaimonShreddingUtils.castShredded;
Expand Down Expand Up @@ -63,4 +67,27 @@ void testRowReaderWithUnshreddedTargetField() {
assertThat(reader.read(shredded, shredded.getBinary(schema.topLevelMetadataIdx)))
.isEqualTo(GenericRow.of(27, null));
}

@Test
void testMapReaderWithNotNullStringKey() {
// The file shreds the object, so the map has to be read from the typed fields.
RowType shreddedType =
RowType.of(
new DataType[] {DataTypes.INT(), DataTypes.INT()}, new String[] {"a", "b"});
VariantSchema schema = buildVariantSchema(variantShreddingSchema(shreddedType));
MapType targetType = new MapType(DataTypes.STRING().notNull(), DataTypes.INT());
InternalRow shredded =
castShredded(GenericVariant.fromJson("{\"a\": 1, \"b\": 2}"), schema);

Map<BinaryString, Integer> expected = new HashMap<>();
expected.put(BinaryString.fromString("a"), 1);
expected.put(BinaryString.fromString("b"), 2);
for (boolean failOnError : new boolean[] {false, true}) {
VariantCastArgs castArgs = new VariantCastArgs(failOnError, ZoneOffset.UTC);
BaseVariantReader reader =
BaseVariantReader.create(schema, targetType, castArgs, false);
assertThat(reader.read(shredded, shredded.getBinary(schema.topLevelMetadataIdx)))
.isEqualTo(new GenericMap(expected));
}
}
}
33 changes: 33 additions & 0 deletions paimon-common/src/test/java/org/apache/paimon/data/variant/GenericVariantTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,13 @@
import org.apache.paimon.data.BinaryString;
import org.apache.paimon.data.Decimal;
import org.apache.paimon.data.GenericArray;
import org.apache.paimon.data.GenericMap;
import org.apache.paimon.data.GenericRow;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.data.Timestamp;
import org.apache.paimon.types.DataType;
import org.apache.paimon.types.DataTypes;
import org.apache.paimon.types.MapType;
import org.apache.paimon.types.RowType;
import org.apache.paimon.utils.InstantiationUtil;

Expand Down Expand Up @@ -282,6 +284,37 @@ public void testVariantGet() {
assertThat(variant.variantGet("$.nullField", DataTypes.BOOLEAN(), castArgs)).isNull();
}

@Test
public void testVariantGetMapWithNotNullStringKey() {
// Spark hands over map key types as NOT NULL, so a STRING key must match whatever its
// nullability is.
Variant variant =
GenericVariant.fromJson("{\"m\": {\"a\": 1, \"b\": 2}, \"arr\": [{\"a\": 1}]}");
VariantCastArgs tryCast = new VariantCastArgs(false, ZoneOffset.UTC);
VariantCastArgs strict = new VariantCastArgs(true, ZoneOffset.UTC);
MapType notNullKey = new MapType(DataTypes.STRING().notNull(), DataTypes.INT());

Map<BinaryString, Integer> expected = new HashMap<>();
expected.put(BinaryString.fromString("a"), 1);
expected.put(BinaryString.fromString("b"), 2);
assertThat(variant.variantGet("$.m", notNullKey, tryCast))
.isEqualTo(new GenericMap(expected));
assertThat(variant.variantGet("$.m", notNullKey, strict))
.isEqualTo(new GenericMap(expected));

Map<BinaryString, Integer> element = new HashMap<>();
element.put(BinaryString.fromString("a"), 1);
GenericArray expectedArray = new GenericArray(new Object[] {new GenericMap(element)});
assertThat(variant.variantGet("$.arr", DataTypes.ARRAY(notNullKey), tryCast))
.isEqualTo(expectedArray);
assertThat(variant.variantGet("$.arr", DataTypes.ARRAY(notNullKey), strict))
.isEqualTo(expectedArray);

// a non string key is still an invalid cast
MapType intKey = new MapType(DataTypes.INT().notNull(), DataTypes.INT());
assertThat(variant.variantGet("$.m", intKey, tryCast)).isNull();
}

@Test
public void testVariantGetDecimalWithScaleAbovePrecision() {
// precision() counts the digits of the unscaled value, so it is below the scale for any
Expand Down
13 changes: 13 additions & 0 deletions paimon-common/src/test/java/org/apache/paimon/data/variant/PaimonShreddingUtilsTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -291,6 +291,19 @@ void testAssembleAllTypes() {
BinaryString.fromString("Apache Paimon"));
}
})));

// cast struct to map whose key is NOT NULL
fields[0] =
buildFieldsToExtract(
new MapType(DataTypes.STRING().notNull(), DataTypes.STRING()),
"$.object",
castArgs,
variantSchema);
Map<BinaryString, BinaryString> expectedMap = new HashMap<>();
expectedMap.put(BinaryString.fromString("age"), BinaryString.fromString("3"));
expectedMap.put(BinaryString.fromString("name"), BinaryString.fromString("Apache Paimon"));
assertThat(assembleVariantStruct(castShredded(v, variantSchema), variantSchema, fields))
.isEqualTo(GenericRow.of(new GenericMap(expectedMap)));
}

@Test
Expand Down
Loading