Fix parsing and syncing NaN Postgres values

BemiHQ · Dec 9, 2024 · 39079e3 · 39079e3
1 parent 924a152
commit 39079e3
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 8 deletions.
diff --git a/scripts/install.sh b/scripts/install.sh
@@ -1,6 +1,6 @@
 #!/bin/bash
 
-VERSION="0.20.0"
+VERSION="0.20.1"
 
 # Detect OS and architecture
 OS=$(uname -s | tr '[:upper:]' '[:lower:]')

diff --git a/scripts/test-data-types.sql b/scripts/test-data-types.sql
@@ -131,7 +131,7 @@ INSERT INTO test_table (
   -9223372036854775807::INT8,               -- int8_column
   NULL,                                     -- xid_column
   NULL,                                     -- xid8_column
-  NULL,                                     -- float4_column
+  'NaN',                                    -- float4_column
   -3.141592653589793::FLOAT8,               -- float8_column
   -12345.00::NUMERIC(10, 2),                -- numeric_column
   NULL,                                     -- date_column

diff --git a/src/init_test.go b/src/init_test.go
@@ -297,7 +297,7 @@ var TEST_LOADED_ROWS = [][]string{
 		"-9223372036854775807",             // int8_column
 		PG_NULL_STRING,                     // xid_column
 		PG_NULL_STRING,                     // xid8_column
-		PG_NULL_STRING,                     // float4_column
+		"NaN",                              // float4_column
 		"-3.141592653589793",               // float8_column
 		"-12345.00",                        // numeric_column
 		PG_NULL_STRING,                     // date_column

diff --git a/src/main.go b/src/main.go
@@ -6,7 +6,7 @@ import (
 	"time"
 )
 
-const VERSION = "0.20.0"
+const VERSION = "0.20.1"
 
 func main() {
 	config := LoadConfig()

diff --git a/src/pg_schema_column.go b/src/pg_schema_column.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"encoding/csv"
+	"math"
 	"strconv"
 	"strings"
 	"time"
@@ -19,6 +20,8 @@ const (
 	PARQUET_SCHEMA_REPETITION_TYPE_REQUIRED = "REQUIRED"
 	PARQUET_SCHEMA_REPETITION_TYPE_OPTIONAL = "OPTIONAL"
 
+	PARQUET_NAN = "NaN"
+
 	// 0000-01-01 00:00:00 +0000 UTC
 	EPOCH_TIME_MS = -62167219200000
 )
@@ -230,11 +233,17 @@ func (pgSchemaColumn *PgSchemaColumn) parquetPrimitiveValue(value string) interf
 	case "float4":
 		floatValue, err := strconv.ParseFloat(value, 32)
 		PanicIfError(err)
+		if math.IsNaN(floatValue) {
+			return PARQUET_NAN
+		}
 		return float32(floatValue)
 	case "float8":
 		floatValue, err := strconv.ParseFloat(value, 64)
 		PanicIfError(err)
-		return float64(floatValue)
+		if math.IsNaN(floatValue) {
+			return PARQUET_NAN
+		}
+		return floatValue
 	case "bool":
 		boolValue, err := strconv.ParseBool(value)
 		PanicIfError(err)

diff --git a/src/query_handler_test.go b/src/query_handler_test.go
@@ -158,13 +158,13 @@ func TestHandleQuery(t *testing.T) {
 			"description": {"xid8_column"},
 			"values":      {""},
 		},
-		"SELECT float4_column FROM public.test_table WHERE float4_column IS NOT NULL": {
+		"SELECT float4_column FROM public.test_table WHERE float4_column = 3.14": {
 			"description": {"float4_column"},
 			"values":      {"3.14"},
 		},
-		"SELECT float4_column FROM public.test_table WHERE float4_column IS NULL": {
+		"SELECT float4_column FROM public.test_table WHERE float4_column != 3.14": {
 			"description": {"float4_column"},
-			"values":      {""},
+			"values":      {"NaN"},
 		},
 		"SELECT float8_column FROM public.test_table WHERE bool_column = TRUE": {
 			"description": {"float8_column"},