Skip to content

Commit d806f3c

Browse files
authored
feat: generic DataFrame (#293)
![image](https://github.com/user-attachments/assets/c528ac63-805c-479b-a3d5-08cf1bd2975c) resolves #292. please check individual commits for detail.
1 parent 70dabf6 commit d806f3c

File tree

9 files changed

+636
-399
lines changed

9 files changed

+636
-399
lines changed

__tests__/dataframe.test.ts

Lines changed: 95 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -756,7 +756,7 @@ describe("dataframe", () => {
756756
expect(actual.columns).toEqual(["foo_new", "bar_new", "ham_new"]);
757757
});
758758
test("replaceAtIdx", () => {
759-
const actual = pl.DataFrame({
759+
const actual: pl.DataFrame = pl.DataFrame({
760760
foo: [1, 2, 3],
761761
bar: [6, 7, 8],
762762
ham: ["a", "b", "c"],
@@ -1315,101 +1315,109 @@ describe("dataframe", () => {
13151315
expect(actual).toFrameEqual(expected);
13161316
});
13171317
test("pivot", () => {
1318-
let df = pl.DataFrame({
1319-
a: pl.Series([1, 2, 3]).cast(pl.Int32),
1320-
b: pl
1321-
.Series([
1322-
[1, 1],
1323-
[2, 2],
1324-
[3, 3],
1325-
])
1326-
.cast(pl.List(pl.Int32)),
1327-
});
1328-
1329-
let expected = pl
1330-
.DataFrame({
1318+
{
1319+
const df = pl.DataFrame({
13311320
a: pl.Series([1, 2, 3]).cast(pl.Int32),
1332-
"1": pl.Series([[1, 1], null, null]).cast(pl.List(pl.Int32)),
1333-
"2": pl.Series([null, [2, 2], null]).cast(pl.List(pl.Int32)),
1334-
"3": pl.Series([null, null, [3, 3]]).cast(pl.List(pl.Int32)),
1335-
})
1336-
.select("a", "1", "2", "3");
1321+
b: pl
1322+
.Series([
1323+
[1, 1],
1324+
[2, 2],
1325+
[3, 3],
1326+
])
1327+
.cast(pl.List(pl.Int32)),
1328+
});
13371329

1338-
let actual = df.pivot("b", {
1339-
index: "a",
1340-
on: "a",
1341-
aggregateFunc: "first",
1342-
sortColumns: true,
1343-
});
1330+
const expected = pl
1331+
.DataFrame({
1332+
a: pl.Series([1, 2, 3]).cast(pl.Int32),
1333+
"1": pl.Series([[1, 1], null, null]).cast(pl.List(pl.Int32)),
1334+
"2": pl.Series([null, [2, 2], null]).cast(pl.List(pl.Int32)),
1335+
"3": pl.Series([null, null, [3, 3]]).cast(pl.List(pl.Int32)),
1336+
})
1337+
.select("a", "1", "2", "3");
13441338

1345-
expect(actual).toFrameEqual(expected, true);
1339+
const actual = df.pivot("b", {
1340+
index: "a",
1341+
on: "a",
1342+
aggregateFunc: "first",
1343+
sortColumns: true,
1344+
});
13461345

1347-
df = pl.DataFrame({
1348-
a: ["beep", "bop"],
1349-
b: ["a", "b"],
1350-
c: ["s", "f"],
1351-
d: [7, 8],
1352-
e: ["x", "y"],
1353-
});
1354-
actual = df.pivot(["a", "e"], {
1355-
index: "b",
1356-
on: ["b"],
1357-
aggregateFunc: "first",
1358-
separator: "|",
1359-
maintainOrder: true,
1360-
});
1346+
expect(actual).toFrameEqual(expected, true);
1347+
}
13611348

1362-
expected = pl.DataFrame({
1363-
b: ["a", "b"],
1364-
"a|a": ["beep", null],
1365-
"a|b": [null, "bop"],
1366-
"e|a": ["x", null],
1367-
"e|b": [null, "y"],
1368-
});
1369-
expect(actual).toFrameEqual(expected, true);
1349+
{
1350+
const df = pl.DataFrame({
1351+
a: ["beep", "bop"],
1352+
b: ["a", "b"],
1353+
c: ["s", "f"],
1354+
d: [7, 8],
1355+
e: ["x", "y"],
1356+
});
1357+
const actual = df.pivot(["a", "e"], {
1358+
index: "b",
1359+
on: ["b"],
1360+
aggregateFunc: "first",
1361+
separator: "|",
1362+
maintainOrder: true,
1363+
});
13701364

1371-
df = pl.DataFrame({
1372-
foo: ["A", "A", "B", "B", "C"],
1373-
N: [1, 2, 2, 4, 2],
1374-
bar: ["k", "l", "m", "n", "o"],
1375-
});
1376-
actual = df.pivot(["N"], {
1377-
index: "foo",
1378-
on: "bar",
1379-
aggregateFunc: "first",
1380-
});
1381-
expected = pl.DataFrame({
1382-
foo: ["A", "B", "C"],
1383-
k: [1, null, null],
1384-
l: [2, null, null],
1385-
m: [null, 2, null],
1386-
n: [null, 4, null],
1387-
o: [null, null, 2],
1388-
});
1389-
expect(actual).toFrameEqual(expected, true);
1365+
const expected = pl.DataFrame({
1366+
b: ["a", "b"],
1367+
"a|a": ["beep", null],
1368+
"a|b": [null, "bop"],
1369+
"e|a": ["x", null],
1370+
"e|b": [null, "y"],
1371+
});
1372+
expect(actual).toFrameEqual(expected, true);
1373+
}
1374+
{
1375+
const df = pl.DataFrame({
1376+
foo: ["A", "A", "B", "B", "C"],
1377+
N: [1, 2, 2, 4, 2],
1378+
bar: ["k", "l", "m", "n", "o"],
1379+
});
1380+
const actual = df.pivot(["N"], {
1381+
index: "foo",
1382+
on: "bar",
1383+
aggregateFunc: "first",
1384+
});
13901385

1391-
df = pl.DataFrame({
1392-
ix: [1, 1, 2, 2, 1, 2],
1393-
col: ["a", "a", "a", "a", "b", "b"],
1394-
foo: [0, 1, 2, 2, 7, 1],
1395-
bar: [0, 2, 0, 0, 9, 4],
1396-
});
1386+
const expected = pl.DataFrame({
1387+
foo: ["A", "B", "C"],
1388+
k: [1, null, null],
1389+
l: [2, null, null],
1390+
m: [null, 2, null],
1391+
n: [null, 4, null],
1392+
o: [null, null, 2],
1393+
});
13971394

1398-
actual = df.pivot(["foo", "bar"], {
1399-
index: "ix",
1400-
on: "col",
1401-
aggregateFunc: "sum",
1402-
separator: "/",
1403-
});
1395+
expect(actual).toFrameEqual(expected, true);
1396+
}
1397+
{
1398+
const df = pl.DataFrame({
1399+
ix: [1, 1, 2, 2, 1, 2],
1400+
col: ["a", "a", "a", "a", "b", "b"],
1401+
foo: [0, 1, 2, 2, 7, 1],
1402+
bar: [0, 2, 0, 0, 9, 4],
1403+
});
14041404

1405-
expected = pl.DataFrame({
1406-
ix: [1, 2],
1407-
"foo/a": [1, 4],
1408-
"foo/b": [7, 1],
1409-
"bar/a": [2, 0],
1410-
"bar/b": [9, 4],
1411-
});
1412-
expect(actual).toFrameEqual(expected, true);
1405+
const actual = df.pivot(["foo", "bar"], {
1406+
index: "ix",
1407+
on: "col",
1408+
aggregateFunc: "sum",
1409+
separator: "/",
1410+
});
1411+
1412+
const expected = pl.DataFrame({
1413+
ix: [1, 2],
1414+
"foo/a": [1, 4],
1415+
"foo/b": [7, 1],
1416+
"bar/a": [2, 0],
1417+
"bar/b": [9, 4],
1418+
});
1419+
expect(actual).toFrameEqual(expected, true);
1420+
}
14131421
});
14141422
});
14151423
describe("join", () => {

0 commit comments

Comments
 (0)