fix: Fix test column generation (#872)

There was a bug in the last implementation causing many types to be missing; it's always that final refactor you do 😄 

This time I added a sanity check for the testdata generation to make sure it can't be catastrophically broken without us knowing about it.

Author: hermanschaaf

schema/testdata.go

@@ -228,24 +228,34 @@ func removeDuplicates(columns []Column) []Column {
 
 func removeColumnsByType(columns []Column, t ...arrow.Type) []Column {
 	var newColumns []Column
-	for _, d := range t {
-		for _, c := range columns {
+	for _, c := range columns {
+		shouldRemove := false
+		for _, d := range t {
 			if c.Type.ID() == d {
-				newColumns = append(newColumns, c)
+				shouldRemove = true
+				break
 			}
 		}
+		if !shouldRemove {
+			newColumns = append(newColumns, c)
+		}
 	}
 	return newColumns
 }
 
 func removeColumnsByDataType(columns []Column, dt ...arrow.DataType) []Column {
 	var newColumns []Column
-	for _, d := range dt {
-		for _, c := range columns {
+	for _, c := range columns {
+		shouldRemove := false
+		for _, d := range dt {
 			if arrow.TypeEqual(c.Type, d) {
-				newColumns = append(newColumns, c)
+				shouldRemove = true
+				break
 			}
 		}
+		if !shouldRemove {
+			newColumns = append(newColumns, c)
+		}
 	}
 	return newColumns
 }

schema/testdata_test.go

@@ -0,0 +1,45 @@
+package schema
+
+import "testing"
+
+func TestTestSourceColumns_Default(t *testing.T) {
+	// basic sanity check for tested columns
+	defaults := TestSourceColumns()
+	if len(defaults) < 100 {
+		t.Fatal("expected at least 100 columns by default")
+	}
+	// test some specific columns
+	checkColumnsExist(t, defaults, []string{"int64", "date32", "timestamp_us", "string", "struct", "string_map", "string_list"})
+}
+
+func TestTestSourceColumns_SkipAll(t *testing.T) {
+	skipAll := ColumnList(TestSourceColumns(
+		WithTestSourceSkipStructs(),
+		WithTestSourceSkipMaps(),
+		WithTestSourceSkipDates(),
+		WithTestSourceSkipTimes(),
+		WithTestSourceSkipTimestamps(),
+		WithTestSourceSkipDurations(),
+		WithTestSourceSkipIntervals(),
+		WithTestSourceSkipLargeTypes(),
+	))
+	// test some specific columns
+	checkColumnsExist(t, skipAll, []string{"int64", "timestamp_us", "string", "string_list"})
+	checkColumnsDontExist(t, skipAll, []string{"date32", "struct", "string_map"})
+}
+
+func checkColumnsExist(t *testing.T, list ColumnList, cols []string) {
+	for _, col := range cols {
+		if list.Get(col) == nil {
+			t.Errorf("expected column %s to be present", col)
+		}
+	}
+}
+
+func checkColumnsDontExist(t *testing.T, list ColumnList, cols []string) {
+	for _, col := range cols {
+		if list.Get(col) != nil {
+			t.Errorf("expected no %s column", col)
+		}
+	}
+}