fix: Support Arrow date64 in Parquet (#410)

Follow-up for #407 Closes cloudquery/cloudquery#16143
cloudquery · Jan 15, 2024 · 27c1c0c · 27c1c0c
1 parent fad9897
commit 27c1c0c
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 37 deletions.
diff --git a/parquet/date.go b/parquet/date.go
@@ -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()
 }
diff --git a/parquet/read.go b/parquet/read.go
@@ -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())

diff --git a/parquet/time.go b/parquet/time.go
@@ -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)
 

diff --git a/parquet/timestamp.go b/parquet/timestamp.go
@@ -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"
@@ -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:

diff --git a/parquet/write_read_test.go b/parquet/write_read_test.go
@@ -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{
@@ -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{
@@ -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{