Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix: Support Arrow date64 in Parquet #410

Merged
merged 2 commits into from
Jan 15, 2024
Merged

fix: Support Arrow date64 in Parquet #410

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
02ffcbe
fix: Support Arrow `date64` in Parquet
candiduslynx Jan 15, 2024
db6282c
lint
candiduslynx Jan 15, 2024
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
28 changes: 12 additions & 16 deletions parquet/date.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,32 +6,28 @@ import (
"github.com/apache/arrow/go/v15/arrow/memory"
)

func reverseTransformDate32(arr *array.Timestamp, toTime toTimeFunc) arrow.Array {
builder := array.NewDate32Builder(memory.DefaultAllocator)
func reverseTransformDate64(arr *array.Date32) *array.Date64 {
builder := array.NewDate64Builder(memory.DefaultAllocator)

for i := 0; i < arr.Len(); i++ {
if arr.IsNull(i) {
builder.AppendNull()
continue
}

builder.Append(arrow.Date32FromTime(toTime(arr.Value(i))))
builder.Append(arrow.Date64FromTime(arr.Value(i).ToTime()))
}

return builder.NewArray()
return builder.NewDate64Array()
}

func reverseTransformDate64(arr *array.Timestamp, toTime toTimeFunc) arrow.Array {
builder := array.NewDate64Builder(memory.DefaultAllocator)

for i := 0; i < arr.Len(); i++ {
if arr.IsNull(i) {
builder.AppendNull()
continue
}

builder.Append(arrow.Date64FromTime(toTime(arr.Value(i))))
func reverseTransformFromDate32(dt arrow.DataType, arr *array.Date32) arrow.Array {
switch dt.(type) {
case *arrow.Date32Type:
return arr
case *arrow.Date64Type:
return reverseTransformDate64(arr)
default:
panic("unsupported " + dt.String() + " type in reverseTransformFromDate32")
}

return builder.NewArray()
}
2 changes: 2 additions & 0 deletions parquet/read.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,8 @@ func reverseTransformArray(dt arrow.DataType, arr arrow.Array) arrow.Array {
return reverseTransformTime32(dt.(*arrow.Time32Type), arr)
case *array.Time64:
return reverseTransformTime64(dt.(*arrow.Time64Type), arr)
case *array.Date32:
return reverseTransformFromDate32(dt, arr)
case *array.Struct:
dt := dt.(*arrow.StructType)
children := make([]arrow.ArrayData, arr.NumField())
Expand Down
4 changes: 0 additions & 4 deletions parquet/time.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,15 +1,11 @@
package parquet

import (
"time"

"github.com/apache/arrow/go/v15/arrow"
"github.com/apache/arrow/go/v15/arrow/array"
"github.com/apache/arrow/go/v15/arrow/memory"
)

type toTimeFunc func(arrow.Timestamp) time.Time

func reverseTransformTime32(dt *arrow.Time32Type, arr *array.Time32) arrow.Array {
builder := array.NewTime32Builder(memory.DefaultAllocator, dt)

Expand Down
11 changes: 0 additions & 11 deletions parquet/timestamp.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
package parquet

import (
"fmt"

"github.com/apache/arrow/go/v15/arrow"
"github.com/apache/arrow/go/v15/arrow/array"
"github.com/apache/arrow/go/v15/arrow/memory"
Expand All @@ -25,16 +23,7 @@ func reverseTransformTimestamp(dt *arrow.TimestampType, arr *array.Timestamp) ar
}

func reverseTransformFromTimestamp(dt arrow.DataType, arr *array.Timestamp) arrow.Array {
toTime, err := arr.DataType().(*arrow.TimestampType).GetToTimeFunc()
if err != nil {
panic(fmt.Errorf("failed GetToTimeFunc: %w", err))
}

switch dt := dt.(type) {
case *arrow.Date32Type:
return reverseTransformDate32(arr, toTime)
case *arrow.Date64Type:
return reverseTransformDate64(arr, toTime)
case *arrow.TimestampType:
return reverseTransformTimestamp(dt, arr)
default:
Expand Down
9 changes: 3 additions & 6 deletions parquet/write_read_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,7 @@ import (
func TestWriteRead(t *testing.T) {
const rows = 10
var b bytes.Buffer
// We can't test DATE64 since arrow changes it to DATE32 see https://github.com/apache/arrow/pull/39460
table := schema.TestTable("test", schema.TestSourceOptions{SkipDates: true})
table := schema.TestTable("test", schema.TestSourceOptions{})
sourceName := "test-source"
syncTime := time.Now().UTC().Round(time.Second)
opts := schema.GenTestDataOptions{
Expand Down Expand Up @@ -66,8 +65,7 @@ func TestWriteRead(t *testing.T) {
func TestWriteReadSliced(t *testing.T) {
const rows = 10
var b bytes.Buffer
// We can't test DATE64 since arrow changes it to DATE32 see https://github.com/apache/arrow/pull/39460
table := schema.TestTable("test", schema.TestSourceOptions{SkipDates: true})
table := schema.TestTable("test", schema.TestSourceOptions{})
sourceName := "test-source"
syncTime := time.Now().UTC().Round(time.Second)
opts := schema.GenTestDataOptions{
Expand Down Expand Up @@ -114,8 +112,7 @@ func TestWriteReadSliced(t *testing.T) {
}

func BenchmarkWrite(b *testing.B) {
// We can't test DATE64 since arrow changes it to DATE32 see https://github.com/apache/arrow/pull/39460
table := schema.TestTable("test", schema.TestSourceOptions{SkipDates: true})
table := schema.TestTable("test", schema.TestSourceOptions{})
sourceName := "test-source"
syncTime := time.Now().UTC().Round(time.Second)
opts := schema.GenTestDataOptions{
Expand Down
Loading