Faster parsing

[michaelkirk]

I'm investigating #223. Recall we currently parse to a serde_json::JsonValue and convert that to our geojson types. This entails parsing every coordinate to a Vec (which, as of #222, then gets converted to a TinyVec). It seems obvious we'd do better to parse directly to TinyVec.

I thought to myself, "Why do we need to have manual deserialization at all? Why can't we derive all of our deserialization with serde?"

It required some structural juggling, but wasn't very hard and is quite concise! Here's a POC I've been messing with:
https://github.com/michaelkirk/geojson_fast

Unfortunately, it's not necessarily any faster, and can even be slower (despite the "_fast" branding in my crate name 😞). This surprised me at first. How could it be slower if we're avoiding the intermediate JSONValue representation (and it's per coordinate Vecs!)

parse GeoJSON Object (geojson crate) (countries.geojson)
                        time:   [1.5592 ms 1.5626 ms 1.5661 ms]

parse GeoJSON Object (geojson_fast) (countries.geojson)                                                                                     
                        time:   [1.8901 ms 1.9085 ms 1.9356 ms]                      

I think it's mostly due to serde-rs/serde#1495

When parsing adjacently tagged enums (e.g. the geojson::GeoJSON enum), (I think) all the fields are effectively buffered in an intermediate representation until deserialization is complete.

My understanding is that the generic serde deserializaer can't know which variant it's encountered until it has completed all the fields in the variant. As a contrived example:

#[serde(tag)]
enum MyEnum {
   Variant1 { coordinates: Vec<TinyVec<f64 ,2>> },
   Variant2 { coordinates: Vec<Vec<f64>> }
}

let input = '{ "coordinates": [[1.0 2.0]], type: "variant2" }';
//                                                       ^________ You have to parse all the way to here 
//                                                                 before you know whether coordinates should have been a
//                                                                 Vec<TinyVec<f64 ,2>> or a Vec<Vec<f64>>

A more concrete example for geojson could be:

{ "coordinates": [[...]], "type": "Polygon" }

## coordinates have same dimensionality
{ "coordinates": [[...]], "type": "MultiLineString" }

Conceptually if we could guarantee the tag came first, we could optimize things with a finite amount of peeking:

{"type": "Polygon" , "coordinates": [[...]]}

## coordinates have same dimensionality
{"type": "MultiLineString", "coordinates": [[...]]}

But we can make no such guarantees, and serde is a generic deserializer, so I don't think we can get any tricks like this for free. It might be possible to write a custom deserializer that is smarter about this. Conceptually it seems like it should be possible to "early terminate" or perhaps provide some ambiguous intermediate structs to help make deserializing geojson optimal.

[michaelkirk]

To add, the serde derived deserialization can be faster if you avoid the enums. E.g. compare these benches which all take the same input and produce the same result

let geojson_str = include_str!("../tests/fixtures/geometry_collection.geojson");
c.bench_function("parse geometry_collection.geojson as Object", |b| {
    b.iter(|| {
        // parsing to the GeoJsonObject enum is slowest
        //     parse geometry_collection.geojson as Object
        //                    time:   [14.319 µs 14.343 µs 14.367 µs]
        let geojson: GeoJsonObject = geojson_str.parse().unwrap();
        let GeoJsonObjectValue::GeometryCollection(geometry_collection) = geojson.value()
        else {
            panic!("expected GeometryCollection");
        };
        assert_eq!(geometry_collection.geometries().len(), 32);
    });
});
c.bench_function("parse geometry_collection.geojson as Geometry", |b| {
    b.iter(|| {
        // parsing to the inner Geometry enum is a bit faster (avoid the Feature/FeatureCollection/Geometry enum)
        //    parse geometry_collection.geojson as Geometry
        //                    time:   [11.463 µs 11.589 µs 11.757 µs]
        let geometry: Geometry = geojson_str.parse().unwrap();
        let Geometry::GeometryCollection(geometry_collection) = geometry else {
            panic!("expected GeometryCollection");
        };
        assert_eq!(geometry_collection.geometries().len(), 32);
    });
});
c.bench_function(
    "parse geometry_collection.geojson as GeometryCollection",
    |b| {
        b.iter(|| {
            // parsing directly to the GeometryCollection is fastest since we skip 2 enum deserializations.
            //    parse geometry_collection.geojson as GeometryCollection
            //                time:   [9.3696 µs 9.4038 µs 9.4614 µs]
            let geometry_collection: GeometryCollection = geojson_str.parse().unwrap();
            assert_eq!(geometry_collection.geometries().len(), 32);
        });
    },
)

So for example, if we did rewrite the deserialization to be based on serde structs, when you are expecting a FeatureCollection, it'd be faster to do:

let feature_collection =  FeatureCollection::from_str(geojson_str)?;

Rather than:

let GeoJson::FeatureCollection(feature_collection) =  GeoJson:from_str(geojson_str)? else {
    return Err("expected feature collection")
} ;

Maybe that's too surprising. I'm still hopeful there's a better way with a custom deserialization strategy.

[michaelkirk]

Maybe a custom buffering type could help? serde-rs/serde#1183

I haven't found the smoking gun yet, but I sort of think this approach is doomed, since we seemingly must deserialize the coordinates in a single pass. We can't "peek" ahead to see how deeply nested our array is — that all happens within a private method call within serde_json.

[michaelkirk]

Ok! I think I've arrived at something...

One false start was that I tried using serde::RawValue to avoid parsing the json into coordinates. It helped a little, but I eventually realized that even though we're not parsing the input into coordinates, serde_json still needs to parse the RawValue to be sure it's valid json. It's obvious in hindsight - how else would it know where the RawValue ends?

So it was a little faster to go from input -> RawValue -> Coordinates rather than input -> serde_json::Value -> Coordinates, but faster yet is obviously input -> Coordinates. I was able to accomplish this in my POC repo with a recursive visitor.

At this point, I think I'm confident enough with the approach that I'm going to try to port the changes over to this crate.

[michaelkirk]

Ah, and some tasty benchmarks:

// currently
parse FeatureCollection countries.geojson/baseline (geojson crate)
                        time:   [2.8144 ms 2.8250 ms 2.8361 ms]
// new impl
parse FeatureCollection countries.geojson/FeatureCollection struct
                        time:   [1.7489 ms 1.7572 ms 1.7660 ms]


// currently
parse argentina.geojson MultiPolygon/baseline (geojson crate)
                        time:   [13.116 µs 13.130 µs 13.144 µs]
// new impl
parse argentina.geojson MultiPolygon/GeoJsonObject enum
                        time:   [8.7915 µs 8.8045 µs 8.8191 µs]


// currently
parse geometry_collection.geojson/baseline (geojson crate)
                        time:   [13.142 µs 13.164 µs 13.190 µs]
// new impl
parse geometry_collection.geojson/GeoJSONObject enum
                        time:   [6.0507 µs 6.0846 µs 6.1223 µs]

Note my POC isn't feature complete — it doesn't support foreign members and I haven't worked on the "deserialize to your own struct", so these numbers might change, but I don't expect them to change much.

[michaelkirk]

To add, the serde derived deserialization can be faster if you avoid the enums. E.g. compare these benches which all take the same input and produce the same result

To be clear, this issue is also resolved. The solution is... more code!

Rather than having adjacently tagged enums, use an intermediate struct that converts to the adjacently tagged enums.

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", try_from = "RawGeometry")]
pub enum Geometry {
    Point(Point),
    LineString(LineString),
    Polygon(Polygon),
    MultiPoint(MultiPoint),
    MultiLineString(MultiLineString),
    MultiPolygon(MultiPolygon),
    GeometryCollection(GeometryCollection),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RawGeometry {
    r#type: GeometryType,
    coordinates: Option<Coordinates>,
    geometries: Option<Vec<Geometry>>,
}

impl TryFrom<RawGeometry> for Geometry {
    type Error = Error;

    fn try_from(value: RawGeometry) -> Result<Self, Self::Error> {
        Ok(match value {
            RawGeometry {
                r#type: GeometryType::Point,
                coordinates: Some(Coordinates::ZeroDimensional(coordinates)),
                geometries: None,
            } => Geometry::Point(Point { coordinates }),
            RawGeometry {
                r#type: GeometryType::LineString,
                coordinates: Some(Coordinates::OneDimensional(coordinates)),
                geometries: None,
            } => Geometry::LineString(LineString { coordinates }),
            RawGeometry {
                r#type: GeometryType::Polygon,
                coordinates: Some(Coordinates::TwoDimensional(coordinates)),
                geometries: None,
            } => Geometry::Polygon(Polygon { coordinates }),
            // Empty coordinates [] deserialize as OneDimensional([])
            RawGeometry {
                r#type: GeometryType::Polygon,
                coordinates: Some(Coordinates::OneDimensional(coords)),
                geometries: None,
            } if coords.is_empty() => Geometry::Polygon(Polygon {
                coordinates: vec![],
            }),
            RawGeometry {
                r#type: GeometryType::MultiPoint,
                coordinates: Some(Coordinates::OneDimensional(coordinates)),
                geometries: None,
            } => Geometry::MultiPoint(MultiPoint { coordinates }),
            RawGeometry {
                r#type: GeometryType::MultiLineString,
                coordinates: Some(Coordinates::TwoDimensional(coordinates)),
                geometries: None,
            } => Geometry::MultiLineString(MultiLineString { coordinates }),
            // Empty coordinates [] deserialize as OneDimensional([])
            RawGeometry {
                r#type: GeometryType::MultiLineString,
                coordinates: Some(Coordinates::OneDimensional(coords)),
                geometries: None,
            } if coords.is_empty() => Geometry::MultiLineString(MultiLineString {
                coordinates: vec![],
            }),
            RawGeometry {
                r#type: GeometryType::MultiPolygon,
                coordinates: Some(Coordinates::ThreeDimensional(coordinates)),
                geometries: None,
            } => Geometry::MultiPolygon(MultiPolygon { coordinates }),
            // Empty coordinates [] deserialize as OneDimensional([])
            RawGeometry {
                r#type: GeometryType::MultiPolygon,
                coordinates: Some(Coordinates::OneDimensional(coords)),
                geometries: None,
            } if coords.is_empty() => Geometry::MultiPolygon(MultiPolygon {
                coordinates: vec![],
            }),
            RawGeometry {
                r#type: GeometryType::GeometryCollection,
                coordinates: None,
                geometries: Some(geometries),
            } => Geometry::GeometryCollection(GeometryCollection { geometries }),
            _ => return Err(Error::InvalidGeometry),
        })
    }
}

It's more code, but pretty straight forward code. The "empty" cases are a little fiddly.

[urschrei]

Where did you get the idea for a recursive visitor?? This is bonkers (positive).

The extra code issue doesn't seem to be a big deal: it's very unlikely to need much / any maintenance once it's done, and that amount is going to have a negligible impact on compile times I imagine.

[michaelkirk]

Coordinates can be:

1. an array of f64 (Point)
2. an array of arrays of f64 (LineString / MultiPoint)
3. an array of arrays of arrays f64 (Polygon / MultiLineString)
4. an array of arrays of arrays of arrays of f64 (MultiPolygon)

So the recursive structure is pretty obvious, and trying to process it recursively seemed natural. It just took me a lot of trial and error to figure out how to express it. I did make heavy use of an LLM for tests and refactoring. I did try to get the LLM to implement various versions of the actual deserialization algorithm, but it was mostly unhelpful (e.g. "ah I know, let me just go back to deserializing the whole thing to an intermediate serde_json::JsonValue, that'll get the tests to pass!").

Serde sort of feels like alien technology to me. It's amazingly powerful. I can easily point it at things and make it go boom, but getting to the point of understanding how it actually works well enough to implement this required days of meditation.

[michaelkirk]

Update: I'm at a pretty good stopping spot with the geojson overhaul, but I could use some help with reviews.

The end result is:

• deserializing is much faster! serializing is a little faster, and converting to geo-types got a big speedup too!
• (Across multiple PRs): 18 files changed, 1621 insertions(+), 1887 deletions(-)
• Parsing errors now include the line/col number.
• Nicer constructors if you are working with the geojson api programatically
  • before: Geometry { value: GeometryValue::Point(vec![1.0, 2.0]), .. }
  • after: Geometry::new_point([1.0, 2.0])

Benches vs v0.24.2:

parse (countries.geojson)
                        time:   [975.45 µs 982.33 µs 990.70 µs]
                        change: [-37.053% -36.531% -35.886%] (p = 0.00 < 0.05)

FeatureReader::features (countries.geojson)
                        time:   [4.4115 ms 4.4268 ms 4.4468 ms]
                        change: [-12.339% -11.986% -11.553%] (p = 0.00 < 0.05)

FeatureReader::deserialize (countries.geojson)
                        time:   [5.6733 ms 5.8691 ms 6.1232 ms]
                        change: [-6.2694% -2.7489% +1.4549%] (p = 0.17 > 0.05)

FeatureReader::deserialize to geo-types (countries.geojson)
                        time:   [5.5962 ms 5.6207 ms 5.6539 ms]
                        change: [-7.5344% -7.1276% -6.5682%] (p = 0.00 < 0.05)

parse (geometry_collection.geojson)
                        time:   [7.8267 µs 7.8426 µs 7.8603 µs]
                        change: [-42.905% -42.647% -42.371%] (p = 0.00 < 0.05)

serialize geojson::FeatureCollection struct (countries.geojson)
                        time:   [409.41 µs 411.05 µs 413.27 µs]
                        change: [+0.1734% +0.5500% +0.9726%] (p = 0.00 < 0.05)

serialize custom struct (countries.geojson)
                        time:   [869.04 µs 870.66 µs 872.33 µs]
                        change: [-1.8347% -1.5441% -1.2567%] (p = 0.00 < 0.05)

serialize custom struct to geo-types (countries.geojson)
                        time:   [920.92 µs 922.36 µs 923.81 µs]
                        change: [-15.454% -15.237% -15.008%] (p = 0.00 < 0.05)

Convert to geo-types    time:   [57.439 µs 57.510 µs 57.577 µs]
                        change: [-75.158% -75.090% -75.023%] (p = 0.00 < 0.05)

The focus of the work was to avoid per-coordinate heap allocations for the common case of people using 2-d geometries.

Most obviously, this meant moving our CoordinateType from Vec<f64> to TinyVec<f64; 2>. 2d coordinates can now be allocated inline. That helped a little but there were a couple other places we were doing per-coordinate heap allocations.

The next required updating our deserialization logic. Our old approach was to parse a serde_json::Object and convert that Into our geojson types. The problem with that approach, is serde_json itself allocates a Vec per coordinate in the serde_json::Object. We're now deserializing directly to our types without an intermediate serde_json representation.

However, one last hitch! GeoJSON has a type column at the same level as it's other fields. We represent the various types (Feature, FeatureCollection, Point, LineString...) in an enum. Parsing one of them, we might have to handle other fields before we get to the type column. It turns out the derived deserialization implementations for "adjacently" or "internally" tagged enums like this are necessarily slow. That's because the derived algorithm is basically a guess-and-check approach. It tries the first variant, and if that fails it tries the next, and so on until a variant is successfully deserialized. Normally, deserializing would consume the input stream, so to be able to "try again" on the next variant, the derived deserialization first persists the entire tokenized input stream before trying to parse any of the variants. So this was yet another place where we were seeing a per-coordinate Vec allocation!

If you're excited about seeing any of this get released, please consider reviewing my remaining pull requests:

• Speed up parsing by deserializing directly to geojson (without intermediate serde_json representation) #269
• Remove methods related to json to/from now that we ser/de directly #270
• Additional ergonomic constructors (for Geometry and FeatureCollection) #271

For context you can take a look at these already merged PRs (thanks @urschrei!): #222, #267, #268