Skip to content

Commit 7729376

Browse files
Bidek56Bidek56
authored
feat: Adding lit_series func (#356)
Adding lit_series func to lit Series into `LiteralValue` to close #351 --------- Co-authored-by: Bidek56 <[email protected]>
1 parent 0c0c45a commit 7729376

File tree

7 files changed

+164
-30
lines changed

7 files changed

+164
-30
lines changed

‎__tests__/lazyframe.test.ts‎

Lines changed: 133 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -276,7 +276,10 @@ describe("lazyframe", () => {
276276
foo: [1],
277277
bar: ["a"],
278278
});
279-
const actual = await df.lazy().select("*").fetch(1);
279+
const actual = await df
280+
.lazy()
281+
.select("*")
282+
.fetch(1, { noOptimization: true });
280283
expect(actual).toFrameEqual(expected);
281284
});
282285
test("fetchSync", () => {
@@ -357,7 +360,46 @@ describe("lazyframe", () => {
357360
});
358361
expect(actual).toFrameEqual(expected);
359362
});
360-
describe("groupby", () => {});
363+
describe("groupby", () => {
364+
test("groupBy", () => {
365+
let actual = pl
366+
.DataFrame({
367+
foo: [1, 2, 3],
368+
ham: ["a", "a", "b"],
369+
})
370+
.lazy()
371+
.groupBy("ham")
372+
.agg(pl.col("foo").sum())
373+
.collectSync();
374+
const expected = pl.DataFrame({
375+
ham: ["a", "b"],
376+
foo: [3, 3],
377+
});
378+
expect(actual).toFrameEqual(expected);
379+
380+
actual = pl
381+
.DataFrame({
382+
foo: [1, 2, 3],
383+
ham: ["a", "a", "b"],
384+
})
385+
.lazy()
386+
.groupBy("ham", true)
387+
.agg(pl.col("foo").sum())
388+
.collectSync();
389+
expect(actual).toFrameEqual(expected);
390+
391+
actual = pl
392+
.DataFrame({
393+
foo: [1, 2, 3],
394+
ham: ["a", "a", "b"],
395+
})
396+
.lazy()
397+
.groupBy("ham", { maintainOrder: true })
398+
.agg(pl.col("foo").sum())
399+
.collectSync();
400+
expect(actual).toFrameEqual(expected);
401+
});
402+
});
361403
test("head", () => {
362404
const actual = pl
363405
.DataFrame({
@@ -1096,6 +1138,90 @@ describe("lazyframe", () => {
10961138
]);
10971139
expect(actual).toFrameEqualIgnoringOrder(expected);
10981140
});
1141+
test("withColumn:series", async () => {
1142+
const actual: pl.DataFrame = pl
1143+
.DataFrame()
1144+
.lazy()
1145+
.withColumn(pl.Series("series1", [1, 2, 3, 4], pl.Int16))
1146+
.collectSync();
1147+
const expected: pl.DataFrame = pl.DataFrame([
1148+
pl.Series("series1", [1, 2, 3, 4], pl.Int16),
1149+
]);
1150+
expect(actual).toFrameEqual(expected);
1151+
});
1152+
test("withColumns:series", async () => {
1153+
const actual: pl.DataFrame = pl
1154+
.DataFrame()
1155+
.lazy()
1156+
.withColumns(
1157+
pl.Series("series1", [1, 2, 3, 4], pl.Int16),
1158+
pl.Series("series2", [1, 2, 3, 4], pl.Int32),
1159+
)
1160+
.collectSync();
1161+
const expected: pl.DataFrame = pl.DataFrame([
1162+
pl.Series("series1", [1, 2, 3, 4], pl.Int16),
1163+
pl.Series("series2", [1, 2, 3, 4], pl.Int32),
1164+
]);
1165+
expect(actual).toFrameEqual(expected);
1166+
});
1167+
test("select:series", async () => {
1168+
let actual: pl.DataFrame = pl
1169+
.DataFrame()
1170+
.lazy()
1171+
.select(
1172+
pl.Series("series1", [1, 2, 3, 4], pl.Int16),
1173+
pl.Series("series2", [1, 2, 3, 4], pl.Int32),
1174+
)
1175+
.collectSync();
1176+
let expected: pl.DataFrame = pl.DataFrame([
1177+
pl.Series("series1", [1, 2, 3, 4], pl.Int16),
1178+
pl.Series("series2", [1, 2, 3, 4], pl.Int32),
1179+
]);
1180+
expect(actual).toFrameEqual(expected);
1181+
actual = pl
1182+
.DataFrame({ text: ["hello"] })
1183+
.lazy()
1184+
.select(pl.Series("series", [1, 2, 3]))
1185+
.collectSync();
1186+
1187+
expected = pl.DataFrame([pl.Series("series", [1, 2, 3])]);
1188+
expect(actual).toFrameEqual(expected);
1189+
1190+
actual = pl
1191+
.DataFrame({ text: ["hello"] })
1192+
.lazy()
1193+
.select("text", pl.Series("series", [1]))
1194+
.collectSync();
1195+
expected = pl.DataFrame({ text: ["hello"], series: [1] });
1196+
expect(actual).toFrameEqual(expected);
1197+
});
1198+
test("select:lit", () => {
1199+
const df = pl.DataFrame({ a: [1] }, { schema: { a: pl.Float32 } });
1200+
let actual = df.lazy().select(pl.col("a"), pl.lit(1)).collectSync();
1201+
const expected = pl.DataFrame({
1202+
a: [1],
1203+
literal: [1],
1204+
});
1205+
expect(actual).toFrameEqual(expected);
1206+
actual = df
1207+
.lazy()
1208+
.select(pl.col("a").mul(2).alias("b"), pl.lit(2))
1209+
.collectSync();
1210+
const expected2 = pl.DataFrame({
1211+
b: [2],
1212+
literal: [2],
1213+
});
1214+
expect(actual).toFrameEqual(expected2);
1215+
});
1216+
test("inspect", () => {
1217+
const actual = pl
1218+
.DataFrame({
1219+
foo: [1, 2, 9],
1220+
bar: [6, 2, 8],
1221+
})
1222+
.lazy();
1223+
expect(actual).toBeDefined();
1224+
});
10991225
test("withColumns", () => {
11001226
const actual = pl
11011227
.DataFrame({
@@ -1120,7 +1246,7 @@ describe("lazyframe", () => {
11201246
bar: [6, 2, 8],
11211247
})
11221248
.lazy()
1123-
.withColumns([pl.lit("a").alias("col_a"), pl.lit("b").alias("col_b")])
1249+
.withColumns(pl.lit("a").alias("col_a"), pl.lit("b").alias("col_b"))
11241250
.collectSync();
11251251
const expected = pl.DataFrame([
11261252
pl.Series("foo", [1, 2, 9], pl.Int16),
@@ -1280,7 +1406,10 @@ describe("lazyframe", () => {
12801406
.lazy();
12811407
await ldf.sinkCSV("./test.csv").collect();
12821408
const newDF: pl.DataFrame = pl.readCSV("./test.csv");
1283-
const actualDf: pl.DataFrame = await ldf.collect({ streaming: true });
1409+
const actualDf: pl.DataFrame = await ldf.collect({
1410+
streaming: true,
1411+
noOptimization: true,
1412+
});
12841413
expect(newDF.sort("foo")).toFrameEqual(actualDf);
12851414
fs.rmSync("./test.csv");
12861415
});

‎__tests__/series.test.ts‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,7 @@ describe("typedArrays", () => {
198198
test("decimal", () => {
199199
const expected = [1n, 2n, 3n];
200200
const expectedDtype = pl.Decimal(10, 2);
201+
expect(expectedDtype.equals(expectedDtype)).toBeTruthy();
201202
const actual = pl.Series("", expected, expectedDtype);
202203
expect(actual.dtype).toEqual(expectedDtype);
203204
try {
@@ -219,6 +220,9 @@ describe("typedArrays", () => {
219220
expect(actual.dtype).toEqual(expectedDtype);
220221
const actualValues = actual.toArray();
221222
expect(actualValues).toEqual(expected);
223+
const lst = expectedDtype.asFixedSizeList();
224+
expect(lst?.inner[0]).toEqual(pl.Float32);
225+
expect(expectedDtype.equals(expectedDtype)).toBeTruthy();
222226
});
223227
});
224228
describe("series", () => {
@@ -949,7 +953,7 @@ describe("series struct", () => {
949953
const expectedFields = [...expectedKeys];
950954
expect(actualFields).toEqual(expectedFields);
951955
});
952-
test.skip("struct:field", () => {
956+
test("struct:field", () => {
953957
const expected = [{ foo: 1, bar: 2, ham: "c" }];
954958
const actual = pl.Series(expected).struct.field("foo").toArray();
955959
expect(actual).toEqual([expected[0]["foo"]]);
@@ -964,15 +968,15 @@ describe("series struct", () => {
964968
});
965969
expect(actual).toFrameEqual(expected);
966970
});
967-
test.skip("struct:renameFields", () => {
971+
test("struct:renameFields", () => {
968972
const expected = [{ foo: 1, bar: 2, ham: "c" }];
969973
const actual = pl
970974
.Series(expected)
971975
.struct.renameFields(["foo", "bar", "ham"])
972976
.toArray();
973977
expect(actual).toEqual(expected);
974978
});
975-
test.skip("struct:nth", () => {
979+
test("struct:nth", () => {
976980
const arr = [
977981
{ foo: 1, bar: 2, ham: "c" },
978982
{ foo: null, bar: 10, ham: null },

‎polars/dataframe.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2853,7 +2853,7 @@ export const _DataFrame = <S extends Schema>(_df: any): DataFrame<S> => {
28532853
);
28542854
}
28552855
return this.lazy()
2856-
.withColumns(columns)
2856+
.withColumns(...columns)
28572857
.collectSync({ noOptimization: true });
28582858
},
28592859
withColumnRenamed(opt, replacement?): any {

‎polars/lazy/dataframe.ts‎

Lines changed: 12 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -478,9 +478,9 @@ export interface LazyDataFrame<S extends Schema = any>
478478
* @see {@link DataFrame.select}
479479
*/
480480
select<U extends keyof S>(...columns: U[]): LazyDataFrame<{ [P in U]: S[P] }>;
481-
select(column: ExprOrString): LazyDataFrame;
482-
select(columns: ExprOrString[]): LazyDataFrame;
483-
select(...columns: ExprOrString[]): LazyDataFrame;
481+
select(column: ExprOrString | Series): LazyDataFrame;
482+
select(columns: (ExprOrString | Series)[]): LazyDataFrame;
483+
select(...columns: (ExprOrString | Series)[]): LazyDataFrame;
484484
/**
485485
* @see {@link DataFrame.shift}
486486
*/
@@ -553,13 +553,12 @@ export interface LazyDataFrame<S extends Schema = any>
553553
* Add or overwrite column in a DataFrame.
554554
* @param expr - Expression that evaluates to column.
555555
*/
556-
withColumn(expr: Expr): LazyDataFrame;
556+
withColumn(expr: Expr | Series): LazyDataFrame;
557557
/**
558558
* Add or overwrite multiple columns in a DataFrame.
559559
* @param exprs - List of Expressions that evaluate to columns.
560560
*
561561
*/
562-
withColumns(exprs: (Expr | Series)[]): LazyDataFrame;
563562
withColumns(...exprs: (Expr | Series)[]): LazyDataFrame;
564563
withColumnRenamed<Existing extends keyof S, New extends string>(
565564
existing: Existing,
@@ -927,15 +926,12 @@ export const _LazyDataFrame = (_ldf: any): LazyDataFrame => {
927926

928927
return _LazyDataFrame(_ldf.filter(predicate));
929928
},
930-
groupBy(opt, maintainOrder: any = true): LazyGroupBy {
931-
if (opt?.by !== undefined) {
932-
const by = selectionToExprList([opt.by], false);
933-
934-
return _LazyGroupBy(_ldf.groupby(by, opt.maintainOrder));
929+
groupBy(by, maintainOrder: any = true): LazyGroupBy {
930+
if (maintainOrder?.maintainOrder !== undefined) {
931+
maintainOrder = maintainOrder.maintainOrder;
935932
}
936-
const by = selectionToExprList([opt], false);
937-
938-
return _LazyGroupBy(_ldf.groupby(by, maintainOrder));
933+
const expr = selectionToExprList([by], false);
934+
return _LazyGroupBy(_ldf.groupby(expr, maintainOrder));
939935
},
940936
groupByRolling({ indexColumn, by, period, offset, closed }) {
941937
offset = offset ?? `-${period}`;
@@ -1165,7 +1161,6 @@ export const _LazyDataFrame = (_ldf: any): LazyDataFrame => {
11651161
},
11661162
select(...exprs) {
11671163
const selections = selectionToExprList(exprs, false);
1168-
11691164
return _LazyDataFrame(_ldf.select(selections));
11701165
},
11711166
shift(periods) {
@@ -1221,12 +1216,11 @@ export const _LazyDataFrame = (_ldf: any): LazyDataFrame => {
12211216
serialize(format) {
12221217
return _ldf.serialize(format);
12231218
},
1224-
withColumn(expr) {
1225-
return _LazyDataFrame(_ldf.withColumn(expr._expr));
1219+
withColumn(column: Expr | Series) {
1220+
return this.withColumns(column);
12261221
},
1227-
withColumns(...columns) {
1222+
withColumns(...columns: (Expr | Series)[]) {
12281223
const exprs = selectionToExprList(columns, false);
1229-
12301224
return _LazyDataFrame(_ldf.withColumns(exprs));
12311225
},
12321226
withColumnRenamed(existing, replacement) {

‎polars/lazy/expr/index.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1944,7 +1944,7 @@ export const exprToLitOrExpr = (expr: any, stringToLit = true): Expr => {
19441944
return expr;
19451945
}
19461946
if (Series.isSeries(expr)) {
1947-
return _Expr(pli.lit((expr as any)._s));
1947+
return _Expr(pli.litSeries((expr as any)._s));
19481948
}
19491949
return _Expr(pli.lit(expr));
19501950
};

‎polars/series/struct.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,17 +47,17 @@ export const SeriesStructFunctions = (_s: any): SeriesStructFunctions => {
4747
},
4848
field(name) {
4949
return DataFrame({})
50-
.select(_Expr(pli.lit(_s).structFieldByName(name)))
50+
.select(_Expr(pli.litSeries(_s).structFieldByName(name)))
5151
.toSeries();
5252
},
5353
renameFields(names) {
5454
return DataFrame({})
55-
.select(_Expr(pli.lit(_s).structRenameFields(names)))
55+
.select(_Expr(pli.litSeries(_s).structRenameFields(names)))
5656
.toSeries();
5757
},
5858
nth(index) {
5959
return DataFrame({})
60-
.select(_Expr(pli.lit(_s).structFieldByIndex(index)))
60+
.select(_Expr(pli.litSeries(_s).structFieldByIndex(index)))
6161
.toSeries();
6262
},
6363
};

‎src/lazy/dsl.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
use crate::conversion::Wrap;
22
use crate::prelude::*;
3+
use crate::series::JsSeries;
34
use crate::utils::reinterpret;
45
use polars::lazy::dsl;
56
use polars::lazy::dsl::Expr;
@@ -1759,6 +1760,12 @@ pub fn arg_sort_by(by: Vec<&JsExpr>, descending: Vec<bool>) -> JsExpr {
17591760
.into()
17601761
}
17611762

1763+
#[napi(catch_unwind)]
1764+
pub fn lit_series(s: &JsSeries) -> JsResult<JsExpr> {
1765+
Ok(s.clone().series.lit().into())
1766+
}
1767+
1768+
17621769
#[napi(catch_unwind)]
17631770
pub fn lit(value: Wrap<AnyValue>) -> JsResult<JsExpr> {
17641771
let lit: LiteralValue = value

0 commit comments

Comments
 (0)