`native_shuffle` crashes for repartition + count

[comphead]

Describe the bug

Valid query

spark.read.parquet("hdfs://location").repartition(50).count()

Crashes with

org.apache.comet.CometNativeException: External error: Execution error: Error in shuffle write: Arrow error: Invalid argument error: must either specify a row count or at least one column

Steps to reproduce

No response

Expected behavior

should work

Additional context

No response

[comphead]

The issue related to native shuffle, reproduce

  test("native datafusion scan - repartition count") {
    withTempPath { dir =>
      withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
        spark
          .range(1000)
          .selectExpr("id", "concat('name_', id) as name")
          .repartition(100)
          .write
          .parquet(dir.toString)
      }
      withSQLConf(
        CometConf.COMET_NATIVE_SCAN_IMPL.key -> CometConf.SCAN_NATIVE_DATAFUSION,
        CometConf.COMET_EXEC_SHUFFLE_WITH_ROUND_ROBIN_PARTITIONING_ENABLED.key -> "true") {
        val count = spark.read.parquet(dir.toString).repartition(10).count()
        assert(count == 1000)
      }
    }
  }

[mbutrovich]

I renamed this. It's a native shuffle issue, not a native_datafusion issue, right? It just manifests when a scan is fully-native?

[comphead]

I renamed this. It's a native shuffle issue, not a native_datafusion issue, right? It just manifests when a scan is fully-native?

Thanks @mbutrovich! upon the investigation, I found the native_datafusion/native_iceberg_compat failing if native shuffle enabled for RR partitioning