From b2d4900d87994cae6cba7137626be9c763bdddd2 Mon Sep 17 00:00:00 2001 From: shayann07 Date: Sun, 4 Oct 2026 23:50:48 +0500 Subject: [PATCH] fix: retain array dtype in filter operations and avoid first() on empty results --- .../kotlinx/multik/api/constructors.kt | 4 +-- .../ndarray/operations/IteratingNDArray.kt | 18 ++++------ .../multik/iterable/IterableNDArrayTest.kt | 35 +++++++++++++++++++ 3 files changed, 43 insertions(+), 14 deletions(-) diff --git a/multik-core/src/commonMain/kotlin/org/jetbrains/kotlinx/multik/api/constructors.kt b/multik-core/src/commonMain/kotlin/org/jetbrains/kotlinx/multik/api/constructors.kt index eb7a09f0..86cc3ddd 100644 --- a/multik-core/src/commonMain/kotlin/org/jetbrains/kotlinx/multik/api/constructors.kt +++ b/multik-core/src/commonMain/kotlin/org/jetbrains/kotlinx/multik/api/constructors.kt @@ -445,8 +445,8 @@ internal fun ndarrayCommon(elements: Collection, shape: In requireElementsWithShape(elements.size, shape.fold(1, Int::times)) val size = shape.reduce { acc, el -> acc * el } - val dtypeFromFirst = DataType.of(elements.first()) - val data = initMemoryView(size, dtype ?: dtypeFromFirst).apply {// TODO: boxing/unboxing!!! Find all usages + val actualDtype = dtype ?: if (elements.isNotEmpty()) DataType.of(elements.first()) else throw IllegalArgumentException("Cannot determine DataType of an empty collection when dtype is not specified.") + val data = initMemoryView(size, actualDtype).apply {// TODO: boxing/unboxing!!! Find all usages var count = 0 for (el in elements) this[count++] = el diff --git a/multik-core/src/commonMain/kotlin/org/jetbrains/kotlinx/multik/ndarray/operations/IteratingNDArray.kt b/multik-core/src/commonMain/kotlin/org/jetbrains/kotlinx/multik/ndarray/operations/IteratingNDArray.kt index 09618b46..9b6ce8c4 100644 --- a/multik-core/src/commonMain/kotlin/org/jetbrains/kotlinx/multik/ndarray/operations/IteratingNDArray.kt +++ b/multik-core/src/commonMain/kotlin/org/jetbrains/kotlinx/multik/ndarray/operations/IteratingNDArray.kt @@ -358,8 +358,7 @@ public inline fun MultiArray.distinctBy(selector: (T if (set.add(key)) list.add(e) } - val dtype = DataType.of(list.first()) - return list.toCommonNDArray(dtype) + return list.toCommonNDArray(this.dtype) } /** @@ -397,8 +396,7 @@ public inline fun MultiArray.dropWhile(predicate: (T) -> Boolean): ND public inline fun MultiArray.filter(predicate: (T) -> Boolean): D1Array { val list = ArrayList() forEach { if (predicate(it)) list.add(it) } - val dtype = DataType.of(list.first()) - return list.toCommonNDArray(dtype) + return list.toCommonNDArray(this.dtype) } /** @@ -408,8 +406,7 @@ public inline fun MultiArray.filter(predicate: (T) -> B public inline fun MultiArray.filterIndexed(predicate: (index: Int, T) -> Boolean): D1Array { val list = ArrayList() forEachIndexed { index, element -> if (predicate(index, element)) list.add(element) } - val dtype = DataType.of(list.first()) - return list.toCommonNDArray(dtype) + return list.toCommonNDArray(this.dtype) } @@ -420,8 +417,7 @@ public inline fun MultiArray.filterIndexed(predicate: (index: Int, T) public inline fun MultiArray.filterMultiIndexed(predicate: (index: IntArray, T) -> Boolean): D1Array { val list = ArrayList() forEachMultiIndexed { index, element -> if (predicate(index, element)) list.add(element) } - val dtype = DataType.of(list.first()) - return list.toCommonNDArray(dtype) + return list.toCommonNDArray(this.dtype) } /** @@ -430,8 +426,7 @@ public inline fun MultiArray.filterMultiIndexed(predica public inline fun MultiArray.filterNot(predicate: (T) -> Boolean): D1Array { val list = ArrayList() for (element in this) if (!predicate(element)) list.add(element) - val dtype = DataType.of(list.first()) - return list.toCommonNDArray(dtype) + return list.toCommonNDArray(this.dtype) } /** @@ -1077,8 +1072,7 @@ public inline fun MultiArray.partition(predicate: (T) - second.add(element) } } - val dtype = DataType.of(first.first()) - return Pair(first.toCommonNDArray(dtype), second.toCommonNDArray(dtype)) + return Pair(first.toCommonNDArray(this.dtype), second.toCommonNDArray(this.dtype)) } /** diff --git a/multik-core/src/commonTest/kotlin/org/jetbrains/kotlinx/multik/iterable/IterableNDArrayTest.kt b/multik-core/src/commonTest/kotlin/org/jetbrains/kotlinx/multik/iterable/IterableNDArrayTest.kt index 59510694..3787c856 100644 --- a/multik-core/src/commonTest/kotlin/org/jetbrains/kotlinx/multik/iterable/IterableNDArrayTest.kt +++ b/multik-core/src/commonTest/kotlin/org/jetbrains/kotlinx/multik/iterable/IterableNDArrayTest.kt @@ -18,8 +18,12 @@ import org.jetbrains.kotlinx.multik.ndarray.operations.distinct import org.jetbrains.kotlinx.multik.ndarray.operations.distinctBy import org.jetbrains.kotlinx.multik.ndarray.operations.drop import org.jetbrains.kotlinx.multik.ndarray.operations.dropWhile +import org.jetbrains.kotlinx.multik.api.zeros +import org.jetbrains.kotlinx.multik.ndarray.data.D1 +import org.jetbrains.kotlinx.multik.ndarray.data.DataType import org.jetbrains.kotlinx.multik.ndarray.operations.filter import org.jetbrains.kotlinx.multik.ndarray.operations.filterIndexed +import org.jetbrains.kotlinx.multik.ndarray.operations.filterMultiIndexed import org.jetbrains.kotlinx.multik.ndarray.operations.filterNot import org.jetbrains.kotlinx.multik.ndarray.operations.find import org.jetbrains.kotlinx.multik.ndarray.operations.findLast @@ -181,6 +185,37 @@ class IterableNDArrayTest { assertEquals(mk.ndarray(expectedList), actual) } + @Test + fun `test_filter_empty_result`() { + val data = mk.arange(10, 30, 1) + val actual = data.filter { it > 100 } + assertEquals(0, actual.size) + assertEquals(DataType.IntDataType, actual.dtype) + + val actualIndexed = data.filterIndexed { _, v -> v > 100 } + assertEquals(0, actualIndexed.size) + assertEquals(DataType.IntDataType, actualIndexed.dtype) + + val actualMultiIndexed = data.filterMultiIndexed { _, v -> v > 100 } + assertEquals(0, actualMultiIndexed.size) + assertEquals(DataType.IntDataType, actualMultiIndexed.dtype) + + val actualNot = data.filterNot { it < 100 } + assertEquals(0, actualNot.size) + assertEquals(DataType.IntDataType, actualNot.dtype) + + val emptyData = mk.zeros(0) + val actualDistinctBy = emptyData.distinctBy { it } + assertEquals(0, actualDistinctBy.size) + assertEquals(DataType.IntDataType, actualDistinctBy.dtype) + + val (firstEmpty, secondAll) = data.partition { it > 100 } + assertEquals(0, firstEmpty.size) + assertEquals(data.size, secondAll.size) + assertEquals(DataType.IntDataType, firstEmpty.dtype) + assertEquals(DataType.IntDataType, secondAll.dtype) + } + @Test fun `test_find`() { val list = listOf(1, 2, 3, 4, 5, 6, 7)