Skip to content

Commit

Permalink
adds support for the Date range aggregations
Browse files Browse the repository at this point in the history
  • Loading branch information
AnyRoad committed May 5, 2024
1 parent d27b35d commit ab22135
Show file tree
Hide file tree
Showing 3 changed files with 114 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -193,7 +193,7 @@ data class RangesBucket(
@Serializable
data class RangesAggregationResult(
override val buckets: List<JsonObject>
) : BucketAggregationResult<TermsBucket>
) : BucketAggregationResult<RangesBucket>

val RangesAggregationResult.parsedBuckets get() = buckets.map { Bucket(it, RangesBucket.serializer()) }

Expand All @@ -207,6 +207,34 @@ fun Aggregations?.rangesResult(name: String, json: Json = DEFAULT_JSON): RangesA
fun Aggregations?.rangesResult(name: Enum<*>, json: Json = DEFAULT_JSON): RangesAggregationResult =
getAggResult(name, json)

@Serializable
data class DateRangesBucket(
val key: String,
@SerialName("doc_count")
val docCount: Long,
val from: Double?,
@SerialName("from_as_string")
val fromAsString: String?,
val to: Double?,
@SerialName("to_as_string")
val toAsString: String?
)

@Serializable
data class DateRangesAggregationResult(
override val buckets: List<JsonObject>
) : BucketAggregationResult<DateRangesBucket>

val DateRangesAggregationResult.parsedBuckets get() = buckets.map { Bucket(it, DateRangesBucket.serializer()) }

fun List<DateRangesBucket>.dateRangeCounts() = this.associate { it.key to it.docCount }

fun Aggregations?.dateRangesResult(name: String, json: Json = DEFAULT_JSON): DateRangesAggregationResult =
getAggResult(name, json)

fun Aggregations?.dateRangesResult(name: Enum<*>, json: Json = DEFAULT_JSON): DateRangesAggregationResult =
getAggResult(name, json)

@Serializable
data class FilterAggregationResult(
@SerialName("doc_count")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,38 @@ class AggQueryTest : SearchTestBase() {
rangeCounts["high"] shouldBe 2
}

@Test
fun shouldDoDateRangesAgg() = coRun {
before()
val aggName = "timestamp_ranges"
val response = repository.search {
resultSize = 0 // we only care about the aggs
agg(aggName, DateRangesAgg(TestDocument::timestamp) {
format = "dd-MM-yyyy"
ranges = listOf(
AggDateRange.create {
key = "before"
to = "now-4d"
},
AggDateRange.create {
key = "after"
from = "now-4d"
}
)
})
}

response.aggregations shouldNotBe null

val rangesAgg = response.aggregations.dateRangesResult(aggName)
rangesAgg shouldNotBe null
val rangeCounts = rangesAgg.parsedBuckets.map { it.parsed }.dateRangeCounts()

rangeCounts.size shouldBe 2
rangeCounts["before"] shouldBe 2
rangeCounts["after"] shouldBe 2
}

@Test
fun nestedAggregationQuery() = coRun {
before()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,59 @@ class RangesAgg(val field: String, block: (RangesAggConfig.() -> Unit)? = null)
}
}

class AggDateRange: JsonDsl() {
/**
* Customizes the key for each range
*/
var key by property<String>()
/**
* Aggregation includes the `from` value
*/
var from by property<String>()
/**
* Aggregation excludes the `to` value
*/
var to by property<String>()

companion object {
fun create(block: AggDateRange.() -> Unit) = AggDateRange().apply(block)
}
}

class DateRangesAggConfig : JsonDsl() {
var field by property<String>()

/**
* Defines how documents that are missing a value should be treated.
* By default, they will be ignored, but it is also possible to treat them as if they had a value.
*/
var missing by property<String>()

/**
* Specifies a date format by which the from and to response fields will be returned.
*/
var format by property<String>()

/**
* Time zones may either be specified as an ISO 8601 UTC offset (e.g. +01:00 or -08:00)
* or as one of the time zone ids from the TZ database.
*/
var timeZone by property<String>()

var ranges by property<List<AggDateRange>>()
}

class DateRangesAgg(val field: String, block: (DateRangesAggConfig.() -> Unit)? = null) : AggQuery("date_range") {
constructor(field: KProperty<*>, block: (DateRangesAggConfig.() -> Unit)? = null) : this(field.name, block)

init {
val config = DateRangesAggConfig()
config.field = field
block?.invoke(config)
put(name, config)
}
}

class DateHistogramAggConfig : JsonDsl() {
var field by property<String>()
var calendarInterval by property<String>("calendar_interval") // can't redefine Map.size sadly
Expand Down

0 comments on commit ab22135

Please sign in to comment.