FIX Issue with Retrieving Columns Named with Colon ':' in Single Selection Using danfojs-node

src/danfojs-base/core/indexing.ts

@@ -22,7 +22,7 @@ const utils = new Utils();
 
 /**
 * Internal function to slice a Series/DataFrame by index based labels
-* @param Object 
+* @param Object
 */
 export function _iloc({ ndFrame, rows, columns }: {
     ndFrame: NDframeInterface
@@ -207,7 +207,7 @@ export function _iloc({ ndFrame, rows, columns }: {
 
 /**
 * Internal function to slice a Series/DataFrame by specified string location based labels
-* @param Object 
+* @param Object
 */
 export function _loc({ ndFrame, rows, columns }: {
     ndFrame: NDframeInterface
@@ -327,9 +327,12 @@ export function _loc({ ndFrame, rows, columns }: {
             throw new Error(`ColumnIndexError: columns parameter must be an array of a string name. For example: columns: ["b"]`)
         }
 
-        if (columns[0].indexOf(":") == -1) { // Input type ==> ["A"] 
+        if (_columnNames.indexOf(columns[0]) !== -1) {
+            // Column exists as a literal name, use it directly
             _columnIndexes = [_columnNames.indexOf(columns[0])]
-
+        } else if (columns[0].indexOf(":") == -1) {
+            // Input type ==> ["A"] but column doesn't exist
+            throw new Error(`ColumnIndexError: Specified column (${columns[0]}) not found`);
         } else { // Input type ==> ["a:b"] or [`"col1":"col5"`]
             const columnSplit = columns[0].split(":")
 
@@ -417,4 +420,4 @@ export function _loc({ ndFrame, rows, columns }: {
 
     }
 
-}
+}

src/danfojs-node/test/core/indexing.test.ts

@@ -7,7 +7,7 @@ describe("Iloc and Loc based Indexing", function () {
         it("throw error for wrong row index value", function () {
             let data = [1, 2, 34, 5, 6];
             let df = new Series(data);
-            
+
             assert.throws(function () { df.iloc(0 as any) }, Error, `rows parameter must be an Array. For example: rows: [1,2] or rows: ["0:10"]`);
         });
 
@@ -428,6 +428,34 @@ describe("Iloc and Loc based Indexing", function () {
 
           });
 
+        it("loc works with column names containing colons", function () {
+            const data = {
+                "my:column": [1, 2, 3, 4],
+                "time:timestamp": ["2023-01-01", "2023-01-02", "2023-01-03", "2023-01-04"],
+                "normal_column": [10, 20, 30, 40]
+            };
+            const df = new DataFrame(data);
+
+            // Test accessing a single column with a colon
+            const singleCol = df.loc({ columns: ["my:column"] });
+            const singleColExpected = [[1], [2], [3], [4]];
+            assert.deepEqual(singleCol.values, singleColExpected);
+            assert.deepEqual(singleCol.columns, ["my:column"]);
+
+            // Test accessing multiple columns with colons
+            const multiCol = df.loc({ columns: ["my:column", "time:timestamp"] });
+            const multiColExpected = [[1, "2023-01-01"], [2, "2023-01-02"], [3, "2023-01-03"], [4, "2023-01-04"]];
+            assert.deepEqual(multiCol.values, multiColExpected);
+            assert.deepEqual(multiCol.columns, ["my:column", "time:timestamp"]);
+
+            // Test accessing a mix of columns with and without colons
+            const mixedCol = df.loc({ columns: ["my:column", "normal_column"] });
+            const mixedColExpected = [[1, 10], [2, 20], [3, 30], [4, 40]];
+            assert.deepEqual(mixedCol.values, mixedColExpected);
+            assert.deepEqual(mixedCol.columns, ["my:column", "normal_column"]);
+
+        });
+
     })
 
-});
+});