#@title Licensed 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
#
# https://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.

"""
Explanation: Copyright 2020 The TensorFlow IO Authors.
End of explanation
"""


!pip install tensorflow-io

"""
Explanation: Avro Dataset API
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://www.tensorflow.org/io/tutorials/avro"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/io/blob/master/docs/tutorials/avro.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/io/blob/master/docs/tutorials/avro.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
      <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/io/docs/tutorials/avro.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png" />Download notebook</a>
  </td>
</table>

Overview
The objective of Avro Dataset API is to load Avro formatted data natively into TensorFlow as <a target="_blank" href="https://www.tensorflow.org/api_docs/python/tf/data/Dataset">TensorFlow dataset</a>. Avro is a data serialization system similiar to Protocol Buffers. It's widely used in Apache Hadoop where it can provide both a serialization format for persistent data, and a wire format for communication between Hadoop nodes. Avro data is a row-oriented, compacted binary data format. It relies on schema which is stored as a separate JSON file. For the spec of Avro format and schema declaration, please refer to <a target="_blank" href="https://avro.apache.org/docs/current/spec.html">the official manual</a>.
Setup package
Install the required tensorflow-io package
End of explanation
"""


import tensorflow as tf
import tensorflow_io as tfio


"""
Explanation: Import packages
End of explanation
"""


print("tensorflow-io version: {}".format(tfio.__version__))
print("tensorflow version: {}".format(tf.__version__))

"""
Explanation: Validate tf and tfio imports
End of explanation
"""


!curl -OL https://github.com/tensorflow/io/raw/master/docs/tutorials/avro/train.avro
!ls -l train.avro

"""
Explanation: Usage
Explore the dataset
For the purpose of this tutorial, let's download the sample Avro dataset. 
Download a sample Avro file:
End of explanation
"""


!curl -OL https://github.com/tensorflow/io/raw/master/docs/tutorials/avro/train.avsc
!ls -l train.avsc

"""
Explanation: Download the corresponding schema file of the sample Avro file:
End of explanation
"""


!pip install avro


"""
Explanation: In the above example, a testing Avro dataset were created based on mnist dataset. The original mnist dataset in TFRecord format is generated from <a target="_blank" href="https://www.tensorflow.org/datasets/api_docs/python/tfds/load">TF named dataset</a>. However, the mnist dataset is too large as a demo dataset. For simplicity purpose, most of it were trimmed and first few records only were kept. Moreover, additional trimming was done for image field in original mnist dataset and mapped it to features field in Avro. So the avro file train.avro has 4 records, each of which has 3 fields: features, which is an array of int, label, an int or null, and dataType, an enum. To view the decoded train.avro (Note <a target="_blank" href="https://github.com/tensorflow/io/raw/master/docs/tutorials/avro/train.avro">the original avro data file</a> is not human readable as avro is a compacted format):
Install the required package to read Avro file:
End of explanation
"""


from avro.io import DatumReader
from avro.datafile import DataFileReader

import json

def print_avro(avro_file, max_record_num=None):
    if max_record_num is not None and max_record_num <= 0:
        return

    with open(avro_file, 'rb') as avro_handler:
        reader = DataFileReader(avro_handler, DatumReader())
        record_count = 0
        for record in reader:
            record_count = record_count+1
            print(record)
            if max_record_num is not None and record_count == max_record_num:
               break

print_avro(avro_file='train.avro')


"""
Explanation: To read and print an Avro file in a human-readable format:
End of explanation
"""


def print_schema(avro_schema_file):
    with open(avro_schema_file, 'r') as handle:
        parsed = json.load(handle)
    print(json.dumps(parsed, indent=4, sort_keys=True))

print_schema('train.avsc')


"""
Explanation: And the schema of train.avro which is represented by train.avsc is a JSON-formatted file.
To view the train.avsc:
End of explanation
"""


features = {
    'features[*]': tfio.experimental.columnar.VarLenFeatureWithRank(dtype=tf.int32),
    'label': tf.io.FixedLenFeature(shape=[], dtype=tf.int32, default_value=-100),
    'dataType': tf.io.FixedLenFeature(shape=[], dtype=tf.string)
}

schema = tf.io.gfile.GFile('train.avsc').read()

dataset = tfio.experimental.columnar.make_avro_record_dataset(file_pattern=['train.avro'],
                                                              reader_schema=schema,
                                                              features=features,
                                                              shuffle=False,
                                                              batch_size=3,
                                                              num_epochs=1)

for record in dataset:
    print(record['features[*]'])
    print(record['label'])
    print(record['dataType'])
    print("--------------------")


"""
Explanation: Prepare the dataset
Load train.avro as TensorFlow dataset with Avro dataset API:
End of explanation
"""


dataset = tfio.experimental.columnar.make_avro_record_dataset(file_pattern=['train.avro'],
                                                              reader_schema=schema,
                                                              features=features,
                                                              shuffle=False,
                                                              batch_size=3,
                                                              drop_final_batch=True,
                                                              num_epochs=1)

for record in dataset:
    print(record)


"""
Explanation: The above example converts train.avro into tensorflow dataset. Each element of the dataset is a dictionary whose key is the feature name, value is the converted sparse or dense tensor. 
E.g, it converts features, label, dataType field to a VarLenFeature(SparseTensor), FixedLenFeature(DenseTensor), and FixedLenFeature(DenseTensor) respectively. Since batch_size is 3, it coerce 3 records from train.avro into one element in the result dataset.
For the first record in train.avro whose label is null, avro reader replaces it with the specified default value(-100).
In this example, there're 4 records in total in train.avro. Since batch size is 3, the result dataset contains 3 elements, last of which's batch size is 1. However user is also able to drop the last batch if the size is smaller than batch size by enabling drop_final_batch. E.g:
End of explanation
"""


dataset = tfio.experimental.columnar.make_avro_record_dataset(file_pattern=['train.avro'],
                                                              reader_schema=schema,
                                                              features=features,
                                                              shuffle=False,
                                                              num_parallel_reads=16,
                                                              batch_size=3,
                                                              drop_final_batch=True,
                                                              num_epochs=1)

for record in dataset:
    print(record)


"""
Explanation: One can also increase num_parallel_reads to expediate Avro data processing by increasing avro parse/read parallelism.
End of explanation
"""


features = {
    'features[*]': tfio.experimental.columnar.VarLenFeatureWithRank(dtype=tf.int32),
    'label': tf.io.FixedLenFeature(shape=[], dtype=tf.int32, default_value=-100),
}


schema = tf.io.gfile.GFile('train.avsc').read()

dataset = tfio.experimental.columnar.make_avro_record_dataset(file_pattern=['train.avro'],
                                                              reader_schema=schema,
                                                              features=features,
                                                              shuffle=False,
                                                              batch_size=1,
                                                              num_epochs=1)


"""
Explanation: For detailed usage of make_avro_record_dataset, please refer to <a target="_blank" href="https://www.tensorflow.org/io/api_docs/python/tfio/experimental/columnar/make_avro_record_dataset">API doc</a>.
Train tf.keras models with Avro dataset
Now let's walk through an end-to-end example of tf.keras model training with Avro dataset based on mnist dataset.
Load train.avro as TensorFlow dataset with Avro dataset API:
End of explanation
"""


def build_and_compile_cnn_model():
    model = tf.keras.Sequential()
    model.compile(optimizer='sgd', loss='mse')
    return model

model = build_and_compile_cnn_model()


"""
Explanation: Define a simple keras model:
End of explanation
"""


def extract_label(feature):
  label = feature.pop('label')
  return tf.sparse.to_dense(feature['features[*]']), label

model.fit(x=dataset.map(extract_label), epochs=1, steps_per_epoch=1, verbose=1)


"""
Explanation: Train the keras model with Avro dataset:
End of explanation
"""


