def _split_row_groups(dataset): if not dataset.metadata or dataset.metadata.num_row_groups == 0: raise NotImplementedError("split_row_groups is only implemented " "if dataset has parquet summary files " "with row group information") # We make a dictionary of how many row groups are in each file in # order to split them. The Parquet Metadata file stores paths as the # relative path from the dataset base dir. row_groups_per_file = dict() for i in range(dataset.metadata.num_row_groups): row_group = dataset.metadata.row_group(i) path = row_group.column(0).file_path row_groups_per_file[path] = row_groups_per_file.get(path, 0) + 1 base_path = os.path.normpath(os.path.dirname(dataset.metadata_path)) split_pieces = [] for piece in dataset.pieces: # Since the pieces are absolute path, we get the # relative path to the dataset base dir to fetch the # number of row groups in the file relative_path = os.path.relpath(piece.path, base_path) # If the path is not in the metadata file, that means there are # no row groups in that file and that file should be skipped if relative_path not in row_groups_per_file: continue for row_group in range(row_groups_per_file[relative_path]): split_piece = compat_make_parquet_piece(piece.path, dataset.fs.open, row_group=row_group, partition_keys=piece.partition_keys) split_pieces.append(split_piece) return split_pieces
def _split_piece(piece, fs_open): metadata = compat_get_metadata(piece, fs_open) return [ compat_make_parquet_piece(piece.path, fs_open, row_group=row_group, partition_keys=piece.partition_keys) for row_group in range(metadata.num_row_groups) ]
def load_row_groups(dataset): """ Load dataset row group pieces from metadata :param dataset: parquet dataset object. :param allow_read_footers: whether to allow reading parquet footers if there is no better way to load row group information :return: splitted pieces, one piece per row group """ # We try to get row group information from metadata file metadata = dataset.metadata common_metadata = dataset.common_metadata if not metadata and not common_metadata: # If we are inferring the schema we allow reading the footers to get the row group information return _split_row_groups_from_footers(dataset) if metadata and metadata.num_row_groups > 0: # If the metadata file exists and has row group information we use it to split the dataset pieces return _split_row_groups(dataset) # If we don't have row groups in the common metadata we look for the old way of loading it dataset_metadata_dict = common_metadata.metadata if ROW_GROUPS_PER_FILE_KEY not in dataset_metadata_dict: raise PetastormMetadataError( 'Could not find row group metadata in _common_metadata file.' ' Use materialize_dataset(..) in petastorm.etl.dataset_metadata.py to generate' ' this file in your ETL code.' ' You can generate it on an existing dataset using petastorm-generate-metadata.py' ) metadata_dict_key = ROW_GROUPS_PER_FILE_KEY row_groups_per_file = json.loads( dataset_metadata_dict[metadata_dict_key].decode()) rowgroups = [] # Force order of pieces. The order is not deterministic since it depends on multithreaded directory # listing implementation inside pyarrow. We stabilize order here, this way we get reproducable order # when pieces shuffling is off. This also enables implementing piece shuffling given a seed sorted_pieces = sorted(dataset.pieces, key=attrgetter('path')) for piece in sorted_pieces: # If we are not using absolute paths, we need to convert the path to a relative path for # looking up the number of row groups. row_groups_key = os.path.relpath(piece.path, dataset.paths) # When reading parquet store directly from an s3 bucket, a separate piece is created for root directory. # This is not a real "piece" and we won't have row_groups_per_file recorded for it. if row_groups_key != ".": for row_group in range(row_groups_per_file[row_groups_key]): rowgroups.append( compat_make_parquet_piece( piece.path, dataset.fs.open, row_group=row_group, partition_keys=piece.partition_keys)) return rowgroups
def _index_columns(piece_info, dataset_url, partitions, indexers, schema, hdfs_driver='libhdfs3'): """ Function build indexes for dataset piece described in piece_info :param piece_info: description of dataset piece :param dataset_url: dataset location :param partitions: dataset partitions :param indexers: list of indexer objects :param schema: dataset schema :param hdfs_driver: A string denoting the hdfs driver to use (if using a dataset on hdfs). Current choices are libhdfs (java through JNI) or libhdfs3 (C++) :return: list of indexers containing index data """ # Resolver in executor context will get hadoop config from environment resolver = FilesystemResolver(dataset_url, hdfs_driver=hdfs_driver) fs = resolver.filesystem() # Create pyarrow piece piece = compat_make_parquet_piece(piece_info.path, fs.open, row_group=piece_info.row_group, partition_keys=piece_info.partition_keys) # Collect column names needed for indexing column_names = set() for indexer in indexers: column_names.update(indexer.column_names) # Read columns needed for indexing column_rows = compat_piece_read(piece, fs.open, columns=list(column_names), partitions=partitions).to_pandas().to_dict('records') # Decode columns values decoded_rows = [utils.decode_row(row, schema) for row in column_rows] if not decoded_rows: raise ValueError('Cannot build index with empty decoded_rows, columns: {}, partitions: {}' .format(column_names, partitions)) # Index columns values for indexer in indexers: indexer.build_index(decoded_rows, piece_info.piece_index) # Indexer objects contain index data, it will be consolidated on reduce phace return indexers
def split_piece(piece): metadata = compat_get_metadata(dataset.pieces[0], dataset.fs.open) return [compat_make_parquet_piece(piece.path, dataset.fs.open, row_group=row_group, partition_keys=piece.partition_keys) for row_group in range(metadata.num_row_groups)]