CASSANALYTICS-171: Avoid Spark 4 partitioning warnings during reads

Spark 4 ignores custom DataSource V2 Partitioning implementations and logs a warning. Cassandra scan partitions are token ranges rather than keyed groups, so report UnknownPartitioning while preserving the input partition count.

Author: liucao-dd

CHANGES.txt

@@ -1,5 +1,6 @@
 0.5.0
 -----
+ * Avoid Spark 4 partitioning warnings during bulk reads (CASSANALYTICS-171)
  * Spark 4.0 Support (CASSANALYTICS-34)
  * Add IAM credential support for S3 storage transport (CASSANALYTICS-155)
  * Make BulkWriterConfig extensible (CASSANALYTICS-168)

cassandra-analytics-core/src/main/spark4/org/apache/cassandra/spark/sparksql/CassandraPartitioning.java

@@ -19,21 +19,16 @@
 
 package org.apache.cassandra.spark.sparksql;
 
-import org.apache.cassandra.spark.data.DataLayer;
-import org.apache.spark.sql.connector.read.partitioning.Partitioning;
+import org.apache.spark.sql.connector.read.partitioning.UnknownPartitioning;
 
-class CassandraPartitioning implements Partitioning
+/**
+ * Cassandra input partitions are token ranges, not groups of rows sharing one partition key value,
+ * so the connector reports Spark's unknown partitioning rather than keyed-group partitioning.
+ */
+class CassandraPartitioning extends UnknownPartitioning
 {
-    final DataLayer dataLayer;
-
-    CassandraPartitioning(DataLayer dataLayer)
-    {
-        this.dataLayer = dataLayer;
-    }
-
-    @Override
-    public int numPartitions()
+    CassandraPartitioning(int numPartitions)
     {
-        return dataLayer.partitionCount();
+        super(numPartitions);
     }
 }

cassandra-analytics-core/src/main/spark4/org/apache/cassandra/spark/sparksql/CassandraScanBuilder.java

@@ -121,7 +121,7 @@ public PartitionReaderFactory createReaderFactory()
     @Override
     public Partitioning outputPartitioning()
     {
-        return new CassandraPartitioning(dataLayer);
+        return new CassandraPartitioning(dataLayer.partitionCount());
     }
 
     private List<PartitionKeyFilter> buildPartitionKeyFilters()

cassandra-analytics-core/src/test/spark4/org/apache/cassandra/spark/sparksql/CassandraScanBuilderTest.java

@@ -0,0 +1,49 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.cassandra.spark.sparksql;
+
+import org.junit.jupiter.api.Test;
+
+import org.apache.cassandra.spark.data.DataLayer;
+import org.apache.spark.sql.connector.read.partitioning.Partitioning;
+import org.apache.spark.sql.connector.read.partitioning.UnknownPartitioning;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.sql.util.CaseInsensitiveStringMap;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+class CassandraScanBuilderTest
+{
+    @Test
+    void outputPartitioningReportsUnknownPartitioningWithPartitionCount()
+    {
+        DataLayer dataLayer = mock(DataLayer.class);
+        when(dataLayer.partitionCount()).thenReturn(7);
+        CassandraScanBuilder builder =
+            new CassandraScanBuilder(dataLayer, new StructType(), CaseInsensitiveStringMap.empty());
+
+        Partitioning partitioning = builder.outputPartitioning();
+
+        assertThat(partitioning).isInstanceOf(UnknownPartitioning.class);
+        assertThat(partitioning.numPartitions()).isEqualTo(7);
+    }
+}