Skip to content

Commit

Permalink
feat: generic Series (#154)
Browse files Browse the repository at this point in the history
  • Loading branch information
scarf005 authored Jul 25, 2024
1 parent 9b49663 commit 95191df
Show file tree
Hide file tree
Showing 4 changed files with 154 additions and 45 deletions.
22 changes: 21 additions & 1 deletion __tests__/series.test.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/* eslint-disable newline-per-chained-call */
import pl from "@polars";
import pl, { DataType } from "@polars";
import Chance from "chance";

describe("from lists", () => {
Expand Down Expand Up @@ -867,3 +867,23 @@ describe("series struct", () => {
expect(actual).toEqual(expected);
});
});
describe("generics", () => {
const series = pl.Series([1, 2, 3]);

test("dtype", () => {
expect(series.dtype).toStrictEqual(DataType.Float64);
});
test("to array", () => {
const arr = series.toArray();
expect<number[]>(arr).toStrictEqual([1, 2, 3]);

const arr2 = [...series];
expect<number[]>(arr2).toStrictEqual([1, 2, 3]);
});
test("to object", () => {
const obj = series.toObject();
expect<{ name: string; datatype: "Float64"; values: number[] }>(
obj,
).toMatchObject({ name: "", datatype: "Float64", values: [1, 2, 3] });
});
});
72 changes: 71 additions & 1 deletion polars/datatypes/datatype.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ import { Field } from "./field";

export abstract class DataType {
get variant() {
return this.constructor.name;
return this.constructor.name as DataTypeName;
}
protected identity = "DataType";
protected get inner(): null | any[] {
Expand Down Expand Up @@ -406,3 +406,73 @@ export namespace DataType {
return DataType[variant](...inner);
}
}

export type DataTypeName =
| "Null"
| "Bool"
| "Int8"
| "Int16"
| "Int32"
| "Int64"
| "UInt8"
| "UInt16"
| "UInt32"
| "UInt64"
| "Float32"
| "Float64"
| "Decimal"
| "Date"
| "Datetime"
| "Utf8"
| "Categorical"
| "List"
| "Struct";

export type JsType = number | boolean | string;
export type JsToDtype<T> = T extends number
? DataType.Float64
: T extends boolean
? DataType.Bool
: T extends string
? DataType.Utf8
: never;
export type DTypeToJs<T> = T extends DataType.Decimal
? bigint
: T extends DataType.Float64
? number
: T extends DataType.Int64
? bigint
: T extends DataType.Int32
? number
: T extends DataType.Bool
? boolean
: T extends DataType.Utf8
? string
: never;
export type DtypeToJsName<T> = T extends DataType.Decimal
? "Decimal"
: T extends DataType.Float64
? "Float64"
: T extends DataType.Float32
? "Float32"
: T extends DataType.Int64
? "Int64"
: T extends DataType.Int32
? "Int32"
: T extends DataType.Int16
? "Int16"
: T extends DataType.Int8
? "Int8"
: T extends DataType.UInt64
? "UInt64"
: T extends DataType.UInt32
? "UInt32"
: T extends DataType.UInt16
? "UInt16"
: T extends DataType.UInt8
? "UInt8"
: T extends DataType.Bool
? "Bool"
: T extends DataType.Utf8
? "Utf8"
: never;
2 changes: 1 addition & 1 deletion polars/lazy/functions.ts
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ import pli from "../internals/polars_internal";
*/
export function col(col: string | string[] | Series | DataType): Expr {
if (Series.isSeries(col)) {
col = col.toArray();
col = col.toArray() as string[];
}
if (Array.isArray(col)) {
return _Expr(pli.cols(col));
Expand Down
103 changes: 61 additions & 42 deletions polars/series/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,54 +12,59 @@ import type {
Comparison,
Cumulative,
Deserialize,
EwmOps,
Rolling,
Round,
Sample,
Serialize,
EwmOps,
} from "../shared_traits";
import { col } from "../lazy/functions";
import type { InterpolationMethod, RankMethod } from "../types";
import type {
DTypeToJs,
DtypeToJsName,
JsToDtype,
JsType,
} from "@polars/datatypes/datatype";

const inspect = Symbol.for("nodejs.util.inspect.custom");
/**
* A Series represents a single column in a polars DataFrame.
*/
export interface Series
extends ArrayLike<any>,
Rolling<Series>,
Arithmetic<Series>,
Comparison<Series>,
Cumulative<Series>,
Round<Series>,
Sample<Series>,
EwmOps<Series>,
export interface Series<T extends DataType = any>
extends ArrayLike<T>,
Rolling<Series<T>>,
Arithmetic<Series<T>>,
Comparison<Series<T>>,
Cumulative<Series<T>>,
Round<Series<T>>,
Sample<Series<T>>,
EwmOps<Series<T>>,
Serialize {
inner(): any;
name: string;
dtype: DataType;
dtype: T;
str: StringNamespace;
lst: ListNamespace;
struct: SeriesStructFunctions;
date: SeriesDateFunctions;
[inspect](): string;
[Symbol.iterator](): IterableIterator<any>;
[Symbol.iterator](): IterableIterator<DTypeToJs<T>>;
// inner(): JsSeries
bitand(other: Series): Series;
bitor(other: Series): Series;
bitxor(other: Series): Series;
bitand(other: Series<T>): Series<T>;
bitor(other: Series<T>): Series<T>;
bitxor(other: Series<T>): Series<T>;
/**
* Take absolute values
*/
abs(): Series;
abs(): Series<T>;
/**
* __Rename this Series.__
*
* @param name - new name
* @see {@link rename}
*
*/
alias(name: string): Series;
alias(name: string): Series<T>;
/**
* __Append a Series to this one.__
* ___
Expand Down Expand Up @@ -113,43 +118,41 @@ export interface Series
argMin(): Optional<number>;
/**
* Get index values where Boolean Series evaluate True.
*
*/
argTrue(): Series;
argTrue(): Series<T>;
/**
* Get unique index as Series.
*/
argUnique(): Series;
argUnique(): Series<T>;
/**
* Index location of the sorted variant of this Series.
* ___
* @param reverse
* @return {SeriesType} indexes - Indexes that can be used to sort this array.
*/
argSort(): Series;
argSort(reverse: boolean): Series;
argSort({ reverse }: { reverse: boolean }): Series;
argSort(): Series<T>;
argSort(reverse: boolean): Series<T>;
argSort({ reverse }: { reverse: boolean }): Series<T>;
/**
* __Rename this Series.__
*
* @param name - new name
* @see {@link rename} {@link alias}
*
*/
as(name: string): Series;
as(name: string): Series<T>;
/**
* Cast between data types.
*/
cast(dtype: DataType, strict?: boolean): Series;
cast<U extends DataType>(dtype: U, strict?: boolean): Series<U>;
/**
* Get the length of each individual chunk
*/
chunkLengths(): Array<any>;
chunkLengths(): Array<T>;
/**
* Cheap deep clones.
*/
clone(): Series;
concat(other: Series): Series;
clone(): Series<T>;
concat(other: Series<T>): Series<T>;

/**
* __Quick summary statistics of a series. __
Expand Down Expand Up @@ -203,14 +206,14 @@ export interface Series
* @param n - number of slots to shift
* @param nullBehavior - `'ignore' | 'drop'`
*/
diff(n: number, nullBehavior: "ignore" | "drop"): Series;
diff(n: number, nullBehavior: "ignore" | "drop"): Series<T>;
diff({
n,
nullBehavior,
}: {
n: number;
nullBehavior: "ignore" | "drop";
}): Series;
}): Series<T>;
/**
* Compute the dot/inner product between two Series
* ___
Expand All @@ -226,7 +229,7 @@ export interface Series
/**
* Create a new Series that copies data from this Series without null values.
*/
dropNulls(): Series;
dropNulls(): Series<T>;
/**
* __Explode a list or utf8 Series.__
*
Expand Down Expand Up @@ -301,7 +304,6 @@ export interface Series
/**
* __Filter elements by a boolean mask.__
* @param {SeriesType} predicate - Boolean mask
*
*/
filter(predicate: Series): Series;
filter({ predicate }: { predicate: Series }): Series;
Expand Down Expand Up @@ -655,7 +657,6 @@ export interface Series
/**
* Count the null values in this Series. --
* _`undefined` values are treated as null_
*
*/
nullCount(): number;
/**
Expand Down Expand Up @@ -744,7 +745,6 @@ export interface Series
* - True -> pl.Int64
* - False -> pl.UInt64
* @see {@link cast}
*
*/
reinterpret(signed?: boolean): Series;
/**
Expand Down Expand Up @@ -1030,7 +1030,6 @@ export interface Series
* ___
* @param mask - Boolean Series
* @param other - Series of same type
*
*/
zipWith(mask: Series, other: Series): Series;

Expand All @@ -1049,7 +1048,7 @@ export interface Series
* true
* ```
*/
toArray(): Array<any>;
toArray(): Array<DTypeToJs<T>>;
/**
* Converts series to a javascript typedArray.
*
Expand All @@ -1060,9 +1059,9 @@ export interface Series

/**
* Get dummy/indicator variables.
* @param separator: str = "_",
* @param separator: str = "_",
* @param dropFirst: bool = False
*
*
* @example
* const s = pl.Series("a", [1, 2, 3])
>>> s.toDummies()
Expand All @@ -1088,7 +1087,7 @@ export interface Series
│ 1 ┆ 0 │
│ 0 ┆ 1 │
└─────┴─────┘
*
*
*/
toDummies(separator?: string, dropFirst?: boolean): DataFrame;

Expand All @@ -1107,7 +1106,11 @@ export interface Series
* }
* ```
*/
toObject(): { name: string; datatype: string; values: any[] };
toObject(): {
name: string;
datatype: DtypeToJsName<T>;
values: DTypeToJs<T>[];
};
toFrame(): DataFrame;
/** compat with `JSON.stringify */
toJSON(): string;
Expand Down Expand Up @@ -1854,6 +1857,7 @@ export interface SeriesConstructor extends Deserialize<Series> {
* ]
* ```
*/
<T extends JsType>(values: ArrayLike<T>): Series<JsToDtype<T>>;
(values: any): Series;
/**
* Create a new named series
Expand All @@ -1871,18 +1875,33 @@ export interface SeriesConstructor extends Deserialize<Series> {
* ]
* ```
*/
<T1 extends JsType>(
name: string,
values: ArrayLike<T1>,
): Series<JsToDtype<T1>>;
<T2 extends DataType>(
name: string,
values: ArrayLike<DTypeToJs<T2>>,
dtype?: T2,
): Series<T2>;
(name: string, values: any[], dtype?): Series;

/**
* Creates an array from an array-like object.
* @param arrayLike — An array-like object to convert to an array.
*/
from<T1 extends JsType>(arrayLike: ArrayLike<T1>): Series<JsToDtype<T1>>;
from<T1>(arrayLike: ArrayLike<T1>): Series;
from<T2 extends JsType>(
name: string,
arrayLike: ArrayLike<T2>,
): Series<JsToDtype<T2>>;
from<T2>(name: string, arrayLike: ArrayLike<T2>): Series;
/**
* Returns a new Series from a set of elements.
* @param items — A set of elements to include in the new Series object.
*/
of<T3 extends JsType>(...items: T3[]): Series<JsToDtype<T3>>;
of<T3>(...items: T3[]): Series;
isSeries(arg: any): arg is Series;
/**
Expand Down

0 comments on commit 95191df

Please sign in to comment.