<jupyter_start><jupyter_text>Dask Geohash Sorted<jupyter_code>import logging
import time
from datetime import datetime
from pathlib import Path
from shapely.geometry import Polygon, box
from polygon_geohasher.polygon_geohasher import polygon_to_geohashes, geohashes_to_polygon
import geohash
from functools import reduce
from math import ceil

import numpy as np
import pandas as pd
import geopandas as gpd
import dask.dataframe as dd
from distributed import LocalCluster, Client
# create local dask cluster
cluster = LocalCluster(#silence_logs=logging.ERROR,
                       dashboard_address=':8790',
                       n_workers=4,
                       threads_per_worker=2,
                       memory_limit='3 GB')
client = Client(cluster)
client
# set up data paths
base_path = Path().cwd().parent.parent
data_dir = base_path.joinpath('data')
# load contiguous us data
df = dd.read_parquet(data_dir.joinpath('contiguous_us_geohash4_sorted.parquet'))
display(df.head(2))
len(df)
# load various size subsets of the zip code data as geodataframes
zips_1 = gpd.read_file(data_dir.joinpath('zip_codes', 'zips_1.geojson')).loc[:, ['geometry']]
zips_10 = gpd.read_file(data_dir.joinpath('zip_codes', 'zips_10.geojson')).loc[:, ['geometry']]
zips_100 = gpd.read_file(data_dir.joinpath('zip_codes', 'zips_100.geojson')).loc[:, ['geometry']]
zips_1000 = gpd.read_file(data_dir.joinpath('zip_codes', 'zips_1000.geojson')).loc[:, ['geometry']]
zips_10000 = gpd.read_file(data_dir.joinpath('zip_codes', 'zips_10000.geojson')).loc[:, ['geometry']]
# filter function
empty_df = pd.DataFrame([], columns=['latitude', 'longitude'])
def spatial_join(large_data_df, zip_codes_gdf):
    """map function for distributed processing"""
    # catch for empty subset dataframes
    if large_data_df.empty:
        return empty_df
    # create geodataframe
    crs = "epsg:4326"
    large_data_gdf = gpd.GeoDataFrame(large_data_df,
                                      geometry=gpd.points_from_xy(large_data_df.longitude,
                                                                  large_data_df.latitude),
                                      crs=crs)
    # perform the spatial join
    rdf = gpd.sjoin(large_data_gdf,
                    zip_codes_gdf,
                    how='inner',
                    op='within').drop(['index_right', 'geometry'], axis=1)
    if rdf.empty:
        return empty_df
    return rdf
%%time
total_points = len(df)
num_partitions = df.npartitions
geohash_precision = 4
num_polygons = []
time_sec = []
num_result_points = []
num_points = len(df.partitions[:num_partitions])

t00 = time.time()
# loop over the subsets of zip codes
for zip_gdf in [zips_1, zips_10, zips_100, zips_1000]:#, zips_10000]:
    num_polygons.append(len(zip_gdf))
    t0 = time.time()
    
    # get unique geohashes from data (could be saved)
    unique_geohashes = df.index.unique().compute()
    
    # convert zip_codes to geohashes
    geohash_df = zip_gdf.geometry.apply(polygon_to_geohashes, 
                                                   precision=geohash_precision,
                                                   inner=False)
    rdfs = []
    for polygon_index, geohash_set in geohash_df.iteritems():
        zip_geohashes = list(geohash_set.intersection(unique_geohashes.values))  # filter out geohashes not in data 
        possible_interior_pts = df.loc[zip_geohashes]
        rdfs.append(possible_interior_pts.map_partitions(spatial_join, zip_codes_gdf=zip_gdf.loc[polygon_index:polygon_index]))
    rdf = dd.concat(rdfs).compute()

    time_sec.append(time.time() - t0)
    num_result_points.append(len(rdf))
    print(f'num_polygons[-1]: {num_polygons[-1]}, time_sec[-1]: {time_sec[-1]:.0f} s')
# save summary info to file
results_df = pd.DataFrame({'num_polygons': num_polygons,
                           'num_points': num_points,
                           'num_result_points': num_result_points,
                           'time_min': np.asarray(time_sec)/60})                       
results_df.to_csv(f'{datetime.now()}_geohash_sorted_results.csv')
results_df
# release the dask workers
cluster.scale(0)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10661236224245328, 'fraction_numerical': 0.017249640632486823, 'mean_word_length': 3.04162633107454, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '533804', 'license': 'MIT', 'n_tokens_mistral': 1415, 'n_tokens_neox': 1353, 'n_words': 264, 'path': '02_compared_cases/sorted_geohash/dask_geohash_sorted.ipynb', 'repo_name': 'jorisvandenbossche/scipy2020_spatial_algorithms_at_scale', 'size': '4174'}