lang/py/test/av_bench.py

#!/usr/bin/env python

# 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.

import sys
import time
from random import sample, choice, randint
from string import lowercase

import avro.datafile
import avro.schema
import avro.io


types = ["A", "CNAME"]

def rand_name():
    return ''.join(sample(lowercase, 15))

def rand_ip():
    return "%s.%s.%s.%s" %(randint(0,255), randint(0,255), randint(0,255), randint(0,255))

def write(n):
    schema_s="""
    { "type": "record",
      "name": "Query",
    "fields" : [
        {"name": "query", "type": "string"},
        {"name": "response", "type": "string"},
        {"name": "type", "type": "string", "default": "A"}
    ]}"""
    out = open("datafile.avr",'w')

    schema = avro.schema.parse(schema_s)
    writer = avro.io.DatumWriter(schema)
    dw = avro.datafile.DataFileWriter(out, writer, schema) #,codec='deflate')
    for _ in xrange(n):
        response = rand_ip()
        query = rand_name()
        type = choice(types)
        dw.append({'query': query, 'response': response, 'type': type})

    dw.close()

def read():
    f = open("datafile.avr")
    reader = avro.io.DatumReader()
    af=avro.datafile.DataFileReader(f,reader)

    x=0
    for _ in af:
        pass

def t(f, *args):
    s = time.time()
    f(*args)
    e = time.time()
    return e-s

if __name__ == "__main__":
    n = int(sys.argv[1])
    print "Write %0.4f" % t(write, n)
    print "Read %0.4f" % t(read)