Repository: snowflakedb/gosnowflake
Pull request #1042 (2024-01-30 01:49:04)

support millisecond and microsecond for arrow batch timestamp
### Description
Address https://github.com/snowflakedb/gosnowflake/issues/910 when driver only needs microsecond or millisecond precision. 

In our use case, the connector passes the arrow record from Snowflake to another service for decoding. If we only want microsecond precision, we can avoid the need to port the `arrowSnowflakeTimestampToTime` logic.


it would be best if we don't need to repack Snowflake response from Snowflake internal format to arrow timestamps in driver. I understand that might require your backend work. In the meanwhile we can have the driver change as a short term solution. 

### Checklist
- [x] Code compiles correctly
- [x] Run ``make fmt`` to fix inconsistent formats
- [x] Run ``make lint`` to get lint errors and fix all of them
- [x] Created tests which fail without the change (if possible)
- [x] All tests passing
- [x] Extended the README / documentation, if necessary


github-actions[bot]: **CLA Assistant Lite bot:** <br/>Thank you for your submission, we really appreciate it. Like many open-source projects, we ask that you sign our [Contributor License Agreement](https://github.com/snowflakedb/CLA/blob/main/README.md) before we can accept your contribution. You can sign the CLA by just posting a Pull Request Comment same as the below format.<br/>
   - - -
   I have read the CLA Document and I hereby sign the CLA
   - - -
   <sub>You can retrigger this bot by commenting **recheck** in this Pull Request</sub>
Yifeng-Sigma: I have read the CLA Document and I hereby sign the CLA
Yifeng-Sigma: recheck
Yifeng-Sigma: @sfc-gh-pfus Could you please trigger the workflow? Thank you!

Patch:
diff --git a/cmd/arrow/batches/arrow_batches.go b/cmd/arrow/batches/arrow_batches.go
index a8746dc..48fc936 100644
--- a/cmd/arrow/batches/arrow_batches.go
+++ b/cmd/arrow/batches/arrow_batches.go
@@ -51,9 +51,9 @@ func main() {
 	}
 
 	ctx :=
-		sf.WithOriginalTimestamp(
+		sf.WithArrowBatchesTimestampOption(
 			sf.WithArrowAllocator(
-				sf.WithArrowBatches(context.Background()), memory.DefaultAllocator))
+				sf.WithArrowBatches(context.Background()), memory.DefaultAllocator), sf.UseOriginalTimestamp)
 
 	query := "SELECT SEQ4(), 'example ' || (SEQ4() * 2), " +
 		" TO_TIMESTAMP_NTZ('9999-01-01 13:13:13.' || LPAD(SEQ4(),9,'0'))  ltz " +
diff --git a/converter.go b/converter.go
index ee6aec7..d4b3a83 100644
--- a/converter.go
+++ b/converter.go
@@ -39,6 +39,21 @@ const (
 	TimeType
 )
 
+type snowflakeArrowBatchesTimestampOption int
+
+const (
+	// UseNanosecondTimestamp uses arrow.Timestamp in nanosecond precision, could cause ErrTooHighTimestampPrecision if arrow.Timestamp cannot fit original timestamp values.
+	UseNanosecondTimestamp snowflakeArrowBatchesTimestampOption = iota
+	// UseMicrosecondTimestamp uses arrow.Timestamp in microsecond precision
+	UseMicrosecondTimestamp
+	// UseMillisecondTimestamp uses arrow.Timestamp in millisecond precision
+	UseMillisecondTimestamp
+	// UseSecondTimestamp uses arrow.Timestamp in second precision
+	UseSecondTimestamp
+	// UseOriginalTimestamp uses original timestamp struct returned by Snowflake. It can be used in case arrow.Timestamp cannot fit original timestamp values.
+	UseOriginalTimestamp
+)
+
 type interfaceArrayBinding struct {
 	hasTimezone       bool
 	tzType            timezoneType
@@ -968,19 +983,22 @@ func higherPrecisionEnabled(ctx context.Context) bool {
 	return ok && d
 }
 
-func originalTimestampEnabled(ctx context.Context) bool {
-	v := ctx.Value(enableOriginalTimestamp)
+func getArrowBatchesTimestampOption(ctx context.Context) snowflakeArrowBatchesTimestampOption {
+	v := ctx.Value(arrowBatchesTimestampOption)
 	if v == nil {
-		return false
+		return UseNanosecondTimestamp
 	}
-	d, ok := v.(bool)
-	return ok && d
+	o, ok := v.(snowflakeArrowBatchesTimestampOption)
+	if !ok {
+		return UseNanosecondTimestamp
+	}
+	return o
 }
 
 func arrowToRecord(ctx context.Context, record arrow.Record, pool memory.Allocator, rowType []execResponseRowType, loc *time.Location) (arrow.Record, error) {
-	useOriginalTimestamp := originalTimestampEnabled(ctx)
+	arrowBatchesTimestampOption := getArrowBatchesTimestampOption(ctx)
 
-	s, err := recordToSchema(record.Schema(), rowType, loc, useOriginalTimestamp)
+	s, err := recordToSchema(record.Schema(), rowType, loc, arrowBatchesTimestampOption)
 	if err != nil {
 		return nil, err
 	}
@@ -1029,27 +1047,49 @@ func arrowToRecord(ctx context.Context, record arrow.Record, pool memory.Allocat
 			}
 			defer newCol.Release()
 		case timestampNtzType, timestampLtzType, timestampTzType:
-			if useOriginalTimestamp {
+			if arrowBatchesTimestampOption == UseOriginalTimestamp {
 				// do nothing - return timestamp as is
 			} else {
+				var unit arrow.TimeUnit
+				switch arrowBatchesTimestampOption {
+				case UseMicrosecondTimestamp:
+					unit = arrow.Microsecond
+				case UseMillisecondTimestamp:
+					unit = arrow.Millisecond
+				case UseSecondTimestamp:
+					unit = arrow.Second
+				case UseNanosecondTimestamp:
+					unit = arrow.Nanosecond
+				}
 				var tb *array.TimestampBuilder
 				if snowflakeType == timestampLtzType {
-					tb = array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: loc.String()})
+					tb = array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: unit, TimeZone: loc.String()})
 				} else {
-					tb = array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Nanosecond})
+					tb = array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: unit})
 				}
 				defer tb.Release()
 
 				for i := 0; i < int(numRows); i++ {
 					ts := arrowSnowflakeTimestampToTime(col, snowflakeType, int(srcColumnMeta.Scale), i, loc)
 					if ts != nil {
-						ar := arrow.Timestamp(ts.UnixNano())
-						// in case of overflow in arrow timestamp return error
-						if ts.Year() != ar.ToTime(arrow.Nanosecond).Year() {
-							return nil, &SnowflakeError{
-								Number:   ErrTooHighTimestampPrecision,
-								SQLState: SQLStateInvalidDataTimeFormat,
-								Message:  fmt.Sprintf("Cannot convert timestamp %v in column %v to Arrow.Timestamp data type due to too high precision. Please use context with WithOriginalTimestamp.", ts.UTC(), srcColumnMeta.Name),
+						var ar arrow.Timestamp
+						switch arrowBatchesTimestampOption {
+						case UseMicrosecondTimestamp:
+							ar = arrow.Timestamp(ts.UnixMicro())
+						case UseMillisecondTimestamp:
+							ar = arrow.Timestamp(ts.UnixMilli())
+						case UseSecondTimestamp:
+							ar = arrow.Timestamp(ts.Unix())
+						case UseNanosecondTimestamp:
+							ar = arrow.Timestamp(ts.UnixNano())
+							// in case of overflow in arrow timestamp return error
+							// this could only happen for nanosecond case
+							if ts.Year() != ar.ToTime(arrow.Nanosecond).Year() {
+								return nil, &SnowflakeError{
+									Number:   ErrTooHighTimestampPrecision,
+									SQLState: SQLStateInvalidDataTimeFormat,
+									Message:  fmt.Sprintf("Cannot convert timestamp %v in column %v to Arrow.Timestamp data type due to too high precision. Please use context with WithOriginalTimestamp.", ts.UTC(), srcColumnMeta.Name),
+								}
 							}
 						}
 						tb.Append(ar)
@@ -1067,7 +1107,7 @@ func arrowToRecord(ctx context.Context, record arrow.Record, pool memory.Allocat
 	return array.NewRecord(s, cols, numRows), nil
 }
 
-func recordToSchema(sc *arrow.Schema, rowType []execResponseRowType, loc *time.Location, useOriginalTimestamp bool) (*arrow.Schema, error) {
+func recordToSchema(sc *arrow.Schema, rowType []execResponseRowType, loc *time.Location, timestampOption snowflakeArrowBatchesTimestampOption) (*arrow.Schema, error) {
 	var fields []arrow.Field
 	for i := 0; i < len(sc.Fields()); i++ {
 		f := sc.Field(i)
@@ -1094,16 +1134,28 @@ func recordToSchema(sc *arrow.Schema, rowType []execResponseRowType, loc *time.L
 		case timeType:
 			t = &arrow.Time64Type{Unit: arrow.Nanosecond}
 		case timestampNtzType, timestampTzType:
-			if useOriginalTimestamp {
+			if timestampOption == UseOriginalTimestamp {
 				// do nothing - return timestamp as is
 				converted = false
+			} else if timestampOption == UseMicrosecondTimestamp {
+				t = &arrow.TimestampType{Unit: arrow.Microsecond}
+			} else if timestampOption == UseMillisecondTimestamp {
+				t = &arrow.TimestampType{Unit: arrow.Millisecond}
+			} else if timestampOption == UseSecondTimestamp {
+				t = &arrow.TimestampType{Unit: arrow.Second}
 			} else {
 				t = &arrow.TimestampType{Unit: arrow.Nanosecond}
 			}
 		case timestampLtzType:
-			if useOriginalTimestamp {
+			if timestampOption == UseOriginalTimestamp {
 				// do nothing - return timestamp as is
 				converted = false
+			} else if timestampOption == UseMicrosecondTimestamp {
+				t = &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: loc.String()}
+			} else if timestampOption == UseMillisecondTimestamp {
+				t = &arrow.TimestampType{Unit: arrow.Millisecond, TimeZone: loc.String()}
+			} else if timestampOption == UseSecondTimestamp {
+				t = &arrow.TimestampType{Unit: arrow.Second, TimeZone: loc.String()}
 			} else {
 				t = &arrow.TimestampType{Unit: arrow.Nanosecond, TimeZone: loc.String()}
 			}
diff --git a/doc.go b/doc.go
index 2dd977d..897bc3b 100644
--- a/doc.go
+++ b/doc.go
@@ -486,9 +486,10 @@ Consequently, Snowflake uses a custom timestamp format in Arrow, which differs o
 
 If you want to use timestamps in Arrow batches, you have two options:
 
- 1. The Go driver can reduce timestamp struct into simple Arrow Timestamp, if all timestamp values fit into Arrow timestamp field.
+ 1. The Go driver can reduce timestamp struct into simple Arrow Timestamp, if you set `WithArrowBatchesTimestampOption` to nanosecond, microsecond, millisecond or second.
+    For nanosecond, some timestamp values might not fit into Arrow timestamp. E.g after year 2262 or before 1677.
  2. You can use native Snowflake values. In that case you will receive complex structs as described above. To transform Snowflake values into the Golang time.Time struct you can use `ArrowSnowflakeTimestampToTime`.
-    To enable this feature, you must use `WithOriginalTimestamp` context.
+    To enable this feature, you must use `WithArrowBatchesTimestampOption` context with value set to`UseOriginalTimestamp`.
 
 # Binding Parameters
 
diff --git a/util.go b/util.go
index 4235028..000a17d 100644
--- a/util.go
+++ b/util.go
@@ -19,18 +19,18 @@ import (
 type contextKey string
 
 const (
-	multiStatementCount     contextKey = "MULTI_STATEMENT_COUNT"
-	asyncMode               contextKey = "ASYNC_MODE_QUERY"
-	queryIDChannel          contextKey = "QUERY_ID_CHANNEL"
-	snowflakeRequestIDKey   contextKey = "SNOWFLAKE_REQUEST_ID"
-	fetchResultByID         contextKey = "SF_FETCH_RESULT_BY_ID"
-	fileStreamFile          contextKey = "STREAMING_PUT_FILE"
-	fileTransferOptions     contextKey = "FILE_TRANSFER_OPTIONS"
-	enableHigherPrecision   contextKey = "ENABLE_HIGHER_PRECISION"
-	arrowBatches            contextKey = "ARROW_BATCHES"
-	arrowAlloc              contextKey = "ARROW_ALLOC"
-	enableOriginalTimestamp contextKey = "ENABLE_ORIGINAL_TIMESTAMP"
-	queryTag                contextKey = "QUERY_TAG"
+	multiStatementCount         contextKey = "MULTI_STATEMENT_COUNT"
+	asyncMode                   contextKey = "ASYNC_MODE_QUERY"
+	queryIDChannel              contextKey = "QUERY_ID_CHANNEL"
+	snowflakeRequestIDKey       contextKey = "SNOWFLAKE_REQUEST_ID"
+	fetchResultByID             contextKey = "SF_FETCH_RESULT_BY_ID"
+	fileStreamFile              contextKey = "STREAMING_PUT_FILE"
+	fileTransferOptions         contextKey = "FILE_TRANSFER_OPTIONS"
+	enableHigherPrecision       contextKey = "ENABLE_HIGHER_PRECISION"
+	arrowBatches                contextKey = "ARROW_BATCHES"
+	arrowAlloc                  contextKey = "ARROW_ALLOC"
+	arrowBatchesTimestampOption contextKey = "ARROW_BATCHES_TIMESTAMP_OPTION"
+	queryTag                    contextKey = "QUERY_TAG"
 )
 
 const (
@@ -111,8 +111,21 @@ func WithArrowAllocator(ctx context.Context, pool memory.Allocator) context.Cont
 // WithOriginalTimestamp in combination with WithArrowBatches returns a context
 // that allows users to retrieve arrow.Record with original timestamp struct returned by Snowflake.
 // It can be used in case arrow.Timestamp cannot fit original timestamp values.
+//
+// Deprecated: please use WithArrowBatchesTimestampOption instead.
 func WithOriginalTimestamp(ctx context.Context) context.Context {
-	return context.WithValue(ctx, enableOriginalTimestamp, true)
+	return context.WithValue(ctx, arrowBatchesTimestampOption, UseOriginalTimestamp)
+}
+
+// WithArrowBatchesTimestampOption in combination with WithArrowBatches returns a context
+// that allows users to retrieve arrow.Record with different timestamp options.
+// UseNanosecondTimestamp: arrow.Timestamp in nanosecond precision, could cause ErrTooHighTimestampPrecision if arrow.Timestamp cannot fit original timestamp values.
+// UseMicrosecondTimestamp: arrow.Timestamp in microsecond precision
+// UseMillisecondTimestamp: arrow.Timestamp in millisecond precision
+// UseSecondTimestamp: arrow.Timestamp in second precision
+// UseOriginalTimestamp: original timestamp struct returned by Snowflake. It can be used in case arrow.Timestamp cannot fit original timestamp values.
+func WithArrowBatchesTimestampOption(ctx context.Context, option snowflakeArrowBatchesTimestampOption) context.Context {
+	return context.WithValue(ctx, arrowBatchesTimestampOption, option)
 }
 
 // WithQueryTag returns a context that will set the given tag as the QUERY_TAG


Test patch:
diff --git a/converter_test.go b/converter_test.go
index bcde6f6..d596bd0 100644
--- a/converter_test.go
+++ b/converter_test.go
@@ -909,17 +909,17 @@ func TestArrowToRecord(t *testing.T) {
 	}
 
 	for _, tc := range []struct {
-		logical  string
-		physical string
-		sc       *arrow.Schema
-		rowType  execResponseRowType
-		values   interface{}
-		error    string
-		origTS   bool
-		nrows    int
-		builder  array.Builder
-		append   func(b array.Builder, vs interface{})
-		compare  func(src interface{}, rec arrow.Record) int
+		logical                     string
+		physical                    string
+		sc                          *arrow.Schema
+		rowType                     execResponseRowType
+		values                      interface{}
+		error                       string
+		arrowBatchesTimestampOption snowflakeArrowBatchesTimestampOption
+		nrows                       int
+		builder                     array.Builder
+		append                      func(b array.Builder, vs interface{})
+		compare                     func(src interface{}, rec arrow.Record) int
 	}{
 		{
 			logical:  "fixed",
@@ -1227,6 +1227,93 @@ func TestArrowToRecord(t *testing.T) {
 				return -1
 			},
 		},
+		// microsecond timestamp_ntz
+		{
+			logical:                     "timestamp_ntz",
+			physical:                    "struct", // timestamp_ntz with scale 4..9 -> int64 + int32
+			values:                      []time.Time{time.Now().Truncate(time.Microsecond), localTime.Truncate(time.Microsecond)},
+			arrowBatchesTimestampOption: UseMicrosecondTimestamp,
+			nrows:                       2,
+			rowType:                     execResponseRowType{Scale: 9},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),
+			builder:                     array.NewStructBuilder(pool, timestampNtzStruct),
+			append: func(b array.Builder, vs interface{}) {
+				sb := b.(*array.StructBuilder)
+				valids = []bool{true, true}
+				sb.AppendValues(valids)
+				for _, t := range vs.([]time.Time) {
+					sb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())
+					sb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))
+				}
+			},
+			compare: func(src interface{}, convertedRec arrow.Record) int {
+				srcvs := src.([]time.Time)
+				for i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {
+					if !srcvs[i].Equal(t.ToTime(arrow.Microsecond)) {
+						return i
+					}
+				}
+				return -1
+			},
+		},
+		// millisecond timestamp_ntz
+		{
+			logical:                     "timestamp_ntz",
+			physical:                    "struct", // timestamp_ntz with scale 4..9 -> int64 + int32
+			values:                      []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond)},
+			arrowBatchesTimestampOption: UseMillisecondTimestamp,
+			nrows:                       2,
+			rowType:                     execResponseRowType{Scale: 9},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),
+			builder:                     array.NewStructBuilder(pool, timestampNtzStruct),
+			append: func(b array.Builder, vs interface{}) {
+				sb := b.(*array.StructBuilder)
+				valids = []bool{true, true}
+				sb.AppendValues(valids)
+				for _, t := range vs.([]time.Time) {
+					sb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())
+					sb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))
+				}
+			},
+			compare: func(src interface{}, convertedRec arrow.Record) int {
+				srcvs := src.([]time.Time)
+				for i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {
+					if !srcvs[i].Equal(t.ToTime(arrow.Millisecond)) {
+						return i
+					}
+				}
+				return -1
+			},
+		},
+		// second timestamp_ntz
+		{
+			logical:                     "timestamp_ntz",
+			physical:                    "struct", // timestamp_ntz with scale 4..9 -> int64 + int32
+			values:                      []time.Time{time.Now().Truncate(time.Second), localTime.Truncate(time.Second)},
+			arrowBatchesTimestampOption: UseSecondTimestamp,
+			nrows:                       2,
+			rowType:                     execResponseRowType{Scale: 9},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),
+			builder:                     array.NewStructBuilder(pool, timestampNtzStruct),
+			append: func(b array.Builder, vs interface{}) {
+				sb := b.(*array.StructBuilder)
+				valids = []bool{true, true}
+				sb.AppendValues(valids)
+				for _, t := range vs.([]time.Time) {
+					sb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())
+					sb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))
+				}
+			},
+			compare: func(src interface{}, convertedRec arrow.Record) int {
+				srcvs := src.([]time.Time)
+				for i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {
+					if !srcvs[i].Equal(t.ToTime(arrow.Second)) {
+						return i
+					}
+				}
+				return -1
+			},
+		},
 		{
 			logical:  "timestamp_ntz",
 			physical: "error",
@@ -1244,14 +1331,14 @@ func TestArrowToRecord(t *testing.T) {
 			compare: func(src interface{}, convertedRec arrow.Record) int { return 0 },
 		},
 		{
-			logical:  "timestamp_ntz",
-			physical: "int64 with original timestamp", // timestamp_ntz with scale 0..3 -> int64
-			values:   []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond), localTimeFarIntoFuture.Truncate(time.Millisecond)},
-			origTS:   true,
-			nrows:    3,
-			rowType:  execResponseRowType{Scale: 3},
-			sc:       arrow.NewSchema([]arrow.Field{{Type: &arrow.Int64Type{}}}, nil),
-			builder:  array.NewInt64Builder(pool),
+			logical:                     "timestamp_ntz",
+			physical:                    "int64 with original timestamp", // timestamp_ntz with scale 0..3 -> int64
+			values:                      []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond), localTimeFarIntoFuture.Truncate(time.Millisecond)},
+			arrowBatchesTimestampOption: UseOriginalTimestamp,
+			nrows:                       3,
+			rowType:                     execResponseRowType{Scale: 3},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: &arrow.Int64Type{}}}, nil),
+			builder:                     array.NewInt64Builder(pool),
 			append: func(b array.Builder, vs interface{}) {
 				for _, t := range vs.([]time.Time) {
 					b.(*array.Int64Builder).Append(t.UnixMilli()) // Millisecond for scale = 3
@@ -1269,14 +1356,14 @@ func TestArrowToRecord(t *testing.T) {
 			},
 		},
 		{
-			logical:  "timestamp_ntz",
-			physical: "struct with original timestamp", // timestamp_ntz with scale 4..9 -> int64 + int32
-			values:   []time.Time{time.Now(), localTime, localTimeFarIntoFuture},
-			origTS:   true,
-			nrows:    3,
-			rowType:  execResponseRowType{Scale: 9},
-			sc:       arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),
-			builder:  array.NewStructBuilder(pool, timestampNtzStruct),
+			logical:                     "timestamp_ntz",
+			physical:                    "struct with original timestamp", // timestamp_ntz with scale 4..9 -> int64 + int32
+			values:                      []time.Time{time.Now(), localTime, localTimeFarIntoFuture},
+			arrowBatchesTimestampOption: UseOriginalTimestamp,
+			nrows:                       3,
+			rowType:                     execResponseRowType{Scale: 9},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),
+			builder:                     array.NewStructBuilder(pool, timestampNtzStruct),
 			append: func(b array.Builder, vs interface{}) {
 				sb := b.(*array.StructBuilder)
 				valids = []bool{true, true, true}
@@ -1347,6 +1434,93 @@ func TestArrowToRecord(t *testing.T) {
 				return -1
 			},
 		},
+		// microsecond timestamp_ltz
+		{
+			logical:                     "timestamp_ltz",
+			physical:                    "struct", // timestamp_ntz with scale 4..9 -> int64 + int32
+			values:                      []time.Time{time.Now().Truncate(time.Microsecond), localTime.Truncate(time.Microsecond)},
+			arrowBatchesTimestampOption: UseMicrosecondTimestamp,
+			nrows:                       2,
+			rowType:                     execResponseRowType{Scale: 9},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),
+			builder:                     array.NewStructBuilder(pool, timestampNtzStruct),
+			append: func(b array.Builder, vs interface{}) {
+				sb := b.(*array.StructBuilder)
+				valids = []bool{true, true}
+				sb.AppendValues(valids)
+				for _, t := range vs.([]time.Time) {
+					sb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())
+					sb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))
+				}
+			},
+			compare: func(src interface{}, convertedRec arrow.Record) int {
+				srcvs := src.([]time.Time)
+				for i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {
+					if !srcvs[i].Equal(t.ToTime(arrow.Microsecond)) {
+						return i
+					}
+				}
+				return -1
+			},
+		},
+		// millisecond timestamp_ltz
+		{
+			logical:                     "timestamp_ltz",
+			physical:                    "struct", // timestamp_ntz with scale 4..9 -> int64 + int32
+			values:                      []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond)},
+			arrowBatchesTimestampOption: UseMillisecondTimestamp,
+			nrows:                       2,
+			rowType:                     execResponseRowType{Scale: 9},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),
+			builder:                     array.NewStructBuilder(pool, timestampNtzStruct),
+			append: func(b array.Builder, vs interface{}) {
+				sb := b.(*array.StructBuilder)
+				valids = []bool{true, true}
+				sb.AppendValues(valids)
+				for _, t := range vs.([]time.Time) {
+					sb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())
+					sb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))
+				}
+			},
+			compare: func(src interface{}, convertedRec arrow.Record) int {
+				srcvs := src.([]time.Time)
+				for i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {
+					if !srcvs[i].Equal(t.ToTime(arrow.Millisecond)) {
+						return i
+					}
+				}
+				return -1
+			},
+		},
+		// second timestamp_ltz
+		{
+			logical:                     "timestamp_ltz",
+			physical:                    "struct", // timestamp_ntz with scale 4..9 -> int64 + int32
+			values:                      []time.Time{time.Now().Truncate(time.Second), localTime.Truncate(time.Second)},
+			arrowBatchesTimestampOption: UseSecondTimestamp,
+			nrows:                       2,
+			rowType:                     execResponseRowType{Scale: 9},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: timestampNtzStruct}}, nil),
+			builder:                     array.NewStructBuilder(pool, timestampNtzStruct),
+			append: func(b array.Builder, vs interface{}) {
+				sb := b.(*array.StructBuilder)
+				valids = []bool{true, true}
+				sb.AppendValues(valids)
+				for _, t := range vs.([]time.Time) {
+					sb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())
+					sb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))
+				}
+			},
+			compare: func(src interface{}, convertedRec arrow.Record) int {
+				srcvs := src.([]time.Time)
+				for i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {
+					if !srcvs[i].Equal(t.ToTime(arrow.Second)) {
+						return i
+					}
+				}
+				return -1
+			},
+		},
 		{
 			logical:  "timestamp_ltz",
 			physical: "error",
@@ -1364,14 +1538,14 @@ func TestArrowToRecord(t *testing.T) {
 			compare: func(src interface{}, convertedRec arrow.Record) int { return 0 },
 		},
 		{
-			logical:  "timestamp_ltz",
-			physical: "int64 with original timestamp", // timestamp_ntz with scale 0..3 -> int64
-			values:   []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond), localTimeFarIntoFuture.Truncate(time.Millisecond)},
-			origTS:   true,
-			nrows:    3,
-			rowType:  execResponseRowType{Scale: 3},
-			sc:       arrow.NewSchema([]arrow.Field{{Type: &arrow.Int64Type{}}}, nil),
-			builder:  array.NewInt64Builder(pool),
+			logical:                     "timestamp_ltz",
+			physical:                    "int64 with original timestamp", // timestamp_ntz with scale 0..3 -> int64
+			values:                      []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond), localTimeFarIntoFuture.Truncate(time.Millisecond)},
+			arrowBatchesTimestampOption: UseOriginalTimestamp,
+			nrows:                       3,
+			rowType:                     execResponseRowType{Scale: 3},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: &arrow.Int64Type{}}}, nil),
+			builder:                     array.NewInt64Builder(pool),
 			append: func(b array.Builder, vs interface{}) {
 				for _, t := range vs.([]time.Time) {
 					b.(*array.Int64Builder).Append(t.UnixMilli()) // Millisecond for scale = 3
@@ -1389,14 +1563,14 @@ func TestArrowToRecord(t *testing.T) {
 			},
 		},
 		{
-			logical:  "timestamp_ltz",
-			physical: "struct with original timestamp", // timestamp_ntz with scale 4..9 -> int64 + int32
-			values:   []time.Time{time.Now(), localTime, localTimeFarIntoFuture},
-			origTS:   true,
-			nrows:    3,
-			rowType:  execResponseRowType{Scale: 9},
-			sc:       arrow.NewSchema([]arrow.Field{{Type: timestampLtzStruct}}, nil),
-			builder:  array.NewStructBuilder(pool, timestampLtzStruct),
+			logical:                     "timestamp_ltz",
+			physical:                    "struct with original timestamp", // timestamp_ntz with scale 4..9 -> int64 + int32
+			values:                      []time.Time{time.Now(), localTime, localTimeFarIntoFuture},
+			arrowBatchesTimestampOption: UseOriginalTimestamp,
+			nrows:                       3,
+			rowType:                     execResponseRowType{Scale: 9},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: timestampLtzStruct}}, nil),
+			builder:                     array.NewStructBuilder(pool, timestampLtzStruct),
 			append: func(b array.Builder, vs interface{}) {
 				sb := b.(*array.StructBuilder)
 				valids = []bool{true, true, true}
@@ -1472,15 +1646,105 @@ func TestArrowToRecord(t *testing.T) {
 				return -1
 			},
 		},
+		// microsecond timestamp_tz
 		{
-			logical:  "timestamp_tz",
-			physical: "struct2 with original timestamp", // timestamp_ntz with scale 0..3 -> int64 + int32
-			values:   []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond), localTimeFarIntoFuture.Truncate(time.Millisecond)},
-			origTS:   true,
-			nrows:    3,
-			rowType:  execResponseRowType{Scale: 3},
-			sc:       arrow.NewSchema([]arrow.Field{{Type: timestampTzStructWithoutFraction}}, nil),
-			builder:  array.NewStructBuilder(pool, timestampTzStructWithoutFraction),
+			logical:                     "timestamp_tz",
+			physical:                    "struct3", // timestamp_tz with scale 4..9 -> int64 + int32 + int32
+			values:                      []time.Time{time.Now().Truncate(time.Microsecond), localTime.Truncate(time.Microsecond)},
+			arrowBatchesTimestampOption: UseMicrosecondTimestamp,
+			nrows:                       2,
+			rowType:                     execResponseRowType{Scale: 9},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: timestampTzStructWithFraction}}, nil),
+			builder:                     array.NewStructBuilder(pool, timestampTzStructWithFraction),
+			append: func(b array.Builder, vs interface{}) {
+				sb := b.(*array.StructBuilder)
+				valids = []bool{true, true}
+				sb.AppendValues(valids)
+				for _, t := range vs.([]time.Time) {
+					sb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())
+					sb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))
+					sb.FieldBuilder(2).(*array.Int32Builder).Append(int32(0)) // timezone index - not important in tests
+				}
+			},
+			compare: func(src interface{}, convertedRec arrow.Record) int {
+				srcvs := src.([]time.Time)
+				for i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {
+					if !srcvs[i].Equal(t.ToTime(arrow.Microsecond)) {
+						return i
+					}
+				}
+				return -1
+			},
+		},
+		// millisecond timestamp_tz
+		{
+			logical:                     "timestamp_tz",
+			physical:                    "struct3", // timestamp_tz with scale 4..9 -> int64 + int32 + int32
+			values:                      []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond)},
+			arrowBatchesTimestampOption: UseMillisecondTimestamp,
+			nrows:                       2,
+			rowType:                     execResponseRowType{Scale: 9},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: timestampTzStructWithFraction}}, nil),
+			builder:                     array.NewStructBuilder(pool, timestampTzStructWithFraction),
+			append: func(b array.Builder, vs interface{}) {
+				sb := b.(*array.StructBuilder)
+				valids = []bool{true, true}
+				sb.AppendValues(valids)
+				for _, t := range vs.([]time.Time) {
+					sb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())
+					sb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))
+					sb.FieldBuilder(2).(*array.Int32Builder).Append(int32(0)) // timezone index - not important in tests
+				}
+			},
+			compare: func(src interface{}, convertedRec arrow.Record) int {
+				srcvs := src.([]time.Time)
+				for i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {
+					if !srcvs[i].Equal(t.ToTime(arrow.Millisecond)) {
+						return i
+					}
+				}
+				return -1
+			},
+		},
+		// second timestamp_tz
+		{
+			logical:                     "timestamp_tz",
+			physical:                    "struct3", // timestamp_tz with scale 4..9 -> int64 + int32 + int32
+			values:                      []time.Time{time.Now().Truncate(time.Second), localTime.Truncate(time.Second)},
+			arrowBatchesTimestampOption: UseSecondTimestamp,
+			nrows:                       2,
+			rowType:                     execResponseRowType{Scale: 9},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: timestampTzStructWithFraction}}, nil),
+			builder:                     array.NewStructBuilder(pool, timestampTzStructWithFraction),
+			append: func(b array.Builder, vs interface{}) {
+				sb := b.(*array.StructBuilder)
+				valids = []bool{true, true}
+				sb.AppendValues(valids)
+				for _, t := range vs.([]time.Time) {
+					sb.FieldBuilder(0).(*array.Int64Builder).Append(t.Unix())
+					sb.FieldBuilder(1).(*array.Int32Builder).Append(int32(t.Nanosecond()))
+					sb.FieldBuilder(2).(*array.Int32Builder).Append(int32(0)) // timezone index - not important in tests
+				}
+			},
+			compare: func(src interface{}, convertedRec arrow.Record) int {
+				srcvs := src.([]time.Time)
+				for i, t := range convertedRec.Column(0).(*array.Timestamp).TimestampValues() {
+					if !srcvs[i].Equal(t.ToTime(arrow.Second)) {
+						return i
+					}
+				}
+				return -1
+			},
+		},
+		{
+			logical:                     "timestamp_tz",
+			physical:                    "struct2 with original timestamp", // timestamp_ntz with scale 0..3 -> int64 + int32
+			values:                      []time.Time{time.Now().Truncate(time.Millisecond), localTime.Truncate(time.Millisecond), localTimeFarIntoFuture.Truncate(time.Millisecond)},
+			arrowBatchesTimestampOption: UseOriginalTimestamp,
+			nrows:                       3,
+			rowType:                     execResponseRowType{Scale: 3},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: timestampTzStructWithoutFraction}}, nil),
+			builder:                     array.NewStructBuilder(pool, timestampTzStructWithoutFraction),
 			append: func(b array.Builder, vs interface{}) {
 				sb := b.(*array.StructBuilder)
 				valids = []bool{true, true, true}
@@ -1502,14 +1766,14 @@ func TestArrowToRecord(t *testing.T) {
 			},
 		},
 		{
-			logical:  "timestamp_tz",
-			physical: "struct3 with original timestamp", // timestamp_ntz with scale 4..9 -> int64 + int32 + int32
-			values:   []time.Time{time.Now(), localTime, localTimeFarIntoFuture},
-			origTS:   true,
-			nrows:    3,
-			rowType:  execResponseRowType{Scale: 9},
-			sc:       arrow.NewSchema([]arrow.Field{{Type: timestampTzStructWithFraction}}, nil),
-			builder:  array.NewStructBuilder(pool, timestampTzStructWithFraction),
+			logical:                     "timestamp_tz",
+			physical:                    "struct3 with original timestamp", // timestamp_ntz with scale 4..9 -> int64 + int32 + int32
+			values:                      []time.Time{time.Now(), localTime, localTimeFarIntoFuture},
+			arrowBatchesTimestampOption: UseOriginalTimestamp,
+			nrows:                       3,
+			rowType:                     execResponseRowType{Scale: 9},
+			sc:                          arrow.NewSchema([]arrow.Field{{Type: timestampTzStructWithFraction}}, nil),
+			builder:                     array.NewStructBuilder(pool, timestampTzStructWithFraction),
 			append: func(b array.Builder, vs interface{}) {
 				sb := b.(*array.StructBuilder)
 				valids = []bool{true, true, true}
@@ -1576,8 +1840,17 @@ func TestArrowToRecord(t *testing.T) {
 			meta.Type = tc.logical
 
 			ctx := context.Background()
-			if tc.origTS {
-				ctx = WithOriginalTimestamp(ctx)
+			switch tc.arrowBatchesTimestampOption {
+			case UseOriginalTimestamp:
+				ctx = WithArrowBatchesTimestampOption(ctx, UseOriginalTimestamp)
+			case UseSecondTimestamp:
+				ctx = WithArrowBatchesTimestampOption(ctx, UseSecondTimestamp)
+			case UseMillisecondTimestamp:
+				ctx = WithArrowBatchesTimestampOption(ctx, UseMillisecondTimestamp)
+			case UseMicrosecondTimestamp:
+				ctx = WithArrowBatchesTimestampOption(ctx, UseMicrosecondTimestamp)
+			default:
+				ctx = WithArrowBatchesTimestampOption(ctx, UseNanosecondTimestamp)
 			}
 
 			transformedRec, err := arrowToRecord(ctx, rawRec, pool, []execResponseRowType{meta}, localTime.Location())
@@ -1706,7 +1979,7 @@ func TestTimestampConversionWithoutArrowBatches(t *testing.T) {
 	})
 }
 
-func TestTimestampConversionWithArrowBatchesFailsForDistantDates(t *testing.T) {
+func TestTimestampConversionWithArrowBatchesNanosecondFailsForDistantDates(t *testing.T) {
 	timestamps := [2]string{
 		"9999-12-12 23:59:59.999999999", // max
 		"0001-01-01 00:00:00.000000000"} // min
@@ -1743,6 +2016,65 @@ func TestTimestampConversionWithArrowBatchesFailsForDistantDates(t *testing.T) {
 	})
 }
 
+// use arrow.Timestamp with microsecond precision and below should not encounter overflow issue.
+func TestTimestampConversionWithArrowBatchesMicrosecondPassesForDistantDates(t *testing.T) {
+	timestamps := [2]string{
+		"9999-12-12 23:59:59.999999999", // max
+		"0001-01-01 00:00:00.000000000"} // min
+	types := [3]string{"TIMESTAMP_NTZ", "TIMESTAMP_LTZ", "TIMESTAMP_TZ"}
+
+	runSnowflakeConnTest(t, func(sct *SCTest) {
+		ctx := WithArrowBatchesTimestampOption(WithArrowBatches(sct.sc.ctx), UseMicrosecondTimestamp)
+
+		pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
+		defer pool.AssertSize(t, 0)
+		ctx = WithArrowAllocator(ctx, pool)
+
+		for _, tsStr := range timestamps {
+			for _, tp := range types {
+				for scale := 0; scale <= 9; scale++ {
+					t.Run(tp+"("+strconv.Itoa(scale)+")_"+tsStr, func(t *testing.T) {
+
+						query := fmt.Sprintf("SELECT '%s'::%s(%v)", tsStr, tp, scale)
+						rows, err := sct.sc.QueryContext(ctx, query, []driver.NamedValue{})
+						if err != nil {
+							t.Fatalf("failed to query: %v", err)
+						}
+						defer rows.Close()
+
+						// getting result batches
+						batches, err := rows.(*snowflakeRows).GetArrowBatches()
+						if err != nil {
+							t.Error(err)
+						}
+
+						rec, err := batches[0].Fetch()
+						if err != nil {
+							t.Error(err)
+						}
+
+						records := *rec
+						r := records[0]
+						defer r.Release()
+						actual := r.Column(0).(*array.Timestamp).TimestampValues()[0]
+						actualYear := actual.ToTime(arrow.Microsecond).Year()
+
+						ts, err := time.Parse("2006-01-02 15:04:05", tsStr)
+						if err != nil {
+							t.Fatalf("failed to parse time: %v", err)
+						}
+						exp := ts.Truncate(time.Duration(math.Pow10(9 - scale)))
+
+						if actualYear != exp.Year() {
+							t.Fatalf("unexpected year in timestamp, expected: %v, got: %v", exp.Year(), actualYear)
+						}
+					})
+				}
+			}
+		}
+	})
+}
+
 func TestTimestampConversionWithArrowBatchesAndWithOriginalTimestamp(t *testing.T) {
 	timestamps := [3]string{
 		"2000-10-10 10:10:10.123456789", // neutral
@@ -1751,7 +2083,7 @@ func TestTimestampConversionWithArrowBatchesAndWithOriginalTimestamp(t *testing.
 	types := [3]string{"TIMESTAMP_NTZ", "TIMESTAMP_LTZ", "TIMESTAMP_TZ"}
 
 	runSnowflakeConnTest(t, func(sct *SCTest) {
-		ctx := WithOriginalTimestamp(WithArrowBatches(sct.sc.ctx))
+		ctx := WithArrowBatchesTimestampOption(WithArrowBatches(sct.sc.ctx), UseOriginalTimestamp)
 		pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
 		defer pool.AssertSize(t, 0)
 		ctx = WithArrowAllocator(ctx, pool)