fio-l1-nyc-csvwkt-addresses

Rerun

# /// script
# requires-python = ">=3.12"
# dependencies = ["geopandas", "shapely", "pyarrow", "pandas"]
# ///

import pandas as pd
import geopandas as gpd
from shapely import wkt
import pyarrow as pa
import pyarrow.parquet as pq

def main():
    input_file = "nyc_addresses.csv"
    output_file = "addresses.geoparquet"

    print(f"Reading {input_file}...")
    df = pd.read_csv(input_file)

    print("Converting geometry...")
    # Convert WKT to geometry
    df['geometry'] = df['geometry_wkt'].apply(wkt.loads)
    
    # Convert to GeoDataFrame
    gdf = gpd.GeoDataFrame(df, geometry='geometry', crs="EPSG:4326")

    print("Processing columns...")
    # 1. recorded_at as timestamp[us]
    gdf['recorded_at'] = pd.to_datetime(gdf['recorded_at']).dt.tz_localize(None)
    
    # 2. unit_count as int32
    gdf['unit_count'] = gdf['unit_count'].astype('int32')

    # 3. All other columns as string
    # Identify columns that are not geometry, recorded_at, or unit_count
    cols_to_stringify = [col for col in gdf.columns if col not in ['geometry', 'recorded_at', 'unit_count']]
    for col in cols_to_stringify:
        gdf[col] = gdf[col].astype(str)

    # Reorder/Clean up: Remove the original wkt column if it exists
    if 'geometry_wkt' in gdf.columns:
        gdf = gdf.drop(columns=['geometry_wkt'])

    print(f"Writing to {output_file}...")
    # GeoParquet writing via geopandas/pyarrow
    # We want to ensure the types are strictly enforced in the parquet schema
    gdf.to_parquet(output_file, index=False)

    # Verify the schema
    print("Verifying schema...")
    table = pq.read_table(output_file)
    print("Schema:")
    print(table.schema)
    
    # Check data types of a few columns
    print("\nColumn types in Parquet:")
    for field in table.schema:
        print(f"{field.name}: {field.type}")

if __name__ == "__main__":
    main()