feature: handle int and float columns

Author: fengelniederhammer

lapis2/docker-compose.yml

@@ -13,7 +13,7 @@ services:
   silo:
     image: ghcr.io/genspectrum/lapis-silo:${SILO_TAG}
     ports:
-      - ":8081"
+      - "8081:8081"
     command:
       - "--api"
       - "--preprocessingConfig=/data/preprocessingConfig.yaml"

lapis2/src/main/kotlin/org/genspectrum/lapis/config/SequenceFilterFields.kt

@@ -32,6 +32,16 @@ private fun mapToSequenceFilterFields(databaseMetadata: DatabaseMetadata) = when
         "${databaseMetadata.name}From" to SequenceFilterFieldType.DateFrom(databaseMetadata.name),
         "${databaseMetadata.name}To" to SequenceFilterFieldType.DateTo(databaseMetadata.name),
     )
+    "int" -> listOf(
+        databaseMetadata.name to SequenceFilterFieldType.Int,
+        "${databaseMetadata.name}From" to SequenceFilterFieldType.IntFrom(databaseMetadata.name),
+        "${databaseMetadata.name}To" to SequenceFilterFieldType.IntTo(databaseMetadata.name),
+    )
+    "float" -> listOf(
+        databaseMetadata.name to SequenceFilterFieldType.Float,
+        "${databaseMetadata.name}From" to SequenceFilterFieldType.FloatFrom(databaseMetadata.name),
+        "${databaseMetadata.name}To" to SequenceFilterFieldType.FloatTo(databaseMetadata.name),
+    )
 
     else -> throw IllegalArgumentException(
         "Unknown field type '${databaseMetadata.type}' for field '${databaseMetadata.name}'",
@@ -53,4 +63,10 @@ sealed class SequenceFilterFieldType(val openApiType: kotlin.String) {
     object VariantQuery : SequenceFilterFieldType("string")
     data class DateFrom(val associatedField: SequenceFilterFieldName) : SequenceFilterFieldType("string")
     data class DateTo(val associatedField: SequenceFilterFieldName) : SequenceFilterFieldType("string")
+    object Int : SequenceFilterFieldType("integer")
+    data class IntFrom(val associatedField: SequenceFilterFieldName) : SequenceFilterFieldType("integer")
+    data class IntTo(val associatedField: SequenceFilterFieldName) : SequenceFilterFieldType("integer")
+    object Float : SequenceFilterFieldType("number")
+    data class FloatFrom(val associatedField: SequenceFilterFieldName) : SequenceFilterFieldType("number")
+    data class FloatTo(val associatedField: SequenceFilterFieldName) : SequenceFilterFieldType("number")
 }

lapis2/src/main/kotlin/org/genspectrum/lapis/model/SiloFilterExpressionMapper.kt

@@ -4,6 +4,10 @@ import org.genspectrum.lapis.config.SequenceFilterFieldType
 import org.genspectrum.lapis.config.SequenceFilterFields
 import org.genspectrum.lapis.silo.And
 import org.genspectrum.lapis.silo.DateBetween
+import org.genspectrum.lapis.silo.FloatBetween
+import org.genspectrum.lapis.silo.FloatEquals
+import org.genspectrum.lapis.silo.IntBetween
+import org.genspectrum.lapis.silo.IntEquals
 import org.genspectrum.lapis.silo.NucleotideSymbolEquals
 import org.genspectrum.lapis.silo.PangoLineageEquals
 import org.genspectrum.lapis.silo.SiloFilterExpression
@@ -35,14 +39,7 @@ class SiloFilterExpressionMapper(
             }
             .groupBy({ it.first }, { it.second })
 
-        if (allowedSequenceFiltersWithType.keys.any { it.second == Filter.VariantQuery } &&
-            allowedSequenceFiltersWithType.keys.any { it.second in variantQueryTypes }
-        ) {
-            throw IllegalArgumentException(
-                "variantQuery filter cannot be used with other variant filters such as: " +
-                    "${variantQueryTypes.joinToString(", ")}",
-            )
-        }
+        crossValidateFilters(allowedSequenceFiltersWithType)
 
         val filterExpressions = allowedSequenceFiltersWithType.map { (key, values) ->
             val (siloColumnName, filter) = key
@@ -52,6 +49,10 @@ class SiloFilterExpressionMapper(
                 Filter.DateBetween -> mapToDateBetweenFilter(siloColumnName, values)
                 Filter.NucleotideSymbolEquals -> mapToNucleotideFilter(values[0].value)
                 Filter.VariantQuery -> mapToVariantQueryFilter(values[0].value)
+                Filter.IntEquals -> mapToIntEqualsFilter(siloColumnName, values)
+                Filter.IntBetween -> mapToIntBetweenFilter(siloColumnName, values)
+                Filter.FloatEquals -> mapToFloatEqualsFilter(siloColumnName, values)
+                Filter.FloatBetween -> mapToFloatBetweenFilter(siloColumnName, values)
             }
         }
 
@@ -70,6 +71,12 @@ class SiloFilterExpressionMapper(
             SequenceFilterFieldType.String -> Pair(key, Filter.StringEquals)
             SequenceFilterFieldType.MutationsList -> Pair(key, Filter.NucleotideSymbolEquals)
             SequenceFilterFieldType.VariantQuery -> Pair(key, Filter.VariantQuery)
+            SequenceFilterFieldType.Int -> Pair(key, Filter.IntEquals)
+            is SequenceFilterFieldType.IntFrom -> Pair(type.associatedField, Filter.IntBetween)
+            is SequenceFilterFieldType.IntTo -> Pair(type.associatedField, Filter.IntBetween)
+            SequenceFilterFieldType.Float -> Pair(key, Filter.FloatEquals)
+            is SequenceFilterFieldType.FloatFrom -> Pair(type.associatedField, Filter.FloatBetween)
+            is SequenceFilterFieldType.FloatTo -> Pair(type.associatedField, Filter.FloatBetween)
 
             null -> throw IllegalArgumentException(
                 "'$key' is not a valid sequence filter key. Valid keys are: " +
@@ -79,6 +86,47 @@ class SiloFilterExpressionMapper(
         return Pair(filterExpressionId, type)
     }
 
+    private fun crossValidateFilters(
+        allowedSequenceFiltersWithType: Map<Pair<SequenceFilterFieldName, Filter>, List<SequenceFilterValue>>,
+    ) {
+        if (allowedSequenceFiltersWithType.keys.any { it.second == Filter.VariantQuery } &&
+            allowedSequenceFiltersWithType.keys.any { it.second in variantQueryTypes }
+        ) {
+            throw IllegalArgumentException(
+                "variantQuery filter cannot be used with other variant filters such as: " +
+                    variantQueryTypes.joinToString(", "),
+            )
+        }
+
+        val intEqualFilters = allowedSequenceFiltersWithType.keys.filter { it.second == Filter.IntEquals }
+        for ((intEqualsColumnName, _) in intEqualFilters) {
+            val intBetweenFilterForSameColumn = allowedSequenceFiltersWithType[
+                Pair(intEqualsColumnName, Filter.IntBetween),
+            ]
+
+            if (intBetweenFilterForSameColumn != null) {
+                throw IllegalArgumentException(
+                    "Cannot filter by exact int field '$intEqualsColumnName' " +
+                        "and by int range field '${intBetweenFilterForSameColumn[0].originalKey}'.",
+                )
+            }
+        }
+
+        val floatEqualFilters = allowedSequenceFiltersWithType.keys.filter { it.second == Filter.FloatEquals }
+        for ((floatEqualsColumnName, _) in floatEqualFilters) {
+            val floatBetweenFilterForSameColumn = allowedSequenceFiltersWithType[
+                Pair(floatEqualsColumnName, Filter.FloatBetween),
+            ]
+
+            if (floatBetweenFilterForSameColumn != null) {
+                throw IllegalArgumentException(
+                    "Cannot filter by exact float field '$floatEqualsColumnName' " +
+                        "and by float range field '${floatBetweenFilterForSameColumn[0].originalKey}'.",
+                )
+            }
+        }
+    }
+
     private fun mapToVariantQueryFilter(variantQuery: String): SiloFilterExpression {
         if (variantQuery.isBlank()) {
             throw IllegalArgumentException(
@@ -123,8 +171,8 @@ class SiloFilterExpressionMapper(
     private inline fun <reified T : SequenceFilterFieldType> findDateOfFilterType(
         dateRangeFilters: List<SequenceFilterValue>,
     ): LocalDate? {
-        val fromFilter = dateRangeFilters.find { (type, _, _) -> type is T }
-        return getAsDate(fromFilter)
+        val filter = dateRangeFilters.find { (type, _, _) -> type is T }
+        return getAsDate(filter)
     }
 
     private fun getAsDate(sequenceFilterValue: SequenceFilterValue?): LocalDate? {
@@ -162,6 +210,88 @@ class SiloFilterExpressionMapper(
         )
     }
 
+    private fun mapToIntEqualsFilter(
+        siloColumnName: SequenceFilterFieldName,
+        values: List<SequenceFilterValue>,
+    ): SiloFilterExpression {
+        val value = values[0].value
+        try {
+            return IntEquals(siloColumnName, value.toInt())
+        } catch (exception: NumberFormatException) {
+            throw IllegalArgumentException(
+                "$siloColumnName '$value' is not a valid integer: ${exception.message}",
+                exception,
+            )
+        }
+    }
+
+    private fun mapToFloatEqualsFilter(
+        siloColumnName: SequenceFilterFieldName,
+        values: List<SequenceFilterValue>,
+    ): SiloFilterExpression {
+        val value = values[0].value
+        try {
+            return FloatEquals(siloColumnName, value.toDouble())
+        } catch (exception: NumberFormatException) {
+            throw IllegalArgumentException(
+                "$siloColumnName '$value' is not a valid float: ${exception.message}",
+                exception,
+            )
+        }
+    }
+
+    private fun mapToIntBetweenFilter(
+        siloColumnName: SequenceFilterFieldName,
+        values: List<SequenceFilterValue>,
+    ): SiloFilterExpression {
+        return IntBetween(
+            siloColumnName,
+            from = findIntOfFilterType<SequenceFilterFieldType.IntFrom>(values),
+            to = findIntOfFilterType<SequenceFilterFieldType.IntTo>(values),
+        )
+    }
+
+    private inline fun <reified T : SequenceFilterFieldType> findIntOfFilterType(
+        dateRangeFilters: List<SequenceFilterValue>,
+    ): Int? {
+        val (_, value, originalKey) = dateRangeFilters.find { (type, _, _) -> type is T } ?: return null
+
+        try {
+            return value.toInt()
+        } catch (exception: NumberFormatException) {
+            throw IllegalArgumentException(
+                "$originalKey '$value' is not a valid integer: ${exception.message}",
+                exception,
+            )
+        }
+    }
+
+    private fun mapToFloatBetweenFilter(
+        siloColumnName: SequenceFilterFieldName,
+        values: List<SequenceFilterValue>,
+    ): SiloFilterExpression {
+        return FloatBetween(
+            siloColumnName,
+            from = findFloatOfFilterType<SequenceFilterFieldType.FloatFrom>(values),
+            to = findFloatOfFilterType<SequenceFilterFieldType.FloatTo>(values),
+        )
+    }
+
+    private inline fun <reified T : SequenceFilterFieldType> findFloatOfFilterType(
+        dateRangeFilters: List<SequenceFilterValue>,
+    ): Double? {
+        val (_, value, originalKey) = dateRangeFilters.find { (type, _, _) -> type is T } ?: return null
+
+        try {
+            return value.toDouble()
+        } catch (exception: NumberFormatException) {
+            throw IllegalArgumentException(
+                "$originalKey '$value' is not a valid float: ${exception.message}",
+                exception,
+            )
+        }
+    }
+
     companion object {
         private var NUCLEOTIDE_MUTATION_REGEX: Regex
 
@@ -188,6 +318,10 @@ class SiloFilterExpressionMapper(
         DateBetween,
         NucleotideSymbolEquals,
         VariantQuery,
+        IntEquals,
+        IntBetween,
+        FloatEquals,
+        FloatBetween,
     }
 
     private val variantQueryTypes = listOf(Filter.PangoLineage, Filter.NucleotideSymbolEquals)

lapis2/src/main/kotlin/org/genspectrum/lapis/silo/SiloQuery.kt

@@ -49,3 +49,11 @@ data class Maybe(val child: SiloFilterExpression) : SiloFilterExpression("Maybe"
 
 data class NOf(val numberOfMatchers: Int, val matchExactly: Boolean, val children: List<SiloFilterExpression>) :
     SiloFilterExpression("N-Of")
+
+data class IntEquals(val column: String, val value: Int) : SiloFilterExpression("IntEquals")
+
+data class IntBetween(val column: String, val from: Int?, val to: Int?) : SiloFilterExpression("IntBetween")
+
+data class FloatEquals(val column: String, val value: Double) : SiloFilterExpression("FloatEquals")
+
+data class FloatBetween(val column: String, val from: Double?, val to: Double?) : SiloFilterExpression("FloatBetween")