Merge pull request #234 from Joshua-Dias-Barreto/numericalColumn

Added #numericalColumns and #numericalColumnNames

Author: jecisc

src/DataFrame-Tests/DataFrameTest.class.st

@@ -2825,6 +2825,37 @@ DataFrameTest >> testNumberOfNils [
 	self assert: df numberOfNils equals: expected
 ]
 
+{ #category : #tests }
+DataFrameTest >> testNumericalColumnNames [
+
+	| dataFrame expected |
+	dataFrame := DataFrame withRows:
+		             #( #( Male 21 Argentina 94 ) #( Female 20 France 97 )
+		                #( Female 22 Spain 95 ) #( Male 24 Portugal 99 ) ).
+	dataFrame columnNames: #( Gender Age Country Score ).
+
+	expected := #( Age Score ) asOrderedCollection.
+
+	self assert: dataFrame numericalColumnNames equals: expected
+]
+
+{ #category : #tests }
+DataFrameTest >> testNumericalColumns [
+
+	| dataFrame expected |
+	dataFrame := DataFrame withRows:
+		             #( #( Male 21 Argentina 94 ) #( Female 20 France 97 )
+		                #( Female 22 Spain 95 ) #( Male 24 Portugal 99 ) ).
+	dataFrame columnNames: #( Gender Age Country Score ).
+
+	expected := DataFrame withRows:
+		            #( #( 21 94 ) #( 20 97 ) #( 22 95 ) #( 24 99 ) ).
+
+	expected columnNames: #( Age Score ).
+
+	self assert: dataFrame numericalColumns equals: expected
+]
+
 { #category : #splitjoin }
 DataFrameTest >> testOuterJoin [
 	| df2 expected |

src/DataFrame/DataFrame.class.st

@@ -1428,6 +1428,21 @@ DataFrame >> numberOfRows [
 	^ contents numberOfRows
 ]
 
+{ #category : #accessing }
+DataFrame >> numericalColumnNames [
+	"Returns the names of all numerical columns of the dataframe"
+
+	^ self columnNames select: [ :columnName |
+		  (self dataTypes at: columnName) includesBehavior: Number ]
+]
+
+{ #category : #accessing }
+DataFrame >> numericalColumns [
+	"Returns all numerical columns of the dataframe"
+
+	^ self columns: self numericalColumnNames
+]
+
 { #category : #splitjoin }
 DataFrame >> outerJoin: aDataFrame [
 	"Performs outer join on aDataFrame with rowNames as keys"