diff --git a/src/danfojs-base/core/indexing.ts b/src/danfojs-base/core/indexing.ts index 5c2b668b..8853b1ce 100644 --- a/src/danfojs-base/core/indexing.ts +++ b/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 }: { } -} \ No newline at end of file +} diff --git a/src/danfojs-node/test/core/indexing.test.ts b/src/danfojs-node/test/core/indexing.test.ts index 1164cf14..8b8f1130 100644 --- a/src/danfojs-node/test/core/indexing.test.ts +++ b/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"]); + + }); + }) -}); \ No newline at end of file +});