fix(arrow/compute): fix scenario where prealloc output is missed

apache · Oct 22, 2024 · 37602e6 · 37602e6
1 parent c124ae4
commit 37602e6
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 5 deletions.
diff --git a/arrow/compute/cast_test.go b/arrow/compute/cast_test.go
@@ -1622,6 +1622,13 @@ func (c *CastSuite) checkCastZeroCopy(from arrow.DataType, json string, to arrow
 	checkCastZeroCopy(c.T(), arr, to, compute.NewCastOptions(to, true))
 }
 
+func (c *CastSuite) TestTimestampToTimestampSimpleTimezone() {
+	c.checkCast(&arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "Etc/UTC"},
+		arrow.FixedWidthTypes.Timestamp_us,
+		`["2023-01-01T19:25:00.123456+00:00", null, "2023-01-01T19:25:00.123456+00:00"]`,
+		`["2023-01-01T19:25:00.123456+00:00", null, "2023-01-01T19:25:00.123456+00:00"]`)
+}
+
 func (c *CastSuite) TestTimestampToTimestamp() {
 	tests := []struct {
 		coarse, fine arrow.DataType

diff --git a/arrow/compute/executor.go b/arrow/compute/executor.go
@@ -587,8 +587,7 @@ func (s *scalarExecutor) executeSpans(data chan<- Datum) (err error) {
 
 	if s.preallocContiguous {
 		// make one big output alloc
-		prealloc := s.prepareOutput(int(s.iterLen))
-		output = *prealloc
+		output := s.prepareOutput(int(s.iterLen))
 
 		output.Offset = 0
 		var resultOffset int64
@@ -598,15 +597,19 @@ func (s *scalarExecutor) executeSpans(data chan<- Datum) (err error) {
 				break
 			}
 			output.SetSlice(resultOffset, input.Len)
-			err = s.executeSingleSpan(&input, &output)
+			err = s.executeSingleSpan(&input, output)
 			resultOffset = nextOffset
 		}
 		if err != nil {
-			prealloc.Release()
+			output.Release()
 			return
 		}
 
-		return s.emitResult(prealloc, data)
+		if output.Offset != 0 {
+			output.SetSlice(0, s.iterLen)
+		}
+
+		return s.emitResult(output, data)
 	}
 
 	// fully preallocating, but not contiguously