This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new f37bbe5980f Map recommender Avro schema by original data type, not
stored type (#18970)
f37bbe5980f is described below
commit f37bbe5980f82424898bd4a8df1b4c8fc426e2da
Author: Xiang Fu <[email protected]>
AuthorDate: Sat Jul 11 03:58:43 2026 +0200
Map recommender Avro schema by original data type, not stored type (#18970)
* Map Avro schema by original data type, not stored type
`AvroSchemaUtil.toAvroSchemaJsonObject` switched on `getStoredType()`,
collapsing logical types to their physical storage: a BOOLEAN field
emitted Avro `"int"` and TIMESTAMP a plain `"long"`, so the generated
Avro schema misrepresented the column and could not round-trip back to
BOOLEAN/TIMESTAMP.
Switch on the original `DataType` instead: BOOLEAN -> Avro `boolean`,
TIMESTAMP -> `timestamp-millis` long, UUID -> `bytes` (unchanged), all
others unchanged. `AvroWriter` coerces the generator's stored int `0`/`1`
for BOOLEAN columns to a `Boolean` so the new schema serializes, with the
boolean columns resolved once per file rather than per row.
The segment-processing converters `AvroUtils.getAvroSchemaFromPinotSchema`
and `SegmentProcessorAvroUtils.convertPinotSchemaToAvroSchema`
intentionally keep the stored type because they serialize physically
stored values; this divergence is documented on `toAvroSchemaJsonObject`.
* Dedupe nullable-union unwrap: reuse AvroRecordAppender.nonNullBranch in
test
Address review nit: expose the appender's union-branch helper as
@VisibleForTesting (returning the branch Schema) and have the test's
nonNullBranch delegate to it instead of re-implementing the union walk.
---
.../recommender/data/writer/AvroWriter.java | 45 ++++++-
.../recommender/data/writer/AvroWriterTest.java | 135 +++++++++++++++++++++
.../plugin/inputformat/avro/AvroSchemaUtil.java | 31 ++++-
.../inputformat/avro/AvroSchemaUtilTest.java | 74 +++++++++++
4 files changed, 280 insertions(+), 5 deletions(-)
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/recommender/data/writer/AvroWriter.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/recommender/data/writer/AvroWriter.java
index 5b82bf10245..04a858035da 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/recommender/data/writer/AvroWriter.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/recommender/data/writer/AvroWriter.java
@@ -21,11 +21,16 @@ package org.apache.pinot.controller.recommender.data.writer;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.google.common.annotations.VisibleForTesting;
import java.io.Closeable;
import java.io.File;
import java.io.IOException;
+import java.util.HashSet;
import java.util.Map;
import java.util.Objects;
+import java.util.Set;
+import org.apache.avro.Schema.Field;
+import org.apache.avro.Schema.Type;
import org.apache.avro.file.DataFileWriter;
import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericDatumWriter;
@@ -93,10 +98,12 @@ public class AvroWriter implements Writer {
class AvroRecordAppender implements Closeable {
private final DataFileWriter<GenericData.Record> _recordWriter;
private final org.apache.avro.Schema _avroSchema;
+ private final Set<String> _booleanColumns;
public AvroRecordAppender(File file, org.apache.avro.Schema avroSchema)
throws IOException {
_avroSchema = avroSchema;
+ _booleanColumns = booleanColumns(avroSchema);
_recordWriter = new DataFileWriter<>(new
GenericDatumWriter<>(_avroSchema));
_recordWriter.create(_avroSchema, file);
}
@@ -104,12 +111,44 @@ class AvroRecordAppender implements Closeable {
public void append(Map<String, Object> record)
throws IOException {
GenericData.Record nextRecord = new GenericData.Record(_avroSchema);
- record.forEach((column, value) -> {
- nextRecord.put(column, record.get(column));
- });
+ record.forEach((column, value) -> nextRecord.put(column,
coerce(_booleanColumns.contains(column), value)));
_recordWriter.append(nextRecord);
}
+ /// The data generator emits Pinot's stored representation (an int `0`/`1`
for BOOLEAN), while the Avro schema now
+ /// declares the logical `boolean` type. Coerce those stored ints to
`Boolean` so they serialize; other values pass
+ /// through unchanged.
+ @VisibleForTesting
+ static Object coerce(boolean booleanColumn, Object value) {
+ if (booleanColumn && value instanceof Number) {
+ return ((Number) value).intValue() != 0;
+ }
+ return value;
+ }
+
+ private static Set<String> booleanColumns(org.apache.avro.Schema avroSchema)
{
+ Set<String> booleanColumns = new HashSet<>();
+ for (Field field : avroSchema.getFields()) {
+ if (nonNullBranch(field.schema()).getType() == Type.BOOLEAN) {
+ booleanColumns.add(field.name());
+ }
+ }
+ return booleanColumns;
+ }
+
+ /// Returns the non-null branch of a nullable union `["null", <type>]`, or
the schema itself if it is not a union.
+ @VisibleForTesting
+ static org.apache.avro.Schema nonNullBranch(org.apache.avro.Schema schema) {
+ if (schema.getType() == Type.UNION) {
+ for (org.apache.avro.Schema member : schema.getTypes()) {
+ if (member.getType() != Type.NULL) {
+ return member;
+ }
+ }
+ }
+ return schema;
+ }
+
@Override
public void close()
throws IOException {
diff --git
a/pinot-controller/src/test/java/org/apache/pinot/controller/recommender/data/writer/AvroWriterTest.java
b/pinot-controller/src/test/java/org/apache/pinot/controller/recommender/data/writer/AvroWriterTest.java
new file mode 100644
index 00000000000..8362c362707
--- /dev/null
+++
b/pinot-controller/src/test/java/org/apache/pinot/controller/recommender/data/writer/AvroWriterTest.java
@@ -0,0 +1,135 @@
+/**
+ * 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.pinot.controller.recommender.data.writer;
+
+import java.io.File;
+import java.nio.file.Files;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import org.apache.avro.Schema;
+import org.apache.avro.file.DataFileReader;
+import org.apache.avro.generic.GenericDatumReader;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.commons.io.FileUtils;
+import org.apache.commons.lang3.IntegerRange;
+import org.apache.pinot.controller.recommender.data.generator.DataGenerator;
+import
org.apache.pinot.controller.recommender.data.generator.DataGeneratorSpec;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.data.FieldSpec.FieldType;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+
+public class AvroWriterTest {
+ private File _baseDir;
+
+ @BeforeMethod
+ public void setUp()
+ throws Exception {
+ _baseDir = Files.createTempDirectory("avroWriterTest").toFile();
+ }
+
+ @AfterMethod
+ public void tearDown()
+ throws Exception {
+ FileUtils.deleteQuietly(_baseDir);
+ }
+
+ /// Guards that logical types round-trip through the generated Avro file:
BOOLEAN is written as an Avro boolean
+ /// (not an int) and TIMESTAMP as a timestamp-millis long. Before the fix
the schema used the stored type, so
+ /// BOOLEAN showed up as "int" and TIMESTAMP as a plain "long".
+ @Test
+ public void testLogicalTypesRoundTrip()
+ throws Exception {
+ List<String> columns = List.of("intCol", "longCol", "boolCol", "tsCol",
"strCol");
+ Map<String, DataType> dataTypes = new HashMap<>();
+ dataTypes.put("intCol", DataType.INT);
+ dataTypes.put("longCol", DataType.LONG);
+ dataTypes.put("boolCol", DataType.BOOLEAN);
+ dataTypes.put("tsCol", DataType.TIMESTAMP);
+ dataTypes.put("strCol", DataType.STRING);
+ Map<String, FieldType> fieldTypes = new HashMap<>();
+ Map<String, Integer> cardinality = new HashMap<>();
+ Map<String, Integer> length = new HashMap<>();
+ for (String column : columns) {
+ fieldTypes.put(column, FieldType.DIMENSION);
+ cardinality.put(column, 5);
+ }
+ length.put("strCol", 6);
+
+ DataGeneratorSpec spec =
+ new DataGeneratorSpec(columns, cardinality, new HashMap<String,
IntegerRange>(), new HashMap<>(),
+ new HashMap<>(), length, dataTypes, fieldTypes, new
HashMap<String, TimeUnit>(), new HashMap<>(),
+ new HashMap<>());
+ DataGenerator generator = new DataGenerator();
+ generator.init(spec);
+
+ int totalDocs = 10;
+ AvroWriterSpec writerSpec = new AvroWriterSpec(generator, _baseDir,
totalDocs, 1, 0);
+ AvroWriter writer = new AvroWriter();
+ writer.init(writerSpec);
+ writer.write();
+
+ // Schema assertions: original types are preserved, not their stored types.
+ Schema avroSchema = AvroWriter.getAvroSchema(writerSpec.getSchema());
+ assertEquals(nonNullBranch(avroSchema, "intCol").getType(),
Schema.Type.INT);
+ assertEquals(nonNullBranch(avroSchema, "longCol").getType(),
Schema.Type.LONG);
+ assertEquals(nonNullBranch(avroSchema, "boolCol").getType(),
Schema.Type.BOOLEAN);
+ Schema tsBranch = nonNullBranch(avroSchema, "tsCol");
+ assertEquals(tsBranch.getType(), Schema.Type.LONG);
+ assertEquals(tsBranch.getProp("logicalType"), "timestamp-millis");
+ assertEquals(nonNullBranch(avroSchema, "strCol").getType(),
Schema.Type.STRING);
+
+ // Value assertions: the written records serialize and read back with the
expected Java types.
+ File avroFile = new File(_baseDir, "part-0.avro");
+ assertTrue(avroFile.exists());
+ int count = 0;
+ try (DataFileReader<GenericRecord> reader = new DataFileReader<>(avroFile,
new GenericDatumReader<>())) {
+ for (GenericRecord record : reader) {
+ assertTrue(record.get("boolCol") instanceof Boolean, "boolCol should
be a Boolean");
+ assertTrue(record.get("tsCol") instanceof Long, "tsCol should be a
Long");
+ assertTrue(record.get("intCol") instanceof Integer, "intCol should be
an Integer");
+ count++;
+ }
+ }
+ assertEquals(count, totalDocs);
+ }
+
+ /// The generator emits Pinot's stored form for BOOLEAN (int 0/1); coercion
must map it to the right Boolean and
+ /// leave non-boolean columns untouched.
+ @Test
+ public void testBooleanCoercion() {
+ assertEquals(AvroRecordAppender.coerce(true, 0), Boolean.FALSE);
+ assertEquals(AvroRecordAppender.coerce(true, 1), Boolean.TRUE);
+ assertEquals(AvroRecordAppender.coerce(true, 5), Boolean.TRUE);
+ // Non-boolean columns pass through unchanged.
+ assertEquals(AvroRecordAppender.coerce(false, 1), 1);
+ assertEquals(AvroRecordAppender.coerce(true, null), null);
+ }
+
+ private static Schema nonNullBranch(Schema recordSchema, String fieldName) {
+ return
AvroRecordAppender.nonNullBranch(recordSchema.getField(fieldName).schema());
+ }
+}
diff --git
a/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtil.java
b/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtil.java
index 021b46f4cd8..52ea2240bd4 100644
---
a/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtil.java
+++
b/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtil.java
@@ -76,10 +76,21 @@ public class AvroSchemaUtil {
}
}
+ /// Builds the Avro schema JSON for a single Pinot field. Used to generate
sample Avro data from a Pinot schema
+ /// (see `AvroWriter`). Each field is emitted as a nullable union `["null",
<type>]`.
+ ///
+ /// The switch is driven by the original (logical) [DataType] rather than
the stored type, so logical types are
+ /// represented faithfully instead of collapsing to their physical storage
type: BOOLEAN maps to Avro `boolean`
+ /// and TIMESTAMP to a `timestamp-millis` long, not a plain `int`/`long`.
+ ///
+ /// This intentionally differs from the segment-processing converters
`AvroUtils.getAvroSchemaFromPinotSchema` and
+ /// `SegmentProcessorAvroUtils.convertPinotSchemaToAvroSchema`, which switch
on the stored type because they
+ /// serialize Pinot's physically-stored values (e.g. an int for BOOLEAN)
directly.
public static ObjectNode toAvroSchemaJsonObject(FieldSpec fieldSpec) {
ObjectNode jsonSchema = JsonUtils.newObjectNode();
jsonSchema.put("name", fieldSpec.getName());
- switch (fieldSpec.getDataType().getStoredType()) {
+ DataType dataType = fieldSpec.getDataType();
+ switch (dataType) {
case INT:
jsonSchema.set("type", convertStringsToJsonArray("null", "int"));
return jsonSchema;
@@ -92,15 +103,31 @@ public class AvroSchemaUtil {
case DOUBLE:
jsonSchema.set("type", convertStringsToJsonArray("null", "double"));
return jsonSchema;
+ case BOOLEAN:
+ jsonSchema.set("type", convertStringsToJsonArray("null", "boolean"));
+ return jsonSchema;
+ case TIMESTAMP:
+ // TIMESTAMP is stored as LONG millis-since-epoch; annotate the long
branch with the timestamp-millis
+ // logical type so the value stays a long but is self-describing as a
timestamp.
+ ObjectNode timestampType = JsonUtils.newObjectNode();
+ timestampType.put("type", "long");
+ timestampType.put("logicalType", "timestamp-millis");
+ ArrayNode timestampUnion = JsonUtils.newArrayNode();
+ timestampUnion.add("null");
+ timestampUnion.add(timestampType);
+ jsonSchema.set("type", timestampUnion);
+ return jsonSchema;
case STRING:
case JSON:
jsonSchema.set("type", convertStringsToJsonArray("null", "string"));
return jsonSchema;
case BYTES:
+ case UUID:
+ // UUID is a logical type stored as fixed-width 16-byte BYTES.
jsonSchema.set("type", convertStringsToJsonArray("null", "bytes"));
return jsonSchema;
default:
- throw new UnsupportedOperationException();
+ throw new UnsupportedOperationException("Unsupported data type: " +
dataType);
}
}
diff --git
a/pinot-plugins/pinot-input-format/pinot-avro-base/src/test/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtilTest.java
b/pinot-plugins/pinot-input-format/pinot-avro-base/src/test/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtilTest.java
new file mode 100644
index 00000000000..924add4bb78
--- /dev/null
+++
b/pinot-plugins/pinot-input-format/pinot-avro-base/src/test/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtilTest.java
@@ -0,0 +1,74 @@
+/**
+ * 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.pinot.plugin.inputformat.avro;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import org.apache.pinot.spi.data.DimensionFieldSpec;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertThrows;
+
+
+public class AvroSchemaUtilTest {
+
+ /// The switch must be driven by the original (logical) data type, not the
stored type. Otherwise BOOLEAN collapses
+ /// to "int" and TIMESTAMP to a plain "long", misrepresenting the column in
the generated Avro schema.
+ @Test
+ public void testToAvroSchemaJsonObjectUsesOriginalType() {
+ assertPrimitiveType(DataType.INT, "int");
+ assertPrimitiveType(DataType.LONG, "long");
+ assertPrimitiveType(DataType.FLOAT, "float");
+ assertPrimitiveType(DataType.DOUBLE, "double");
+ assertPrimitiveType(DataType.STRING, "string");
+ assertPrimitiveType(DataType.JSON, "string");
+ assertPrimitiveType(DataType.BYTES, "bytes");
+ // UUID is a logical type stored as BYTES; it must keep mapping to Avro
bytes (not fall through to default).
+ assertPrimitiveType(DataType.UUID, "bytes");
+ // Logical types must not collapse to their stored INT/LONG type.
+ assertPrimitiveType(DataType.BOOLEAN, "boolean");
+
+ JsonNode type = typeOf(DataType.TIMESTAMP);
+ assertEquals(type.get(0).asText(), "null");
+ JsonNode timestampBranch = type.get(1);
+ assertEquals(timestampBranch.get("type").asText(), "long");
+ assertEquals(timestampBranch.get("logicalType").asText(),
"timestamp-millis");
+ }
+
+ /// Types with no Avro mapping (e.g. BIG_DECIMAL) must be rejected rather
than silently mishandled.
+ @Test
+ public void testToAvroSchemaJsonObjectRejectsUnsupportedType() {
+ assertThrows(UnsupportedOperationException.class,
+ () -> AvroSchemaUtil.toAvroSchemaJsonObject(new
DimensionFieldSpec("col", DataType.BIG_DECIMAL, true)));
+ }
+
+ private static void assertPrimitiveType(DataType dataType, String
expectedAvroType) {
+ JsonNode type = typeOf(dataType);
+ assertEquals(type.get(0).asText(), "null");
+ assertEquals(type.get(1).asText(), expectedAvroType);
+ }
+
+ private static JsonNode typeOf(DataType dataType) {
+ ObjectNode jsonSchema = AvroSchemaUtil.toAvroSchemaJsonObject(new
DimensionFieldSpec("col", dataType, true));
+ assertEquals(jsonSchema.get("name").asText(), "col");
+ return jsonSchema.get("type");
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
