Skip to content

Commit 56b794f

Browse files
authored
fix(arrow/compute): fix scenario where prealloc output is missed (#167)
Fixes #52
1 parent 6843412 commit 56b794f

File tree

2 files changed

+15
-5
lines changed

2 files changed

+15
-5
lines changed

arrow/compute/cast_test.go

+7
Original file line numberDiff line numberDiff line change
@@ -1622,6 +1622,13 @@ func (c *CastSuite) checkCastZeroCopy(from arrow.DataType, json string, to arrow
16221622
checkCastZeroCopy(c.T(), arr, to, compute.NewCastOptions(to, true))
16231623
}
16241624

1625+
func (c *CastSuite) TestTimestampToTimestampSimpleTimezone() {
1626+
c.checkCast(&arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "Etc/UTC"},
1627+
arrow.FixedWidthTypes.Timestamp_us,
1628+
`["2023-01-01T19:25:00.123456+00:00", null, "2023-01-01T19:25:00.123456+00:00"]`,
1629+
`["2023-01-01T19:25:00.123456+00:00", null, "2023-01-01T19:25:00.123456+00:00"]`)
1630+
}
1631+
16251632
func (c *CastSuite) TestTimestampToTimestamp() {
16261633
tests := []struct {
16271634
coarse, fine arrow.DataType

arrow/compute/executor.go

+8-5
Original file line numberDiff line numberDiff line change
@@ -587,8 +587,7 @@ func (s *scalarExecutor) executeSpans(data chan<- Datum) (err error) {
587587

588588
if s.preallocContiguous {
589589
// make one big output alloc
590-
prealloc := s.prepareOutput(int(s.iterLen))
591-
output = *prealloc
590+
output := s.prepareOutput(int(s.iterLen))
592591

593592
output.Offset = 0
594593
var resultOffset int64
@@ -598,15 +597,19 @@ func (s *scalarExecutor) executeSpans(data chan<- Datum) (err error) {
598597
break
599598
}
600599
output.SetSlice(resultOffset, input.Len)
601-
err = s.executeSingleSpan(&input, &output)
600+
err = s.executeSingleSpan(&input, output)
602601
resultOffset = nextOffset
603602
}
604603
if err != nil {
605-
prealloc.Release()
604+
output.Release()
606605
return
607606
}
608607

609-
return s.emitResult(prealloc, data)
608+
if output.Offset != 0 {
609+
output.SetSlice(0, s.iterLen)
610+
}
611+
612+
return s.emitResult(output, data)
610613
}
611614

612615
// fully preallocating, but not contiguously

0 commit comments

Comments
 (0)