Permalink
e197edf Sep 13, 2018
1 contributor

Users who have contributed to this file

57 lines (42 sloc) 1.87 KB
# Copyright (c) 2017-2018 Uber Technologies, Inc.
#
# 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
#
# 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.
"""Minimal example of how to read samples from a dataset generated by `generate_hello_world_dataset.py`
using pyspark"""
from __future__ import print_function
from pyspark.sql import SparkSession
from examples.hello_world.generate_hello_world_dataset import HelloWorldSchema
from petastorm.spark_utils import dataset_as_rdd
def pyspark_hello_world(dataset_url='file:///tmp/hello_world_dataset'):
spark = SparkSession \
.builder \
.master('local[1]') \
.getOrCreate()
# dataset_as_rdd creates an rdd of named tuples.
rdd = dataset_as_rdd(dataset_url, spark, [HelloWorldSchema.id, HelloWorldSchema.image1])
print('An id in the dataset: ', rdd.first().id)
# Create a dataframe object from a parquet file
dataframe = spark.read.parquet(dataset_url)
# Show a schema
dataframe.printSchema()
# Count all
dataframe.count()
# Show just some columns
dataframe.select('id').show()
# This is how you can use a standard SQL to query a dataset. Note that the data is not decoded in this case.
number_of_rows = spark.sql(
'SELECT count(id) '
'from parquet.`{}` '.format(dataset_url)).collect()
print('Number of rows in the dataset: {}'.format(number_of_rows[0][0]))
if __name__ == '__main__':
pyspark_hello_world()