fix(go/adbc/driver/snowflake): properly handle time fields (#1021)

Fixes #1019
apache · Sep 1, 2023 · d772fd1 · d772fd1
1 parent 1bb507f
commit d772fd1
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 3 deletions.
diff --git a/go/adbc/driver/snowflake/driver_test.go b/go/adbc/driver/snowflake/driver_test.go
@@ -333,13 +333,25 @@ func (suite *SnowflakeTests) TestSqlIngestTimestamp() {
 	sc := arrow.NewSchema([]arrow.Field{{
 		Name: "col", Type: arrow.FixedWidthTypes.Timestamp_us,
 		Nullable: true,
-	}}, nil)
+	}, {
+		Name: "col2", Type: arrow.FixedWidthTypes.Time64us,
+		Nullable: true,
+	}, {
+		Name: "col3", Type: arrow.PrimitiveTypes.Int64,
+		Nullable: true,
+	},
+	}, nil)
 
 	bldr := array.NewRecordBuilder(memory.DefaultAllocator, sc)
 	defer bldr.Release()
 
 	tbldr := bldr.Field(0).(*array.TimestampBuilder)
 	tbldr.AppendValues([]arrow.Timestamp{0, 0, 42}, []bool{false, true, true})
+	tmbldr := bldr.Field(1).(*array.Time64Builder)
+	tmbldr.AppendValues([]arrow.Time64{420000, 0, 86000}, []bool{true, false, true})
+	ibldr := bldr.Field(2).(*array.Int64Builder)
+	ibldr.AppendValues([]int64{-1, 25, 0}, []bool{true, true, false})
+
 	rec := bldr.NewRecord()
 	defer rec.Release()
 

diff --git a/go/adbc/driver/snowflake/record_reader.go b/go/adbc/driver/snowflake/record_reader.go
@@ -110,9 +110,15 @@ func getTransformer(sc *arrow.Schema, ld gosnowflake.ArrowStreamLoader) (*arrow.
 				}
 			}
 		case "TIME":
-			f.Type = arrow.FixedWidthTypes.Time64ns
+			var dt arrow.DataType
+			if srcMeta.Scale < 6 {
+				dt = &arrow.Time32Type{Unit: arrow.TimeUnit(srcMeta.Scale / 3)}
+			} else {
+				dt = &arrow.Time64Type{Unit: arrow.TimeUnit(srcMeta.Scale / 3)}
+			}
+			f.Type = dt
 			transformers[i] = func(ctx context.Context, a arrow.Array) (arrow.Array, error) {
-				return compute.CastArray(ctx, a, compute.SafeCastOptions(f.Type))
+				return compute.CastArray(ctx, a, compute.SafeCastOptions(dt))
 			}
 		case "TIMESTAMP_NTZ":
 			dt := &arrow.TimestampType{Unit: arrow.TimeUnit(srcMeta.Scale / 3)}