feat: do not count extra partition for custom queries with limits

Author: fbiville

common/src/main/scala/org/neo4j/spark/service/SchemaService.scala

@@ -568,15 +568,7 @@ class SchemaService(
         Seq(PartitionPagination.EMPTY)
       } else {
         val partitionSize = Math.ceil(count.toDouble / options.partitions).toInt
-        val partitions = options.query.queryType match {
-          case QueryType.QUERY => if (options.queryMetadata.queryCount.nonEmpty) {
-              options.partitions // for custom query count we overfetch
-            } else {
-              options.partitions - 1
-            }
-          case _ => options.partitions - 1
-        }
-        (0 to partitions)
+        (0 until options.partitions)
           .map(index => PartitionPagination(index, index * partitionSize, TopN(partitionSize, Array.empty)))
       }
     }

spark-3/src/test/scala/org/neo4j/spark/DataSourceReaderTSE.scala

@@ -998,9 +998,9 @@ class DataSourceReaderTSE extends SparkConnectorScalaBaseTSE {
       )
       .load()
 
-    assertEquals(6, partitionedQueryCountDf.rdd.getNumPartitions)
+    assertEquals(5, partitionedQueryCountDf.rdd.getNumPartitions)
     assertEquals(50, partitionedQueryCountDf.collect().map(_.getAs[String]("person")).toSet.size)
-    assertEquals(50, partitionedQueryCountDf.collect().map(_.getAs[String]("person")).size)
+    assertEquals(50, partitionedQueryCountDf.collect().map(_.getAs[String]("person")).length)
 
     val partitionedQueryCountLiteralDf = ss.read.format(classOf[DataSource].getName)
       .option("url", SparkConnectorScalaSuiteIT.server.getBoltUrl)
@@ -1014,9 +1014,9 @@ class DataSourceReaderTSE extends SparkConnectorScalaBaseTSE {
       .option("query.count", "50")
       .load()
 
-    assertEquals(6, partitionedQueryCountLiteralDf.rdd.getNumPartitions)
+    assertEquals(5, partitionedQueryCountLiteralDf.rdd.getNumPartitions)
     assertEquals(50, partitionedQueryCountLiteralDf.collect().map(_.getAs[String]("person")).toSet.size)
-    assertEquals(50, partitionedQueryCountLiteralDf.collect().map(_.getAs[String]("person")).size)
+    assertEquals(50, partitionedQueryCountLiteralDf.collect().map(_.getAs[String]("person")).length)
   }
 
   @Test