Skip to content

Commit

Permalink
Add type stub information
Browse files Browse the repository at this point in the history
  • Loading branch information
wence- committed Nov 19, 2024
1 parent f2a905e commit 010c1da
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 1 deletion.
46 changes: 45 additions & 1 deletion python/pylibcudf/pylibcudf/io/parquet.pyi
Original file line number Diff line number Diff line change
@@ -1,7 +1,20 @@
# Copyright (c) 2024, NVIDIA CORPORATION.

from collections.abc import Mapping
from typing import Self

from pylibcudf.expressions import Expression
from pylibcudf.io.types import SourceInfo, TableWithMetadata
from pylibcudf.io.types import (
CompressionType,
DictionaryPolicy,
PartitionInfo,
SinkInfo,
SourceInfo,
StatisticsFreq,
TableInputMetadata,
TableWithMetadata,
)
from pylibcudf.table import Table

class ChunkedParquetReader:
def __init__(
Expand Down Expand Up @@ -34,3 +47,34 @@ def read_parquet(
# reader_column_schema: ReaderColumnSchema = *,
# timestamp_type: DataType = *
) -> TableWithMetadata: ...

class ParquetWriterOptions:
def __init__(self): ...
@staticmethod
def builder(
sink: SinkInfo, table: Table
) -> ParquetWriterOptionsBuilder: ...
def set_partitions(self, partitions: list[PartitionInfo]) -> None: ...
def set_column_chunks_file_paths(self, file_paths: list[str]) -> None: ...
def set_row_group_size_bytes(self, size_bytes: int) -> None: ...
def set_row_group_size_rows(self, size_rows: int) -> None: ...
def set_max_page_size_bytes(self, size_bytes: int) -> None: ...
def set_max_page_size_rows(self, size_rows: int) -> None: ...
def set_max_dictionary_size(self, size_rows: int) -> None: ...

class ParquetWriterOptionsBuilder:
def __init__(self): ...
def metadata(self, metadata: TableInputMetadata) -> Self: ...
def key_value_metadata(
self, metadata: list[Mapping[str, str]]
) -> Self: ...
def compression(self, compression: CompressionType) -> Self: ...
def stats_level(self, sf: StatisticsFreq) -> Self: ...
def int96_timestamps(self, enabled: bool) -> Self: ...
def write_v2_headers(self, enabled: bool) -> Self: ...
def dictionary_policy(self, val: DictionaryPolicy) -> Self: ...
def utc_timestamps(self, enabled: bool) -> Self: ...
def write_arrow_schema(self, enabled: bool) -> Self: ...
def build(self) -> ParquetWriterOptions: ...

def write_parquet(options: ParquetWriterOptions) -> memoryview: ...
6 changes: 6 additions & 0 deletions python/pylibcudf/pylibcudf/io/types.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,12 @@ class QuoteStyle(IntEnum):
ColumnNameSpec: TypeAlias = tuple[str, list[ColumnNameSpec]]
ChildNameSpec: TypeAlias = Mapping[str, ChildNameSpec]

class PartitionInfo:
def __init__(self, start_row: int, num_rows: int): ...

class TableInputMetadata:
def __init__(self, table: Table): ...

class TableWithMetadata:
tbl: Table
def __init__(
Expand Down
2 changes: 2 additions & 0 deletions python/pylibcudf/pylibcudf/io/types.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -42,10 +42,12 @@ __all__ = [
"CompressionType",
"DictionaryPolicy",
"JSONRecoveryMode",
"PartitionInfo",
"QuoteStyle",
"SinkInfo",
"SourceInfo",
"StatisticsFreq",
"TableInputMetadata",
"TableWithMetadata",
]

Expand Down

0 comments on commit 010c1da

Please sign in to comment.