xiangfu0 commented on code in PR #18970:
URL: https://github.com/apache/pinot/pull/18970#discussion_r3562723599
##########
pinot-controller/src/main/java/org/apache/pinot/controller/recommender/data/writer/AvroWriter.java:
##########
@@ -93,23 +98,55 @@ public void cleanup() {
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);
}
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) {
Review Comment:
Right — `DataGenerator.buildSpec` sets `setSingleValueField(true)`, so the
recommender never emits MV here, and (as you note) MV output is unhandled for
every type in this writer, not just BOOLEAN. Kept the coercion scalar-only to
match. Thanks for flagging.
##########
pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtil.java:
##########
@@ -92,15 +103,31 @@ public static ObjectNode toAvroSchemaJsonObject(FieldSpec
fieldSpec) {
case DOUBLE:
jsonSchema.set("type", convertStringsToJsonArray("null", "double"));
return jsonSchema;
+ case BOOLEAN:
+ jsonSchema.set("type", convertStringsToJsonArray("null", "boolean"));
+ return jsonSchema;
+ case TIMESTAMP:
Review Comment:
Agreed. Closing that loop means teaching `extractFieldDataType`/`valueOf`
about the `timestamp-millis` logical type, which changes the shared Avro→Pinot
inference path — existing Avro files with `timestamp-millis` would start
inferring `TIMESTAMP` instead of `LONG`, a back-compat behavior change beyond
this recommender fix. Better as a separate follow-up. Data-wise, as you said,
`AvroRecordExtractor` already honors it.
##########
pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroSchemaUtil.java:
##########
@@ -92,15 +103,31 @@ public static ObjectNode toAvroSchemaJsonObject(FieldSpec
fieldSpec) {
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:
Review Comment:
Right, unreachable via the generator today. I kept it grouped with `BYTES`
purely so the public `getAvroSchema`/`toAvroSchemaJsonObject` contract does not
regress from the old stored-type behavior (UUID stored type is BYTES).
##########
pinot-controller/src/test/java/org/apache/pinot/controller/recommender/data/writer/AvroWriterTest.java:
##########
@@ -0,0 +1,143 @@
+/**
+ * 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) {
Review Comment:
Done in 9b58463c13 — exposed `AvroRecordAppender.nonNullBranch` as
`@VisibleForTesting` (returning the branch `Schema`) and the test helper now
delegates to it instead of re-walking the union.
--
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]
