################################################################################
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
################################################################################
"""
Conversion utilities for creating DataFrame from Python objects.
"""
import builtins
from typing import Any, List, Mapping, Optional, Sequence, Tuple, Union, TYPE_CHECKING
from pyflink.dataframe.context import get_or_create_table_environment
from pyflink.dataframe.dataframe import DataFrame
from pyflink.table import Schema
from pyflink.table.types import (
_create_converter,
_create_type_verifier,
_infer_schema_from_data,
DataType,
DataTypes,
from_arrow_type,
LocalZonedTimestampType,
RowField,
RowType,
TimestampType,
)
if TYPE_CHECKING:
from pyflink.table import Table
__all__ = ["from_dict", "from_records", "from_table", "from_pandas", "from_arrow", "range"]
def _validate_watermark(watermark: Optional[Tuple[str, str]]) -> Optional[Tuple[str, str]]:
if watermark is None:
return None
if not isinstance(watermark, tuple) or len(watermark) != 2:
raise TypeError("watermark must be a tuple of (column, expression)")
column_name, watermark_expr = watermark
if (
not isinstance(column_name, str)
or not column_name.strip()
or not isinstance(watermark_expr, str)
or not watermark_expr.strip()
):
raise TypeError("watermark column and expression must be non-empty strings")
return watermark
def _build_table_schema(row_type: RowType,
watermark: Optional[Tuple[str, str]]) -> Optional[Schema]:
watermark = _validate_watermark(watermark)
if watermark is None:
return None
row_type = _normalize_watermark_row_type(row_type, watermark)
builder = Schema.new_builder().from_row_data_type(row_type)
builder.watermark(*watermark)
return builder.build()
def _normalize_watermark_row_type(row_type: RowType, watermark: Tuple[str, str]) -> RowType:
column_name = watermark[0]
fields = []
changed = False
for field in row_type.fields:
data_type = field.data_type
if (
field.name == column_name
and isinstance(data_type, (TimestampType, LocalZonedTimestampType))
and data_type.precision > 3
):
data_type = type(data_type)(3, data_type._nullable)
changed = True
fields.append(RowField(field.name, data_type, field.description))
return RowType(fields, row_type._nullable) if changed else row_type
def _from_rows(
rows: Sequence[Any],
schema: Optional[Union[RowType, DataType, List[str], Tuple[str, ...]]],
watermark: Optional[Tuple[str, str]],
row_type: Optional[RowType] = None,
) -> DataFrame:
t_env = get_or_create_table_environment()
if row_type is None:
if isinstance(schema, RowType):
row_type = schema
verify_func = _create_type_verifier(row_type)
def verify_obj(obj):
verify_func(obj)
return obj
elif isinstance(schema, DataType):
data_type = schema
row_type = RowType().add("value", schema)
verify_func = _create_type_verifier(data_type, name="field value")
def verify_obj(obj):
verify_func(obj)
return obj
else:
def verify_obj(obj):
return obj
if row_type is None:
row_type = _infer_schema_from_data(rows, names=schema)
converter = _create_converter(row_type)
rows = list(map(converter, rows))
elif not isinstance(row_type, RowType):
raise TypeError(
"schema should be RowType, list, tuple or None, but got: %s" % schema)
else:
verify_func = _create_type_verifier(row_type)
def verify_obj(obj):
verify_func(obj)
return obj
rows = list(map(verify_obj, rows))
sql_rows = [row_type.to_sql_type(row) for row in rows]
table_schema = _build_table_schema(row_type, watermark)
return DataFrame(t_env._from_elements(sql_rows, row_type, table_schema))
def _row_type_from_pandas(pdf: Any, schema: Optional[Any]) -> RowType:
import pyarrow as pa
import pyarrow_hotfix # noqa # pylint: disable=unused-import
arrow_schema = pa.Schema.from_pandas(pdf, preserve_index=False)
if schema is not None:
if isinstance(schema, RowType):
return schema
elif isinstance(schema, (list, tuple)) and schema and isinstance(schema[0], str):
return RowType(
[RowField(field_name, from_arrow_type(field.type, field.nullable))
for field_name, field in zip(schema, arrow_schema)])
elif isinstance(schema, (list, tuple)) and schema and isinstance(schema[0], DataType):
return RowType(
[RowField(field_name, field_type) for field_name, field_type
in zip(arrow_schema.names, schema)])
else:
raise TypeError("Unsupported schema type, it could only be of RowType, a "
"list of str or a list of DataType, got %s" % schema)
return RowType([RowField(field.name, from_arrow_type(field.type, field.nullable))
for field in arrow_schema])
def _pandas_scalar_to_python(value: Any) -> Any:
if value is None:
return None
import pandas as pd
try:
if bool(pd.isna(value)):
return None
except (TypeError, ValueError):
pass
to_pydatetime = getattr(value, "to_pydatetime", None)
if to_pydatetime is not None:
return to_pydatetime()
to_pytimedelta = getattr(value, "to_pytimedelta", None)
if to_pytimedelta is not None:
return to_pytimedelta()
item = getattr(value, "item", None)
if item is not None:
try:
return item()
except (AttributeError, ValueError):
return value
return value
def _pandas_rows_to_python(pdf: Any) -> List[Tuple[Any, ...]]:
return [
tuple(_pandas_scalar_to_python(value) for value in row)
for row in pdf.itertuples(index=False, name=None)
]
[docs]def from_table(table: "Table") -> DataFrame:
"""
Create a DataFrame from a PyFlink Table.
Args:
table: The PyFlink Table to wrap.
Returns:
A new DataFrame instance.
Example::
>>> import pyflink.dataframe as pf
>>> table = t_env.from_elements([(1, 'a')], ['id', 'name'])
>>> df = pf.from_table(table)
"""
return DataFrame(table)
[docs]def from_pandas(
pdf: Any,
schema: Optional[List[str]] = None,
watermark: Optional[Tuple[str, str]] = None,
) -> DataFrame:
"""
Create a DataFrame from a Pandas DataFrame.
Args:
pdf: The Pandas DataFrame.
schema: Optional list of column names. If None, uses pandas column names.
watermark: Optional ``(column, expression)`` event-time watermark definition.
Returns:
A new DataFrame instance.
Example::
>>> import pandas as pd
>>> import pyflink.dataframe as pf
>>> pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"]})
>>> df = pf.from_pandas(pdf)
"""
t_env = get_or_create_table_environment()
if watermark is None:
table = t_env.from_pandas(pdf, schema)
return DataFrame(table)
_validate_watermark(watermark)
import pandas as pd
if not isinstance(pdf, pd.DataFrame):
raise TypeError("Unsupported type, expected pandas.DataFrame, got %s" % type(pdf))
row_type = _row_type_from_pandas(pdf, schema)
rows = _pandas_rows_to_python(pdf)
return _from_rows(rows, schema, watermark, row_type)
[docs]def from_arrow(
table: Any,
schema: Optional[List[str]] = None,
watermark: Optional[Tuple[str, str]] = None,
) -> DataFrame:
"""
Create a DataFrame from a PyArrow Table.
Args:
table: The PyArrow Table.
schema: Optional list of column names. If None, uses Arrow column names.
watermark: Optional ``(column, expression)`` event-time watermark definition.
Returns:
A new DataFrame instance.
"""
import pyarrow as pa
import pyarrow_hotfix # noqa # pylint: disable=unused-import
if not isinstance(table, pa.Table):
raise TypeError("Unsupported type, expected pyarrow.Table, got %s" % type(table))
return from_pandas(table.to_pandas(), schema=schema, watermark=watermark)
[docs]def from_dict(
data: Mapping[str, Sequence[Any]],
schema: Optional[List[str]] = None,
watermark: Optional[Tuple[str, str]] = None,
) -> DataFrame:
"""
Create a DataFrame from a dictionary of sequences.
Args:
data: A dictionary where keys are column names and values are
sequences of column values. All sequences must have the same length.
schema: Optional list of column names. If provided, only these columns
will be used, in this order. If None, uses all keys from data.
watermark: Optional ``(column, expression)`` event-time watermark definition.
Returns:
A new DataFrame.
Example::
>>> import pyflink.dataframe as pf
>>> df = pf.from_dict({"a": [1, 2, 3], "b": ["x", "y", "z"]})
"""
t_env = get_or_create_table_environment()
_validate_watermark(watermark)
# Validate that all columns have the same length
if not data:
raise ValueError("data dictionary cannot be empty")
lengths = {key: len(values) for key, values in data.items()}
if len(set(lengths.values())) > 1:
raise ValueError(
f"All columns must have the same length. Got lengths: {lengths}"
)
# Convert to list of rows for from_elements
num_rows = len(next(iter(data.values())))
columns = schema if schema else list(data.keys())
# Validate schema columns exist in data
for col_name in columns:
if col_name not in data:
raise ValueError(f"Column '{col_name}' not found in data")
rows = []
for i in builtins.range(num_rows):
row = tuple(data[col][i] for col in columns)
rows.append(row)
if watermark is not None:
return _from_rows(rows, columns, watermark)
table = t_env.from_elements(rows, columns)
return DataFrame(table)
[docs]def from_records(
data: Sequence[Union[Sequence[Any], Mapping[str, Any]]],
schema: Optional[List[str]] = None,
watermark: Optional[Tuple[str, str]] = None,
) -> DataFrame:
"""
Create a DataFrame from a sequence of records.
Args:
data: A sequence of records. Each record can be:
- A sequence (tuple/list) of values
- A dictionary with column names as keys
schema: Column names. Required when data contains sequences.
Ignored when data contains dictionaries (keys become column names).
watermark: Optional ``(column, expression)`` event-time watermark definition.
Returns:
A new DataFrame.
Example::
>>> import pyflink.dataframe as pf
>>> # From list of dicts
>>> df = pf.from_records([{"a": 1, "b": "x"}, {"a": 2, "b": "y"}])
>>> # From list of tuples with schema
>>> df = pf.from_records([(1, "x"), (2, "y")], schema=["a", "b"])
"""
if not data:
raise ValueError("data cannot be empty")
t_env = get_or_create_table_environment()
_validate_watermark(watermark)
# Check the type of the first record
first_record = data[0]
if isinstance(first_record, Mapping):
# Data is list of dicts
if schema is None:
# Infer schema from keys of the first record
schema = list(first_record.keys())
# Convert dicts to tuples
rows = []
records: Sequence[Any] = data
for record in records:
row = tuple(record.get(col) for col in schema)
rows.append(row)
if watermark is not None:
return _from_rows(rows, schema, watermark)
table = t_env.from_elements(rows, schema)
return DataFrame(table)
elif isinstance(first_record, (list, tuple)):
# Data is list of sequences
if schema is None:
raise ValueError(
"schema is required when data contains sequences (list/tuple). "
"Please provide column names via schema parameter."
)
# Validate each row has the same length as schema
expected_len = len(schema)
for i, record in enumerate(data):
if len(record) != expected_len:
raise ValueError(
f"Record at index {i} has {len(record)} elements, "
f"expected {expected_len} (length of schema)"
)
if watermark is not None:
return _from_rows(data, schema, watermark)
table = t_env.from_elements(data, schema)
return DataFrame(table)
else:
raise TypeError(
f"Unsupported record type: {type(first_record)}. "
"Expected dict, list, or tuple."
)
[docs]def range(start_or_end: int, end: Optional[int] = None, step: int = 1) -> DataFrame:
"""
Create a DataFrame containing a single ``id`` column over an integer range.
"""
if not isinstance(start_or_end, int):
raise TypeError("start_or_end must be an integer")
if end is not None and not isinstance(end, int):
raise TypeError("end must be an integer")
if not isinstance(step, int):
raise TypeError("step must be an integer")
if step == 0:
raise ValueError("step must not be zero")
if end is None:
start = 0
stop = start_or_end
else:
start = start_or_end
stop = end
rows = [(value,) for value in builtins.range(start, stop, step)]
if not rows:
row_type = DataTypes.ROW([DataTypes.FIELD("id", DataTypes.BIGINT())])
return _from_rows(rows, row_type, None, row_type)
return from_records(rows, schema=["id"])