yashmayya commented on code in PR #18970:
URL: https://github.com/apache/pinot/pull/18970#discussion_r3562200148
##########
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:
nit: this won't handle the `List<Integer>` a multi-value `BooleanGenerator`
emits, but MV output was already broken for every type here (a `List` fails
union resolution against these scalar unions, and `DataGenerator.buildSpec`
forces single-value anyway), so just noting for the record.
##########
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:
Note that the schema-inference direction still doesn't close the loop for
TIMESTAMP: `getPinotSchemaFromAvroSchema` goes through
`AvroSchemaUtil.valueOf(Schema.Type)`, which ignores `logicalType`, so
inferring a Pinot schema back from the generated file yields LONG (BOOLEAN does
round-trip now). Data-wise it's fine since `AvroRecordExtractor` honors
timestamp-millis. Might be worth a follow-up to teach inference about timestamp
logical types.
##########
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:
FWIW this case is unreachable through the data generator today: there's no
UUID generator (`NumberGenerator` throws at init), and `BytesGenerator` emits a
hex String that wouldn't serialize against `bytes` either. Keeping it for
parity with the old stored-type behavior makes sense though, since
`getAvroSchema` is public.
##########
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:
nit: duplicates `AvroRecordAppender.unwrapNullableType` — could expose that
as @VisibleForTesting and reuse. Fine as-is for a test.
--
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]
