Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions python/pyspark/sql/tests/arrow/test_arrow.py
Original file line number Diff line number Diff line change
Expand Up @@ -1871,6 +1871,41 @@ def test_toArrow_with_compression_codec_large_dataset(self):
self.assertEqual(t.num_rows, 10000)
self.assertEqual(t.column_names, ["id", "str_col", "mod_col"])

def test_toPandas_double_nested_array_empty_outer(self):
schema = StructType([StructField("data", ArrayType(ArrayType(StringType())))])
df = self.spark.createDataFrame([Row(data=[])], schema=schema)
pdf = df.toPandas()
self.assertEqual(len(pdf), 1)
self.assertEqual(len(pdf["data"][0]), 0)

def test_toPandas_array_of_map_empty_outer(self):
schema = StructType([StructField("data", ArrayType(MapType(StringType(), StringType())))])
df = self.spark.createDataFrame([Row(data=[])], schema=schema)
pdf = df.toPandas()
self.assertEqual(len(pdf), 1)
self.assertEqual(len(pdf["data"][0]), 0)

def test_toPandas_triple_nested_array_empty_outer(self):
# SPARK-55056: This used to trigger SIGSEGV before the upstream arrow-java fix.
# When the outer array is empty, the second-level ArrayWriter is never
# invoked, so its count stays 0. Arrow format requires ListArray offset
# buffer to have N+1 entries even when N=0, but getBufferSizeFor(0)
# returns 0 and the buffer is omitted in IPC serialization.
schema = StructType([StructField("data", ArrayType(ArrayType(ArrayType(StringType()))))])
df = self.spark.createDataFrame([Row(data=[])], schema=schema)
pdf = df.toPandas()
self.assertEqual(len(pdf), 1)
self.assertEqual(len(pdf["data"][0]), 0)

def test_toPandas_nested_array_with_map_empty_outer(self):
schema = StructType(
[StructField("data", ArrayType(ArrayType(MapType(StringType(), StringType()))))]
)
df = self.spark.createDataFrame([Row(data=[])], schema=schema)
pdf = df.toPandas()
self.assertEqual(len(pdf), 1)
self.assertEqual(len(pdf["data"][0]), 0)


@unittest.skipIf(
not have_pandas or not have_pyarrow,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -875,4 +875,45 @@ class ArrowWriterSuite extends SparkFunSuite {
assert(map2.keyArray().array().mkString(",") == Array(1).mkString(","))
assert(stringRepr(map2) == Array("bob", "40").mkString(","))
}

test("SPARK-55056: triple nested array with empty outer array") {
// Schema: array<array<array<string>>>
// This used to trigger SIGSEGV before the upstream arrow-java fix.
// When the outer array is empty, the second-level ArrayWriter is never
// invoked, so its count stays 0. Arrow format requires ListArray offset
// buffer to have N+1 entries even when N=0, but getBufferSizeFor(0)
// returns 0 and the buffer is omitted.
val schema = new StructType()
.add("data", ArrayType(ArrayType(ArrayType(StringType))))
val writer = ArrowWriter.create(schema, null)
assert(writer.schema === schema)

// Write a row with an empty outer array
writer.write(InternalRow(ArrayData.toArrayData(Array.empty)))
writer.finish()

val reader = new ArrowColumnVector(writer.root.getFieldVectors().get(0))
val array0 = reader.getArray(0)
assert(array0.numElements() === 0)

writer.root.close()
}

test("SPARK-55056: nested array with map inside empty outer array") {
// Schema: array<array<map<string, string>>>
val schema = new StructType()
.add("data", ArrayType(ArrayType(MapType(StringType, StringType))))
val writer = ArrowWriter.create(schema, null)
assert(writer.schema === schema)

// Write a row with an empty outer array
writer.write(InternalRow(ArrayData.toArrayData(Array.empty)))
writer.finish()

val reader = new ArrowColumnVector(writer.root.getFieldVectors().get(0))
val array0 = reader.getArray(0)
assert(array0.numElements() === 0)

writer.root.close()
}
}