Support multiple Avro schema version in Pulsar SQL

pulsar-client/src/main/java/org/apache/pulsar/client/impl/schema/StructSchema.java

@@ -115,12 +115,12 @@ protected static org.apache.avro.Schema createAvroSchema(SchemaDefinition schema
         }
     }
 
-    protected static org.apache.avro.Schema parseAvroSchema(String schemaJson) {
+    public static org.apache.avro.Schema parseAvroSchema(String schemaJson) {
         final Parser parser = new Parser();
         return parser.parse(schemaJson);
     }
 
-    protected static <T> SchemaInfo parseSchemaInfo(SchemaDefinition<T> schemaDefinition, SchemaType schemaType) {
+    public static <T> SchemaInfo parseSchemaInfo(SchemaDefinition<T> schemaDefinition, SchemaType schemaType) {
         return SchemaInfo.builder()
                 .schema(createAvroSchema(schemaDefinition).toString().getBytes(UTF_8))
                 .properties(schemaDefinition.getProperties())

pulsar-common/src/main/java/org/apache/pulsar/common/api/raw/RawMessage.java

@@ -101,4 +101,12 @@ public interface RawMessage {
      * @return the key of the message
      */
     Optional<String> getKey();
+
+    /**
+     * Get the schema verison of the message
+     *
+     * @return the schema version of the message
+     */
+    byte[] getSchemaVersion();
+
 }

pulsar-common/src/main/java/org/apache/pulsar/common/api/raw/RawMessageImpl.java

@@ -38,6 +38,7 @@ public class RawMessageImpl implements RawMessage {
     private ReferenceCountedObject<MessageMetadata> msgMetadata;
     private PulsarApi.SingleMessageMetadata.Builder singleMessageMetadata;
     private ByteBuf payload;
+    private MessageMetadata.Builder msgMetadataBuilder;
 
     private static final Recycler<RawMessageImpl> RECYCLER = new Recycler<RawMessageImpl>() {
         @Override
@@ -71,6 +72,7 @@ public static RawMessage get(ReferenceCountedObject<MessageMetadata> msgMetadata
             ByteBuf payload,
             long ledgerId, long entryId, long batchIndex) {
         RawMessageImpl msg = RECYCLER.get();
+        msg.msgMetadataBuilder = MessageMetadata.newBuilder(msgMetadata.get());
         msg.msgMetadata = msgMetadata;
         msg.msgMetadata.retain();
         msg.singleMessageMetadata = singleMessageMetadata;
@@ -140,4 +142,13 @@ public Optional<String> getKey() {
             return Optional.empty();
         }
     }
+
+    @Override
+    public byte[] getSchemaVersion() {
+        if (msgMetadataBuilder != null && msgMetadataBuilder.hasSchemaVersion()) {
+            return msgMetadataBuilder.getSchemaVersion().toByteArray();
+        } else {
+            return null;
+        }
+    }
 }

pulsar-sql/presto-pulsar/src/main/java/org/apache/pulsar/sql/presto/AvroSchemaHandler.java

@@ -18,86 +18,101 @@
  */
 package org.apache.pulsar.sql.presto;
 
+import com.google.common.annotations.VisibleForTesting;
 import io.airlift.log.Logger;
 
 import io.netty.buffer.ByteBuf;
-import io.netty.buffer.ByteBufAllocator;
-import io.netty.util.ReferenceCountUtil;
 import io.netty.util.concurrent.FastThreadLocal;
-import org.apache.avro.Schema;
-import org.apache.avro.generic.GenericDatumReader;
-import org.apache.avro.generic.GenericRecord;
-import org.apache.avro.io.BinaryDecoder;
-import org.apache.avro.io.DatumReader;
-import org.apache.avro.io.DecoderFactory;
-
-import java.io.IOException;
+
+import org.apache.pulsar.client.api.PulsarClientException;
+import org.apache.pulsar.client.impl.schema.generic.GenericAvroRecord;
+import org.apache.pulsar.client.impl.schema.generic.GenericAvroSchema;
+import org.apache.pulsar.common.api.raw.RawMessage;
+import org.apache.pulsar.common.naming.TopicName;
+import org.apache.pulsar.common.schema.SchemaInfo;
+
 import java.util.List;
 
 public class AvroSchemaHandler implements SchemaHandler {
 
-    private final DatumReader<GenericRecord> datumReader;
-
     private final List<PulsarColumnHandle> columnHandles;
 
-    private static final FastThreadLocal<BinaryDecoder> decoders =
-            new FastThreadLocal<>();
+    private final GenericAvroSchema genericAvroSchema;
+
+    private final SchemaInfo schemaInfo;
+
+    private static final FastThreadLocal<byte[]> tmpBuffer = new FastThreadLocal<byte[]>() {
+        @Override
+        protected byte[] initialValue() {
+            return new byte[1024];
+        }
+    };
 
     private static final Logger log = Logger.get(AvroSchemaHandler.class);
 
-    public AvroSchemaHandler(Schema schema, List<PulsarColumnHandle> columnHandles) {
-        this.datumReader = new GenericDatumReader<>(schema);
+    public AvroSchemaHandler(TopicName topicName, PulsarConnectorConfig pulsarConnectorConfig,
+                             SchemaInfo schemaInfo, List<PulsarColumnHandle> columnHandles) throws PulsarClientException {
+        this.schemaInfo = schemaInfo;
+        this.genericAvroSchema = new GenericAvroSchema(schemaInfo);
+        this.genericAvroSchema
+                .setSchemaInfoProvider(new PulsarSqlSchemaInfoProvider(topicName, pulsarConnectorConfig.getPulsarAdmin()));
         this.columnHandles = columnHandles;
     }
 
-    @Override
-    public Object deserialize(ByteBuf payload) {
-
-        ByteBuf heapBuffer = null;
-        try {
-            BinaryDecoder decoderFromCache = decoders.get();
-
-            // Make a copy into a heap buffer, since Avro cannot deserialize directly from direct memory
-            int size = payload.readableBytes();
-            heapBuffer = ByteBufAllocator.DEFAULT.heapBuffer(size, size);
-            heapBuffer.writeBytes(payload);
+    AvroSchemaHandler(PulsarSqlSchemaInfoProvider pulsarSqlSchemaInfoProvider, SchemaInfo schemaInfo, List<PulsarColumnHandle> columnHandles) {
+        this.schemaInfo = schemaInfo;
+        this.genericAvroSchema = new GenericAvroSchema(schemaInfo);
+        this.genericAvroSchema.setSchemaInfoProvider(pulsarSqlSchemaInfoProvider);
+        this.columnHandles = columnHandles;
+    }
 
-            BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(heapBuffer.array(), heapBuffer.arrayOffset(),
-                    heapBuffer.readableBytes(), decoderFromCache);
-            if (decoderFromCache==null) {
-                decoders.set(decoder);
-            }
-            return this.datumReader.read(null, decoder);
-        } catch (IOException e) {
-            log.error(e);
-        } finally {
-            ReferenceCountUtil.safeRelease(heapBuffer);
+    @Override
+    public Object deserialize(RawMessage rawMessage) {
+        ByteBuf payload = rawMessage.getData();
+        int size = payload.readableBytes();
+        byte[] buffer = tmpBuffer.get();
+        if (buffer.length < size) {
+            // If the thread-local buffer is not big enough, replace it with
+            // a bigger one
+            buffer = new byte[size * 2];
+            tmpBuffer.set(buffer);
+        }
+        payload.readBytes(buffer, 0, size);
+        if (rawMessage.getSchemaVersion() != null) {
+            return genericAvroSchema.decode(buffer, rawMessage.getSchemaVersion());
+        } else {
+            return genericAvroSchema.decode(buffer);
         }
-        return null;
     }
 
     @Override
     public Object extractField(int index, Object currentRecord) {
         try {
-            GenericRecord record = (GenericRecord) currentRecord;
+            GenericAvroRecord record = (GenericAvroRecord) currentRecord;
             PulsarColumnHandle pulsarColumnHandle = this.columnHandles.get(index);
-            Integer[] positionIndices = pulsarColumnHandle.getPositionIndices();
-            Object curr = record.get(positionIndices[0]);
-            if (curr == null) {
-                return null;
-            }
-            if (positionIndices.length > 0) {
-                for (int i = 1 ; i < positionIndices.length; i++) {
-                    curr = ((GenericRecord) curr).get(positionIndices[i]);
-                    if (curr == null) {
-                        return null;
-                    }
+            String[] names = pulsarColumnHandle.getFieldNames();
+
+            if (names.length == 1) {
+                return record.getField(pulsarColumnHandle.getName());
+            } else {
+                for (int i = 0 ; i < names.length - 1; i++) {
+                    record = (GenericAvroRecord) record.getField(names[i]);
                 }
+                return record.getField(names[names.length - 1]);
             }
-            return curr;
         } catch (Exception ex) {
             log.debug(ex,"%s", ex);
         }
         return null;
     }
+
+    @VisibleForTesting
+    GenericAvroSchema getSchema() {
+        return this.genericAvroSchema;
+    }
+
+    @VisibleForTesting
+    SchemaInfo getSchemaInfo() {
+        return schemaInfo;
+    }
 }

pulsar-sql/presto-pulsar/src/main/java/org/apache/pulsar/sql/presto/JSONSchemaHandler.java

@@ -30,6 +30,7 @@
 
 import io.netty.buffer.ByteBuf;
 import io.netty.util.concurrent.FastThreadLocal;
+import org.apache.pulsar.common.api.raw.RawMessage;
 
 public class JSONSchemaHandler implements SchemaHandler {
 
@@ -51,9 +52,10 @@ public JSONSchemaHandler(List<PulsarColumnHandle> columnHandles) {
     }
 
     @Override
-    public Object deserialize(ByteBuf payload) {
+    public Object deserialize(RawMessage rawMessage) {
         // Since JSON deserializer only works on a byte[] we need to convert a direct mem buffer into
         // a byte[].
+        ByteBuf payload = rawMessage.getData();
         int size = payload.readableBytes();
         byte[] buffer = tmpBuffer.get();
         if (buffer.length < size) {

pulsar-sql/presto-pulsar/src/main/java/org/apache/pulsar/sql/presto/PulsarColumnMetadata.java

@@ -22,7 +22,6 @@
 import com.facebook.presto.spi.type.Type;
 
 import java.util.Arrays;
-import java.util.List;
 
 public class PulsarColumnMetadata extends ColumnMetadata {
 

pulsar-sql/presto-pulsar/src/main/java/org/apache/pulsar/sql/presto/PulsarPrimitiveSchemaHandler.java

@@ -18,15 +18,14 @@
  */
 package org.apache.pulsar.sql.presto;
 
-import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufUtil;
 import java.sql.Time;
 import java.sql.Timestamp;
 import java.util.Date;
 
-import io.netty.util.concurrent.FastThreadLocal;
 import org.apache.pulsar.client.api.Schema;
 import org.apache.pulsar.client.impl.schema.AutoConsumeSchema;
+import org.apache.pulsar.common.api.raw.RawMessage;
 import org.apache.pulsar.common.schema.SchemaInfo;
 
 /**
@@ -43,8 +42,8 @@ public PulsarPrimitiveSchemaHandler(SchemaInfo schemaInfo) {
     }
 
     @Override
-    public Object deserialize(ByteBuf byteBuf) {
-        byte[] data = ByteBufUtil.getBytes(byteBuf);
+    public Object deserialize(RawMessage rawMessage) {
+        byte[] data = ByteBufUtil.getBytes(rawMessage.getData());
         Object currentRecord = schema.decode(data);
         switch (schemaInfo.getType()) {
             case DATE:

pulsar-sql/presto-pulsar/src/main/java/org/apache/pulsar/sql/presto/PulsarRecordCursor.java

@@ -138,10 +138,9 @@ private void initialize(List<PulsarColumnHandle> columnHandles, PulsarSplit puls
         this.readOffloaded = pulsarConnectorConfig.getManagedLedgerOffloadDriver() != null;
         this.pulsarConnectorConfig = pulsarConnectorConfig;
 
-        this.schemaHandler = PulsarSchemaHandlers.newPulsarSchemaHandler(
-                pulsarSplit.getSchemaInfo(),
-                columnHandles
-        );
+        this.schemaHandler = PulsarSchemaHandlers
+                .newPulsarSchemaHandler(this.topicName,
+                        this.pulsarConnectorConfig, pulsarSplit.getSchemaInfo(), columnHandles);
 
         log.info("Initializing split with parameters: %s", pulsarSplit);
 
@@ -401,7 +400,7 @@ public boolean advanceNextPosition() {
         //start time for deseralizing record
         metricsTracker.start_RECORD_DESERIALIZE_TIME();
 
-        currentRecord = this.schemaHandler.deserialize(this.currentMessage.getData());
+        currentRecord = this.schemaHandler.deserialize(this.currentMessage);
         metricsTracker.incr_NUM_RECORD_DESERIALIZED();
 
         // stats for time spend deserializing
@@ -457,7 +456,11 @@ public long getLong(int field) {
         } else if (type.equals(TIME)) {
             return ((Number) record).longValue();
         } else if (type.equals(TIMESTAMP)) {
-            return ((Number) record).longValue();
+            if (record instanceof String) {
+                return Long.parseLong((String) record);
+            } else {
+                return ((Number) record).longValue();
+            }
         } else if (type.equals(TIMESTAMP_WITH_TIME_ZONE)) {
             return packDateTimeWithZone(((Number) record).longValue(), 0);
         } else if (type.equals(TINYINT)) {
@@ -540,4 +543,9 @@ private void checkFieldType(int field, Class<?> expected) {
         Class<?> actual = getType(field).getJavaType();
         checkArgument(actual == expected, "Expected field %s to be type %s but is %s", field, expected, actual);
     }
+
+    @VisibleForTesting
+    SchemaHandler getSchemaHandler() {
+        return this.schemaHandler;
+    }
 }

pulsar-sql/presto-pulsar/src/main/java/org/apache/pulsar/sql/presto/PulsarSchemaHandlers.java

@@ -19,35 +19,40 @@
 package org.apache.pulsar.sql.presto;
 
 import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
-import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.facebook.presto.spi.PrestoException;
 import java.util.List;
+
+import org.apache.pulsar.client.api.PulsarClientException;
+import org.apache.pulsar.common.naming.TopicName;
 import org.apache.pulsar.common.schema.SchemaInfo;
 
 class PulsarSchemaHandlers {
 
-    static SchemaHandler newPulsarSchemaHandler(SchemaInfo schemaInfo,
-                                                List<PulsarColumnHandle> columnHandles) {
+    static SchemaHandler newPulsarSchemaHandler(TopicName topicName, PulsarConnectorConfig pulsarConnectorConfig, SchemaInfo schemaInfo,
+                                                List<PulsarColumnHandle> columnHandles) throws RuntimeException{
         if (schemaInfo.getType().isPrimitive()) {
             return new PulsarPrimitiveSchemaHandler(schemaInfo);
         } else if (schemaInfo.getType().isStruct()) {
-            switch (schemaInfo.getType()) {
-                case JSON:
-                    return new JSONSchemaHandler(columnHandles);
-                case AVRO:
-                    return new AvroSchemaHandler(PulsarConnectorUtils
-                            .parseSchema(new String(schemaInfo.getSchema(), UTF_8)
-                    ), columnHandles);
-                default:
-                    throw new PrestoException(NOT_SUPPORTED, "Not supported schema type: " + schemaInfo.getType());
-            }
 
+            try {
+                switch (schemaInfo.getType()) {
+                    case JSON:
+                        return new JSONSchemaHandler(columnHandles);
+                    case AVRO:
+                        return new AvroSchemaHandler(topicName, pulsarConnectorConfig, schemaInfo, columnHandles);
+                    default:
+                        throw new PrestoException(NOT_SUPPORTED, "Not supported schema type: " + schemaInfo.getType());
+                }
+            } catch (PulsarClientException e) {
+                throw new RuntimeException(new Throwable("PulsarAdmin gets version schema fail, topicName : " + topicName.toString(), e));
+            }
         } else {
             throw new PrestoException(
                     NOT_SUPPORTED,
                     "Schema `" + schemaInfo.getType() + "` is not supported by presto yet : " + schemaInfo);
         }
+
     }
 
 }