Skip to content

improve wfs size estimation #277

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
utas-raymondng merged 1 commit into from
May 26, 2026
Merged

improve wfs size estimation #277

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 64 additions & 19 deletions ...ain/java/au/org/aodn/ogcapi/server/core/service/geoserver/wfs/DownloadWfsDataService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,9 @@
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cache.annotation.Cacheable;
import org.springframework.context.annotation.Lazy;
import org.springframework.http.*;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
Expand All @@ -27,6 +29,10 @@ public class DownloadWfsDataService {
protected final ObjectMapper objectMapper;
protected static final int SAMPLES_SIZE = 500; // A not too small sample for download size estimation

@Autowired
@Lazy
protected DownloadWfsDataService self;

public DownloadWfsDataService(
WfsServer wfsServer,
RestTemplate restTemplate,
Expand Down Expand Up @@ -88,13 +94,61 @@ public String prepareWfsRequestUrl(
}

/**
* We just need to estimate the download size, the way we do it is issue two query:
* a. Issue a query and get the number or record hit
* b. Issue a query with data download but then limit the records size, and do a liner interpolation
* Unfiltered total feature count for a layer
* Cached per (uuid, layerName)
*/
@Cacheable(CacheConfig.DOWNLOADABLE_SIZE)
public BigInteger getUnfilteredRecordCount(String uuid, String layerName) {
String countUrl = prepareWfsRequestUrl(
uuid, null, null, null, null, layerName, "application/json", 1L, false
);

ResponseEntity<String> response = restTemplate.exchange(countUrl, HttpMethod.GET, pretendUserEntity, String.class);
Comment thread
NekoLyn marked this conversation as resolved.
Dismissed
if (response.getStatusCode().is2xxSuccessful() && response.getBody() != null) {
try {
JsonNode root = objectMapper.readTree(response.getBody());
if (!root.has("totalFeatures")) {
throw new RuntimeException("GeoServer GeoJSON response missing totalFeatures field");
}
return BigInteger.valueOf(root.get("totalFeatures").asLong());
} catch (IOException e) {
log.error("Failed to parse unfiltered count response for {}/{}", uuid, layerName, e);
}
}
return null;
}

/**
* Average bytes per record for the layer in the requested output format.
* Issues an unfiltered sample download so the result can be reused across calls
* Cached per (uuid, layerName, outputFormat).
*/
@Cacheable(CacheConfig.DOWNLOADABLE_SIZE)
public BigInteger getBytesPerRecord(String uuid, String layerName, String outputFormat) {
BigInteger totalCount = self.getUnfilteredRecordCount(uuid, layerName);
if (totalCount == null || totalCount.equals(BigInteger.ZERO)) {
return BigInteger.ZERO;
}

long sampleSize = totalCount.longValue() < SAMPLES_SIZE ? totalCount.longValue() : SAMPLES_SIZE;

String sampleUrl = prepareWfsRequestUrl(
uuid, null, null, null, null, layerName, outputFormat, sampleSize, false
);

ResponseEntity<byte[]> bytes = restTemplate.exchange(sampleUrl, HttpMethod.GET, pretendUserEntity, byte[].class);
Comment thread
NekoLyn marked this conversation as resolved.
Dismissed
if (bytes.getStatusCode().is2xxSuccessful() && bytes.getBody() != null) {
return BigInteger.valueOf(bytes.getBody().length).divide(BigInteger.valueOf(sampleSize));
}
return null;
}

/**
* Estimate download size for the user's subset. Runs the inherently subset-dependent
* count query, then multiplies by the cached bytes-per-record sample.
*
* @return The estimated file size
*/
@Cacheable(CacheConfig.DOWNLOADABLE_SIZE)
public BigInteger estimateDownloadSize(
String uuid,
String layerName,
Expand All @@ -104,7 +158,7 @@ public BigInteger estimateDownloadSize(
List<String> fields,
String outputFormat) throws IllegalArgumentException {

// Get total feature count via GeoJSON response
// Subset-filtered count — not cacheable here because the subset would explode the key space.
String countUrl = prepareWfsRequestUrl(
uuid, startDate, endDate, multiPolygon, fields, layerName, "application/json", 1L, false
);
Expand All @@ -118,26 +172,17 @@ public BigInteger estimateDownloadSize(
throw new RuntimeException("GeoServer GeoJSON response missing totalFeatures field");
}
BigInteger featureCount = BigInteger.valueOf(root.get("totalFeatures").asLong());
log.debug("Total record hits {}", featureCount);
log.debug("Subset record hits {}", featureCount);

if (featureCount.equals(BigInteger.ZERO)) {
return BigInteger.ZERO;
}

// In case the records we have is smaller than our predefined SAMPLES_SIZE, we use smaller one.
long sampleSize = featureCount.longValue() < SAMPLES_SIZE ? featureCount.longValue() : SAMPLES_SIZE;

// Download a small sample to measure bytes per record in the requested output format
String sampleUrl = prepareWfsRequestUrl(
uuid, startDate, endDate, multiPolygon, fields, layerName, outputFormat, sampleSize, false
);

ResponseEntity<byte[]> bytes = restTemplate.exchange(sampleUrl, HttpMethod.GET, pretendUserEntity, byte[].class);
if (bytes.getStatusCode().is2xxSuccessful() && bytes.getBody() != null) {
return featureCount
.multiply(BigInteger.valueOf(bytes.getBody().length))
.divide(BigInteger.valueOf(sampleSize));
BigInteger bytesPerRecord = self.getBytesPerRecord(uuid, layerName, outputFormat);
if (bytesPerRecord == null) {
return null;
}
return featureCount.multiply(bytesPerRecord);
} catch (IOException e) {
log.error("Fail to get feature count for estimate", e);
}
Expand Down
111 changes: 91 additions & 20 deletions ...java/au/org/aodn/ogcapi/server/core/service/geoserver/wfs/DownloadWfsDataServiceTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,8 @@ public void setUp() {
downloadWfsDataService = new DownloadWfsDataService(
wfsServer, restTemplate, pretendUserEntity, 16384, new ObjectMapper()
);

downloadWfsDataService.self = downloadWfsDataService;
}

/**
Expand Down Expand Up @@ -268,6 +270,7 @@ public void verifyRequestUrlGenerateCorrect() {
);
assertEquals("https://test.com/geoserver/wfs?VERSION=1.0.0&typeName=test:layer&SERVICE=WFS&REQUEST=GetFeature&outputFormat=shape-zip&cql_filter=((timestamp DURING 2024-01-01T00:00:00Z/2024-12-31T23:59:59Z))", result, "Correct url 1");
}

/**
* Make sure the url generated contains the correct polygon
*
Expand Down Expand Up @@ -302,6 +305,7 @@ public void verifyRequestUrlGenerateCorrectWithPolygon() throws JsonProcessingEx
result,
"Correct url 1");
}

/**
* Verify estimate size on success request
*/
Expand All @@ -315,25 +319,28 @@ void shouldReturnEstimatedSizeWhenBothRequestsSucceed() {
List<String> fields = List.of("name", "area");
String format = "application/json";

// 1. Count response: GeoJSON with totalFeatures (1 record requested, but totalFeatures = full count)
// 1. Count response: GeoJSON with totalFeatures (1 record requested, but totalFeatures = full count).
// Returned for BOTH the subset-filtered count and the unfiltered count probe inside
// getBytesPerRecord — both URLs use maxFeatures=1.
String countJson = "{\"totalFeatures\": 227193, \"features\": []}";
ResponseEntity<String> countResponse = new ResponseEntity<>(countJson, HttpStatus.OK);

// 2. Sample response (small payload in requested format)
byte[] sampleBytes = "fake data".getBytes();
// 2. Sample response. Use a payload >= SAMPLES_SIZE so bytesPerRecord = sampleBytes / 500
// yields a non-zero integer (10 bytes/record here).
byte[] sampleBytes = new byte[DownloadWfsDataService.SAMPLES_SIZE * 10];
ResponseEntity<byte[]> sampleResponse = new ResponseEntity<>(sampleBytes, HttpStatus.OK);

doReturn(countResponse)
.when(restTemplate).exchange(
argThat((String url) -> url != null && url.contains("maxFeatures=1")),
eq(HttpMethod.GET),
any(HttpEntity.class),
eq(String.class));
argThat((String url) -> url != null && url.contains("maxFeatures=1")),
eq(HttpMethod.GET),
any(HttpEntity.class),
eq(String.class));

doReturn(sampleResponse)
.when(restTemplate).exchange(
argThat((String url) -> url != null && url.contains("maxFeatures=" + DownloadWfsDataService.SAMPLES_SIZE)),
eq(HttpMethod.GET), any(), eq(byte[].class));
argThat((String url) -> url != null && url.contains("maxFeatures=" + DownloadWfsDataService.SAMPLES_SIZE)),
eq(HttpMethod.GET), any(), eq(byte[].class));

doReturn(Optional.of("http://dummy.com/wfs"))
.when(wfsServer).getFeatureServerUrl(eq(uuid), anyString());
Expand All @@ -350,26 +357,45 @@ void shouldReturnEstimatedSizeWhenBothRequestsSucceed() {
BigInteger size = downloadWfsDataService.estimateDownloadSize(
uuid, layer, start, end, multiPolygon, fields, format);

// Should call with maxFeatures=1 to get totalFeatures count via GeoJSON
// Subset-filtered count (carries the cql_filter built from start/end dates)
verify(restTemplate).exchange(
argThat((String url) -> url != null && url.contains("maxFeatures=1") && url.contains("outputFormat=application")),
argThat((String url) -> url != null
&& url.contains("maxFeatures=1")
&& url.contains("outputFormat=application")
&& url.contains("cql_filter")),
eq(HttpMethod.GET),
any(),
eq(String.class)
);

// Should also call with maxFeatures=500 to sample bytes for size interpolation
// Unfiltered count probe issued inside getBytesPerRecord — same maxFeatures=1
// pattern but without cql_filter. Acceptance criterion: sample/count path ignores subsetting.
verify(restTemplate).exchange(
argThat((String url) -> url != null && url.contains("maxFeatures=" + DownloadWfsDataService.SAMPLES_SIZE)),
argThat((String url) -> url != null
&& url.contains("maxFeatures=1")
&& url.contains("outputFormat=application")
&& !url.contains("cql_filter")),
eq(HttpMethod.GET),
any(),
eq(String.class)
);

// Sample download with maxFeatures=500, also without subset params.
verify(restTemplate).exchange(
argThat((String url) -> url != null
&& url.contains("maxFeatures=" + DownloadWfsDataService.SAMPLES_SIZE)
&& !url.contains("cql_filter")),
eq(HttpMethod.GET),
any(),
eq(byte[].class)
);

// totalFeatures=227193, sampleBytes=9 bytes, SAMPLES_SIZE=500
long expected = 227193L * sampleBytes.length / DownloadWfsDataService.SAMPLES_SIZE;
// bytesPerRecord = sampleBytes.length / SAMPLES_SIZE; total = featureCount * bytesPerRecord
long bytesPerRecord = sampleBytes.length / DownloadWfsDataService.SAMPLES_SIZE;
long expected = 227193L * bytesPerRecord;
assertEquals(BigInteger.valueOf(expected), size, "Size match");
}

@Test
void shouldReturnZeroWhenTotalFeaturesIsZero() {
String uuid = "lyr-123";
Expand All @@ -385,10 +411,10 @@ void shouldReturnZeroWhenTotalFeaturesIsZero() {

doReturn(countResponse)
.when(restTemplate).exchange(
argThat((String url) -> url != null && url.contains("maxFeatures=1")),
eq(HttpMethod.GET),
any(HttpEntity.class),
eq(String.class));
argThat((String url) -> url != null && url.contains("maxFeatures=1")),
eq(HttpMethod.GET),
any(HttpEntity.class),
eq(String.class));

doReturn(Optional.of("http://dummy.com/wfs"))
.when(wfsServer).getFeatureServerUrl(eq(uuid), anyString());
Expand Down Expand Up @@ -490,8 +516,53 @@ void returnsNullWhenParserThrowsException() {
.when(wfsServer).getDownloadableFields(eq(uuid), any(WfsServer.WfsFeatureRequest.class));

BigInteger size = downloadWfsDataService.estimateDownloadSize(
uuid, layer, start, end, multiPolygon, fields, format);
uuid, layer, start, end, multiPolygon, fields, format);

assertNull(size, "Size should be null when JSON parsing fails");
}

@Test
void sampleRequestIgnoresSubsetFilter() throws JsonProcessingException {
String uuid = "lyr-123";
String layer = "test:layer";
String start = "2024-01-01";
String end = "2024-12-31";
Object multiPolygon = new ObjectMapper().readValue(
"{ \"type\": \"MultiPolygon\", \"coordinates\": [[[[0,0],[1,0],[1,1],[0,1],[0,0]]]] }",
HashMap.class
);
List<String> fields = List.of("name", "area");
String format = "text/csv";

String countJson = "{\"totalFeatures\": 1000, \"features\": []}";
ResponseEntity<String> countResponse = new ResponseEntity<>(countJson, HttpStatus.OK);
byte[] sampleBytes = new byte[DownloadWfsDataService.SAMPLES_SIZE * 4];
ResponseEntity<byte[]> sampleResponse = new ResponseEntity<>(sampleBytes, HttpStatus.OK);

doReturn(countResponse)
.when(restTemplate).exchange(
argThat((String url) -> url != null && url.contains("maxFeatures=1")),
eq(HttpMethod.GET), any(HttpEntity.class), eq(String.class));
doReturn(sampleResponse)
.when(restTemplate).exchange(
argThat((String url) -> url != null
&& url.contains("maxFeatures=" + DownloadWfsDataService.SAMPLES_SIZE)),
eq(HttpMethod.GET), any(), eq(byte[].class));

doReturn(Optional.of("http://dummy.com/wfs"))
.when(wfsServer).getFeatureServerUrl(eq(uuid), anyString());
doReturn(createTestWFSFieldModel())
.when(wfsServer).getDownloadableFields(eq(uuid), any(WfsServer.WfsFeatureRequest.class));

downloadWfsDataService.estimateDownloadSize(uuid, layer, start, end, multiPolygon, fields, format);

// Sample URL must NOT carry the subset filter (no cql_filter, no DURING, no INTERSECTS).
verify(restTemplate).exchange(
argThat((String url) -> url != null
&& url.contains("maxFeatures=" + DownloadWfsDataService.SAMPLES_SIZE)
&& !url.contains("cql_filter")
&& !url.contains("DURING")
&& !url.contains("INTERSECTS")),
eq(HttpMethod.GET), any(), eq(byte[].class));
}
}
Loading